diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index bf7e62e295..7e8d52525c 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -990,6 +990,7 @@ jobs: # First seen on transformers >=5,<6; each represents a slow # or recursive source-rewriter path the zoo can address. "beit": "TimeoutError: compile exceeds per-model budget", + "deepseek_ocr2": "TimeoutError: compile exceeds per-model budget", "sam": "TimeoutError: compile exceeds per-model budget", "sam_hq": "TimeoutError: compile exceeds per-model budget", "deepseek_ocr2": "TimeoutError: compile exceeds per-model budget", diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml index c1297b84cc..3a4c76a2bb 100644 --- a/.github/workflows/studio-inference-smoke.yml +++ b/.github/workflows/studio-inference-smoke.yml @@ -20,7 +20,7 @@ # enable_tools / enabled_tools, and enable_thinking on/off. # # 3. JSON, images -# gemma-4-E2B-it UD-IQ3_XXS (~2.4 GiB) + mmproj-F16 (~986 MiB). +# Qwen3-VL-2B-Instruct UD-IQ2_XXS (~570 MiB) + mmproj-F16 (~780 MiB). # response_format JSON-schema decoding and OpenAI image_url # (data URI) plus Anthropic source/base64 image inputs. # @@ -791,9 +791,9 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 30 env: - GGUF_REPO: unsloth/gemma-4-E2B-it-GGUF - GGUF_VARIANT: UD-IQ3_XXS - GGUF_FILE: gemma-4-E2B-it-UD-IQ3_XXS.gguf + GGUF_REPO: unsloth/Qwen3-VL-2B-Instruct-GGUF + GGUF_VARIANT: UD-IQ2_XXS + GGUF_FILE: Qwen3-VL-2B-Instruct-UD-IQ2_XXS.gguf MMPROJ_FILE: mmproj-F16.gguf STUDIO_PORT: '18890' HF_HOME: ${{ github.workspace }}/hf-cache @@ -888,13 +888,23 @@ jobs: -H 'content-type: application/json' \ -d "{\"username\":\"unsloth\",\"password\":\"$NEW\"}" | jq -r .access_token) echo "API_KEY=$TOKEN" >> "$GITHUB_ENV" - # Load the GGUF (mmproj is auto-detected via the HF repo - # lookup, the cached file is pulled out of HF_HOME). - curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \ - -H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \ - --max-time 900 \ - -d "{\"model_path\":\"$GGUF_REPO\",\"gguf_variant\":\"$GGUF_VARIANT\",\"is_lora\":false,\"max_seq_length\":2048}" \ - | jq '{status, display_name, is_vision}' + # Retry: llama-server startup can race process teardown after a + # failed attempt. Keep curl out of a pipe so HTTP failures are not + # masked by jq. + LOAD_OK=0 + for attempt in 1 2 3; do + HTTP=$(curl -s -o /tmp/load.json -w '%{http_code}' \ + -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \ + -H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \ + --max-time 900 \ + -d "{\"model_path\":\"$GGUF_REPO\",\"gguf_variant\":\"$GGUF_VARIANT\",\"is_lora\":false,\"max_seq_length\":2048}") + if [ "$HTTP" = "200" ]; then LOAD_OK=1; break; fi + echo "::warning::/api/inference/load attempt $attempt returned $HTTP; response:" + cat /tmp/load.json || true + sleep 10 + done + [ "$LOAD_OK" = "1" ] || { echo "::error::/api/inference/load failed 3 attempts"; exit 22; } + jq '{status, display_name, is_vision}' /tmp/load.json - name: JSON schema decoding + image input env: diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index bbe6b9e33a..6148856016 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -13,7 +13,7 @@ # 2. Tool calling Tests # Qwen3.5-2B UD-Q4_K_XL (~890 MiB). # 3. JSON, images -# gemma-4-E2B-it UD-Q4_K_XL + mmproj-F16 (~3.4 GiB total). +# Qwen3-VL-2B-Instruct UD-IQ2_XXS + mmproj-F16 (~1.4 GiB total). # Within the 14 GB windows-latest SSD budget. name: Windows Studio GGUF CI @@ -843,9 +843,9 @@ jobs: run: shell: bash env: - GGUF_REPO: unsloth/gemma-4-E2B-it-GGUF - GGUF_VARIANT: UD-Q4_K_XL - GGUF_FILE: gemma-4-E2B-it-UD-Q4_K_XL.gguf + GGUF_REPO: unsloth/Qwen3-VL-2B-Instruct-GGUF + GGUF_VARIANT: UD-IQ2_XXS + GGUF_FILE: Qwen3-VL-2B-Instruct-UD-IQ2_XXS.gguf MMPROJ_FILE: mmproj-F16.gguf STUDIO_PORT: '18899' HF_HOME: ${{ github.workspace }}/hf-cache @@ -1123,7 +1123,7 @@ jobs: ) data_uri = f"data:image/png;base64,{PNG_64X64_RED_B64}" - # On Windows + the gemma-4-E2B mmproj, llama.cpp's vision + # On Windows + the Qwen3-VL mmproj, llama.cpp's vision # path runs on CPU (no Metal involvement). The wrapper is # kept for resilience but the vision path is expected to # work on Windows; an exception here is a real regression. diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index e018803a5f..56d32c9312 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.15 + rev: v0.15.16 hooks: - id: ruff args: diff --git a/scripts/check_frontend_dep_removal.py b/scripts/check_frontend_dep_removal.py index 3ec4e9037f..b95c4ca7f6 100644 --- a/scripts/check_frontend_dep_removal.py +++ b/scripts/check_frontend_dep_removal.py @@ -43,7 +43,7 @@ DEP_FIELDS = ( "optionalDependencies", ) -# Sources where seeing a package name does NOT count as usage. +# Files where seeing a package name does NOT count as usage. EXPECTED_NOISE_FILES = { "studio/frontend/package.json", "studio/frontend/package-lock.json", @@ -51,17 +51,15 @@ EXPECTED_NOISE_FILES = { "studio/backend/core/data_recipe/oxc-validator/package-lock.json", } -# Only quoted-string occurrences in these file types can be module specifiers. +# File types where a quoted string can be a module specifier. JS_LIKE_EXT = re.compile(r"\.(ts|tsx|js|jsx|mjs|cjs|html|htm|css|scss|sass|json|jsonc)$") -# Files where JS-syntactic import patterns (static/dynamic/require/re-export) -# could be a real module reference. Markdown gets a separate gate (.mdx is +# Files where JS import patterns could be a real module reference (.mdx is # real ESM; .md code fences are not). SCRIPT_LIKE_EXT = re.compile(r"\.(ts|tsx|js|jsx|mjs|cjs|mdx)$") STYLE_EXT = re.compile(r"\.(css|scss|sass)$") HTML_EXT = re.compile(r"\.(html|htm)$") TS_LIKE_EXT = re.compile(r"\.(ts|tsx|mts|cts|mdx)$") -# Files where a removed package's CLI binary could be invoked (npx, bunx, -# yarn dlx, pnpm exec, or a bare `pkg --flag` shell call). +# Files where a removed package's CLI binary could be invoked. COMMAND_LIKE_EXT = re.compile(r"(\.(ya?ml|sh|ps1|bat)$|(^|/)Dockerfile[^/]*$)") GREP_INCLUDES = [ @@ -102,7 +100,7 @@ GREP_EXCLUDES = [ "--exclude-dir=venv", ] -# A pip-installed playwright reference is the PyPI package, not npm. +# A pip-installed playwright ref is the PyPI package, not npm. PIP_PLAYWRIGHT = re.compile( r"(pip\s+install\s+['\"]?playwright" r"|python\s+-m\s+playwright" @@ -153,9 +151,8 @@ def all_decl_names(pkg: dict) -> set[str]: def _resolve_install_path(parent_path: str, name: str, pkgs: dict) -> str | None: - """Walk up the nested node_modules chain from `parent_path` to find - where `name` actually resolves. Mirrors Node module resolution. - """ + """Walk up the nested node_modules chain from `parent_path` to find where + `name` resolves, mirroring Node module resolution.""" parts = parent_path.split("/node_modules/") for i in range(len(parts), 0, -1): prefix = "/node_modules/".join(parts[:i]) @@ -168,11 +165,8 @@ def _resolve_install_path(parent_path: str, name: str, pkgs: dict) -> str | None def _deps_of(meta: dict) -> dict: - """Deps npm actually installs. Optional peers are skipped: npm only - installs them when another package declares the same dep, so for the - purpose of "is this package still reachable" they cannot keep a - removed top-level dep alive on their own. - """ + """Deps npm actually installs. Optional peers are skipped: they can't keep + a removed top-level dep reachable on their own.""" out = {} for field in ("dependencies", "optionalDependencies"): out.update(meta.get(field) or {}) @@ -185,10 +179,8 @@ def _deps_of(meta: dict) -> dict: def reachable_from_head(head_pkg: dict, lock: dict) -> set[str]: - """BFS the lockfile dep graph starting from `head_pkg`'s top-level - declared deps. Returns the set of lockfile install paths that survive. - Stale lockfile entries (orphaned by the new package.json) are excluded. - """ + """BFS the lockfile dep graph from `head_pkg`'s top-level deps. Returns the + surviving install paths, excluding stale (orphaned) lockfile entries.""" pkgs = lock.get("packages", {}) if not pkgs: return set() @@ -215,23 +207,17 @@ def reachable_from_head(head_pkg: dict, lock: dict) -> set[str]: def classify(pkg: str, file: str, content: str) -> str | None: """Return why `content` references `pkg`, or None. - `content` may span multiple lines (for multi-line imports/exports); - each pattern uses re.DOTALL where it matters. The bare-spec - regexes use a word-boundary check on the package name so that - `foobar` does not match `foo`. - - File-type gating: JS-syntactic patterns only fire on .ts/.tsx/.js/.jsx/ - .mjs/.cjs/.mdx files, so an `import x from "pkg"` snippet inside a - Python test fixture or a Markdown code block is not mistaken for a - real npm usage. CSS patterns only fire on .css/.scss/.sass. HTML - patterns only fire on .html/.htm. + `content` may span multiple lines (multi-line imports/exports use re.DOTALL). + Bare-spec regexes word-boundary the package name so `foobar` doesn't match + `foo`. File-type gating restricts JS patterns to .ts/.tsx/.js/.jsx/.mjs/ + .cjs/.mdx, CSS to .css/.scss/.sass, HTML to .html/.htm, so a snippet inside + a Python fixture or Markdown code block isn't mistaken for real npm usage. """ if file in EXPECTED_NOISE_FILES: return None esc = re.escape(pkg) - # Subpath gate: after the package name, the next char must be either - # the closing quote, `/`, or end-of-string. Prevents foo matching foobar. + # Subpath gate: pkg must be followed by quote, `/`, or end-of-string. sub = r"(?:/[^'\"`]*)?" flags_dotall = re.DOTALL | re.MULTILINE @@ -241,30 +227,22 @@ def classify(pkg: str, file: str, content: str) -> str | None: is_html = bool(HTML_EXT.search(file)) is_ts = bool(TS_LIKE_EXT.search(file)) - # If the file is none of script / style / html / json (which is the - # quoted-string fallback surface) and is not an mdx file, no classify - # rule applies. This is what gates out Python fixtures, Markdown code - # blocks, shell snippets, etc. + # Gate out Python fixtures, Markdown code blocks, shell snippets, etc. is_json = file.endswith(".json") or file.endswith(".jsonc") if not (is_script or is_style or is_html or is_json): return None - # CSS @import is checked first so it does not collide with the - # side-effect-import regex below. + # CSS @import first so it doesn't collide with side-effect-import below. if is_style and re.search(rf"@import\s+['\"]{esc}{sub}['\"]", content): return "css_import" - # Static imports: handle multi-line `import { ... } from "pkg"` by - # allowing arbitrary content (newlines included) between `import` - # and `from`. The non-greedy match plus the required `from` keeps - # this scoped to a single statement. + # Static imports, including multi-line `import { ... } from "pkg"`. if is_script and re.search( rf"(? str | None: # require / require.resolve if is_script and re.search(rf"\brequire(?:\.resolve)?\(\s*['\"]{esc}{sub}['\"]\s*\)", content): return "require" - # Re-exports: `export * from "pkg"`, `export { x } from "pkg"`, - # `export type { Foo } from "pkg"`. Multi-line supported. + # Re-exports: `export * from`, `export { x } from`, `export type { Foo } from`. if is_script and re.search( rf"\bexport\b[^;'\"]*?\bfrom\s+['\"]{esc}{sub}['\"]", content, flags_dotall, ): return "re_export" - # HTML script / link. Match the package name as a complete path - # segment bounded by a quote / `#` / `?` or a subpath `/`, so - # `/node_modules/foo-extra/...` is NOT treated as usage of `foo`. + # HTML script / link. Match pkg as a complete path segment so + # `/node_modules/foo-extra/...` is not treated as usage of `foo`. html_pkg = rf"{esc}(?:/[^'\"#?]*)?(?=['\"#?])" if is_html and re.search(rf"]*src\s*=\s*['\"][^'\"]*/{html_pkg}", content): return "html_script" @@ -295,8 +271,7 @@ def classify(pkg: str, file: str, content: str) -> str | None: # new URL("pkg/...", import.meta.url) if is_script and re.search(rf"\bnew\s+URL\(\s*['\"]{esc}{sub}['\"]", content): return "new_url" - # CSS url(...). Accept quoted ("pkg/x") AND unquoted (pkg/x) variants, - # bounded by a path-segment lookahead so `pkg-extra` does not match. + # CSS url(...), quoted and unquoted, bounded so `pkg-extra` doesn't match. if is_style and re.search( rf"\burl\(\s*['\"]?(?:[^)'\"\s]+/)?{esc}(?:/[^)'\"`]*)?['\"]?\s*\)", content, @@ -309,20 +284,18 @@ def classify(pkg: str, file: str, content: str) -> str | None: if is_script and re.search(rf"@import\(\s*['\"]{esc}{sub}['\"]\s*\)", content): return "jsdoc_import" # Bare quoted-string fallback (config plugin lists, vite aliases, - # tsconfig paths, biome config plugin arrays, shadcn registries). + # tsconfig paths, biome plugin arrays, shadcn registries). if not JS_LIKE_EXT.search(file): return None - # Boundary: pkg must be followed by `'`, `"`, or `/` to avoid - # matching `foo` inside `foobar`. + # pkg must be followed by `'`, `"`, or `/` so `foo` doesn't match `foobar`. if re.search(rf"['\"]{esc}(?:['\"]|/)", content): return "string_literal" return None def lockfile_root_sync(head_pkg: dict, head_lock: dict) -> list[str]: - """Return a list of warnings if package-lock.json's dep map - disagrees with package.json (i.e., npm install was not re-run). - """ + """Warn if package-lock.json's dep map disagrees with package.json + (i.e. npm install was not re-run).""" warnings = [] if not head_lock: return warnings @@ -350,10 +323,8 @@ def lockfile_root_sync(head_pkg: dict, head_lock: dict) -> list[str]: def types_orphan_warnings(head_pkg: dict) -> list[str]: - """Flag @types/ deps where is no longer declared anywhere - in package.json. Removing X without also dropping @types/X leaves - dangling type packages. - """ + """Flag @types/ deps where is no longer declared in package.json, + which leaves dangling type packages.""" decl = set() for f in DEP_FIELDS: decl.update((head_pkg.get(f) or {}).keys()) @@ -361,9 +332,7 @@ def types_orphan_warnings(head_pkg: dict) -> list[str]: for name in decl: if not name.startswith("@types/"): continue - # @types/foo provides types for `foo` - # @types/foo-bar provides types for `foo-bar` - # @types/scope__pkg provides types for `@scope/pkg` + # @types/scope__pkg provides types for @scope/pkg. target = name[len("@types/") :] if "__" in target: scope, sub = target.split("__", 1) @@ -386,8 +355,7 @@ _PKG_JSON_SKIP_KEYS = { "bundledDependencies", } -# Top-level fields whose contents are never package references. We walk -# everything else recursively. +# Top-level fields whose contents are never package references. _PKG_JSON_OPAQUE_KEYS = { "browserslist", # browser queries "keywords", # free-form strings @@ -429,20 +397,12 @@ _PKG_JSON_OPAQUE_KEYS = { def package_json_extra_refs(pkg: dict, target: str) -> list[str]: - """Walk every key/value in package.json EXCEPT the dep declaration - blocks, and return citations for string values or dict keys that - equal `target` (or `target/subpath`). + """Walk package.json (except dep declaration blocks) and return citations + for string values or dict keys equal to `target` (or `target/subpath`). - Catches the patterns the public dep-checker tools commonly miss: - - `overrides` / `resolutions` / `pnpm.overrides` keys - - `pnpm.patchedDependencies` keys - - `peerDependenciesMeta` keys - - `prettier`: "@my/prettier-config" - - `eslintConfig.extends`: ["..."] / "..." - - `stylelint.extends` / `stylelint.plugins` - - `babel.presets` / `babel.plugins` - - `jest.preset` / `jest.setupFiles` / `jest.transform` - - `commitlint.extends`, `renovate.extends`, `remarkConfig.plugins` + Catches refs that public dep-checkers commonly miss: overrides/resolutions/ + pnpm.overrides keys, pnpm.patchedDependencies, peerDependenciesMeta, + prettier, eslintConfig.extends, stylelint, babel, jest, commitlint, etc. """ target_sub = target + "/" cites: list[str] = [] @@ -453,14 +413,11 @@ def package_json_extra_refs(pkg: dict, target: str) -> list[str]: def walk(obj: object, path: str) -> None: if isinstance(obj, dict): for k, v in obj.items(): - # Skip top-level dep declaration fields entirely. if path == "" and k in _PKG_JSON_SKIP_KEYS: continue - # Top-level fields whose contents are never package refs. if path == "" and k in _PKG_JSON_OPAQUE_KEYS: continue - # Inside `overrides` / `resolutions` / etc., the KEY itself - # is a package reference. + # Inside overrides/resolutions/etc., the KEY is a package ref. if matches(k): cites.append(f"{path}.{k}" if path else k) walk(v, f"{path}.{k}" if path else k) @@ -476,9 +433,8 @@ def package_json_extra_refs(pkg: dict, target: str) -> list[str]: def build_bin_to_pkg(head_lock: dict) -> dict[str, str]: - """Map a binary name (e.g. 'vite', 'tsc', 'eslint') to the package - that provides it. Built from each lockfile entry's `bin` field. - """ + """Map a binary name (e.g. 'vite', 'eslint') to its providing package, + from each lockfile entry's `bin` field.""" out: dict[str, str] = {} if not head_lock: return out @@ -497,42 +453,29 @@ def build_bin_to_pkg(head_lock: dict) -> dict[str, str]: _SCRIPT_TOKENIZE = re.compile(r"\s*(?:&&|\|\||;|\|(?!\|))\s*") -# Wrappers that delegate to a real CLI in the same shell word list. -# After stripping env prefixes and (optionally) `npx`/`pnpm exec`/`yarn dlx`/ -# `bunx`, if the leading token is one of these we advance past the -# wrapper's own flags and any further env-prefix tokens, then re-check. -# `cross-env` is the common one; `dotenv-cli` / `dotenvx` use `--` as a -# separator. Wrappers that operate on named npm-scripts (concurrently, -# npm-run-all, run-s, run-p, wireit, turbo, nx) intentionally aren't -# here -- they reference script names, not bin names, so the real bin -# is in the *target* script's chunk which we already tokenize. +# Wrappers that delegate to a real CLI in the same shell word list; we skip +# past them and their flags to find the wrapped bin. Script-name wrappers +# (concurrently, npm-run-all, turbo, nx) are excluded: they reference script +# names, so the real bin lives in the target script's chunk we already tokenize. _SCRIPT_WRAPPERS = {"cross-env", "dotenv", "dotenvx", "env-cmd"} _ENV_PREFIX_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=") def _next_real_bin(words: list[str], idx: int) -> str | None: - """Walk `words` from `idx`, peeling env-prefix tokens, the leading - package-manager runner (`npx`, `pnpm exec`, etc.), and the known - wrapper bins. Return the next token that looks like the real CLI - binary, or None if the chunk has nothing to look up. - - Recursion depth is bounded by the chunk's word count, so the loop - cannot run away on a pathological wrapper chain. - """ + """Walk `words` from `idx`, peeling env-prefix tokens, the package-manager + runner (npx, pnpm exec, etc.), and known wrapper bins. Return the next + real CLI binary, or None. Bounded by the chunk's word count.""" seen_wrappers: set[str] = set() while idx < len(words): - # 1. env-prefix run: `FOO=bar BAZ="a b" cmd ...`. shlex has - # already collapsed quoted values into one word, so this - # tokenizer is safe for them. + # 1. env-prefix run `FOO=bar BAZ="a b" cmd ...` (shlex pre-collapsed). while idx < len(words) and _ENV_PREFIX_RE.match(words[idx]): idx += 1 if idx >= len(words): return None first = words[idx] - # 2. Package-manager runner: `npx args`, `pnpm exec `, - # `yarn dlx `, `bunx `. Strip and continue (so the - # wrapped command goes through the same unwrap loop). + # 2. Package-manager runner (npx/pnpm exec/yarn dlx/bunx): strip and + # continue so the wrapped command re-enters the unwrap loop. if first in {"npx", "pnpx", "bunx"} and idx + 1 < len(words): idx += 1 continue @@ -540,21 +483,17 @@ def _next_real_bin(words: list[str], idx: int) -> str | None: idx += 2 continue - # 3. Wrapper bin (cross-env, dotenv, etc.). Skip the wrapper's - # own flags and any subsequent env-prefix tokens, then re-loop. + # 3. Wrapper bin (cross-env, dotenv): skip its flags and env prefixes. bin_token = first.removeprefix("./node_modules/.bin/").removeprefix("node_modules/.bin/") if bin_token in _SCRIPT_WRAPPERS and bin_token not in seen_wrappers: seen_wrappers.add(bin_token) idx += 1 - # cross-env / env-cmd: no flags; just more env-prefix tokens. - # dotenv / dotenvx: skip `-e ` style flags and the - # optional `--` separator before the wrapped command. + # dotenv/dotenvx use `-e ` flags and an optional `--`. while idx < len(words): tok = words[idx] if tok.startswith("-") and tok != "--": idx += 1 - # `-e .env` style: also skip the flag's argument - # when it does not look like another flag. + # `-e .env`: also skip the flag's argument. if ( idx < len(words) and not words[idx].startswith("-") @@ -572,18 +511,12 @@ def _next_real_bin(words: list[str], idx: int) -> str | None: def scripts_bin_refs(head_pkg: dict, bin_to_pkg: dict[str, str]) -> dict[str, list[str]]: - """Return `{package_name: ['scripts.X: cmd', ...]}` listing every - package referenced via its bin name in package.json scripts. + """Return `{package_name: ['scripts.X: cmd', ...]}` for every package + referenced via its bin name in package.json scripts. - Each script value is split on shell separators (`&&`, `||`, `;`, - `|`). Within each chunk, `_next_real_bin()` unwraps env prefixes, - package-manager runners (`npx` / `pnpm exec` / `yarn dlx` / `bunx`), - and wrapper bins like `cross-env` / `dotenv` so that - `cross-env CI=1 biome check` correctly credits `biome` to its - declaring package. - - Tokenization uses shlex.split so quoted env values - (`FOO="a b" biome`) survive unbroken. + Each script is split on shell separators; `_next_real_bin()` unwraps env + prefixes, package-manager runners, and wrapper bins so `cross-env CI=1 + biome check` credits `biome`. Uses shlex.split so quoted env values survive. """ import shlex @@ -599,7 +532,7 @@ def scripts_bin_refs(head_pkg: dict, bin_to_pkg: dict[str, str]) -> dict[str, li try: words = shlex.split(chunk, posix = True) except ValueError: - # Unbalanced quotes -- fall back to plain split. + # Unbalanced quotes: fall back to plain split. words = chunk.split() if not words: continue @@ -613,11 +546,8 @@ def scripts_bin_refs(head_pkg: dict, bin_to_pkg: dict[str, str]) -> dict[str, li def tsconfig_compiler_types_refs() -> set[str]: - """Read studio/frontend/tsconfig*.json and return the set of - package names referenced in compilerOptions.types arrays. These are - implicitly loaded by tsc and count as a real use even though they - have no explicit import. - """ + """Return package names in tsconfig*.json compilerOptions.types arrays. + These are implicitly loaded by tsc and count as real uses.""" out: set[str] = set() base = REPO_ROOT / "studio/frontend" for name in ("tsconfig.json", "tsconfig.app.json", "tsconfig.node.json"): @@ -635,26 +565,16 @@ def tsconfig_compiler_types_refs() -> set[str]: for t in types: if not isinstance(t, str): continue - # `vite/client` resolves to `vite` package. + # `vite/client` resolves to the `vite` package. pkg = t.split("/", 1)[0] if not t.startswith("@") else "/".join(t.split("/", 2)[:2]) out.add(pkg) return out def enumerate_dep_usage(head_pkg: dict, head_lock: dict) -> dict[str, list]: - """For every declared dep, classify whether it appears used. Returns - a dict with these categories: - - used: has at least one detected usage in src/, - config files, scripts.bin, package.json - field refs, or tsconfig types - - unused: no detected usage anywhere - - type_pkg_kept: @types/X where X is still declared - - type_pkg_orphan: @types/X where X is no longer declared - (or X is removed) -- candidate for removal - - Each entry is the package name. The categorisation is opinionated; - `unused` is a CANDIDATE list, not a guarantee. The caller should - verify before deletion. + """For every declared dep, classify usage into a dict of package-name lists: + used, unused, type_pkg_kept (@types/X with X declared), type_pkg_orphan + (@types/X with X gone). `unused` is a CANDIDATE list; verify before deletion. """ decl = all_decl_names(head_pkg) bin_to_pkg = build_bin_to_pkg(head_lock) if head_lock else {} @@ -683,11 +603,8 @@ def enumerate_dep_usage(head_pkg: dict, head_lock: dict) -> dict[str, list]: # Real-source-usage check hits = find_usage(name) used = bool(hits) - # CLI usage in shell / workflow / Dockerfile surfaces. Skip for - # `@types/*` packages because they never expose a CLI binary and - # the unscoped-tail bin name candidate would scan workflow files - # for the bare runtime name (a removed `@types/foo` would look - # for invocations of `foo`). + # CLI usage in shell/workflow/Dockerfile. Skipped for @types/* (no CLI + # binary; the bare-name bin candidate would false-match the runtime). if not used and not name.startswith("@types/") and find_command_usage(name): used = True # Bin scripts @@ -707,28 +624,15 @@ def enumerate_dep_usage(head_pkg: dict, head_lock: dict) -> dict[str, list]: def find_imports_without_decl(head_pkg: dict) -> list[tuple[str, int, str]]: - """Reverse check: find bare-specifier imports in studio/frontend/src - that don't correspond to any declared package.json dep. Catches the - case where someone adds an import but forgets the dep declaration. - Returns (file, line, spec) tuples. - - Match shapes covered: - import "pkg" - import Foo from "pkg" - import { Foo } from "pkg" - import type { Foo } from "pkg" - const x = require("pkg") - const x = await import("pkg") + """Reverse check: find bare-specifier imports in studio/frontend/src with + no matching package.json dep (import added but dep declaration forgotten). + Covers import/require/dynamic-import shapes. Returns (file, line, spec). """ decl = set() for f in DEP_FIELDS: decl.update((head_pkg.get(f) or {}).keys()) - # Also: anything tsconfig path-aliases (just '@/...' here) is internal. - # The capture group is the specifier; the leading alternation accepts - # any of: `from "..."`, bare side-effect `import "..."`, - # `import("..."), or `require("...")`. We exclude relative paths and - # the `@/` alias prefix by requiring the first char of the specifier - # to be neither `.` nor `/`. + # Exclude relative paths and the `@/` alias by requiring the specifier's + # first char to be neither `.` nor `/`. Capture group is the specifier. pattern = ( r"(?:\bfrom\s+|" r"\bimport\s+(?:\(\s*)?|" @@ -754,7 +658,7 @@ def find_imports_without_decl(head_pkg: dict) -> list[tuple[str, int, str]]: file, ln, content = m.group(1), int(m.group(2)), m.group(3) for spec_match in re.finditer(pattern, content): spec = spec_match.group(1) - # Resolve to package name (strip subpath) + # Resolve to package name (strip subpath). if spec.startswith("@"): parts = spec.split("/", 2) pkg_name = "/".join(parts[:2]) if len(parts) >= 2 else spec @@ -762,7 +666,7 @@ def find_imports_without_decl(head_pkg: dict) -> list[tuple[str, int, str]]: pkg_name = spec.split("/", 1)[0] if pkg_name in decl: continue - # Internal aliases like '@/foo' or starts with builtin names + # Internal aliases like '@/foo' or builtin names. if pkg_name == "@": continue if pkg_name in { @@ -807,11 +711,10 @@ def _read_file(path: str) -> list[str]: def find_usage(pkg: str) -> list[Hit]: - """Return real usages of `pkg`. Filters pip-playwright separately. + """Return real usages of `pkg` (pip-playwright filtered separately). - For each filename returned by grep, also feed a multi-line window - around the matching line into classify() so multi-line imports - (`import {\n a\n} from "pkg"`) get picked up. + For each grep hit, also feed a multi-line window into classify() so + multi-line imports get picked up. """ rows = grep_repo(re.escape(pkg)) hits = [] @@ -822,10 +725,8 @@ def find_usage(pkg: str) -> list[Hit]: # Try the single-line classify first. kind = classify(pkg, file, content) if not kind: - # Multi-line window: a generous 25 lines above + the line + - # 25 below so Prettier's one-import-per-line formatting for - # 12-20+ named imports still includes the `import` keyword - # in the same window as the `from "pkg"` clause. + # Multi-line window (25 lines each side) so Prettier's + # one-import-per-line formatting still pairs `import` with `from`. lines = _read_file(file) lo = max(0, lineno - 26) hi = min(len(lines), lineno + 25) @@ -841,28 +742,21 @@ def find_usage(pkg: str) -> list[Hit]: def _candidate_bin_names(pkg: str) -> set[str]: - """Names a removed package's CLI could be invoked under in shell - scripts and workflow files. Most npm CLIs use the package name - (`vite`, `eslint`, `playwright`); scoped CLI packages commonly - expose an unscoped binary name (`@biomejs/biome` -> `biome`). - """ + """Bin names a removed package's CLI could be invoked under. Most npm CLIs + use the package name; scoped ones expose an unscoped bin (@biomejs/biome -> + biome).""" return {pkg, pkg.rsplit("/", 1)[-1]} def find_command_usage(pkg: str) -> list[Hit]: - """Find package CLI invocations in shell / workflow / Dockerfile - surfaces: `npx pkg`, `bunx pkg`, `pnpm exec pkg`, `yarn dlx pkg`, - or a bare `pkg --flag`. Returns Hit("command_bin"). - - Detection is bounded to COMMAND_LIKE_EXT files so a JS string that - happens to contain `npx foo` inside a TS test fixture is not - mistaken for a real invocation. + """Find package CLI invocations in shell/workflow/Dockerfile surfaces (npx, + bunx, pnpm exec, yarn dlx, or bare `pkg --flag`). Bounded to + COMMAND_LIKE_EXT so `npx foo` in a TS fixture isn't mistaken for real use. """ bins = sorted(_candidate_bin_names(pkg), key = len, reverse = True) esc_bins = "|".join(re.escape(b) for b in bins) - # grep ERE pattern (POSIX classes for whitespace/word boundaries). - # Build without f-strings to avoid f-string-vs-{} confusion with the - # POSIX `[[:space:]]` literals and trailing `})}` boundary class. + # grep ERE pattern. Built without f-strings to avoid clashing with the + # POSIX `[[:space:]]` literals. grep_pat = ( r"(^|[[:space:]:;&|(\[])" r"(npx[[:space:]]+|pnpm[[:space:]]+exec[[:space:]]+" @@ -894,10 +788,8 @@ def find_command_usage(pkg: str) -> list[Hit]: def types_target_name(pkg: str) -> str | None: - """Strip `@types/` prefix and decode the npm scope-encoding so the - return value matches the runtime package name. `@types/foo` -> `foo`, - `@types/foo__bar` -> `@foo/bar`. Returns None for non-@types packages. - """ + """Strip `@types/` and decode scope-encoding to the runtime package name + (`@types/foo__bar` -> `@foo/bar`). None for non-@types packages.""" if not pkg.startswith("@types/"): return None target = pkg[len("@types/") :] @@ -908,11 +800,8 @@ def types_target_name(pkg: str) -> str | None: def find_types_runtime_usage(pkg: str, tsc_types: set[str]) -> list[Hit]: - """For a removed `@types/X`, find usages of `X` itself: explicit - `/// `, `tsconfig.compilerOptions.types: ["X"]`, - and runtime `import "X"` shapes. The whole point of `@types/X` is to - type one of those; if any are present, the type package must stay. - """ + """For a removed `@types/X`, find usages of `X` itself (triple-slash + reference, tsconfig types, runtime import). If any exist, @types/X stays.""" target = types_target_name(pkg) if target is None: return [] @@ -997,10 +886,9 @@ def main() -> int: return 2 head_lock = read_pkg_file(head_lock_path) - # Base lockfile is best-effort. We use it only to recover the - # bin -> package mapping for packages the PR is removing -- so a - # `scripts.biome:check` cite still fires when `@biomejs/biome` is - # being dropped and the head lockfile no longer has it. + # Base lockfile is best-effort: only used to recover the bin -> package + # mapping for packages the PR removes, so a scripts.biome cite still fires + # when @biomejs/biome is dropped from the head lockfile. if args.base_lock: base_lock_path = Path(args.base_lock) base_lock = read_pkg_file(base_lock_path) if base_lock_path.exists() else {} @@ -1011,9 +899,8 @@ def main() -> int: head_names = all_decl_names(head_pkg) removed = sorted(base_names - head_names) - # All hygiene checks compute up front so they can run on both the - # removal-present and removal-empty paths (so `--strict` actually - # fails when only hygiene issues exist). + # Hygiene checks compute up front so they run on both the removal-present + # and removal-empty paths (so --strict fails on hygiene-only issues). sync_warns = lockfile_root_sync(head_pkg, head_lock) types_warns = types_orphan_warnings(head_pkg) missing_imports = find_imports_without_decl(head_pkg) @@ -1074,12 +961,9 @@ def main() -> int: print() reachable_paths = reachable_from_head(head_pkg, head_lock) if head_lock else set() - # bin -> package map: start from the head lockfile, then layer the - # base lockfile's entries on top for packages this PR is removing. - # A correct removal updates the head lockfile to drop node_modules/foo, - # so build_bin_to_pkg(head_lock) loses the mapping; we recover it - # from the base lockfile so `scripts.biome:check` still flags as a - # usage when `@biomejs/biome` is being dropped. + # bin -> package map from the head lockfile, layering base-lockfile entries + # for removed packages so scripts.biome still flags when @biomejs/biome is + # dropped (head lockfile no longer maps it). bin_to_pkg = build_bin_to_pkg(head_lock) if head_lock else {} base_bin_to_pkg = build_bin_to_pkg(base_lock) if base_lock else {} removed_set = set(removed) @@ -1091,9 +975,8 @@ def main() -> int: def reachable_install_paths(name: str) -> tuple[str | None, list[str]]: """Return (top_level_path, nested_paths). top_level is what bare - `import "name"` from src/ actually resolves to; nested copies are - only visible inside the parent package that nested them. - """ + `import "name"` resolves to; nested copies are only visible inside + their parent package.""" top = f"node_modules/{name}" top_path = top if top in reachable_paths else None nested = sorted( @@ -1106,8 +989,7 @@ def main() -> int: hits = find_usage(name) # CLI invocations in shell scripts / workflows / Dockerfiles. hits.extend(find_command_usage(name)) - # @types/X is "used" if X is referenced as a type or as a - # runtime import elsewhere in the repo. + # @types/X is "used" if X is referenced as a type or runtime import. hits.extend(find_types_runtime_usage(name, tsc_types)) for cite in script_refs.get(name, []): hits.append(Hit("studio/frontend/package.json", 0, "script_bin", cite)) @@ -1115,9 +997,8 @@ def main() -> int: hits.append(Hit("studio/frontend/package.json", 0, "pkg_json_field", cite)) top, nested = reachable_install_paths(name) importable_top_level = top is not None - # Source imports of bare specifier `name` resolve ONLY to top-level - # node_modules/. Nested copies under another package are - # invisible to src/ files. + # Bare specifier `name` resolves ONLY to top-level node_modules/; + # nested copies are invisible to src/ files. if hits and not importable_top_level: status = "FAIL" elif hits and importable_top_level: diff --git a/scripts/check_new_install_scripts.py b/scripts/check_new_install_scripts.py index 2beaada0bf..604d9b9f90 100644 --- a/scripts/check_new_install_scripts.py +++ b/scripts/check_new_install_scripts.py @@ -4,33 +4,18 @@ """Diff two `package-lock.json` files and flag NEW install-script deps. -A package with `"hasInstallScript": true` runs `preinstall` / `install` / -`postinstall` lifecycle hooks every time `npm ci` lays it down. Every -npm supply-chain compromise of the last 18 months (Shai-Hulud, -TanStack, axios-style, ArmorCode hijacks) leveraged exactly this lever: -the attacker publishes a new malicious version of a dep we already -trust, and the post-install hook runs the next time CI installs. +A `"hasInstallScript": true` package runs preinstall/install/postinstall +hooks on every `npm ci` -- the lever behind recent npm supply-chain +compromises (attacker publishes a malicious version of a trusted dep). +This refuses to land a newly-introduced install-script dep without a +maintainer eyeball; pre-existing ones are not re-flagged. -This scanner refuses to allow a newly-introduced install-script dep to -land without a maintainer eyeball on the lifecycle script body. -Existing install-script deps are NOT re-flagged -- if `node-gyp` has -been in the lockfile since day one, it's not part of this PR's threat -model. Only new entries are surfaced. +Supports lockfileVersion 1 (recursive `dependencies`) and 2/3 (flat +`packages` with `node_modules/.../node_modules/...` nesting). For each +new entry we best-effort fetch the registry metadata to recover the +postinstall command body; the finding is still emitted if unreachable. -Supports lockfileVersion 1 (`dependencies` key, recursive), 2 and 3 -(flat `packages` key with `node_modules//node_modules/` nesting -for transitive entries). For each NEW install-script package we -attempt a stdlib-only fetch of -`https://registry.npmjs.org//` to recover the actual -postinstall command body. If the network is blocked we still emit the -finding -- the lifecycle command body is informational, not -load-bearing. - -Exit codes -========== - 0 no newly-added install-script deps - 1 one or more newly-added install-script deps; listed on stderr - 2 internal error (missing lockfile, malformed JSON, etc.) +Exit codes: 0 = none; 1 = one or more (on stderr); 2 = internal error. """ from __future__ import annotations @@ -68,21 +53,14 @@ class Finding: ) -# ───────────────────────────────────────────────────────────────────── # Lockfile parsing. -# ───────────────────────────────────────────────────────────────────── def _strip_nm_prefix(key: str) -> str: - """Convert a v2/v3 `packages` key into a bare package name. - - `node_modules/foo` -> `foo`; `node_modules/foo/node_modules/bar` -> - `bar`. The empty key (`""`) is the project root and returns "". - """ + """Convert a v2/v3 `packages` key into a bare package name (leaf after last `node_modules/`).""" if not key: return "" - # Use the LAST `node_modules/` segment so transitives map to their - # leaf name, matching how npm install resolves a postinstall. + # LAST node_modules/ segment so transitives map to their leaf name. marker = "node_modules/" idx = key.rfind(marker) if idx == -1: @@ -91,14 +69,9 @@ def _strip_nm_prefix(key: str) -> str: def _collect_install_script_entries(lock: dict) -> dict[str, str]: - """Walk a parsed lockfile and return {package_name: version} for - every entry with `hasInstallScript: true` (v2/v3) OR a - non-empty `scripts.preinstall|install|postinstall` (v1). + """Return {name@version: name} for entries with hasInstallScript (v2/v3) or a lifecycle script (v1). - The same package may appear at multiple versions in a single - lockfile (de-duplicated copies under different parents); we key by - `name@version` so we don't lose either copy. Returns a dict keyed - by `name@version` -> the same string for convenience. + Keyed by name@version so dup copies at different versions aren't lost. """ seen: dict[str, str] = {} version = lock.get("lockfileVersion") @@ -118,10 +91,7 @@ def _collect_install_script_entries(lock: dict) -> dict[str, str]: ver = entry.get("version") or "" seen[f"{name}@{ver}"] = name - # v1 also embeds a `dependencies` tree; v2/v3 carry both for - # backwards-compat but `packages` is canonical for them. For v1 - # there is no `hasInstallScript` flag, so look for a non-empty - # `scripts.preinstall|install|postinstall` directly. + # v1 has no hasInstallScript flag; detect lifecycle scripts directly. def _walk_v1(deps: dict, depth: int = 0) -> None: if depth > 64 or not isinstance(deps, dict): return @@ -133,8 +103,6 @@ def _collect_install_script_entries(lock: dict) -> dict[str, str]: isinstance(scripts, dict) and scripts.get(hook) for hook in ("preinstall", "install", "postinstall") ) - # v1 also sets `requires` only on the parent, no flag, so - # the lifecycle-script presence is the only signal. if lifecycle: ver = entry.get("version") or "" seen[f"{name}@{ver}"] = name @@ -155,19 +123,11 @@ def _load_lockfile(path: Path) -> dict: raise ValueError(f"{path}: not valid JSON: {exc}") from exc -# ───────────────────────────────────────────────────────────────────── # Registry lookup for the postinstall command body (best-effort). -# ───────────────────────────────────────────────────────────────────── def _fetch_registry_scripts(name: str, version: str) -> dict[str, str] | None: - """Return {hook: command} for any of preinstall / install / - postinstall published in the registry metadata for this name@ver. - - Returns None on any error (network blocked, 404, malformed JSON). - Never raises; the caller treats absence as "could not enrich, emit - finding anyway". - """ + """Return {hook: command} for lifecycle hooks in registry metadata; None on any error (never raises).""" safe_name = urllib.parse.quote(name, safe = "@/") url = f"{REGISTRY_BASE}{safe_name}/{urllib.parse.quote(version)}" try: @@ -190,9 +150,7 @@ def _fetch_registry_scripts(name: str, version: str) -> dict[str, str] | None: return keep or None -# ───────────────────────────────────────────────────────────────────── # Diff. -# ───────────────────────────────────────────────────────────────────── def diff_new_install_scripts(base_lock: dict, head_lock: dict) -> list[Finding]: @@ -203,7 +161,6 @@ def diff_new_install_scripts(base_lock: dict, head_lock: dict) -> list[Finding]: if key in base: continue # pre-existing install-script dep; not in scope name = head[key] - # key is "name@version"; rsplit("@", 1) handles scoped names. version = key[len(name) + 1 :] if key.startswith(name + "@") else "" scripts = _fetch_registry_scripts(name, version) if scripts: @@ -226,9 +183,7 @@ def diff_new_install_scripts(base_lock: dict, head_lock: dict) -> list[Finding]: return findings -# ───────────────────────────────────────────────────────────────────── # CLI. -# ───────────────────────────────────────────────────────────────────── def main(argv: list[str] | None = None) -> int: diff --git a/scripts/enforce_kwargs_spacing.py b/scripts/enforce_kwargs_spacing.py index fdef950d7f..dc4a6d6821 100755 --- a/scripts/enforce_kwargs_spacing.py +++ b/scripts/enforce_kwargs_spacing.py @@ -20,13 +20,8 @@ from pathlib import Path def _atomic_write_text(path: Path, data: str, encoding: str) -> None: - """Write ``data`` to ``path`` atomically. - - Stages a tmp file in the same directory (so it's on the same - filesystem as the destination), fsyncs, then `os.replace`s into - place. A crash mid-write therefore leaves either the previous - content or the fully new content -- never a truncated source file. - """ + """Write ``data`` to ``path`` atomically via same-dir tmp + fsync + os.replace, + so a crash mid-write leaves either the old or full new content, never a truncation.""" dirpath = str(path.parent) or "." fd, tmp_path = tempfile.mkstemp(prefix=".kwargs_fix.", dir=dirpath) try: @@ -142,7 +137,7 @@ def remove_redundant_passes(text: str) -> tuple[str, bool]: lines[start] = segment if segment.strip() else "" continue - # Defensive fall-back for unexpected multi-line 'pass'. + # Fall-back for unexpected multi-line 'pass'. prefix = lines[start][: node.col_offset] lines[start] = prefix if prefix.strip() else "" for idx in range(start + 1, end): @@ -164,15 +159,12 @@ def remove_redundant_passes(text: str) -> tuple[str, bool]: def remove_blank_after_short_import(text: str) -> tuple[str, bool]: - """Drop blank line(s) after an import block in a *small* nested suite. + """Drop blank line(s) after an import block in a small nested suite. - Inside an indented suite of <= 3 statements (function/try/if/with/etc., never - module level), when a run of consecutive ``import`` / ``from ... import`` - statements is directly followed -- across one or more blank lines and nothing - else -- by another statement in the same suite, remove those blank lines so - the import sits next to the code that uses it. A comment in the gap blocks the - rule. Removing blank lines never changes the AST, so this is always - semantics-preserving. + In an indented suite of <= 3 statements (never module level), when consecutive + imports are followed across blank lines (nothing else) by another statement, + remove those blanks. A comment in the gap blocks the rule. Removing blank lines + never changes the AST. """ try: tree = ast.parse(text) @@ -226,12 +218,11 @@ _DEF_MIN_PARAMS_FOR_MULTILINE = 3 # signatures with < this many params stay one def _def_specs_by_line(tree: ast.AST) -> dict[int, tuple[int, bool]]: - """Map the line of each def keyword to (param count, has-any-default). + """Map each def keyword line to (param count, has-any-default). - One def per line, so the line is a stable key. ``*`` / ``/`` markers are not - parameters and are not counted. A default exists if any positional default - is present or any keyword-only default is not ``None`` (a ``None`` entry in - ``kw_defaults`` means a required keyword-only arg). + ``*`` / ``/`` markers aren't counted. A default exists if any positional default + is present or any keyword-only default is not ``None`` (``None`` in ``kw_defaults`` + means a required keyword-only arg). """ out: dict[int, tuple[int, bool]] = {} for node in ast.walk(tree): @@ -250,20 +241,12 @@ def _def_specs_by_line(tree: ast.AST) -> dict[int, tuple[int, bool]]: def normalize_def_trailing_comma(text: str) -> tuple[str, bool]: - """Force a def / async-def signature one-per-line iff it has >= 3 parameters - AND at least one default value; otherwise keep it collapsible. + """Force a def signature one-per-line iff >= 3 params AND a default; else collapsible. - Rationale: signatures with defaults read better one parameter per line, but - only once they are non-trivial (< 3 params always stay on one line). A - signature with >= 3 params and a default gets a magic trailing comma added - (ruff then wraps it one-per-line regardless of length); every other - signature has its trailing comma stripped so ruff collapses it onto one line - when it fits (and wraps a genuinely long one by length alone). - - Function-definition parameter lists only, never call sites or collection - literals. Parameter counts and defaults come from the AST. Run BEFORE ruff - format. Adding or removing a def trailing comma never changes the AST, which - is re-checked before returning. + A qualifying signature gets a magic trailing comma added (ruff wraps it + one-per-line); every other signature has its trailing comma stripped so ruff + collapses it when it fits. Def parameter lists only, never call sites or + collection literals. Run BEFORE ruff format. Never changes the AST (re-checked). """ try: tree = ast.parse(text) @@ -332,12 +315,10 @@ def normalize_def_trailing_comma(text: str) -> tuple[str, bool]: def _split_string_token(s: str) -> tuple[str, str, str] | None: - """Split a string literal's source into (prefix, quote, body). + """Split a string literal source into (prefix, quote, body). - ``prefix`` is the letters before the opening quote (``r``/``f``/``b``/``u`` - in any case/order), ``quote`` is the opening delimiter (``'``, ``"``, - ``'''`` or ``\"\"\"``) and ``body`` is everything between the delimiters. - Returns ``None`` if ``s`` is not a recognizable string literal. + ``prefix`` is the letters before the opening quote, ``quote`` the delimiter, + ``body`` everything between. ``None`` if not a recognizable string literal. """ i = 0 while i < len(s) and s[i] not in ("'", '"'): @@ -393,15 +374,12 @@ def _string_pieces( def _merge_string_run(pieces: list[tuple[str, str]]) -> str | None: """Merge a run of adjacent string pieces into one literal's source text. - ``pieces`` is a list of ``(kind, raw_source)`` where kind is ``"str"`` or - ``"f"``. Rules: bytes are left side-by-side (return ``None``); a run with no - f-string merges plain/raw/unicode pieces sharing one prefix+quote by simple - body concatenation; a run mixing an f-string with at least one plain string - (and no bytes, no raw) folds into a single f-string -- f pieces keep their - bodies verbatim and plain pieces have their braces escaped (``{`` -> ``{{``). - Runs of only f-strings are left side-by-side. The caller re-checks the file - AST and drops the change if it differs, so any subtle case (e.g. ``\\N{...}``) - that this would mis-handle is caught and skipped. + ``pieces`` is ``(kind, raw_source)`` with kind ``"str"`` or ``"f"``. Bytes are + left side-by-side (``None``); a run with no f-string merges plain/raw/unicode + sharing one prefix+quote by body concatenation; a run mixing an f-string with a + plain string (no bytes, no raw) folds into one f-string with plain braces escaped. + Runs of only f-strings are left alone. Caller re-checks the AST and drops a + differing change, so subtle cases are caught. """ parsed = [] for kind, raw in pieces: @@ -455,13 +433,10 @@ def _fold_collapses( ) -> bool: """Whether an f-string fold at ``row[c0:c1]`` -> ``merged`` is safe to apply. - Only ``assert`` statements wrap awkwardly when a message is folded: ruff - parenthesizes the *condition* once ``assert cond, msg`` no longer fits on one - line. For every other construct (call argument, ``raise``, assignment, ...) a - folded long message wraps acceptably, so the fold is always allowed. For an - ``assert`` the fold is allowed only when the statement is already one physical - line, or its estimated one-line length after folding fits the line length; - otherwise the message is left side-by-side. + Only ``assert`` wraps awkwardly when a message folds (ruff parenthesizes the + condition once it no longer fits one line); every other construct wraps + acceptably so is always allowed. An ``assert`` fold is allowed only if already + one line, or its estimated folded one-line length fits the line length. """ stmt = _enclosing_stmt(tree, row) if not isinstance(stmt, ast.Assert): @@ -483,15 +458,12 @@ def _fold_collapses( def merge_adjacent_string_literals(text: str) -> tuple[str, bool]: - """Merge a run of adjacent string literals on ONE physical line into a single - literal (the ``"a" "b"`` form ruff emits when it collapses an implicit - concatenation). Plain/raw/unicode runs merge by concatenation; a run mixing - an f-string with a plain string folds into one f-string (plain parts' braces - escaped) -- but only when the statement still fits on one line, so a long - message is left side-by-side rather than forcing the statement to re-wrap. - Runs of only f-strings, and bytes, are left side-by-side. The whole file's AST - is re-checked and the change is dropped if it would differ, so the transform - can never change meaning. + """Merge adjacent string literals on ONE physical line into a single literal. + + Plain/raw/unicode runs merge by concatenation; an f-string + plain string folds + into one f-string (plain braces escaped) only while the statement still fits one + line. Runs of only f-strings, and bytes, are left side-by-side. The file AST is + re-checked and a differing change dropped, so meaning never changes. """ try: toks = list(tokenize.generate_tokens(io.StringIO(text).readline)) @@ -552,17 +524,11 @@ def merge_adjacent_string_literals(text: str) -> tuple[str, bool]: def collapse_short_asserts(text: str) -> tuple[str, bool]: """Collapse a multi-line ``assert`` onto one line when it would fit. - An ``assert`` is often kept multi-line only by a magic trailing comma inside - a collection / call / tuple-message (``assert x == {\"a\": 1,}`` written - across lines). When the whole statement's estimated one-line length fits the - line length, strip those trailing commas (the comma before a ``)`` / ``]`` / - ``}``) so ruff joins it back onto one line on the following format pass. - - Run BEFORE ruff format. Skips any assert that contains a comment (a comment - forces ruff to keep it multi-line, which would oscillate). Stripping a - trailing comma is non-semantic except for a one-element tuple ``(x,)``; the - file AST is re-checked and any assert whose strip would change it is left - alone. + When the statement's estimated one-line length fits, strip the magic trailing + commas (comma before a closer) holding it open so ruff rejoins it. Run BEFORE + ruff format. Skips asserts with a comment (would oscillate). Stripping is + non-semantic except for a one-element tuple; AST is re-checked and changing + asserts left alone. """ try: tree = ast.parse(text) diff --git a/scripts/lint_workflow_triggers.py b/scripts/lint_workflow_triggers.py index 34868b7d68..0688f6c65c 100644 --- a/scripts/lint_workflow_triggers.py +++ b/scripts/lint_workflow_triggers.py @@ -4,37 +4,19 @@ """Refuse dangerous GitHub Actions trigger patterns at PR time. -Two patterns are banned outright, both of which powered the TanStack -GHSA-g7cv-rxg3-hmpx supply-chain compromise: +Bans patterns behind the TanStack GHSA-g7cv-rxg3-hmpx compromise: -1. `pull_request_target` -- runs a fork's workflow YAML against the - BASE repository's secrets and permissions. The fork can inject - arbitrary code into the base context. The TanStack worm used this - to land base-context execution from a fork PR. There is essentially - no safe use of this trigger for a public open-source project; - `pull_request` is the safe alternative. +1. `pull_request_target` -- runs a fork's workflow against the base + repo's secrets/permissions; use `pull_request` instead. +2. `workflow_run` chained to a PR-triggered workflow -- same trust + boundary problem one hop later (poisoned artifacts/caches run with + elevated permissions). +3. Cache keys shared between PR-triggered and publish/release/push + workflows -- a fork PR could poison a cache the publish workflow + restores. Partition the key namespaces. -2. `workflow_run` chained to a PR-triggered workflow -- carries the - same trust boundary problem one hop later. If a PR-triggered - workflow can poison artifacts/caches and a `workflow_run` trigger - fires off the result with elevated permissions, the attacker still - reaches the trusted context. - -3. Shared cache keys between PR-triggered workflows and publish / - release / push-triggered workflows. The TanStack worm poisoned the - Actions cache from a fork PR and the legitimate release workflow - then restored the poisoned cache. Cache keys must be partitioned - so that nothing a PR can write is ever read by a workflow that - holds secrets. - -Exit codes -========== - - 0 no findings - 1 one or more findings; stderr lists each with file path - -Run from repo root: - python3 scripts/lint_workflow_triggers.py +Exit codes: 0 = no findings, 1 = findings (listed on stderr). +Run from repo root: python3 scripts/lint_workflow_triggers.py """ from __future__ import annotations diff --git a/scripts/lockfile_supply_chain_audit.py b/scripts/lockfile_supply_chain_audit.py index 5b503dcffe..66b48c094d 100644 --- a/scripts/lockfile_supply_chain_audit.py +++ b/scripts/lockfile_supply_chain_audit.py @@ -5,64 +5,21 @@ """Lockfile supply-chain audit for the Studio frontend and Tauri shell. Runs BEFORE `npm ci` / `cargo fetch` in CI. Refuses to proceed when a -lockfile contains patterns that indicate the kind of supply-chain -injection seen in the npm Shai-Hulud waves and the cargo -crates.io brand-squat attempts. +lockfile contains patterns indicating supply-chain injection (npm +Shai-Hulud waves, cargo crates.io brand-squats). -What it checks -============== +Checks package-lock.json (lockfileVersion 2/3): `resolved` URL must be +the npm registry (direct git/github/file refs are the injection vector); +`integrity` SHA must be present; known IOC substrings grepped from the +body. Checks Cargo.lock: `source` must be the crates.io registry index; +known cargo IOC substrings. -studio/frontend/package-lock.json (lockfileVersion 2 or 3): +Exit codes: 0 = clean (or skip env var set to a justification >=5 chars, +not '1'/'true'); 1 = findings; 2 = internal error. - 1. `resolved` URL origin. Every entry must resolve through - `https://registry.npmjs.org/`. Direct GitHub-hosted dependencies - (`git+ssh://`, `git+https://`, `github:owner/repo#sha`, - `file:`, `http://`) are refused -- npm's TanStack incident used - exactly this vector to land an unaudited GitHub commit hash as - an optional dependency. - - 2. `integrity` field presence. Every non-workspace entry must carry - an `integrity` SHA. A missing integrity means the registry can - swap the tarball after lockfile generation and CI will not - notice. - - 3. Known IOC strings. A hardcoded set of indicator-of-compromise - substrings is grepped across the entire lockfile body (file - names, dependency keys, URLs). The list is updated as new - campaigns surface. Catching one means the local install was - about to pull a publicly-known malicious release. - -studio/src-tauri/Cargo.lock: - - 4. `source` field origin. Every entry with a `source` must point at - `registry+https://github.com/rust-lang/crates.io-index`. Direct - git sources (`git+https://...`) and `path+...` for cross-crate - paths warrant manual review and are flagged. - - 5. Known cargo IOC strings. Same idea as (3), separate list. - -Exit codes -========== - - 0 no findings, or an opt-out env var (UNSLOTH_LOCKFILE_AUDIT_SKIP) - is set to a justification string (>=5 chars, not '1'/'true'/etc). - A value like '1' or 'true' is now REJECTED loudly and the audit - runs normally - 1 one or more findings; stderr lists them with file path and line - number where derivable - 2 internal error (missing dependency, malformed JSON, etc.) - -Operational stance -================== - -This scanner only PARSES the lockfiles -- it never executes anything -in them, never resolves anything against the network. Safe to run -ahead of every `npm ci`. The IOC list is short by design; this -complements (not replaces) `npm audit`, OSV-Scanner, and the -advisory-DB pipeline in `.github/workflows/security-audit.yml`. The -shape of the catch is "we refuse to proceed because the lockfile -itself is shaped wrong", which fires before any third-party install -script gets a chance to run on the runner. +Only PARSES the lockfiles, never executes or networks. Complements (not +replaces) `npm audit` / OSV-Scanner / the advisory-DB pipeline. Fires +before any third-party install script runs on the runner. """ from __future__ import annotations @@ -77,14 +34,9 @@ from pathlib import Path REPO_ROOT = Path(__file__).resolve().parents[1] -# ───────────────────────────────────────────────────────────────────── -# Known IOC strings (case-sensitive substring match). -# ───────────────────────────────────────────────────────────────────── -# -# Keep these short and FACTUAL. Each entry is tied to a public advisory -# and is the literal string an attacker would have to embed for the -# attack to work. Adding speculative or generic patterns here would -# generate false positives on dependency upgrades. +# Known IOC strings (case-sensitive substring match). Each is tied to a +# public advisory; speculative/generic patterns would false-positive on +# upgrades. NPM_IOC_STRINGS: tuple[str, ...] = ( # Shai-Hulud TanStack wave -- May 11, 2026 (GHSA-g7cv-rxg3-hmpx). "router_init.js", @@ -328,36 +280,22 @@ BLOCKED_NPM_VERSIONS: dict[str, set[str]] = { } CARGO_IOC_STRINGS: tuple[str, ...] = ( - # Reserved for future cargo-side incidents. Empty by default -- - # `source` origin check below catches the structural pattern. + # Empty by default; the `source` origin check catches the structural + # pattern. Reserved for future cargo-side incidents. ) -# ───────────────────────────────────────────────────────────────────── # Allowed lockfile origins. -# ───────────────────────────────────────────────────────────────────── NPM_REGISTRY_PREFIX = "https://registry.npmjs.org/" - -# Tarballs are also fetched from this mirror on some GH Actions cached -# runs (npm rewrites the resolved URL on cache hit). Allow either. NPM_REGISTRY_PREFIXES_ALLOWED: tuple[str, ...] = (NPM_REGISTRY_PREFIX,) CARGO_REGISTRY_SOURCE = "registry+https://github.com/rust-lang/crates.io-index" -# ───────────────────────────────────────────────────────────────────── -# Cargo non-registry source allowlist. -# ───────────────────────────────────────────────────────────────────── -# -# Each entry is `(crate_name, exact_source_string)`. The crate must -# match by name AND the source must match the full pinned-SHA string -# verbatim. Bumping the commit SHA forces a re-review here: the -# scanner fires until the new SHA is appended. -# -# Studio's Tauri shell pulls `fix-path-env` directly from -# tauri-apps/fix-path-env-rs because the crate is not published to -# crates.io. The pinned commit (c4c45d5) was reviewed at the time it -# landed; future bumps need explicit approval. +# Cargo non-registry source allowlist: `(crate_name, exact_source_string)`. +# Both must match verbatim; bumping the pinned SHA forces a re-review. +# Studio's Tauri shell pulls `fix-path-env` from git because it is not +# published to crates.io; commit c4c45d5 was reviewed when it landed. CARGO_SOURCE_ALLOWLIST: tuple[tuple[str, str], ...] = ( ( "fix-path-env", @@ -367,11 +305,6 @@ CARGO_SOURCE_ALLOWLIST: tuple[tuple[str, str], ...] = ( ) -# ───────────────────────────────────────────────────────────────────── -# Finding container. -# ───────────────────────────────────────────────────────────────────── - - class Finding: __slots__ = ("path", "package", "kind", "detail") @@ -390,27 +323,19 @@ class Finding: def _gha_escape(text: str) -> str: - """Escape a string for use in a GitHub Actions `::warning::` / - `::error::` workflow command message. GH Actions truncates - annotation messages at the first newline unless `\\n` is - escaped as `%0A`; carriage returns and the percent sign need - matching escapes per the workflow-commands spec. Order matters: - `%` must be replaced first so the subsequent `%0A` / `%0D` - sequences are not double-encoded. + """Escape a string for a GH Actions `::warning::`/`::error::` message. + + GH Actions truncates at the first newline unless `\\n`/`\\r` are + escaped as `%0A`/`%0D`. `%` must be replaced first to avoid + double-encoding the subsequent escapes. """ return text.replace("%", "%25").replace("\r", "%0D").replace("\n", "%0A") -# ───────────────────────────────────────────────────────────────────── -# package-lock.json audit. -# ───────────────────────────────────────────────────────────────────── - - def audit_npm_lockfile(path: Path) -> list[Finding]: findings: list[Finding] = [] if not path.exists(): - # A missing requested lockfile is a config error, not a clean - # audit; surface it so a deleted default cannot pass silently. + # Missing lockfile is a config error, not a clean audit. findings.append( Finding( path = str(path), @@ -427,8 +352,7 @@ def audit_npm_lockfile(path: Path) -> list[Finding]: try: raw = path.read_text(encoding = "utf-8") except OSError as exc: - # Permission denied, is-a-directory, broken-pipe etc. -- surface - # as a finding instead of crashing CI with a raw traceback. + # Surface as a finding instead of crashing CI with a traceback. findings.append( Finding( path = str(path), @@ -464,9 +388,7 @@ def audit_npm_lockfile(path: Path) -> list[Finding]: packages = lock.get("packages") or {} for key, entry in packages.items(): - # The empty key "" is the project root; workspace entries use - # keys like "node_modules/foo" or "studio/frontend/sub-pkg". - # Skip the project root (it has no `resolved`). + # Empty key "" is the project root (no `resolved`); skip it. if key == "": continue if entry.get("link"): @@ -474,12 +396,8 @@ def audit_npm_lockfile(path: Path) -> list[Finding]: continue resolved = entry.get("resolved") - # Entries living inside another package's `node_modules/` - # tree are bundled fold-ins -- the parent's tarball ships - # their source verbatim and the parent's `integrity` covers - # the whole subtree. npm represents them in lockfileVersion 3 - # as nested entries with no `resolved` and no `integrity` of - # their own. Treat them as transparent to this audit. + # Entries nested in another package's node_modules are bundled + # fold-ins covered by the parent's integrity; treat as transparent. nested = key.count("/node_modules/") >= 1 # 1. resolved-URL origin. @@ -545,12 +463,10 @@ def audit_npm_lockfile(path: Path) -> list[Finding]: ) ) - # 4. Known IOC strings: scan the raw file body so we hit fields the - # structural pass above doesn't enumerate (scripts, optional - # dependencies, etc.). Cheap and complete. + # 4. Known IOC strings: scan the raw body to catch fields the + # structural pass doesn't enumerate (scripts, optional deps, etc.). for ioc in NPM_IOC_STRINGS: if ioc in raw: - # Best-effort line number lookup. line_no = _first_line_containing(raw, ioc) findings.append( Finding( @@ -575,14 +491,7 @@ def _first_line_containing(text: str, needle: str) -> int | None: return None -# ───────────────────────────────────────────────────────────────────── -# Cargo.lock audit. -# ───────────────────────────────────────────────────────────────────── - - -# Cargo.lock is TOML; parse with stdlib tomllib (Python 3.11+). The -# studio's Tauri shell already requires a modern toolchain so this is -# always available where CI runs. +# Cargo.lock is TOML; parsed with stdlib tomllib (Python 3.11+). _PACKAGE_HEADER = re.compile(r"^\[\[package\]\]\s*$") @@ -706,24 +615,15 @@ def audit_cargo_lockfile(path: Path) -> list[Finding]: return findings -# ───────────────────────────────────────────────────────────────────── -# CLI. -# ───────────────────────────────────────────────────────────────────── - - # Finding kinds split into BLOCKING vs ADVISORY for the default run mode. -# Blocking findings come from public supply-chain attack indicators (a -# version we know is malicious, a string an attacker would have to embed -# for an attack to work). Advisory findings are structural lockfile -# anomalies (missing integrity, non-default registry, etc.) -- they -# WARN the maintainer but do not block merges. Pass --strict to make -# every finding blocking (PR-5479-style behavior for opt-in adopters). +# Blocking = public attack indicators (known-malicious version, IOC +# string). Advisory = structural anomalies that warn but don't block. +# --strict makes every finding blocking. BLOCKING_KINDS: frozenset[str] = frozenset( { "blocked-known-malicious", "known-ioc-string", - # Internal-failure kinds: a structurally broken lockfile MIGHT - # be hiding a real attack, so we keep these blocking too. + # A structurally broken lockfile might hide a real attack. "malformed-lockfile", "missing-lockfile", "unreadable-lockfile", @@ -779,15 +679,9 @@ def main(argv: list[str] | None = None) -> int: ) args = parser.parse_args(argv) - # SF4: require a real justification (e.g. JIRA ticket id) for the - # skip env var. Treat the trivially-set values ("1", "true", "yes", - # "on", empty) as INVALID -- they look like accidental flips and - # silently bypassed the supply-chain audit. A valid value is a - # non-empty string >=5 chars after stripping that does not match - # any of the boolean-shaped tokens above. An invalid value emits a - # loud GitHub Actions warning to stderr and FALLS THROUGH to run - # the audit normally (fail-safe). A valid value emits a warning - # naming the reason and skips with rc=0 (compat). + # Require a real justification (>=5 chars, not a boolean-shaped token) + # for the skip env var. An invalid value warns and falls through to + # run the audit (fail-safe); a valid one warns and skips with rc=0. _skip_raw = os.environ.get("UNSLOTH_LOCKFILE_AUDIT_SKIP") if _skip_raw is not None: _skip = _skip_raw.strip() @@ -809,8 +703,7 @@ def main(argv: list[str] | None = None) -> int: return 0 root = Path(args.root).resolve() - # Explicit --npm-lockfile/--cargo-lockfile scopes the scan to those - # paths; defaults apply only to the no-args CI invocation. + # Explicit flags scope the scan; defaults apply only to no-args CI. _user_explicit = args.npm_lockfile is not None or args.cargo_lockfile is not None if _user_explicit: npm_paths = [root / p for p in (args.npm_lockfile or ())] @@ -835,11 +728,9 @@ def main(argv: list[str] | None = None) -> int: ) return 0 - # Split findings into blocking (known-malicious / IOC / structurally - # broken) and advisory (everything else, e.g. missing integrity on a - # registry-published tarball). In default mode advisory findings are - # printed but do not change the exit code; --strict treats every - # finding as blocking. + # Split into blocking (known-malicious / IOC / structurally broken) + # and advisory (everything else). Default mode prints advisories + # without changing the exit code; --strict makes all blocking. blocking = [f for f in all_findings if f.kind in BLOCKING_KINDS] advisory = [f for f in all_findings if f.kind not in BLOCKING_KINDS] @@ -854,12 +745,8 @@ def main(argv: list[str] | None = None) -> int: file = sys.stderr, ) for f in advisory: - # Surface in GitHub Actions UI as a warning annotation when run - # under Actions; harmless prefix elsewhere. GH Actions - # truncates annotation messages at the first newline unless - # newlines are escaped as `%0A`, so the full multi-line - # Finding (kind + path + package + detail) only renders in - # the UI after _gha_escape collapses it onto one line. + # GH Actions warning annotation; _gha_escape collapses the + # multi-line Finding onto one line so it renders fully in the UI. print(f"::warning::{_gha_escape(str(f))}", file = sys.stderr) print(file = sys.stderr) @@ -876,9 +763,7 @@ def main(argv: list[str] | None = None) -> int: file = sys.stderr, ) for f in blocking: - # Same %-encoding rationale as the advisory branch above: the - # GH Actions annotation is truncated at the first newline - # unless the message is escaped. + # Same %-encoding rationale as the advisory branch above. print(f"::error::{_gha_escape(str(f))}", file = sys.stderr) print(file = sys.stderr) print( diff --git a/scripts/notebook_to_python.py b/scripts/notebook_to_python.py index cc9dc0acda..4b64123d6a 100644 --- a/scripts/notebook_to_python.py +++ b/scripts/notebook_to_python.py @@ -22,19 +22,14 @@ import urllib.parse from pathlib import Path -# Hosts we are willing to fetch raw notebook JSON from. Anything else -# is rejected before `urlopen` so a typoed / hostile URL cannot pull -# code from arbitrary infrastructure. +# Allowlist of hosts for raw notebook fetches; anything else rejected before urlopen. _ALLOWED_NOTEBOOK_HOSTS = { "raw.githubusercontent.com", "gist.githubusercontent.com", } -# Shell metacharacters that imply the cell's `!cmd` line cannot be -# parsed as a flat argv. If any of these appears, `shlex.split` would -# either fail or, worse, silently strip the operator -- so we keep -# `shell=True` for that command and emit a review marker. +# Metacharacters that mean a `!cmd` line can't be a flat argv -> keep shell=True + review marker. _SHELL_METACHARS_RE = re.compile(r"\$\(|`|\|\||\||&&|>>?|< bool: def github_blob_to_raw(url: str) -> str: """Convert GitHub blob URL to raw URL.""" - # https://github.com/user/repo/blob/branch/path - # -> https://raw.githubusercontent.com/user/repo/branch/path - # Compare the parsed host exactly (not as a substring) so a URL - # like https://attacker.example.com/github.com/blob/... does NOT - # get rewritten to a github raw URL. Closes CodeQL alert - # py/incomplete-url-substring-sanitization. + # github.com/user/repo/blob/branch/path -> raw.githubusercontent.com/user/repo/branch/path + # Exact host match (not substring) so attacker.example.com/github.com/blob/... is not rewritten. parsed = urllib.parse.urlparse(url) if parsed.netloc != "github.com" or "/blob/" not in parsed.path: return url @@ -63,18 +54,12 @@ def github_blob_to_raw(url: str) -> str: def download_notebook(url: str) -> tuple[str, str]: """Download notebook from URL. Returns (content, filename).""" - # Convert blob URL to raw if needed raw_url = github_blob_to_raw(url) - # Extract filename from URL parsed = urllib.parse.urlparse(raw_url) filename = os.path.basename(urllib.parse.unquote(parsed.path)) - # Host allowlist. Refuse to fetch from anywhere the campaign IOC - # tables flag (or just anywhere we don't recognise). The blob->raw - # conversion above only emits `raw.githubusercontent.com`, so a - # rejection here means the caller hand-typed a URL pointing - # somewhere we don't trust. + # Host allowlist: refuse to fetch from anything we don't recognise. host = parsed.hostname if host not in _ALLOWED_NOTEBOOK_HOSTS: raise ValueError( @@ -82,7 +67,6 @@ def download_notebook(url: str) -> tuple[str, str]: f"{sorted(_ALLOWED_NOTEBOOK_HOSTS)}" ) - # Download print(f"Downloading {url}...") with urllib.request.urlopen(raw_url, timeout = 60) as response: content = response.read().decode("utf-8") @@ -97,29 +81,18 @@ def is_url(path: str) -> bool: def replace_colab_paths(source: str) -> str: """Replace Colab-specific /content/ paths with current working directory.""" - # Replace /content/ with f-string using _WORKING_DIR source = source.replace('"/content/', 'f"{_WORKING_DIR}/') source = source.replace("'/content/", "f'{_WORKING_DIR}/") return source def _emit_shell_command(indent: str, full_cmd: str, *, allow_shell: bool) -> list[str]: - """Render a `!cmd` notebook line as one or more Python statements. + """Render a `!cmd` notebook line as Python statements. - When the command body is f-string-interpolated, contains shell - metacharacters, or spans multiple lines, falling back to - `shell=True` is the only correct option -- `shlex.split` would - either drop operators or fail outright. We surface that with a - `# WARNING: shell=True; reviewed for hostile input` comment so a - reviewer cannot miss it. - - Otherwise we emit `subprocess.run(shlex.split(cmd), shell=False)` - so the converted script is not a re-injection vector if the - notebook ever interpolates user-controlled data. - - `allow_shell` defaults to True at the CLI for backwards - compatibility. Setting it to False makes `shell=True` emission a - hard error (no surprise behaviour). + f-string interpolation, shell metacharacters, or multiline force + shell=True (shlex.split would drop operators), flagged with a + WARNING comment. Otherwise emit shell=False argv form. allow_shell + False makes shell=True emission a hard error. """ needs_f = needs_fstring(full_cmd) has_meta = bool(_SHELL_METACHARS_RE.search(full_cmd)) @@ -144,7 +117,6 @@ def _emit_shell_command(indent: str, full_cmd: str, *, allow_shell: bool) -> lis stmt = f"{indent}subprocess.run({f_prefix}{full_cmd!r}, shell=True)" return [warn, stmt] - # Shell-safe argv form. return [f"{indent}subprocess.run(shlex.split({full_cmd!r}), shell=False)"] @@ -159,12 +131,10 @@ def convert_cell_to_python(source: str, *, allow_shell: bool = True) -> str: stripped = line.strip() indent = line[: len(line) - len(line.lstrip())] - # Skip %%capture if stripped.startswith("%%capture"): i += 1 continue - # Handle %%file magic if stripped.startswith("%%file "): filename = stripped[7:].strip() file_lines = [] @@ -178,7 +148,6 @@ def convert_cell_to_python(source: str, *, allow_shell: bool = True) -> str: result.append(f'{indent} _f.write("""{file_content}""")') continue - # Handle ! shell commands if stripped.startswith("!"): cmd_lines = [stripped[1:]] while cmd_lines[-1].rstrip().endswith("\\") and i + 1 < len(lines): @@ -308,21 +277,16 @@ def convert_notebook_to_script( content = f.read() source_name = source - # Generate output filename output_filename = filename.replace(".ipynb", ".py") - # Clean up filename output_filename = output_filename.replace("(", "").replace(")", "").replace("-", "_") - # Add output directory if specified if output_dir: output_path = os.path.join(output_dir, output_filename) else: output_path = output_filename - # Convert script = convert_notebook(content, source_name, allow_shell = allow_shell) - # Write output with open(output_path, "w", encoding = "utf-8") as f: f.write(script) @@ -349,11 +313,7 @@ Examples: ) parser.add_argument("notebooks", nargs = "+", help = "Notebook files or URLs to convert.") parser.add_argument("-o", "--output", dest = "output_dir", default = ".", help = "Output directory.") - # Default True for backwards compatibility: existing Colab notebooks - # routinely use pipes / redirection / interpolation in `!cmd` lines - # and the converted script needs to keep working. Operators who - # convert untrusted notebooks should pass --no-allow-shell to force - # a hard error on every metacharacter-bearing cell. + # Default True for backwards compat; pass --no-allow-shell for untrusted notebooks. parser.add_argument( "--allow-shell", dest = "allow_shell", @@ -371,14 +331,9 @@ Examples: args = parser.parse_args() - # Create output directory if needed os.makedirs(args.output_dir, exist_ok = True) - # SF2: track per-notebook failures so a CI invocation that converts - # 10 notebooks but silently fails on 3 is no longer reported as - # success. Each failure is collected and the loop continues so the - # caller sees the full set; final exit status is 1 if anything - # failed. + # Track per-notebook failures; continue the loop and exit 1 if any failed. failures: list[tuple[str, str]] = [] ok = 0 total = len(args.notebooks) diff --git a/scripts/notebook_validator.py b/scripts/notebook_validator.py index 74bfaa8d41..c1be7a63a4 100644 --- a/scripts/notebook_validator.py +++ b/scripts/notebook_validator.py @@ -49,12 +49,9 @@ from typing import Any, Iterable, Iterator def _atomic_write_bytes(path: pathlib.Path, data: bytes) -> None: - """Atomic write helper. See `scripts/scan_packages.py::update_req_file`. - - A crash between `mkstemp` and `os.replace` leaves the prior file - untouched, so a half-downloaded PyPI metadata cache file cannot - poison subsequent runs of the validator. - """ + """Atomic write (see scripts/scan_packages.py::update_req_file). A crash + between mkstemp and os.replace leaves the prior file intact, so a + half-downloaded cache file can't poison later runs.""" path.parent.mkdir(parents = True, exist_ok = True) dirpath = str(path.parent) or "." fd, tmp_path = tempfile.mkstemp(prefix = ".nb_val.", dir = dirpath) @@ -81,12 +78,10 @@ COLAB_PIP_FREEZE_URL = ( ) COLAB_FALLBACK_FILE = DATA_DIR / "colab_pip_freeze.gpu.txt" -# Oracle files we snapshot from googlecolab/backend-info. The diff -# subcommand fetches each, compares against the committed snapshot, -# and surfaces NEW / REMOVED / CHANGED entries so upstream Colab base -# image rotations land in CI within ~24h instead of when a notebook -# breaks. Every rule in this validator that resolves against the -# Colab preinstall (R-INST-002/003/004/005) gets earlier signal. +# Oracle files snapshotted from googlecolab/backend-info. The colab-diff +# subcommand surfaces NEW/REMOVED/CHANGED entries so upstream Colab base +# image rotations land in CI within ~24h, giving R-INST-002/003/004/005 +# earlier signal. COLAB_ORACLE_FILES: dict[str, str] = { "pip-freeze.gpu.txt": "colab_pip_freeze.gpu.txt", "apt-list-gpu.txt": "colab_apt_list.gpu.txt", @@ -145,9 +140,8 @@ class Finding: def iter_notebooks( notebooks_dir: pathlib.Path, include_templates: bool = False ) -> Iterator[pathlib.Path]: - """Yield user-facing .ipynb files under nb/ and kaggle/. Pass - include_templates=True to also walk original_template/ (used by the - convert subcommand which doesn't lint install cells).""" + """Yield user-facing .ipynb files under nb/ and kaggle/. + include_templates=True also walks original_template/ (for convert).""" subs = ("nb", "kaggle") if include_templates: subs = ("nb", "kaggle", "original_template") @@ -198,11 +192,8 @@ def install_cells(nb: dict[str, Any]) -> list[tuple[int, str]]: return out -# Notebook target environment. The Colab oracle (pip-freeze.gpu.txt) only -# applies to notebooks that actually run on Colab; AMD-Dev-Cloud, -# Kaggle, HuggingFace-Course, and DGX-Spark notebooks have their own -# preinstalled environments and the Colab-vs-cell rules are not -# applicable to them. +# Colab oracle only applies to notebooks that run on Colab; AMD, Kaggle, +# DGX-Spark have their own preinstalls and the Colab-vs-cell rules don't apply. def target_environment(notebook_name: str) -> str: parts = pathlib.PurePath(notebook_name).parts base = parts[-1] if parts else notebook_name @@ -466,19 +457,12 @@ def constraint_satisfied(version: str, ops: list[tuple[str, str]]) -> bool: def resolved_set(install_cell: str, colab: dict[str, str]) -> dict[str, str]: - """Merge install-cell explicit constraints with Colab pip-freeze. Cell - wins. + """Merge install-cell constraints with Colab pip-freeze (cell wins). - Resolution order per package, when more than one form is present: - 1. Exact `==V` pin in any install line (definitive). - 2. Upper-bound `<=V` constraint (pip picks the highest - allowed; that's V). - 3. Colab pip-freeze fallback. - - The lower-bound `>=V` is intentionally NOT reflected here — a `>=V` - by itself doesn't change the resolved version when a higher - Colab-preinstalled version is already in scope. (R-INST-003 calls - `_install_cell_lower_bound` separately to model that case.) + Resolution order per package: (1) exact `==V` pin, (2) upper-bound `<=V` + (pip picks the highest allowed = V), (3) Colab fallback. Lower-bound `>=V` + is intentionally NOT reflected (it doesn't lower an already-higher Colab + version); R-INST-003 models that via `_install_cell_lower_bound`. """ out = dict(colab) pinned: set[str] = set() @@ -543,8 +527,7 @@ def rule_inst_002_no_deps_transitive( v = explicit_pin(sp) if v is None: continue - # Check transitive constraints on a curated short list of pkgs we - # care about (transformers/peft/trl/accelerate/torchao/torchcodec). + # Check transitive constraints on a curated short list of pkgs. for target in ( "tokenizers", "torchao", @@ -575,10 +558,9 @@ def rule_inst_002_no_deps_transitive( def _install_cell_lower_bound(install_cell: str, target: str) -> str | None: - """Return the highest LOWER bound that any install line places on `target`, - or None if no constraint is present. Treats `==V` as both lower and upper. - Used by R-INST-003: a `pip install torchao>=0.16.0` line is enough to - satisfy a `torchao>=0.16.0` floor even though it's not a `==` pin.""" + """Return the highest lower bound any install line places on `target` + (treating `==V` as both bounds), or None. Used by R-INST-003 so a + `torchao>=0.16.0` line satisfies the floor without a `==` pin.""" best: str | None = None for inv in iter_pip_invocations(install_cell): for raw in inv.packages: @@ -652,20 +634,17 @@ def rule_inst_004_torchcodec_torch( def rule_inst_005_transformers_tokenizers( install_cell: str, colab: dict[str, str], file: str, cell_idx: int ) -> list[Finding]: - """Fires only when transformers is installed with `--no-deps`. Without - `--no-deps`, pip resolves the correct tokenizers transitively, so the - rule would be a false positive (this is the case for older notebooks - that pin `transformers==4.51.3` but rely on pip's transitive resolver). - The rule targets the exact pattern PR #261b / #264 fixed: - `pip install --no-deps transformers==X` next to a Colab preinstall - `tokenizers` outside transformers's window.""" + """Fires only when transformers is installed with `--no-deps` (otherwise + pip resolves tokenizers transitively and flagging would be a false + positive). Targets the PR #261b/#264 pattern: `--no-deps transformers==X` + next to a Colab `tokenizers` outside transformers's window.""" findings: list[Finding] = [] res = resolved_set(install_cell, colab) tf = res.get("transformers") tok = res.get("tokenizers") if not tf or tok is None: return findings - # Find the install line that pins transformers and check for --no-deps. + # Find the transformers pin and check for --no-deps. transformers_line_no_deps = False for inv in iter_pip_invocations(install_cell): for raw in inv.packages: @@ -724,11 +703,9 @@ def rule_inst_006_double_bang(install_cell: str, file: str, cell_idx: int) -> li class _APIScanner(ast.NodeVisitor): """Scan user-facing code cells for known deprecated patterns. R-API-001 - (`for_training`/`for_inference`) is intentionally absent: those helpers - are still part of the live unsloth surface as of 2026-05; PR #221 removed - the calls cosmetically from Vision notebooks but did not deprecate the - methods. R-API-004 (live API surface diff) catches actual removals - dynamically without us hand-coding them.""" + (`for_training`/`for_inference`) is intentionally absent: those helpers are + still live as of 2026-05 (PR #221 removed them cosmetically, not as a + deprecation). R-API-004 catches actual removals dynamically.""" def __init__(self, file: str, cell_idx: int): self.file = file @@ -736,14 +713,10 @@ class _APIScanner(ast.NodeVisitor): self.findings: list[Finding] = [] def visit_Call(self, node: ast.Call) -> None: - # SFTConfig with suboptimal optim (R-API-003). - # NOTE: PR #221 also stripped `gradient_checkpointing` / - # `gradient_checkpointing_kwargs` from a handful of vision notebooks, - # but those kwargs are still accepted by live TRL (verified against - # trl==0.25.1 in the unsloth workspace) so removing them was - # cosmetic, not a deprecation. We do NOT flag them. R-API-004 (live - # API surface diff in the api subcommand) is the right way to catch - # actual TRL signature drift. + # SFTConfig with suboptimal optim (R-API-003). + # NOTE: PR #221 also stripped gradient_checkpointing kwargs from some + # vision notebooks, but they're still accepted by live TRL (trl==0.25.1) + # so that was cosmetic. We don't flag them; R-API-004 catches real drift. if isinstance(node.func, ast.Name) and node.func.id == "SFTConfig": for kw in node.keywords: if ( @@ -799,13 +772,9 @@ POLICY_CLAUSES_DEFAULT = [ def extract_policy_clauses(update_script: pathlib.Path) -> list[tuple[str, re.Pattern[str], Any]]: - """Best-effort: scan update_all_notebooks.py for canonical phrases used by - multiple templates. Falls back to POLICY_CLAUSES_DEFAULT. - - Today we use POLICY_CLAUSES_DEFAULT directly; the regex form is - intentionally permissive so a template-side reword (e.g. comment changes) - doesn't cause false positives. New clauses become 1-line PRs to this list. - """ + """Best-effort scan of update_all_notebooks.py for canonical phrases; + falls back to POLICY_CLAUSES_DEFAULT (which we use directly today). The + permissive regexes avoid false positives on template rewords.""" return list(POLICY_CLAUSES_DEFAULT) @@ -868,14 +837,9 @@ def cmd_drift(args: argparse.Namespace) -> int: check = False, capture_output = True, ) - # SF3: the restore MUST run even on SystemExit / KeyboardInterrupt / - # segfault-propagated exception, otherwise the user's working tree - # silently stays rolled back into the stash. A bare try/finally - # (NOT try/except/finally) preserves the original exception and - # still runs the cleanup. The pre-existing try/except around - # `subprocess.run` of the updater is folded inside the new outer - # try so its early returns still happen, but the stash pop is - # protected. + # The restore MUST run even on SystemExit/KeyboardInterrupt, else the + # working tree stays rolled back into the stash. A bare try/finally keeps + # the original exception while still running the cleanup (stash pop). findings: list[Finding] = [] rc: int try: @@ -920,8 +884,7 @@ def cmd_drift(args: argparse.Namespace) -> int: ) rc = 0 if not findings else 1 finally: - # Restore the working tree. Both commands MUST run regardless of - # how the try block exited (including SystemExit/KeyboardInterrupt). + # Restore the working tree (both commands run regardless of exit path). subprocess.run( ["git", "-C", str(nbdir), "checkout", "."], check = False, @@ -1004,19 +967,16 @@ def cmd_lint(args: argparse.Namespace) -> int: continue rel = str(path.relative_to(nbdir)) env = target_environment(rel) - # The Colab oracle is the source of truth ONLY for Colab notebooks. - # Other targets (amd / kaggle / dgx_spark) have their own runtime - # preinstall sets that aren't tracked here yet, so we apply the - # environment-agnostic rules and skip the Colab-specific ones. + # Colab oracle applies only to Colab notebooks; other targets get the + # environment-agnostic rules only (their preinstalls aren't tracked). oracle = colab if env == "colab" else {} cells = install_cells(nb) - # Per-cell rules: forbid-pattern checks scoped to a single line. + # Per-cell forbid-pattern checks. for idx, cell in cells: findings += rule_inst_001_git_plus(cell, rel, idx) findings += rule_inst_006_double_bang(cell, rel, idx) - # Whole-notebook rules: a notebook's install steps are sometimes split - # across multiple cells (initial install + post-install bumps). Merge - # all install cells before resolving compat against Colab. + # Whole-notebook rules: install steps may span multiple cells, so merge + # before resolving compat against Colab. merged = "\n".join(c for _, c in cells) if env == "colab" and merged: first_cell = cells[0][0] if cells else None @@ -1147,8 +1107,7 @@ def _parse_apt_lines(text: str) -> dict[str, str]: def _parse_os_lines(text: str) -> dict[str, str]: - """Free-form ` ` lines. Skip comments. The key is the - first token lower-cased; the value is the rest of the line.""" + """Free-form ` ` lines -> {tool_lower: rest}.""" out: dict[str, str] = {} for line in text.splitlines(): line = line.strip() @@ -1185,10 +1144,9 @@ def _diff_oracle( def cmd_colab_diff(args: argparse.Namespace) -> int: - """Fetch every Colab oracle file in COLAB_ORACLE_FILES, diff against - the committed snapshot, and print NEW / REMOVED / CHANGED. Advisory - by default (rc=0); --strict promotes any diff to rc=1 so the daily - cron can fail loudly when upstream rotates.""" + """Diff each Colab oracle file against its committed snapshot and print + NEW/REMOVED/CHANGED. Advisory (rc=0) by default; --strict makes any diff + rc=1 so the daily cron fails loudly on upstream rotation.""" snapshot_dir = pathlib.Path(args.snapshot_dir).resolve() any_diff = False for upstream_name, snapshot_name in COLAB_ORACLE_FILES.items(): diff --git a/scripts/run_ruff_format.py b/scripts/run_ruff_format.py index 2f0a6c6bba..817565eddf 100755 --- a/scripts/run_ruff_format.py +++ b/scripts/run_ruff_format.py @@ -19,10 +19,8 @@ def main(argv: list[str]) -> int: spacing_script = HERE / "enforce_kwargs_spacing.py" - # Pre-ruff: normalize def-signature trailing commas (>=3 params with a - # default -> one-per-line; everything else collapsible) and strip the magic - # comma from a short multi-line assert, so ruff wraps signatures accordingly - # and joins the assert back onto one line. + # Pre-ruff: normalize def-signature magic commas and strip the magic comma + # from short multi-line asserts so ruff wraps/joins accordingly. pre_cmd = [sys.executable, str(spacing_script), "--pre", *files] pre_proc = subprocess.run(pre_cmd) if pre_proc.returncode != 0: diff --git a/scripts/scan_npm_packages.py b/scripts/scan_npm_packages.py index 14a3cd6560..fe90afa7e6 100644 --- a/scripts/scan_npm_packages.py +++ b/scripts/scan_npm_packages.py @@ -7,73 +7,32 @@ """scan_npm_packages.py -- npm-side content scanner. -Counterpart to scripts/scan_packages.py for the pip ecosystem. Reads +npm counterpart to scripts/scan_packages.py. Reads studio/frontend/package-lock.json, downloads each resolved tarball DIRECTLY from registry.npmjs.org (never via `npm install` -- no -lifecycle scripts ever run), verifies the lockfile integrity hash, -unpacks each tarball into a sandboxed temp dir behind size / count / -path-escape / symlink guards, and pattern-scans the extracted file -contents for the signatures common to npm supply-chain attacks: +lifecycle scripts run), verifies the lockfile integrity hash, unpacks +each into a sandboxed temp dir behind size/count/path-escape/symlink +guards, and pattern-scans extracted contents for npm supply-chain +attack signatures: malicious lifecycle scripts, C2 / exfil hosts, +credential-stealing references, known IOC filenames, and obfuscation +shapes. - - Lifecycle (preinstall / install / postinstall / prepare) scripts - in any package.json that fetch + execute external code. - - C2 / exfiltration hosts (getsession.org, AWS IMDS endpoints, - Kubernetes ServiceAccount token paths, GitHub Actions OIDC, - HashiCorp Vault endpoints). - - Credential-stealing references (~/.npmrc, ~/.aws/credentials, - GITHUB_TOKEN / NPM_TOKEN in JS sources). - - Known IOC filenames from public advisories - (router_init.js, tanstack_runner.js, router_runtime.js). - - Obfuscation shapes (large single JS in package root with a low - whitespace ratio + Function/eval against a base64-decoded blob). - -Safety stance -============= - -This script ingests attacker-controlled archives. Every parse path -assumes the worst: - - 1. Downloads ONLY from `registry.npmjs.org`. Any tarball URL with a - different hostname is refused without fetching. - 2. Tarball download is size-capped (HARD_MAX_TARBALL_BYTES default - 64 MiB). HEAD-style probe via the Content-Length response header - plus a chunked read that aborts on overflow. +Safety stance (ingests attacker-controlled archives; assumes worst): + 1. Downloads ONLY from registry.npmjs.org; other hosts refused. + 2. Tarball download size-capped via Content-Length probe + chunked + read that aborts on overflow. 3. SHA-512 integrity verified against the lockfile entry BEFORE the - tarball is even opened. A mismatch aborts that package -- the - scanner does not "fall back" to the registry-published hash. - 4. tar extraction goes through `safe_extract`: - - rejects symbolic links (`SYMTYPE`, `LNKTYPE`) - - rejects absolute paths, `..` traversal, paths outside the - extract root after resolution - - rejects character / block / FIFO devices - - per-file uncompressed size cap (HARD_MAX_FILE_BYTES, default - 8 MiB) AND cumulative cap (HARD_MAX_TOTAL_BYTES, default - 128 MiB) AND member-count cap (HARD_MAX_MEMBERS, default - 50_000) - - tar reads happen via `tarfile.open(mode='r|gz')` streaming - so an oversized file is detected before write - 5. NOTHING from the extracted tree is ever executed. Files are read - as raw bytes, decoded with `errors='replace'`, and grepped. We - never call `node`, `eval`, `compile`, `subprocess.run`, - `os.system`, or anything that would touch the tarball's - declared scripts. - 6. Tempdir is created with `tempfile.mkdtemp(prefix='npm-scan-')`, - fully resolved with .resolve(), and registered with atexit to be - wiped on every termination path. - 7. Stdlib only. No third-party deps -- adding one would itself be a - supply-chain liability. + tarball is opened; mismatch aborts that package (no fallback). + 4. tar extraction via `safe_extract`: rejects symlinks, absolute / + `..` paths, device files; enforces per-file, cumulative, and + member-count caps; streams (`r|gz`) so oversize is caught early. + 5. NOTHING extracted is executed -- files are read as bytes and + grepped only. + 6. Tempdir resolved and atexit-wiped on every termination path. + 7. Stdlib only (a dep would be a supply-chain liability itself). -Exit codes -========== - - 0 no findings of severity HIGH or higher - 1 one or more HIGH/CRITICAL findings (or pre-scan structural - anomalies -- non-registry resolved URL, missing integrity) - 2 internal error (lockfile missing, integrity mismatch on - download, malformed tarball, etc.) - -The script is meant to be run in CI on every PR that touches -package-lock.json and on a nightly schedule. +Exit codes: 0 = no HIGH+ findings; 1 = HIGH/CRITICAL or pre-scan +structural anomaly; 2 = internal error. Run in CI per-PR and nightly. """ from __future__ import annotations @@ -565,12 +524,9 @@ CRED_HOST_NEEDS_CONTEXT: tuple[tuple[str, str], ...] = ( ), ) -# Credentials a frontend package should NEVER need to read. Bare -# substring match is too noisy (object-treeify ships a `docker` dev -# script that mounts ~/.npmrc -- legitimate dev tooling, never run -# at install time). We instead surface these only when they appear -# inside a LIFECYCLE script (preinstall / install / postinstall / -# prepare), which is the only path that runs automatically on +# Credentials a frontend package should never read. Bare substring +# match is too noisy (legit dev tooling mounts ~/.npmrc), so we flag +# these only inside lifecycle scripts -- the only auto-run path on # `npm ci`. See `scan_package_json` below. CRED_PATH_SUBSTRINGS: tuple[tuple[str, str], ...] = ( ("/.npmrc", "npm credentials file"), @@ -603,9 +559,8 @@ _JS_FETCH_EVAL = re.compile( """, ) -# `process.env.GITHUB_TOKEN` / `NPM_TOKEN` / `AWS_*` access in -# top-level / install-time code is suspicious. We also catch -# `os.environ["GITHUB_TOKEN"]` for the rare Python-in-npm postinstall. +# Token env access in install-time code; also catches os.environ[...] +# for the rare Python-in-npm postinstall. _JS_ENV_TOKEN = re.compile( r"""(process\.env\.|os\.environ\[?['"])(?: GITHUB_TOKEN | GH_TOKEN | NPM_TOKEN | NODE_AUTH_TOKEN @@ -616,11 +571,9 @@ _JS_ENV_TOKEN = re.compile( re.VERBOSE, ) -# Suspicious lifecycle-script payloads. Anything in a package.json -# `scripts` field that wgets/curls an external resource and executes -# it. We do NOT block ALL curl/wget in scripts (some legit packages -# fetch test fixtures into devDependencies), but we DO block the -# fetch+exec chain. +# Lifecycle-script fetch+exec chain: curl/wget an external resource +# and run it. Bare curl/wget is allowed (legit fixture fetches); only +# the fetch+exec chain is blocked. _LIFECYCLE_FETCH_EXEC = re.compile( r"""(?xs) (?:curl|wget|fetch|http\.get|axios\.get)\s+ # fetch verb @@ -634,9 +587,8 @@ _LIFECYCLE_FETCH_EXEC = re.compile( """, ) -# Obfuscation: large JS file that is mostly one line of base64-ish -# blob with a Function() / eval() bookend. Tuned against the -# router_init.js shape (2.3 MB obfuscated single-blob). +# Obfuscation: large single-line base64-ish blob behind Function()/ +# eval(). Tuned against the router_init.js shape (2.3 MB blob). _OBFUSC_BLOB = re.compile( r"""(?xs) (?:Function|eval)\s*\(\s*['"`]? @@ -653,11 +605,9 @@ _OBFUSC_BLOB = re.compile( def parse_lockfile(path: Path) -> tuple[list[PackageEntry], list[Finding]]: """Return (entries, structural_findings). - Structural findings here are HIGH-severity refusals that should - short-circuit the scan -- a lockfile with non-registry resolved - URLs is itself a finding (covered by scripts/lockfile_supply_chain - _audit.py in detail; we surface a summary here so this scanner is - standalone-runnable). + Structural findings are HIGH-severity refusals that short-circuit + the scan (e.g. non-registry resolved URLs). A summary is surfaced + here so this scanner is standalone-runnable. """ entries: list[PackageEntry] = [] findings: list[Finding] = [] @@ -701,9 +651,8 @@ def parse_lockfile(path: Path) -> tuple[list[PackageEntry], list[Finding]]: resolved = entry.get("resolved") if not resolved: continue - # Strict registry origin check. lockfile_supply_chain_audit - # already catches this; double-defend here so this scanner - # cannot be tricked into fetching from an attacker-chosen URL. + # Strict registry origin check so this scanner can't be tricked + # into fetching from an attacker-chosen URL. parsed = urllib.parse.urlparse(resolved) if parsed.scheme != "https" or parsed.hostname != ALLOWED_DOWNLOAD_HOST: findings.append( @@ -775,16 +724,12 @@ def download_tarball( timeout: float = HARD_HTTP_TIMEOUT_S, max_bytes: int = HARD_MAX_TARBALL_BYTES, ) -> tuple[Path, str | None]: - """Stream-download entry.resolved to dest. Verify SRI integrity. + """Stream-download entry.resolved to dest and verify SRI integrity. - Returns (downloaded_path, error_or_none). On any error the - returned path may not exist. Network access is restricted to - https://{ALLOWED_DOWNLOAD_HOST}/ -- the caller passes a Request - we already validated. + Returns (downloaded_path, error_or_none); on error the path may not + exist. Network access is restricted to ALLOWED_DOWNLOAD_HOST. """ - # Re-assert hostname; the entry was validated at parse time but a - # defence-in-depth check here means a future refactor cannot - # accidentally bypass it. + # Re-assert hostname (defence-in-depth against a future refactor). parsed = urllib.parse.urlparse(entry.resolved) if parsed.scheme != "https" or parsed.hostname != ALLOWED_DOWNLOAD_HOST: return dest, (f"refused download from non-allowlisted URL {entry.resolved!r}") @@ -873,8 +818,7 @@ def safe_extract( total = 0 count = 0 try: - # Open in streaming mode so we never seek backwards in the - # input. `r|gz` rejects malformed gzip frames immediately. + # Streaming mode (no backward seeks); `r|gz` rejects bad gzip. with tarfile.open(tarball_path, mode = "r|gz") as tf: for member in tf: count += 1 @@ -889,9 +833,8 @@ def safe_extract( return f"refused link member {name!r} (sym/lnk)" if member.isdev() or member.isfifo(): return f"refused special member {name!r}" - # Cumulative cap is checked against DECLARED size up - # front to short-circuit obvious bombs without reading - # the body. + # Check declared size up front to short-circuit bombs + # without reading the body. declared = max(member.size, 0) if declared > HARD_MAX_BINARY_FILE_BYTES: return ( @@ -903,9 +846,8 @@ def safe_extract( f"cumulative bytes {total + declared} > cap " f"{max_total_bytes} at {name!r}" ) - # Strip leading "package/" -- the npm convention. We do - # NOT trust npm to be right, so we explicitly resolve - # the destination and refuse anything that escapes. + # Resolve destination and refuse anything escaping root + # (don't trust the npm "package/" convention). dest = extract_root / name if not _is_within(extract_root, dest): return f"refused escape: {name!r} resolved outside root" @@ -919,10 +861,8 @@ def safe_extract( src = tf.extractfile(member) if src is None: continue - # Sniff first 16 bytes to classify text vs binary. - # Text-cap members get the tight 16 MiB limit; binary - # members (executables, .node, .wasm, native libs) - # get the generous binary cap. We bound BOTH cases. + # Sniff first 16 bytes to classify text vs binary; + # each gets its own cap (both are bounded). header = src.read(16) is_binary = _looks_binary(name, header) file_cap = HARD_MAX_BINARY_FILE_BYTES if is_binary else HARD_MAX_TEXT_FILE_BYTES @@ -941,8 +881,7 @@ def safe_extract( f"({'binary' if is_binary else 'text'})" ) total += len(data) - # Write with restrictive mode (rw-r--r--) so even if - # someone runs the extract dir nothing is executable. + # Restrictive mode (rw-r--r--): nothing executable. with open(dest, "wb") as out: out.write(data) os.chmod(dest, 0o644) @@ -1008,10 +947,8 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: ), ) ) - # Credential file paths inside a lifecycle script are - # exfiltration prep -- npm runs these scripts automatically - # on `npm ci`. Manual `scripts.*` entries (like a `docker` - # dev script) are out of scope: npm does not run them. + # Cred file paths in a lifecycle script are exfil prep (npm + # auto-runs these on `npm ci`); manual scripts are out of scope. for path_substr, why in CRED_PATH_SUBSTRINGS: if path_substr in body: findings.append( @@ -1071,20 +1008,12 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: def _host_in_outbound_context(text: str, host: str) -> bool: - """True if `host` appears in a way consistent with an outbound call. + """True if `host` appears consistent with an outbound call. - A bare `"169.254.169.254"` array literal (defensive blocklist) is - safe; a `fetch("http://169.254.169.254/...")` is not. The signal - is co-occurrence with either an HTTP URL scheme or a fetch verb - within a short window. - - A defensive blocklist looks like: - const CLOUD_METADATA_IPS = ["169.254.169.254", "169.254.170.2"]; - An exfil call looks like: - fetch("http://169.254.169.254/latest/meta-data/...") - http.request({ host: "169.254.169.254", path: "/..." }) + A bare array literal (defensive blocklist) is safe; co-occurrence + with an HTTP URL scheme or a fetch verb in a short window is not. """ - # Esc for use in a regex (IPs contain dots). + # Escape for regex (IPs contain dots). host_re = re.escape(host) # 1. URL form: http://host or https://host or //host/ or //host" url_form = re.compile( @@ -1127,8 +1056,7 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: ) ) - # Credential surfaces. Tier 1: hosts with no legitimate use, - # bare substring is enough. + # Cred surfaces, tier 1: hosts with no legit use; bare substring. for needle, why in CRED_HOST_ALWAYS_BAD: if needle in text: findings.append( @@ -1145,8 +1073,8 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: ) ) - # Credential surfaces. Tier 2: hosts that do appear in defensive - # code; require co-occurrence with a fetch verb or URL prefix. + # Cred surfaces, tier 2: hosts that appear in defensive code too; + # require co-occurrence with a fetch verb or URL prefix. for needle, why in CRED_HOST_NEEDS_CONTEXT: if needle in text and _host_in_outbound_context(text, needle): findings.append( @@ -1164,11 +1092,8 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: ) ) - # Credential PATHS are deliberately not scanned here; they have - # too high a false-positive rate at file scope (defensive code, - # docker mounts, AWS SDK docs strings). `scan_package_json` - # catches the malicious case -- credential paths inside a - # lifecycle script run automatically on `npm ci`. + # Credential PATHS aren't scanned here (too many FPs at file + # scope); scan_package_json catches them inside lifecycle scripts. # JS-specific regex. if _JS_FETCH_EVAL.search(text): @@ -1211,9 +1136,8 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: return findings -# Filename suffix decides which scanners run. We deliberately treat -# *.cjs/*.mjs/*.ts the same as *.js -- attackers use whichever -# extension the consumer's bundler / loader resolves. +# Filename suffix decides which scanners run; .cjs/.mjs/.ts are +# treated like .js (attackers use whichever the loader resolves). _TEXT_SUFFIXES = ( ".js", ".mjs", @@ -1241,12 +1165,9 @@ def scan_extracted_tree(pkg: PackageEntry, root: Path) -> list[Finding]: rel = path.relative_to(root).as_posix() lower = rel.lower() if not lower.endswith(_TEXT_SUFFIXES): - # Skip native binaries entirely -- regex over compiled - # machine code is just noise (false positives in WASM - # opcodes, .node BSS segments, image pixel data). Use - # content-magic detection so extensionless executables - # (eg `package/biome`) and versioned shared libraries - # are also skipped. + # Skip native binaries (regex over machine code is noise); + # content-magic detection also skips extensionless + # executables and versioned shared libraries. try: if path.stat().st_size > HARD_MAX_TEXT_FILE_BYTES: continue @@ -1288,12 +1209,12 @@ def scan_extracted_tree(pkg: PackageEntry, root: Path) -> list[Finding]: def scan_one(pkg: PackageEntry, workspace: Path) -> tuple[list[Finding], str | None]: - """Download + extract + scan a single package. Cleans up its dir. + """Download + extract + scan a single package; cleans up its dir. Returns (findings, error). `error` is non-None only on hard - failures (download error, integrity mismatch, malformed tarball); - on a clean run with findings the error is None and the caller - decides exit code based on severity. + failures (download, integrity mismatch, malformed tarball); on a + clean run with findings, error is None and the caller decides the + exit code from severity. """ pkg_dir = workspace / f"{pkg.name.replace('/', '_')}-{pkg.version}" pkg_dir.mkdir(parents = True, exist_ok = True) diff --git a/scripts/scan_packages.py b/scripts/scan_packages.py index d753a17e49..861b35617b 100644 --- a/scripts/scan_packages.py +++ b/scripts/scan_packages.py @@ -56,27 +56,22 @@ from dataclasses import dataclass, field from pathlib import Path -# --------------------------------------------------------------------------- # Severity -# --------------------------------------------------------------------------- CRITICAL = "CRITICAL" HIGH = "HIGH" MEDIUM = "MEDIUM" SEVERITY_ORDER = {CRITICAL: 0, HIGH: 1, MEDIUM: 2} -# Hard pin-blocks for publicly confirmed malicious PyPI versions. -# Source: Socket.dev 2026-05-12 disclosure (Mini Shai-Hulud May-12 wave) and -# earlier Semgrep / Endor reports for the `lightning` entries. +# Hard pin-blocks for confirmed malicious PyPI versions (Socket.dev 2026-05-12 +# Mini Shai-Hulud wave; earlier Semgrep/Endor reports for `lightning`). BLOCKED_PYPI_VERSIONS: dict[str, set[str]] = { "guardrails-ai": {"0.10.1"}, "mistralai": {"2.4.6"}, "lightning": {"2.6.2", "2.6.3"}, } -# --------------------------------------------------------------------------- # Pattern definitions -# --------------------------------------------------------------------------- # Subprocess / OS exec patterns RE_SUBPROCESS = re.compile( @@ -312,10 +307,8 @@ RE_C2_POLLING = re.compile( re.DOTALL, ) -# Developer-tool persistence hooks. The PyTorch Lightning 2.6.x compromise -# planted SessionStart hooks into Claude Code, VS Code tasks, and Cursor -# settings so the payload re-attached on every editor open. Catches any -# package writing into a known dev-tool config that supports auto-run. +# Developer-tool persistence hooks. Lightning 2.6.x planted SessionStart hooks +# into Claude Code / VS Code / Cursor so the payload re-attached on editor open. RE_DEV_TOOL_HIJACK = re.compile( r"\.claude/settings\.json" r"|\.cursor/.*hooks" @@ -326,9 +319,8 @@ RE_DEV_TOOL_HIJACK = re.compile( r"|\bautomator\b.*\.workflow\b", ) -# Hard-coded credential / API-token regexes embedded in source. Packages -# that ship regexes for OTHER people's secrets are nearly always -# stealers (litellm 1.82.7, elementary-data 0.23.3, Shai-Hulud). +# Hard-coded credential / API-token regexes embedded in source. Packages that +# ship regexes for OTHER people's secrets are nearly always stealers. RE_TOKEN_REGEX = re.compile( r"\bgh[psoru]_[A-Za-z0-9_]{20,}" # GitHub PAT/OAuth/etc. r"|\bgithub_pat_[A-Za-z0-9_]{20,}" @@ -342,20 +334,16 @@ RE_TOKEN_REGEX = re.compile( r"|\bglpat-[0-9A-Za-z_-]{20,}", # GitLab PAT ) -# Mini Shai-Hulud May-12 2026 wave indicators. The dropper artifact name -# `transformers.pyz` is high-confidence (no legit PyPI package ships a `.pyz` -# named after `transformers`); the host + slogans are CRITICAL. +# Mini Shai-Hulud May-12 2026 wave indicators. `transformers.pyz` dropper name +# is high-confidence; the host + slogans are CRITICAL. RE_MAY12_IOC = re.compile( r"(git-tanstack\.com|/tmp/transformers\.pyz|transformers\.pyz" r"|With Love TeamPCP|We've been online over 2 hours)", re.IGNORECASE, ) -# JavaScript-side obfuscation. The npm chalk/debug compromise and the -# Lightning router_runtime.js use the same minifier-style hex-var name -# pattern; a bundle full of `_0x1f2e3d` identifiers is a near-universal -# tell for a malicious npm payload (and very rare in legit minified code -# that ships in PyPI wheels). +# JavaScript-side obfuscation. A bundle full of `_0x1f2e3d` hex-var identifiers +# is a near-universal tell for a malicious npm payload, rare in legit wheels. RE_JS_OBFUSCATION = re.compile( r"_0x[a-f0-9]{4,6}\s*=\s*function" r"|var\s+_0x[a-f0-9]{4,6}\b" @@ -363,9 +351,8 @@ RE_JS_OBFUSCATION = re.compile( r"|String\.fromCharCode\s*\(\s*\d+\s*(?:,\s*\d+\s*){10,}\)", ) -# Web3 / wallet-hijack pattern. The Qix npm phish overrode fetch / -# XMLHttpRequest and attached a `window.ethereum` listener that -# Levenshtein-swapped recipient addresses on the way to the network. +# Web3 / wallet-hijack pattern. The Qix npm phish overrode fetch/XMLHttpRequest +# and swapped recipient addresses via a `window.ethereum` listener. RE_WEB3_HIJACK = re.compile( r"\bwindow\.ethereum\b" r"|\bweb3\.eth\.\w+\s*\(" @@ -374,11 +361,9 @@ RE_WEB3_HIJACK = re.compile( r"|TronWeb|solanaWeb3", ) -# Self-propagating supply-chain worms (Shai-Hulud, ForceMemo) plant -# their own GitHub workflow in every repo they can reach, and lean on -# trufflehog/gitleaks for credential discovery. The combo of any of -# these strings inside a *package payload* is overwhelming evidence of -# repo-takeover intent. +# Self-propagating worms (Shai-Hulud, ForceMemo) plant their own GitHub workflow +# in every repo they reach and use trufflehog/gitleaks for credential discovery. +# Any of these strings in a package payload is strong repo-takeover evidence. RE_WORKFLOW_INJECT = re.compile( r"\.github/workflows/[^\"\']*\.ya?ml" r"|\btrufflehog\b|\bgitleaks\b" @@ -388,9 +373,8 @@ RE_WORKFLOW_INJECT = re.compile( re.IGNORECASE | re.DOTALL, ) -# Shell-side patterns specific to install.sh / postinstall scripts that -# pipe remote code into a shell. `curl ... | sh` and friends are the -# canonical npm postinstall dropper. +# install.sh / postinstall scripts piping remote code into a shell. +# `curl ... | sh` is the canonical npm postinstall dropper. RE_SHELL_DROPPER = re.compile( r"\bcurl\b[^\n|]*\|\s*(?:sh|bash|zsh)\b" r"|\bwget\b[^\n|]*-O-\s*\|\s*(?:sh|bash|zsh)\b" @@ -400,9 +384,6 @@ RE_SHELL_DROPPER = re.compile( ) -# --------------------------------------------------------------------------- -# Finding dataclass -# --------------------------------------------------------------------------- @dataclass class Finding: severity: str @@ -412,9 +393,7 @@ class Finding: evidence: str = "" -# --------------------------------------------------------------------------- # Checkers -# --------------------------------------------------------------------------- def check_pth_file(content: str, filename: str, package: str) -> list[Finding]: @@ -425,7 +404,7 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]: """ findings = [] - # Only care about .pth files that have import lines (executable) + # Only .pth files with import lines are executable import_lines = [line for line in content.splitlines() if RE_PTH_IMPORT.match(line)] if not import_lines: return findings # Pure path entries, inert @@ -470,7 +449,7 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]: ) ) - # Large base64 blob (special handling for blob size) + # Large base64 blob if RE_LARGE_BLOB.search(content): blob = RE_LARGE_BLOB.search(content).group() findings.append( @@ -483,7 +462,7 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]: ) ) - # Catch-all: any import line at all in .pth (if nothing else triggered) + # Catch-all: any import line in .pth if nothing else triggered if not findings and import_lines: evidence = "\n".join(import_lines[:5]) if len(import_lines) > 5: @@ -521,7 +500,7 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: is_setup = basename in ("setup.py", "setup.cfg") is_init = basename == "__init__.py" - # Pre-compute all pattern matches + # Pre-compute pattern matches has_network = bool(RE_NETWORK.search(content)) has_subprocess = bool(RE_SUBPROCESS.search(content)) has_base64 = bool(RE_BASE64.search(content)) @@ -546,9 +525,7 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: has_c2_polling = bool(RE_C2_POLLING.search(content)) has_may12_ioc = bool(RE_MAY12_IOC.search(content)) - # --------------------------------------------------------------- # CRITICAL: combination patterns that strongly indicate malice - # --------------------------------------------------------------- # base64 decode + subprocess execution (staged payload) if has_base64 and has_subprocess: @@ -740,9 +717,7 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: ) ) - # --------------------------------------------------------------- # HIGH: single strong signals or weaker combinations - # --------------------------------------------------------------- # Obfuscated payload: base64 + exec/eval + large blob if has_base64 and has_exec_eval and has_blob: @@ -879,9 +854,7 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: ) ) - # --------------------------------------------------------------- # MEDIUM: standalone signals (informational, may be legitimate) - # --------------------------------------------------------------- # base64 + exec/eval without blob if has_base64 and has_exec_eval and not has_blob: @@ -978,23 +951,18 @@ def _extract_evidence( return " | ".join(matches) if matches else "" -# --------------------------------------------------------------------------- # Non-Python checkers -# --------------------------------------------------------------------------- -# Several recent PyPI compromises (PyTorch Lightning 2.6.x, ForceMemo) -# carried the active payload in a bundled .js / .sh / workflow yaml so -# the Python imports looked clean on first glance. These checkers scan -# those file types when they appear inside a Python wheel/sdist. +# Recent PyPI compromises (Lightning 2.6.x, ForceMemo) carried the payload in a +# bundled .js / .sh / workflow yaml so the Python imports looked clean. These +# checkers scan those file types when they appear inside a wheel/sdist. def check_js_file(content: str, filename: str, package: str) -> list[Finding]: """Run JS-side checks. Triggered by .js / .mjs / .cjs / .ts.""" findings = [] - # A JS file *inside a Python wheel* that's larger than 100 KB is - # itself anomalous (legit Python packages don't ship hand-written - # JS bundles). Combined with ANY of the other JS heuristics it is - # CRITICAL; standalone it is HIGH. + # A >100 KB JS file inside a Python wheel is anomalous: CRITICAL combined + # with any other JS heuristic, HIGH standalone. is_large = len(content) > 100 * 1024 has_obf = bool(RE_JS_OBFUSCATION.search(content)) has_web3 = bool(RE_WEB3_HIJACK.search(content)) @@ -1119,10 +1087,8 @@ def check_shell_file(content: str, filename: str, package: str) -> list[Finding] def check_workflow_file(content: str, filename: str, package: str) -> list[Finding]: """Run GitHub-Actions workflow checks. Triggered by .github/workflows/*.yml.""" findings = [] - # A GitHub workflow file inside a *PyPI package* is itself - # suspicious (Shai-Hulud's whole MO is to plant `shai-hulud.yml` - # in every repo it can write to). Anything matching the workflow - # injection signature gets flagged CRITICAL. + # A workflow file inside a PyPI package is suspicious (Shai-Hulud plants + # `shai-hulud.yml` everywhere); injection-signature matches are CRITICAL. if RE_WORKFLOW_INJECT.search(content): findings.append( Finding( @@ -1166,15 +1132,11 @@ def check_workflow_file(content: str, filename: str, package: str) -> list[Findi return findings -# --------------------------------------------------------------------------- # Archive handling -# --------------------------------------------------------------------------- # Tarbomb caps, mirrored from scripts/scan_npm_packages.py::safe_extract. -# Refuses zip-of-death / tar-of-death archives so a hostile sdist or -# wheel cannot exhaust memory or fill the temp dir before content -# scanning even starts. Keep these constants in sync with the npm side; -# we duplicate rather than import to keep `scan_packages.py` standalone. +# Refuses zip/tar-of-death so a hostile archive cannot exhaust memory before +# scanning. Keep in sync with the npm side; duplicated to stay standalone. HARD_MAX_FILE_BYTES = 64 * 1024 * 1024 # 64 MiB per member HARD_MAX_TOTAL_BYTES = 512 * 1024 * 1024 # 512 MiB cumulative HARD_MAX_MEMBERS = 50_000 # entries per archive @@ -1183,11 +1145,8 @@ HARD_MAX_MEMBERS = 50_000 # entries per archive def _refuse_unsafe_member_name(name: str) -> str | None: """Return a refusal reason for a member name, or None if safe. - Mirrors `scan_npm_packages.py::safe_extract` semantics: no absolute - paths, no `..` traversal segments. The caller is responsible for - checking the resolved path lands inside the extract root, but for - iter_archive_files we never write to disk so the name-shape check - plus the in-memory size cap is sufficient. + Mirrors `safe_extract`: no absolute paths, no `..` traversal. We never write + to disk, so the name-shape check plus the in-memory size cap is sufficient. """ if name.startswith("/") or ".." in Path(name).parts: return f"unsafe member name {name!r}" @@ -1197,9 +1156,8 @@ def _refuse_unsafe_member_name(name: str) -> str | None: def iter_archive_files(archive_path: str): """Yield (filename, text_content) for every file in a wheel/sdist. - Streams members with size + count caps applied at the member level - so a tarbomb / zipbomb cannot blow up the scanner's memory budget. - On cap breach we emit a `[WARN]` log and short-circuit the archive. + Streams members with per-member size + count caps so a tarbomb/zipbomb can't + blow the memory budget. On cap breach, emits a `[WARN]` and short-circuits. """ path = Path(archive_path) @@ -1225,7 +1183,7 @@ def iter_archive_files(archive_path: str): file = sys.stderr, ) continue - # Declared (uncompressed) size cap. + # Declared (uncompressed) size cap if info.file_size > HARD_MAX_FILE_BYTES: print( f" [WARN] {path.name}: skipped {info.filename!r} " @@ -1262,9 +1220,8 @@ def iter_archive_files(archive_path: str): file = sys.stderr, ) return - # Refuse symlinks / hardlinks / devices outright -- the - # scanner never writes them anyway, but tar parsers - # have historically dereferenced them on extract. + # Refuse symlinks/hardlinks/devices: tar parsers have + # historically dereferenced them on extract. if member.issym() or member.islnk(): print( f" [WARN] {path.name}: refused link member " f"{member.name!r}", @@ -1305,8 +1262,7 @@ def iter_archive_files(archive_path: str): f = tf.extractfile(member) if f is None: continue - # Bound the read so a tar header that lies about - # size cannot OOM us. + # Bound the read: a tar header may lie about size data = f.read(HARD_MAX_FILE_BYTES + 1) if len(data) > HARD_MAX_FILE_BYTES: print( @@ -1327,11 +1283,9 @@ def iter_archive_files(archive_path: str): def scan_archive(archive_path: str, package: str) -> list[Finding]: """Scan all files in an archive for malicious patterns. - A corrupted archive container (truncated wheel, bad gzip header, - etc.) used to be silently skipped by an ``except Exception: continue`` - inside ``iter_archive_files``. Per the silent-failure hardening - (SF1) it now emits a CRITICAL ``archive_corrupted`` finding so the - main loop counts and surfaces it rather than reporting "0 findings". + A corrupted archive container (truncated wheel, bad gzip header, etc.) emits + a CRITICAL ``archive_corrupted`` finding rather than being silently skipped + and reported as "0 findings" (silent-failure hardening SF1). """ findings: list[Finding] = [] try: @@ -1342,22 +1296,17 @@ def scan_archive(archive_path: str, package: str) -> list[Finding]: elif lower.endswith(".py"): findings.extend(check_py_file(content, filename, package)) elif lower.endswith((".js", ".mjs", ".cjs", ".ts")): - # Lightning 2.6.x hid its real payload in a 14.8 MB - # router_runtime.js inside a Python wheel. Without this - # branch we'd have only seen the small Python loader. + # Lightning 2.6.x hid its payload in a 14.8 MB router_runtime.js; + # without this branch we'd only see the small Python loader. findings.extend(check_js_file(content, filename, package)) elif lower.endswith((".sh", ".bash")): findings.extend(check_shell_file(content, filename, package)) elif "/.github/workflows/" in lower and lower.endswith((".yml", ".yaml")): - # Shai-Hulud / ForceMemo plant their own GHA workflow. - # A workflow file inside a *PyPI package* is on its own - # already a yellow flag; pattern-match the worm signatures. + # Shai-Hulud/ForceMemo plant their own GHA workflow findings.extend(check_workflow_file(content, filename, package)) except (zipfile.BadZipFile, tarfile.TarError, EOFError, OSError) as exc: - # The archive cannot be opened or is structurally broken. A - # benign wheel/sdist always opens; a malformed one is either a - # transport corruption (treat as scan failure) or a deliberate - # attempt to bypass scanners that swallow archive errors. + # Archive cannot be opened / is structurally broken: either transport + # corruption or a deliberate attempt to bypass error-swallowing scanners. findings.append( Finding( CRITICAL, @@ -1370,9 +1319,7 @@ def scan_archive(archive_path: str, package: str) -> list[Finding]: return findings -# --------------------------------------------------------------------------- # Download packages -# --------------------------------------------------------------------------- _RE_PYPI_SPEC_VERSION = re.compile(r"==\s*([A-Za-z0-9_.\-+!]+)") @@ -1382,10 +1329,8 @@ def _check_blocked_pypi_versions(specs: list[str]) -> tuple[list[str], list[Find """Filter ``specs`` against ``BLOCKED_PYPI_VERSIONS``. Returns ``(safe_specs, findings)``. Each blocked spec emits a CRITICAL - ``Finding`` and is removed from the returned spec list so the caller - never fetches the malicious tarball. Specs without an ``==X.Y.Z`` pin - pass through unchanged -- pip will resolve them at download time and - the existing scanners will catch the payload via the IOC regexes. + ``Finding`` and is dropped so the malicious tarball is never fetched. Specs + without an ``==X.Y.Z`` pin pass through; the IOC regexes catch them later. """ safe: list[str] = [] findings: list[Finding] = [] @@ -1407,7 +1352,7 @@ def _check_blocked_pypi_versions(specs: list[str]) -> tuple[list[str], list[Find f"{name}=={version} is on the BLOCKED_PYPI_VERSIONS list", ) ) - # Drop the spec; do not download. + # Drop the spec; do not download continue safe.append(spec) return safe, findings @@ -1416,14 +1361,11 @@ def _check_blocked_pypi_versions(specs: list[str]) -> tuple[list[str], list[Find def _pip_download_env() -> dict[str, str]: """Return a scrubbed environment for invoking `pip download`. - Hostile shells / CI configs can override the index with PIP_INDEX_URL, - PIP_EXTRA_INDEX_URL, or a user `pip.conf`. We strip every PIP_* - override and route the resolver explicitly at PyPI. PIP_CONFIG_FILE - is forced to /dev/null so a stray ~/.pip/pip.conf with an - extra-index-url cannot bypass the pin. + Strips every PIP_* override and forces the resolver at PyPI; PIP_CONFIG_FILE + is /dev/null so a stray pip.conf extra-index-url cannot bypass the pin. """ env = {**os.environ} - # Drop any user override. + # Drop any user override for key in [k for k in env if k.startswith("PIP_")]: env.pop(key, None) env["PIP_INDEX_URL"] = "https://pypi.org/simple" @@ -1433,10 +1375,8 @@ def _pip_download_env() -> dict[str, str]: return env -# Pip resolver flags shared by both download branches. Pinning the -# index URL on the CLI is belt + braces with the env scrub above. -# `--no-build-isolation` is deliberately NOT set; we never invoke -# setup.py at all because of `--only-binary :all:`. +# Pip resolver flags shared by both download branches. CLI index-URL pin is +# belt + braces with the env scrub; `--only-binary :all:` avoids running setup.py. _PIP_DOWNLOAD_PIN_FLAGS = [ "--index-url", "https://pypi.org/simple", @@ -1445,9 +1385,8 @@ _PIP_DOWNLOAD_PIN_FLAGS = [ ] -# Strip any character that could escape `dest` via `os.path.join`. This -# is the last line of defence before `pkg_dir = os.path.join(dest, ...)` -# so a spec like `../../etc/foo==1.0` cannot land outside the temp tree. +# Strip characters that could escape `dest` via `os.path.join`, so a spec like +# `../../etc/foo==1.0` cannot land outside the temp tree. _RE_PKG_NAME_SANITIZE = re.compile(r"[^A-Za-z0-9._-]") @@ -1459,27 +1398,21 @@ def download_packages( ) -> tuple[list[tuple[str, str]], list[str]]: """Download packages to dest using pip download. NEVER installs. - Returns ``(results, download_errors)`` where ``results`` is a list of - ``(spec_or_name, filepath)`` for every downloaded archive and - ``download_errors`` is a list of one-line transport-failure summaries. - A non-empty ``download_errors`` MUST cause the caller to exit non-zero - even if no findings were produced; a silent ``0 findings, scan - incomplete`` is the bug class this return-shape was widened to fix. + Returns ``(results, download_errors)``: ``results`` is ``(spec_or_name, + filepath)`` per archive; ``download_errors`` is one-line transport-failure + 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". - When with_deps=True, downloads the full transitive dependency tree - in a single pip invocation (all archives land in one flat dir). - When with_deps=False (default), downloads each spec individually - with --no-deps. + 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. """ results: list[tuple[str, str]] = [] download_errors: list[str] = [] env = _pip_download_env() if with_deps: - # Single pip download call for all specs + their transitive deps. - # `--only-binary :all:` refuses sdists so we never execute a - # setup.py just to learn dependency metadata; combined with the - # scrubbed env, pip is wired hard at pypi.org. + # 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, @@ -1495,7 +1428,7 @@ def download_packages( cmd, capture_output = True, text = True, - timeout = 600, # transitive resolution can be slow + timeout = 600, # transitive resolution is slow env = env, ) if proc.returncode != 0: @@ -1511,14 +1444,12 @@ def download_packages( for fname in sorted(os.listdir(dest)): fpath = os.path.join(dest, fname) if os.path.isfile(fpath): - # Derive package name from filename pkg_name = fname.split("-")[0].replace("_", "-").lower() results.append((pkg_name, fpath)) else: for spec in specs: raw_name = _extract_pkg_name(spec) - # Sanitize before joining into `dest` so a hostile spec - # cannot path-traverse out of the destination directory. + # Sanitize before joining into `dest` to prevent path traversal safe_name = _RE_PKG_NAME_SANITIZE.sub("_", raw_name) or "_pkg" pkg_dir = os.path.join(dest, safe_name) os.makedirs(pkg_dir, exist_ok = True) @@ -1552,7 +1483,6 @@ def download_packages( download_errors.append(msg) continue - # Find downloaded file(s) for fname in os.listdir(pkg_dir): fpath = os.path.join(pkg_dir, fname) if os.path.isfile(fpath): @@ -1560,9 +1490,7 @@ def download_packages( return results, download_errors -# --------------------------------------------------------------------------- # Parse requirements files -# --------------------------------------------------------------------------- _RE_NAME = re.compile(r"^([A-Za-z0-9]([A-Za-z0-9._-]*[A-Za-z0-9])?)") @@ -1591,7 +1519,7 @@ def parse_requirements(req_files: list[str]) -> list[dict]: if not line or line.startswith("#") or line.startswith("-"): continue is_git = line.startswith("git+") or "git+" in line.split("#")[0] - # Strip inline comments and environment markers for spec + # Strip inline comments and env markers spec = line.split("#")[0].strip() spec = spec.split(";")[0].strip() if not spec: @@ -1624,7 +1552,7 @@ def get_downloaded_version(archive_path: str) -> str | None: parts = basename[:-4].split("-") if len(parts) >= 2: return parts[1] - # Sdist: name-version.tar.gz / .tar.bz2 / .zip + # Sdist: name-version. for ext in (".tar.gz", ".tar.bz2", ".tar.xz", ".tar", ".zip"): if basename.endswith(ext): stem = basename[: -len(ext)] @@ -1634,9 +1562,7 @@ def get_downloaded_version(archive_path: str) -> str | None: return None -# --------------------------------------------------------------------------- # Display -# --------------------------------------------------------------------------- def severity_color(sev: str) -> str: @@ -1652,7 +1578,6 @@ def print_findings(findings: list[Finding]) -> None: print("\n All clean. No suspicious patterns found.") return - # Sort by severity findings.sort(key = lambda f: SEVERITY_ORDER.get(f.severity, 99)) print(f"\n {'=' * 72}") @@ -1682,9 +1607,7 @@ def print_findings(findings: list[Finding]) -> None: print(f" Summary: {', '.join(parts)}") -# --------------------------------------------------------------------------- # PyPI version queries and --fix logic -# --------------------------------------------------------------------------- def version_sort_key(v: str) -> tuple: @@ -1708,15 +1631,13 @@ def version_sort_key(v: str) -> tuple: base = v_clean[0] suffix = v[len(base) :] - # Parse numeric parts parts = [] for seg in base.split("."): try: parts.append(int(seg)) except ValueError: parts.append(0) - # Pad to at least 3 parts - while len(parts) < 3: + while len(parts) < 3: # pad to at least 3 parts parts.append(0) # Suffix ordering: dev < alpha < beta < rc < (none) < post @@ -1732,7 +1653,7 @@ def version_sort_key(v: str) -> tuple: elif suffix_lower.startswith("post"): suffix_rank = 1 else: - suffix_rank = 0 # stable release + suffix_rank = 0 # stable return (epoch, tuple(parts), suffix_rank, suffix) @@ -1772,11 +1693,10 @@ def find_safe_version( print(f" [WARN] No versions found on PyPI for {name}", file = sys.stderr) return None - # Find index of bad version try: bad_idx = versions.index(bad_ver) except ValueError: - # bad_ver might have been resolved to a different string; search by sort key + # bad_ver may resolve to a different string; search by sort key bad_key = version_sort_key(bad_ver) bad_idx = None for i, v in enumerate(versions): @@ -1820,7 +1740,6 @@ def find_safe_version( print(f" {ver} -- CRITICAL finding(s), skipping") break - # Clean up scan dir for this version shutil.rmtree(scan_dir, ignore_errors = True) if clean: @@ -1850,8 +1769,7 @@ def update_req_line(raw_line: str, safe_ver: str, old_ver: str | None) -> str: code_part, marker = code_part.split(";", 1) marker = ";" + marker - # Replace version specifier - # Match patterns like ==1.2.3, >=1.2, ~=1.0, <=2.0, !=1.1, or bare name + # Replace version specifier (==1.2.3, >=1.2, ~=1.0, !=1.1, or bare name) rewritten = re.sub( r"([A-Za-z0-9._-]+)\s*(?:[><=!~]=?[^;#,\s]*(?:\s*,\s*[><=!~]=?[^;#,\s]*)*)?", lambda m: f"{m.group(1)}=={safe_ver}", @@ -1870,12 +1788,8 @@ def update_req_file(filepath: str, updates: dict[int, str]) -> None: updates: {line_num (1-indexed): new_line_text} - Writes atomically: stage in a sibling tmp file on the same - filesystem, fsync, then `os.replace` over the original. A SIGKILL - or power loss mid-write therefore either leaves the original - intact or leaves the fully new file -- never a half-written - requirements file (which would silently re-introduce a malicious - pin). + Writes atomically (sibling tmp file, fsync, os.replace) so a crash mid-write + never leaves a half-written file that re-introduces a malicious pin. """ with open(filepath) as f: lines = f.readlines() @@ -1883,8 +1797,7 @@ def update_req_file(filepath: str, updates: dict[int, str]) -> None: for line_num, new_text in updates.items(): idx = line_num - 1 if 0 <= idx < len(lines): - # Preserve original line ending - ending = "\n" if lines[idx].endswith("\n") else "" + ending = "\n" if lines[idx].endswith("\n") else "" # preserve line ending lines[idx] = new_text + ending dirpath = os.path.dirname(os.path.abspath(filepath)) or "." @@ -1909,7 +1822,7 @@ def update_req_file(filepath: str, updates: dict[int, str]) -> None: def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> None: """Run the --fix flow: find safe versions, update requirements files.""" - # Map package names to their entries for source tracking + # Map package names to entries for source tracking pkg_entries: dict[str, list[dict]] = {} for e in entries: norm = e["name"].lower().replace("-", "_").replace(".", "_") @@ -1931,8 +1844,7 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N changes_summary.append(f" SKIP {pkg_name} (git URL)") continue - # Get the currently resolved version - # Try to extract from the spec (e.g. name==1.2.3) + # Resolved version: try to extract from the spec (name==1.2.3) current_ver = None for e in related: spec = e["spec"] @@ -1947,7 +1859,6 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N downloaded = download_packages([pkg_name], dl_dir) if downloaded: current_ver = get_downloaded_version(downloaded[0][1]) - # Delete resolution download immediately shutil.rmtree(dl_dir, ignore_errors = True) if not current_ver: @@ -1986,7 +1897,6 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N for filepath, updates in file_updates.items(): update_req_file(filepath, updates) - # Print summary print(f"\n {'=' * 72}") print(f" FIX SUMMARY") print(f" {'=' * 72}") @@ -1995,9 +1905,7 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N print(f"\n Re-run without --fix to verify the scan is clean.") -# --------------------------------------------------------------------------- # Directory scanning -# --------------------------------------------------------------------------- def _find_requirements_files(root: str) -> list[str]: @@ -2014,7 +1922,7 @@ def _find_requirements_files(root: str) -> list[str]: skip_dirs = {"__pycache__", "node_modules", "venv", ".venv", "site-packages"} results = [] for dirpath, dirnames, filenames in os.walk(root): - # Skip hidden dirs and known non-requirement dirs + # Skip hidden and known non-requirement dirs dirnames[:] = [ d for d in dirnames @@ -2024,18 +1932,15 @@ def _find_requirements_files(root: str) -> list[str]: for fname in sorted(filenames): if not fname.endswith(".txt"): continue - # Match requirements*.txt anywhere if fnmatch.fnmatch(fname.lower(), "requirements*.txt"): results.append(os.path.join(dirpath, fname)) - # Match *.txt inside a directory named "requirements" + # *.txt inside a directory named "requirements" elif dirname == "requirements": results.append(os.path.join(dirpath, fname)) return sorted(results) -# --------------------------------------------------------------------------- # Main -# --------------------------------------------------------------------------- def main() -> int: @@ -2134,7 +2039,7 @@ def main() -> int: all_findings: list[Finding] = [] - # Hard pin-block: refuse to download known-malicious PyPI versions. + # Hard pin-block: refuse to download known-malicious PyPI versions specs, blocked_findings = _check_blocked_pypi_versions(specs) all_findings.extend(blocked_findings) @@ -2172,10 +2077,8 @@ def main() -> int: ) _run_fix(critical_pkgs, entries, args.max_search) - # Surface any pip-download failures BEFORE the scan-result exit code so - # an empty / partial download cannot mask itself as "0 findings, all - # clean". This is item (4) of the silent-failure hardening: an - # unresolvable spec or PyPI timeout used to print to stderr and exit 0. + # Surface pip-download failures BEFORE the exit code so a partial download + # can't masquerade as "0 findings, all clean" (silent-failure hardening 4). if download_errors: print( f"\n {'=' * 72}\n" diff --git a/scripts/stamp_studio_release.py b/scripts/stamp_studio_release.py index a3c8712a6e..7dab35ea8a 100644 --- a/scripts/stamp_studio_release.py +++ b/scripts/stamp_studio_release.py @@ -22,11 +22,8 @@ def _atomic_write_text( data: str, encoding: str = "utf-8", ) -> None: - """Atomic version of ``Path.write_text``. - - A crash or signal mid-write leaves the prior file intact; the - Studio build never reads a partial ``_studio_release_build.py``. - """ + """Atomic ``Path.write_text``: a crash mid-write leaves the prior file + intact, so the build never reads a partial ``_studio_release_build.py``.""" dirpath = str(path.parent) or "." path.parent.mkdir(parents = True, exist_ok = True) fd, tmp_path = tempfile.mkstemp(prefix = ".stamp_studio.", dir = dirpath) diff --git a/scripts/verify_comment_only_diff.py b/scripts/verify_comment_only_diff.py index b661f76859..23d06b85df 100644 --- a/scripts/verify_comment_only_diff.py +++ b/scripts/verify_comment_only_diff.py @@ -57,9 +57,7 @@ def _git_show(rev: str, path: str) -> str: def _strip_docstrings(tree: ast.AST) -> ast.AST: - """Remove every string-literal docstring (Module / FunctionDef / - AsyncFunctionDef / ClassDef). Empty body becomes ``pass`` so - ast.unparse stays valid.""" + """Remove docstrings; empty bodies become ``pass`` so unparse stays valid.""" for node in ast.walk(tree): if isinstance( node, @@ -87,9 +85,8 @@ def _normalize_py(src: str) -> str: def _strip_shell_comments(s: str) -> str: - """Strip pure-comment lines and inline trailing comments from a shell - snippet, then collapse runs of blank lines. Heuristic only: leaves a - line untouched if it has an odd quote count (open string).""" + """Strip shell comments and collapse blank lines. Heuristic: skips lines + with an odd quote count (open string).""" out = [] for line in s.splitlines(): stripped = line.lstrip() @@ -116,9 +113,7 @@ def _strip_shell_comments(s: str) -> str: def _normalize_yaml_run_strings(obj: Any) -> Any: - """Walk the parsed YAML object; for any multi-line string (i.e. a - ``run: |`` script body), strip shell comments. Returns a normalised - copy.""" + """Strip shell comments from any multi-line string (``run: |`` body).""" if isinstance(obj, dict): return {k: _normalize_yaml_run_strings(v) for k, v in obj.items()} if isinstance(obj, list): diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index 6f5b452f3c..b8cb0573fe 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -122,16 +122,13 @@ class _Builder(ast.NodeVisitor): def __init__(self): self.module = Scope("module", "", None) - self.uses: list[tuple[Scope, str, int]] = [] # (scope, name, lineno) hard loads - self.soft_uses: list[tuple[Scope, str, int]] = [] # annotations: count as "used" - # but never as "unresolved" - # (forward refs / string annos) + self.uses: list[tuple[Scope, str, int]] = [] # hard loads + # annotations: count as "used" but never as "unresolved" (forward refs) + self.soft_uses: list[tuple[Scope, str, int]] = [] def _visit_annotation(self, node, scope: Scope) -> None: - """Annotation context: with `from __future__ import annotations` these are - never evaluated (strings), and even otherwise they routinely contain forward - references. Record contained names as SOFT uses so an import used only in an - annotation still counts as used, but a forward-ref name is never 'unresolved'.""" + """Record annotation names as SOFT uses: an import used only in an annotation + counts as used, but a forward-ref name is never 'unresolved'.""" if node is None: return for n in ast.walk(node): @@ -189,7 +186,7 @@ class _Builder(ast.NodeVisitor): child = Scope("function", f"{scope.qualname}.{node.name}", scope) self._bind_type_params(node, child) self._bind_args(node.args, child) - # arg + return annotations: soft uses (may be strings / forward refs) + # arg + return annotations: soft uses for a in self._all_args(node.args): self._visit_annotation(a.annotation, child) self._visit_annotation(getattr(node, "returns", None), child) @@ -355,7 +352,7 @@ class _Builder(ast.NodeVisitor): if isinstance(node, (ast.ListComp, ast.SetComp, ast.GeneratorExp, ast.DictComp)): child = Scope("comp", f"{scope.qualname}.", scope) for i, gen in enumerate(node.generators): - # first iterable is evaluated in the enclosing scope + # first iterable evaluates in the enclosing scope self._visit_expr(gen.iter, scope if i == 0 else child) self._bind_targets(child, gen.target) for cond in gen.ifs: @@ -394,9 +391,8 @@ def _any_star(scope: Scope) -> bool: def _resolve(scope: Scope, name: str): - """LEGB resolution. Returns (status, bindings) where status in + """LEGB resolution. Returns (status, bindings); status in {'local','import','other','builtin','star','unresolved'}.""" - # global / nonlocal redirection start = scope if name in scope.globals: chain = [_module_of(scope)] @@ -441,7 +437,7 @@ def _legb_chain(scope: Scope) -> list[Scope]: chain = [scope] p = scope.parent while p is not None: - if p.kind != "class" or p.parent is None: # module-level class never happens; keep module + if p.kind != "class" or p.parent is None: # skip class scopes, keep module if p.kind != "class": chain.append(p) p = p.parent @@ -455,7 +451,7 @@ def _analyze(src: str): tree = ast.parse(src) b = _Builder() b.run(tree) - # Per-scope: unresolved load names, and import targets it resolves to. + # Per-scope: unresolved load names + import targets it resolves to. unresolved: dict[str, set[str]] = {} targets_by_scope: dict[str, set[str]] = {} target_by_use: dict[tuple[str, str], set[str]] = {} @@ -467,7 +463,7 @@ def _analyze(src: str): tids = {bd.target for bd in binds if bd.target} targets_by_scope.setdefault(scope.qualname, set()).update(tids) target_by_use.setdefault((scope.qualname, name), set()).update(tids) - # soft uses (annotations): only contribute to "used", never to "unresolved" + # soft uses (annotations): contribute to "used" only, never "unresolved" for scope, name, _ln in b.soft_uses: status, binds = _resolve(scope, name) if status == "import": @@ -495,7 +491,7 @@ def _analyze(src: str): x.kind not in ("import", "importfrom") for x in bs ): ambiguous.setdefault(scope.qualname, set()).add(n) - # scope tree isn't stored; rebuild via uses is hard. We approximate with module only. + # scope tree isn't stored; approximate with module only. walk_scopes(module) return { @@ -524,14 +520,9 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] Blocker signals (precise, no relocation false-positives): UNRESOLVED-NEW - a load became undefined (dangling alias / removed import). - NEW-UNUSED-HOIST - a module-level import added by THIS change is resolved by - NO load. A correct hoist always wires its new import to a - reference; if the alias was left un-normalized OR renamed - to the wrong name, the hoisted import ends up unused. This - single signal catches BOTH user-described failure modes and - does NOT fire for code merely relocated to another file - (that removes the import, it doesn't add an unused one). - TARGET-CHANGED - the same (scope, name) load resolves to a different import + NEW-UNUSED-HOIST - a module-level import added by this change is resolved by + NO load (un-normalized alias or wrong rename target). + TARGET-CHANGED - same (scope, name) load resolves to a different import target before vs after (a same-name re-point). """ a = _analyze(before_src) @@ -566,14 +557,13 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] ) ) - # 2. HOISTED-IMPORT-UNUSED (the core botched-hoist / wrong-rename signal) - # A module-level import in AFTER that NO load resolves to, and which was - # either newly added by this change OR was actually used before. Excludes: - # - relocation (the import is REMOVED, so it's not in after at all) - # - stable pre-existing re-exports (unused before AND after, not newly added) + # 2. HOISTED-IMPORT-UNUSED (core botched-hoist / wrong-rename signal) + # A module-level import in AFTER that NO load resolves to, that was either + # newly added by this change OR actually used before. Excludes relocation + # (import removed) and stable pre-existing re-exports. for n, tids in b["module_import_targets"].items(): if tids & after_used: - continue # resolved by something -> fine + continue # resolved -> fine newly_added = bool(tids - before_module_targets) was_used_before = bool(tids & before_used) if newly_added or was_used_before: @@ -619,8 +609,7 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] findings.append(("WARN", f"{path}: AMBIGUOUS-BIND '{n}' import+non-import in {scope}")) # 6. TARGET-MISSING (informational): a scope stopped resolving to an import - # target. Real bugs are already covered above; remaining cases are code - # relocated to another file (e.g. a moved helper). Shown for transparency. + # target. Real bugs are covered above; remaining cases are relocated code. for scope, tbefore in a["targets_by_scope"].items(): tafter = b["targets_by_scope"].get(scope, set()) for t in sorted(tbefore - tafter): @@ -667,13 +656,13 @@ _SELF_TESTS = { "BLOCKER", ), "local_var_clash": ( - # _b renamed to b, but b is a LOCAL variable in f -> import silently unused + # _b renamed to b, but b is a LOCAL var in f -> import silently unused "def f(b):\n import re as _b\n return _b.compile(b)\n", "import re\ndef f(b):\n return b.compile(b)\n", # 'b' is the param, not the module "BLOCKER", ), "substring_safe": ( - # correct _copy->copy rename while a config_copy var exists: NO false positive + # correct _copy->copy rename while config_copy var exists: NO false positive "def f(config):\n" " import copy as _copy\n" " config_copy = _copy.deepcopy(config)\n" @@ -728,10 +717,9 @@ def _pyflakes_undefined(path: str) -> set[str] | None: def audit_files(paths: list[str]) -> int: - """Single-version robustness audit. For every file: confirm the analyzer does - not crash, then cross-check its 'unresolved' names against pyflakes. Any name - the resolver flags that pyflakes does NOT call undefined is a tool FALSE - POSITIVE (a resolver gap to fix).""" + """Single-version robustness audit: confirm the analyzer doesn't crash, then + cross-check its 'unresolved' names against pyflakes. A name the resolver flags + that pyflakes accepts is a tool false positive.""" n_files = n_err = n_fp = n_syntax = 0 fp_detail: dict[str, set[str]] = {} err_detail: dict[str, str] = {} diff --git a/studio/backend/_platform_compat.py b/studio/backend/_platform_compat.py index 490e1820d7..013667d5ca 100644 --- a/studio/backend/_platform_compat.py +++ b/studio/backend/_platform_compat.py @@ -4,17 +4,16 @@ """ Compatibility shim for Anaconda/conda-forge Python builds. -Anaconda modifies sys.version to include distributor metadata between pipe -characters, e.g. '3.12.4 | packaged by Anaconda, Inc. | (main, ...) [MSC ...]'. -Python's platform._sys_version() has a hardcoded regex that cannot parse this, -raising ValueError. CPython closed this as "not planned" (cpython#102396). +Anaconda puts distributor metadata between pipes in sys.version, e.g. +'3.12.4 | packaged by Anaconda, Inc. | (main, ...) [MSC ...]'. The regex in +platform._sys_version() can't parse this and raises ValueError (cpython#102396, +closed as "not planned"). -This module seeds platform._sys_version_cache so the stdlib parser never sees -the problematic string, fixing the import chain: +We seed platform._sys_version_cache so the stdlib parser never sees the bad +string, fixing the import chain: structlog -> rich.pretty -> attrs._compat -> platform.python_implementation() -Import this module before any library imports that may trigger the above chain. -Safe to import multiple times (no-op if cache is already seeded or no pipes). +Import before any library that may trigger that chain. Idempotent. """ import platform @@ -23,18 +22,17 @@ import sys def _seed_sys_version_cache() -> None: - """One-shot cache prime: parse a cleaned sys.version and seed the cache.""" + """Parse a cleaned sys.version and seed the cache once.""" raw = sys.version # Strip paired |...| segments (Anaconda, conda-forge metadata) cleaned = re.sub(r"\s*\|[^|]*\|\s*", " ", raw).strip() - # Format B: "ver (build) | label | (build_dup) \n[compiler]" - # After pipe-strip, two consecutive (...) groups remain; drop the second. + # Pipe-strip can leave two consecutive (...) groups; drop the second. cleaned = re.sub(r"(\([^)]*\))\s+\([^)]*\)", r"\1", cleaned) if "|" in cleaned: - # Unpaired pipe remaining -- keep version + everything from "(" onward + # Unpaired pipe left: keep version + everything from "(" onward m = re.match(r"([\w.+]+)\s*", cleaned) p = cleaned.find("(") if m and p > 0: @@ -43,13 +41,12 @@ def _seed_sys_version_cache() -> None: if cleaned == raw: return # Nothing to fix - # Parse the cleaned string through the real stdlib parser try: result = platform._sys_version(cleaned) except ValueError: - return # Cleaning didn't produce a parseable string; don't make things worse + return # Still unparsable; don't make things worse - # Seed the cache so future calls with the raw string skip parsing entirely + # Seed the cache so future calls with the raw string skip parsing cache = getattr(platform, "_sys_version_cache", None) if isinstance(cache, dict): cache[raw] = result diff --git a/studio/backend/auth/__init__.py b/studio/backend/auth/__init__.py index b3e1a8a9c0..a63d5303d4 100644 --- a/studio/backend/auth/__init__.py +++ b/studio/backend/auth/__init__.py @@ -1,9 +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 -""" -Authentication module for JWT-based auth with SQLite storage. -""" +"""Authentication module for JWT-based auth with SQLite storage.""" from .authentication import ( create_access_token, diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py index bb5b873e65..9dd56489eb 100644 --- a/studio/backend/auth/authentication.py +++ b/studio/backend/auth/authentication.py @@ -58,7 +58,7 @@ def create_access_token( """ Create a signed JWT for the given subject (e.g. username). - Tokens are valid across restarts because the signing secret is stored in SQLite. + Valid across restarts: the signing secret is stored in SQLite. """ to_encode = {"sub": subject} if desktop: @@ -100,7 +100,7 @@ def create_refresh_token(subject: str, *, desktop: bool = False) -> str: """ Create a random refresh token, store its hash in SQLite, and return it. - Refresh tokens are opaque (not JWTs) and expire after REFRESH_TOKEN_EXPIRE_DAYS. + Refresh tokens are opaque (not JWTs); expire after REFRESH_TOKEN_EXPIRE_DAYS. """ token = secrets.token_urlsafe(48) expires_at = datetime.now(timezone.utc) + timedelta(days = REFRESH_TOKEN_EXPIRE_DAYS) @@ -112,8 +112,8 @@ def refresh_access_token(refresh_token: str) -> Tuple[Optional[str], Optional[st """ Validate a refresh token and issue a new access token. - The refresh token itself is NOT consumed — it stays valid until expiry. - Returns a new access_token or None if the refresh token is invalid/expired. + The refresh token is NOT consumed; it stays valid until expiry. + Returns a new access_token, or None if the refresh token is invalid/expired. """ verified = verify_refresh_token(refresh_token) if verified is None: @@ -128,7 +128,7 @@ def refresh_access_token(refresh_token: str) -> Tuple[Optional[str], Optional[st def reload_secret() -> None: """ - Keep legacy API compatibility for callers expecting auth storage init. + Legacy API compat for callers expecting auth storage init. Auth now resolves the current signing secret directly from SQLite. """ @@ -156,15 +156,7 @@ async def get_current_subject_allow_password_change( async def _get_current_subject( credentials: HTTPAuthorizationCredentials, *, allow_password_change: bool ) -> str: - """ - FastAPI dependency to validate the JWT and return the subject. - - Use this as a dependency on routes that should be protected, e.g.: - - @router.get("/secure") - async def secure_endpoint(current_subject: str = Depends(get_current_subject)): - ... - """ + """FastAPI dependency: validate the JWT and return the subject. Use on protected routes.""" token = credentials.credentials # --- API key path (sk-unsloth-...) --- diff --git a/studio/backend/auth/hashing.py b/studio/backend/auth/hashing.py index 873381db04..0a1b0c15be 100644 --- a/studio/backend/auth/hashing.py +++ b/studio/backend/auth/hashing.py @@ -23,7 +23,7 @@ def hash_password(password: str, salt: str | None = None) -> Tuple[str, str]: "sha256", password.encode("utf-8"), salt.encode("utf-8"), - 100_000, # 100k iterations + 100_000, ) return salt, dk.hex() diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py index 0e34a0cf28..fa5b985513 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -1,9 +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 -""" -SQLite storage for authentication data (user credentials + JWT secret). -""" +"""SQLite storage for auth data (user credentials + JWT secret).""" import hashlib import os @@ -17,42 +15,40 @@ from utils.paths import auth_db_path, ensure_dir DB_PATH = auth_db_path() DEFAULT_ADMIN_USERNAME = "unsloth" -# Plaintext bootstrap password file — lives beside auth.db, deleted on -# first password change so the credential never lingers on disk. +# Plaintext bootstrap password file beside auth.db, deleted on first password +# change so the credential never lingers on disk. _BOOTSTRAP_PW_PATH = DB_PATH.parent / ".bootstrap_password" -# In-process cache so we don't re-read the file on every HTML serve. +# In-process cache to avoid re-reading the file on every HTML serve. _bootstrap_password: Optional[str] = None def generate_bootstrap_password() -> str: """Generate a 4-word diceware passphrase and persist it to disk. - The passphrase is written to ``_BOOTSTRAP_PW_PATH`` so that it - survives server restarts (the DB only stores the *hash*). On - subsequent calls / restarts, the persisted value is returned. + Persisted (the DB stores only the hash) so it survives restarts; later + calls return the persisted value. """ global _bootstrap_password - # 1. Already cached in this process? + # Cached in this process? if _bootstrap_password is not None: return _bootstrap_password - # 2. Already persisted from a previous run? + # Persisted from a previous run? if _BOOTSTRAP_PW_PATH.is_file(): _bootstrap_password = _BOOTSTRAP_PW_PATH.read_text().strip() if _bootstrap_password: return _bootstrap_password - # 3. First-ever startup — generate a fresh passphrase. + # First startup: generate a fresh passphrase. import diceware _bootstrap_password = diceware.get_passphrase( options = diceware.handle_options(args = ["-n", "4", "-d", "", "-c"]) ) - # Persist so the *same* passphrase is used if the server restarts - # before the user changes the password. + # Persist so the same passphrase survives restarts until password change. ensure_dir(_BOOTSTRAP_PW_PATH.parent) _BOOTSTRAP_PW_PATH.write_text(_bootstrap_password) try: @@ -88,21 +84,12 @@ def clear_bootstrap_password() -> None: def _hash_token(token: str) -> str: - """SHA-256 hash helper used for refresh token storage. + """SHA-256 hash helper for refresh token storage. - Plain SHA-256 is intentional here: refresh tokens are high-entropy - random strings from ``secrets.token_urlsafe(48)`` (384 bits of - entropy), so a slow KDF (Argon2 / bcrypt / PBKDF2) provides zero - additional security — no attacker can brute-force 2^384 regardless - of hash speed — while adding tens of ms of CPU to every refresh. - See the OWASP Password Storage Cheat Sheet on fast-vs-slow hashing - of high-entropy inputs. - - API keys use the separate ``_pbkdf2_api_key`` helper below, which - runs PBKDF2-HMAC-SHA256 with a persistent server-side salt — not - for cryptographic reasons (128-bit random tokens don't need slow - hashing), but because CodeQL's ``py/weak-sensitive-data-hashing`` - query mislabels API keys as passwords and demands a KDF. + Plain SHA-256 is intentional: refresh tokens are 384-bit random strings, so + a slow KDF adds no security while costing per-refresh latency. API keys use + the separate ``_pbkdf2_api_key`` helper, only to satisfy CodeQL's + ``py/weak-sensitive-data-hashing`` query, not for crypto reasons. """ return hashlib.sha256(token.encode("utf-8")).hexdigest() @@ -176,22 +163,19 @@ def get_connection() -> sqlite3.Connection: # ── API-key PBKDF2 salt ──────────────────────────────────────────────── # -# Module-level cache for the persistent API-key PBKDF2 salt. Populated -# lazily on first use via ``_get_or_create_api_key_pbkdf2_salt``. Not -# protected by a lock because (a) the ``INSERT OR IGNORE`` provides -# atomicity at the SQLite layer and (b) concurrent populations converge -# on the same value, so the worst case is a harmless duplicate read on -# startup. +# Module-level cache for the persistent API-key PBKDF2 salt, populated lazily +# via ``_get_or_create_api_key_pbkdf2_salt``. No lock needed: (a) ``INSERT OR +# IGNORE`` is atomic at the SQLite layer and (b) concurrent populations +# converge on the same value, so the worst case is a harmless duplicate read +# on startup. _api_key_pbkdf2_salt_cache: Optional[bytes] = None def _get_or_create_api_key_pbkdf2_salt() -> bytes: """Return the persistent API-key PBKDF2 salt, generating it once if missing. - Stored as a hex-encoded 32-byte random value in the ``app_secrets`` - table under key ``"api_key_pbkdf2_salt"``. Regenerated only if the row - is missing (i.e. fresh install, or operator manually deleted the row - and accepts invalidating existing API keys). + Hex-encoded 32-byte random value in ``app_secrets``. Regenerated only when + the row is missing (fresh install, or operator deleted it). """ global _api_key_pbkdf2_salt_cache if _api_key_pbkdf2_salt_cache is not None: @@ -233,22 +217,10 @@ _DESKTOP_SECRET_CREATED_AT_KEY = "desktop_secret_created_at" def _pbkdf2_api_key(raw_key: str) -> str: """PBKDF2-HMAC-SHA256 an API key with a persistent server-side salt. - Used for API-key storage ONLY, not refresh tokens. Matches the - PBKDF2 algorithm + iteration count used by the password hasher in - ``auth/hashing.py`` so the codebase is consistent on which KDF it - uses for credential storage. - - Notes on why a slow KDF here is *only* a CodeQL appeasement and - *not* a cryptographic requirement: API keys are cryptographically - random 128-bit tokens (via ``secrets.token_hex``), so brute force - against 2^128 is infeasible regardless of hash speed. CodeQL's - ``py/weak-sensitive-data-hashing`` query mislabels these tokens as - "password" sensitive data and then demands a KDF from its - allowlist (Argon2 / scrypt / bcrypt / PBKDF2). Per the query's - own recommendation page we use PBKDF2. The persistent salt is - still loaded from ``app_secrets`` so an attacker dumping the - ``api_keys`` table alone cannot derive hashes for candidate - tokens without also obtaining the salt row. + For API-key storage ONLY, not refresh tokens. The slow KDF is only to + appease CodeQL's ``py/weak-sensitive-data-hashing`` query, not a crypto + requirement (API keys are random 128-bit tokens). The salt lives in + ``app_secrets`` so dumping ``api_keys`` alone can't derive hashes. """ salt = _get_or_create_api_key_pbkdf2_salt() dk = hashlib.pbkdf2_hmac( @@ -513,7 +485,7 @@ def verify_refresh_token(token: str) -> Optional[Tuple[str, bool]]: token_hash = _hash_token(token) conn = get_connection() try: - # Clean up any expired tokens while we're here + # Opportunistically clean up expired tokens conn.execute( "DELETE FROM refresh_tokens WHERE expires_at < ?", (datetime.now(timezone.utc).isoformat(),), diff --git a/studio/backend/colab.py b/studio/backend/colab.py index c3a1e03fbe..ba46c52a6a 100644 --- a/studio/backend/colab.py +++ b/studio/backend/colab.py @@ -2,15 +2,13 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """ -Colab-specific helpers for running Unsloth Studio. -Uses Colab's built-in proxy - no external tunneling needed! +Colab helpers for Unsloth Studio. Uses Colab's built-in proxy. """ from pathlib import Path import sys -# Fix for Anaconda/conda-forge Python: seed platform._sys_version_cache before -# any library imports that trigger attrs -> rich -> structlog -> platform crash. +# Seed platform._sys_version_cache before attrs->rich->structlog->platform crash on conda Python. # See: https://github.com/python/cpython/issues/102396 _backend_dir = str(Path(__file__).parent) if _backend_dir not in sys.path: @@ -25,11 +23,10 @@ logger = get_logger(__name__) def get_colab_url(port: int = 8888) -> str: """ - Get the actual Colab proxy URL for a port. + Get the Colab proxy URL for a port. - Retries up to 3 times and validates that the result is a real HTTPS Colab - URL before returning. Falls back to http://localhost:{port} only when all - attempts fail. + Retries up to 3 times, validating the result is a real HTTPS Colab URL. + Falls back to http://localhost:{port} only when all attempts fail. """ import time as _time @@ -43,7 +40,7 @@ def get_colab_url(port: int = 8888) -> str: for attempt in range(3): try: url = eval_js(f"google.colab.kernel.proxyPort({port})", timeout_sec = 10) - # A valid Colab proxy URL starts with https:// and embeds the port. + # Valid proxy URL is https:// and embeds the port. if url and isinstance(url, str) and url.startswith("https://") and str(port) in url: return url.rstrip("/") except Exception as e: @@ -61,16 +58,13 @@ def get_colab_url(port: int = 8888) -> str: def show_link(port: int = 8888, *, _url: "str | None" = None): """Display a styled clickable link to the UI. - *_url* is an optional pre-fetched Colab proxy URL. When omitted, - ``get_colab_url(port)`` is called internally. Pass it from - ``_show_and_embed`` to avoid a second ``eval_js`` round-trip. + *_url* is an optional pre-fetched proxy URL; pass it to avoid a second eval_js round-trip. """ from IPython.display import display, HTML url = _url if _url is not None else get_colab_url(port) - # Build a truncated display URL. Wrap in try/except so an unexpected URL - # shape never prevents the link from rendering. + # Truncated display URL; try/except so an odd URL shape still renders the link. try: port_prefix = f"{port}-" idx = url.index(port_prefix) @@ -79,8 +73,7 @@ def show_link(port: int = 8888, *, _url: "str | None" = None): except (ValueError, IndexError): short_url = url - # Also emit a plain-text line so the URL is visible even if HTML display - # is suppressed or fails. + # Plain-text line so the URL shows even if HTML display fails. logger.info(f"🌐 Unsloth Studio URL: {url}") html = f""" @@ -123,12 +116,8 @@ def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool: def _show_and_embed(port: int): """Embed the Studio inline for *port* with a branded header bar. - Fetches the Colab proxy URL once (registering the port with Colab's - reverse-proxy at the same time) then renders a header bar + full-height - iframe as a single HTML block. - - Falls back to ``serve_kernel_port_as_iframe`` if ``IPython.display.HTML`` - is unavailable for any reason. + Fetches the proxy URL once (registering the port), then renders header bar + + iframe. Falls back to serve_kernel_port_as_iframe if IPython HTML is unavailable. """ url = get_colab_url(port) logger.info(f"🌐 Unsloth Studio URL: {url}") @@ -138,7 +127,7 @@ def _show_and_embed(port: int): iframe_id = f"unsloth-studio-{port}" - # Truncated URL shown in the header — best-effort, falls back to full URL. + # Truncated header URL — best-effort, falls back to full URL. try: port_prefix = f"{port}-" idx = url.index(port_prefix) @@ -167,7 +156,7 @@ def _show_and_embed(port: int): """) ) except Exception: - # Fallback: Colab's built-in helper (less control, but always works) + # Fallback: Colab's built-in helper. try: from google.colab import output as colab_output colab_output.serve_kernel_port_as_iframe(port, height = 900, width = "100%") @@ -187,9 +176,8 @@ def start(port: int = 8888): logger.info("🦥 Starting Unsloth Studio...") - # --- Fast path: Studio is already running (cell re-run) --- - # Re-launching would either collide on the port or silently shift to a new - # port and confuse the user. Just re-show the link and iframe instead. + # Fast path: Studio already running (cell re-run). Re-launching would collide on + # the port, so just re-show the link and iframe. if _is_studio_healthy(port): logger.info(f" Studio is already running on port {port} — reusing existing server.") _show_and_embed(port) @@ -222,16 +210,14 @@ def start(port: int = 8888): logger.error(f"❌ Unsloth Studio failed to start: {exc}") return - # run_server auto-increments the port when the requested one is already in - # use (e.g. Jupyter occupying 8888). Read back the actual bound port so the - # Colab proxy URL and iframe always point at the right place. + # run_server auto-increments the port if in use; read back the bound port so the + # proxy URL and iframe point at the right place. actual_port: int = getattr(getattr(app, "state", None), "server_port", None) or port logger.info(f" Server started on port {actual_port}!") - # Poll health endpoint to confirm the server is truly reachable before - # showing the link and registering the iframe — avoids the race where - # ready_event fires but the process hasn't finished binding. + # Poll health endpoint before showing the link — avoids the race where ready_event + # fires but the process hasn't finished binding. import urllib.request server_ready = False @@ -252,9 +238,8 @@ def start(port: int = 8888): _show_and_embed(actual_port) - # Keep kernel alive so the daemon server thread stays running. - # Handle KeyboardInterrupt cleanly so the user gets a readable message - # rather than a raw traceback when they interrupt the cell. + # Keep kernel alive so the daemon server thread runs; handle KeyboardInterrupt + # cleanly so interrupting the cell gives a readable message. try: for _ in range(10000): time.sleep(300) diff --git a/studio/backend/core/__init__.py b/studio/backend/core/__init__.py index 2ba9d5c65c..30ecb27f81 100644 --- a/studio/backend/core/__init__.py +++ b/studio/backend/core/__init__.py @@ -4,18 +4,16 @@ """ Unified core module for Unsloth backend -Imports are LAZY (via __getattr__) so that training subprocesses can -import core.training.worker without pulling in heavy ML dependencies -like unsloth, transformers, or torch before the version activation -code has a chance to run. +Imports are LAZY (via __getattr__) so training subprocesses can import +core.training.worker without pulling in heavy ML deps (unsloth, transformers, +torch) before the version-activation code runs. """ import sys from pathlib import Path -# Ensure the backend directory is on sys.path so that bare "from utils.*" -# imports used throughout the backend work when core is imported as a package -# (e.g. from the CLI: "from studio.backend.core import ModelConfig"). +# Add backend dir to sys.path so bare "from utils.*" imports work when core +# is imported as a package. _backend_dir = str(Path(__file__).resolve().parent.parent) if _backend_dir not in sys.path: sys.path.insert(0, _backend_dir) @@ -69,7 +67,7 @@ def __getattr__(name): globals()["TrainingProgress"] = TrainingProgress return globals()[name] - # Config (from utils.models) + # Config (utils.models) if name in ( "is_vision_model", "ModelConfig", diff --git a/studio/backend/core/_torchao_stub.py b/studio/backend/core/_torchao_stub.py index 217c25fda3..6336954bd5 100644 --- a/studio/backend/core/_torchao_stub.py +++ b/studio/backend/core/_torchao_stub.py @@ -3,18 +3,11 @@ """Shared torchao Windows-ROCm import stub. -torchao (pulled in by transformers.quantizers) imports -torch.distributed._functional_collectives at module level, which imports -distributed_c10d.py unconditionally — that file crashes on Windows ROCm because -torch._C._distributed_c10d (the RCCL backend) is absent. -torch/distributed/__init__.py itself is guarded by `if is_available()` so -`import torch.distributed` alone is safe; the crash only comes via torchao's -import chain. Stubbing torchao short-circuits it entirely. -_StubSubpackageFinder handles any depth of torchao.xxx.yyy imports. - -This logic used to be duplicated inline inside run_export_process() and -run_training_process(); it now lives here so both worker subprocesses call the -single `install_torchao_windows_rocm_stub()` entrypoint before importing +torchao (pulled in by transformers.quantizers) imports distributed_c10d.py +unconditionally, which crashes on Windows ROCm because the RCCL backend +(torch._C._distributed_c10d) is absent. Stubbing torchao short-circuits its +import chain; _StubSubpackageFinder handles any depth of torchao.xxx.yyy. +Worker subprocesses call install_torchao_windows_rocm_stub() before importing transformers / unsloth_zoo. """ @@ -28,12 +21,9 @@ import importlib.machinery _STUB_SENTINEL = object() -# Metaclass for stub types so that isinstance(x, StubClass) returns False -# instead of raising TypeError ("arg 2 must be a type"). -# peft/tuners/lora/torchao.py does: -# from torchao.dtypes import AffineQuantizedTensor, LinearActivationQuantizedTensor -# isinstance(weight, (AffineQuantizedTensor, LinearActivationQuantizedTensor)) -# If those names resolve to stub modules rather than types, isinstance() raises. +# Metaclass for stub types so isinstance(x, StubClass) returns False instead of +# raising TypeError -- peft's lora/torchao.py does isinstance() against torchao +# types, which fails if those names resolve to stub modules rather than types. class _StubTypeMeta(type): def __instancecheck__(cls, instance): return False @@ -71,8 +61,7 @@ def _make_mod_stub(mod_name): ): if attr.startswith("__"): raise AttributeError(attr) - # Return a stub CLASS (not a module) so that isinstance(x, attr) - # works and returns False instead of raising TypeError. + # Return a stub CLASS (not module) so isinstance() returns False, not TypeError. child = _make_stub_type(f"{_n}.{attr}") setattr(_m, attr, child) return child @@ -114,15 +103,12 @@ class _StubSubpackageFinder(importlib.abc.MetaPathFinder): def install_torchao_windows_rocm_stub() -> None: """Pre-stub torchao on Windows ROCm so transformers/peft imports don't crash. - No-op on every other platform (Windows CUDA included — there torchao is real - and shadowing it would break torchao-based quantization paths). Must run - before any import of transformers / unsloth_zoo. Safe to call once per worker - process. + No-op elsewhere (incl. Windows CUDA, where torchao is real). Must run before + importing transformers / unsloth_zoo. Safe to call once per worker. """ - # Gate on the active torch runtime, not env-var presence -- HIP_PATH / - # ROCM_PATH stay set after a user installs the HIP SDK and reverts to a - # CUDA torch wheel. AMD SDK / Radeon ROCm wheels may not set torch.version.hip - # but still encode "rocm" in torch.__version__, so accept either. + # Gate on the active torch runtime, not env-var presence -- HIP_PATH/ROCM_PATH + # persist after reverting to a CUDA wheel. Some ROCm wheels lack + # torch.version.hip but still encode "rocm" in __version__, so accept either. _is_win32_rocm = False if sys.platform == "win32": try: @@ -135,8 +121,7 @@ def install_torchao_windows_rocm_stub() -> None: except Exception: pass if _is_win32_rocm: - # Register the finder only on Windows ROCm -- on other platforms there - # are no stub modules seeded, so appending is a pure accumulation. + # Register the finder only on Windows ROCm. sys.meta_path.append(_StubSubpackageFinder()) # Seed torchao top-level + key submodules; the finder handles the rest. for _tao_name in ( diff --git a/studio/backend/core/data_recipe/__init__.py b/studio/backend/core/data_recipe/__init__.py index 239e406300..ebc2798375 100644 --- a/studio/backend/core/data_recipe/__init__.py +++ b/studio/backend/core/data_recipe/__init__.py @@ -1,9 +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 -""" -Data Recipe core (DataDesigner wrapper + job runner). -""" +"""Data Recipe core (DataDesigner wrapper + job runner).""" from .jobs import JobManager, get_job_manager diff --git a/studio/backend/core/data_recipe/huggingface.py b/studio/backend/core/data_recipe/huggingface.py index d5a6db6baf..7a1219b2c3 100644 --- a/studio/backend/core/data_recipe/huggingface.py +++ b/studio/backend/core/data_recipe/huggingface.py @@ -95,8 +95,7 @@ def publish_recipe_dataset( tags = None, ) card.text = card.text.replace(_DATA_DESIGNER_FOOTER, _UNSLOTH_STUDIO_FOOTER) - # Data Designer currently drops the explicit token when pushing the - # dataset card. Push it ourselves so auth stays request-local. + # Data Designer drops the explicit token, so push the card ourselves to keep auth request-local. card.push_to_hub(repo_id, token = hf_token, repo_type = "dataset") client._upload_main_dataset_files( diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py index 75dc1efb9c..32523e469e 100644 --- a/studio/backend/core/data_recipe/jobs/manager.py +++ b/studio/backend/core/data_recipe/jobs/manager.py @@ -132,10 +132,9 @@ class JobManager: ) -> str: """Spawn the job subprocess (one at a time, no cap). - ``internal_api_key_id`` is the row id of a workflow-scoped - sk-unsloth-* key minted by the route layer for local providers. - JobManager revokes it when the job reaches a terminal state so the - key's live window is no longer than the run. + ``internal_api_key_id`` is a workflow-scoped sk-unsloth-* key row id + minted by the route layer; revoked on terminal state so the key's + live window is no longer than the run. """ llm_columns = recipe.get("columns") or [] llm_column_count = 0 @@ -202,7 +201,7 @@ class JobManager: return True def get_status(self, job_id: str) -> dict | None: - """UI friendly snapshot that we need. Alternative to sse kinda of and structured""" + """UI-friendly structured snapshot; an alternative to SSE.""" with self._lock: if self._job is None or self._job.job_id != job_id: return None @@ -537,15 +536,14 @@ class JobManager: def _retire_workflow_key(self, job: Job) -> None: """Revoke the workflow-scoped sk-unsloth-* key, if one was minted. - Best-effort: revocation failures are swallowed. The key would - expire on its own after 24h, so a missed revoke is a latency - concern, not a correctness one. + Best-effort: failures are swallowed. The key expires after 24h, so a + missed revoke is a latency, not correctness, concern. """ key_id = getattr(job, "internal_api_key_id", None) if not key_id: return try: - from auth import storage # deferred: avoids circular import + from auth import storage # deferred: avoid circular import storage.revoke_internal_api_key(int(key_id)) except Exception: pass diff --git a/studio/backend/core/data_recipe/jobs/parse.py b/studio/backend/core/data_recipe/jobs/parse.py index 8ca3702edd..3be830d0e4 100644 --- a/studio/backend/core/data_recipe/jobs/parse.py +++ b/studio/backend/core/data_recipe/jobs/parse.py @@ -46,7 +46,7 @@ class ParsedUpdate: source_progress: SourceProgress | None = None -# kinda of a bummber but currently only option, Best effort parser from data-designer logs -> structured status for UI. +# Best-effort parser from data-designer logs -> structured status for UI. _RE_SAMPLERS = re.compile( r"Preparing samplers to generate (?P\d+) records across (?P\d+) columns" ) @@ -327,7 +327,7 @@ def apply_update(job: Job, update: ParsedUpdate) -> None: _apply_source_progress(job, update.source_progress) if update.stage in USAGE_RESET_STAGES: - # usage summary is a short block so we reset once we move into the next stage. + # Usage summary is a short block; reset on the next stage. job._in_usage_summary = False if update.usage_section_start is not None: diff --git a/studio/backend/core/data_recipe/jobs/types.py b/studio/backend/core/data_recipe/jobs/types.py index 3d3ddb974e..c30606b2d1 100644 --- a/studio/backend/core/data_recipe/jobs/types.py +++ b/studio/backend/core/data_recipe/jobs/types.py @@ -90,9 +90,8 @@ class Job: progress_columns_total: int | None = None source_progress_estimated_total: int | None = None completed_columns: list[str] = field(default_factory = list) - # Id of the internal sk-unsloth-* API key minted for a local-model - # workflow. Revoked when the job terminates so the key's live window - # matches the run rather than its 24h TTL. + # Id of the internal sk-unsloth-* API key minted for a local-model workflow. + # Revoked when the job ends so the key's window matches the run, not its 24h TTL. internal_api_key_id: int | None = None _current_usage_model: str | None = None _in_usage_summary: bool = False diff --git a/studio/backend/core/data_recipe/jobs/worker.py b/studio/backend/core/data_recipe/jobs/worker.py index f59538769e..4073288cb3 100644 --- a/studio/backend/core/data_recipe/jobs/worker.py +++ b/studio/backend/core/data_recipe/jobs/worker.py @@ -73,13 +73,10 @@ def _build_dataset_name(*, run_name: str | None, job_id: str, artifact_root: Pat def run_job_process(*, event_queue, recipe: dict[str, Any], run: dict[str, Any]) -> None: - """ - Subprocess entrypoint. - Sends events to `event_queue`. - """ + """Subprocess entrypoint. Sends events to `event_queue`.""" import os - os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports + os.environ["PYTHONWARNINGS"] = "ignore" # suppress C-level warnings before imports import warnings from loggers.config import LogConfig @@ -115,8 +112,8 @@ def run_job_process(*, event_queue, recipe: dict[str, Any], run: dict[str, Any]) builder = build_config_builder(recipe) designer = create_data_designer(recipe, artifact_path = str(_ARTIFACT_ROOT)) - # DataDesigner configures root logging in DataDesigner.__init__. - # Attach queue logger directly to `data_designer` so parser events survive root resets. + # DataDesigner resets root logging in __init__; attach the queue handler + # to the named loggers directly so parser events survive. handler = _QueueLogHandler(event_queue) handler.setLevel(logging.INFO) for logger_name in ( diff --git a/studio/backend/core/data_recipe/local_callable_validators.py b/studio/backend/core/data_recipe/local_callable_validators.py index a40261fa92..fbe847f9ce 100644 --- a/studio/backend/core/data_recipe/local_callable_validators.py +++ b/studio/backend/core/data_recipe/local_callable_validators.py @@ -184,7 +184,7 @@ def _build_oxc_validation_function(lang: str, validation_mode: str, code_shape: normalized_code_shape = code_shape if code_shape in _OXC_CODE_SHAPES else "auto" def _validator(df): - import pandas as pd # imported lazily for local callable runtime + import pandas as pd # lazy import for local callable runtime row_count = int(len(df.index)) if row_count == 0: diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py index cb5901f04e..9d8ca5cfcc 100644 --- a/studio/backend/core/data_recipe/service.py +++ b/studio/backend/core/data_recipe/service.py @@ -168,8 +168,7 @@ def build_mcp_providers(recipe: dict[str, Any]) -> list: from data_designer.config.mcp import LocalStdioMCPProvider, MCPProvider # pyright: ignore[reportMissingImports] # Same gate as the chat MCP path: stdio providers spawn a local subprocess, - # so only build them when this host allows it (desktop / explicit opt-in). - # Skip them otherwise so a recipe carried onto a hosted host cannot spawn. + # so build them only when this host allows it (desktop / explicit opt-in). from core.inference.mcp_client import stdio_mcp_enabled stdio_allowed = stdio_mcp_enabled() @@ -258,7 +257,7 @@ def build_config_builder(recipe: dict[str, Any]): ) # DataDesignerConfigBuilder.from_config currently skips processors. - # Re-attach explicitly so drop_columns/schema_transform survive API payload. + # Re-attach so drop_columns/schema_transform survive the API payload. for processor in recipe_core.get("processors") or []: if not isinstance(processor, dict): continue @@ -282,9 +281,8 @@ def create_data_designer(recipe: dict[str, Any], *, artifact_path: str | None = model_providers = build_model_providers(recipe) _validate_recipe_runtime_support(recipe, model_providers) - # DataDesigner requires at least one model provider in its registry even - # when the pipeline contains no LLM columns. Supply a lightweight stub - # so sampler/expression-only recipes can run without a real provider. + # DataDesigner requires >=1 model provider even with no LLM columns; stub + # one so sampler/expression-only recipes run without a real provider. if not model_providers: from data_designer.config.models import ModelProvider # pyright: ignore[reportMissingImports] model_providers = [ diff --git a/studio/backend/core/export/__init__.py b/studio/backend/core/export/__init__.py index 16fba368ef..f7241cd925 100644 --- a/studio/backend/core/export/__init__.py +++ b/studio/backend/core/export/__init__.py @@ -1,12 +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 -""" -Export submodule - Model export operations +"""Export submodule - model export operations. -The default get_export_backend() returns an ExportOrchestrator that -delegates to a subprocess. The original ExportBackend runs inside -the subprocess and can be imported directly from .export when needed. +get_export_backend() returns an ExportOrchestrator that delegates to a +subprocess. The original ExportBackend runs inside the subprocess and can be +imported directly from .export when needed. """ from .orchestrator import ExportOrchestrator, get_export_backend diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index e74dc13e8a..d5a3de06df 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -1,10 +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 -# backend/export.py -""" -Export backend - handles model exporting in various formats -""" +"""Export backend - exports models in various formats.""" import glob import json @@ -46,10 +43,8 @@ def _is_wsl(): def _apply_wsl_sudo_patch(): """On WSL, monkey-patch do_we_need_sudo() to return False. - WSL doesn't have passwordless sudo, and do_we_need_sudo() runs - `sudo apt-get update` which hangs waiting for a stdin password - inside a non-interactive subprocess. setup.sh pre-installs the - build dependencies on WSL, so sudo is not needed at runtime. + WSL lacks passwordless sudo and do_we_need_sudo()'s `sudo apt-get update` + hangs on a stdin password; setup.sh pre-installs the build deps anyway. """ if not _is_wsl(): return @@ -110,18 +105,15 @@ class ExportBackend: try: logger.info("Starting memory cleanup...") - # Unload all models from inference backend model_names = list(self.inference_backend.models.keys()) for model_name in model_names: self.inference_backend.unload_model(model_name) - # Clear current export state self.current_model = None self.current_tokenizer = None self.current_checkpoint = None self._audio_type = None - # Clear GPU memory cache (handles gc + backend-specific cleanup) clear_gpu_cache() logger.info("Memory cleanup completed successfully") @@ -137,8 +129,7 @@ class ExportBackend: """ Scan outputs folder for training runs and their checkpoints. - Returns: - List of tuples: [(model_name, [(display_name, checkpoint_path), ...]), ...] + Returns: [(model_name, [(display_name, checkpoint_path), ...]), ...] """ from utils.models.checkpoints import scan_checkpoints return scan_checkpoints(outputs_dir = outputs_dir) @@ -159,12 +150,11 @@ class ExportBackend: try: logger.info(f"Loading checkpoint: {checkpoint_path}") - # First, cleanup existing models self.cleanup_memory() checkpoint_path_obj = Path(checkpoint_path) - # Determine the model identity for type detection + # Model identity for type detection adapter_config = checkpoint_path_obj / "adapter_config.json" base_model = None if adapter_config.exists(): @@ -174,11 +164,9 @@ class ExportBackend: model_id = base_model or checkpoint_path - # Detect audio type and vision self._audio_type = detect_audio_type(model_id) self.is_vision = not self._audio_type and is_vision_model(model_id) - # Load model based on type if self._audio_type == "csm": from unsloth import FastModel from transformers import CsmForConditionalGeneration @@ -246,7 +234,7 @@ class ExportBackend: load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, ) - tokenizer = processor # For vision models, processor acts as tokenizer + tokenizer = processor # vision: processor acts as tokenizer else: logger.info("Loading as text model...") @@ -258,14 +246,12 @@ class ExportBackend: trust_remote_code = trust_remote_code, ) - # Check if PEFT / LoRA model if _IS_MLX: # MLX doesn't use PeftModel — detect LoRA via adapter_config.json self.is_peft = adapter_config.exists() else: self.is_peft = isinstance(model, (PeftModel, PeftModelForCausalLM)) - # Store loaded model self.current_model = model self.current_tokenizer = tokenizer self.current_checkpoint = checkpoint_path @@ -349,7 +335,6 @@ class ExportBackend: else: save_method = "merged_16bit" - # Save locally if requested if save_directory: save_directory = str(resolve_export_dir(save_directory)) logger.info(f"Saving merged model locally to: {save_directory}") @@ -370,7 +355,6 @@ class ExportBackend: logger.info(f"Model saved successfully to {save_directory}") output_path = str(Path(save_directory).resolve()) - # Push to hub if requested if push_to_hub: if not repo_id or not hf_token: return ( @@ -451,7 +435,6 @@ class ExportBackend: output_path: Optional[str] = None try: - # Save locally if requested if save_directory: save_directory = str(resolve_export_dir(save_directory)) logger.info(f"Saving base model locally to: {save_directory}") @@ -459,7 +442,7 @@ class ExportBackend: if _IS_MLX: # MLX: save_pretrained_merged handles non-LoRA models too - # (fuse() is a no-op when there are no LoRA layers) + # (fuse() is a no-op without LoRA layers) self.current_model.save_pretrained_merged( save_directory, self.current_tokenizer, @@ -474,7 +457,6 @@ class ExportBackend: logger.info(f"Model saved successfully to {save_directory}") output_path = str(Path(save_directory).resolve()) - # Push to hub if requested if push_to_hub: if not repo_id or not hf_token: return ( @@ -509,12 +491,11 @@ class ExportBackend: private = private, ) else: - # Get base model name from request or model config + # Base model name from request or model config base_model = ( base_model_id or self.current_model.config._name_or_path or "unknown" ) - # Create repo hf_api = HfApi(token = hf_token) repo_id = PushToHubMixin._create_repo( PushToHubMixin, @@ -524,7 +505,6 @@ class ExportBackend: ) username = repo_id.split("/")[0] - # Create and push model card content = MODEL_CARD.format( username = username, base_model = base_model, @@ -535,7 +515,6 @@ class ExportBackend: card = ModelCard(content) card.push_to_hub(repo_id, token = hf_token, commit_message = "Unsloth Model Card") - # Upload model files if save_directory: hf_api.upload_folder( folder_path = save_directory, @@ -585,13 +564,11 @@ class ExportBackend: output_path: Optional[str] = None try: - # Convert quantization method to lowercase for unsloth + # unsloth expects lowercase quant method quant_method = quantization_method.lower() - # Pin convert_hf_to_gguf.py to the same llama.cpp ref as the - # llama-quantize binary (Studio installs at a tagged ref via - # setup.sh) so it can't drift past the pinned binary's gguf API. - # Set before both branches; hub-only export has save_directory == "". + # Pin convert_hf_to_gguf.py to setup.sh's tagged llama.cpp ref so it + # can't drift past the pinned llama-quantize binary's gguf API. global _LLAMA_CPP_SCRIPTS_WARNING_EMITTED try: from unsloth_zoo.llama_cpp import ( @@ -610,30 +587,23 @@ class ExportBackend: ) _LLAMA_CPP_SCRIPTS_WARNING_EMITTED = True - # Save locally if requested if save_directory: save_directory = str(resolve_export_dir(save_directory)) - # Resolve to absolute path so unsloth's relative-path internals - # (check_llama_cpp, use_local_gguf, _download_convert_hf_to_gguf) - # all resolve against the repo root cwd, NOT the export directory. + # Absolute path so unsloth's relative-path internals resolve + # against the repo root cwd, not the export directory. abs_save_dir = os.path.abspath(save_directory) logger.info(f"Saving GGUF model locally to: {abs_save_dir}") - # Create the directory if it doesn't exist ensure_dir(Path(abs_save_dir)) # On WSL, patch out sudo check before llama.cpp build _apply_wsl_sudo_patch() - # Snapshot existing .gguf files in cwd before conversion. - # unsloth's convert_to_gguf writes output files relative to - # cwd (repo root), so we diff afterwards and relocate them. + # convert_to_gguf writes output relative to cwd (repo root); + # snapshot existing .gguf so we can diff and relocate afterwards. cwd = os.getcwd() pre_existing_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf"))) - # Pass absolute path — no os.chdir needed. - # unsloth saves intermediate HF model files into model_save_path. - # unsloth-zoo's check_llama_cpp() uses ~/.unsloth/llama.cpp by default. model_save_path = os.path.join(abs_save_dir, "model") self.current_model.save_pretrained_gguf( model_save_path, @@ -641,18 +611,14 @@ class ExportBackend: quantization_method = quant_method, ) - # Relocate GGUF artifacts into the export directory. - # convert_to_gguf writes .gguf files to cwd (repo root) - # because --outfile is a relative path like "model.Q4_K_M.gguf". + # Relocate the .gguf that convert_to_gguf wrote to cwd (repo root). new_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf"))) - pre_existing_ggufs for src in sorted(new_ggufs): dest = os.path.join(abs_save_dir, os.path.basename(src)) shutil.move(src, dest) logger.info(f"Relocated GGUF: {os.path.basename(src)} → {abs_save_dir}/") - # Flatten any .gguf files from subdirectories into abs_save_dir. - # save_pretrained_gguf may create subdirs (e.g. model_gguf/) - # with a name different from model_save_path. + # Flatten any .gguf from subdirs (e.g. model_gguf/) into abs_save_dir. for sub in list(Path(abs_save_dir).iterdir()): if not sub.is_dir(): continue @@ -660,13 +626,11 @@ class ExportBackend: dest = os.path.join(abs_save_dir, src.name) shutil.move(str(src), dest) logger.info(f"Relocated GGUF: {src.name} → {abs_save_dir}/") - # Clean up the subdirectory (intermediate HF files, etc.) shutil.rmtree(str(sub), ignore_errors = True) logger.info(f"Cleaned up subdirectory: {sub.name}") - # For non-PEFT models, save_pretrained_gguf redirects to the - # checkpoint path, leaving a *_gguf directory in outputs/. - # Relocate any GGUFs from there and clean it up. + # For non-PEFT models, save_pretrained_gguf leaves a *_gguf dir at + # the checkpoint path; relocate its GGUFs and clean it up. if self.current_checkpoint: ckpt = Path(self.current_checkpoint) gguf_dir = ckpt.parent / f"{ckpt.name}_gguf" @@ -686,8 +650,6 @@ class ExportBackend: # Write export metadata so the Chat page can identify the base model self._write_export_metadata(abs_save_dir) - # Log final file locations (after relocation) so it's clear - # where the GGUF files actually ended up. final_ggufs = sorted(glob.glob(os.path.join(abs_save_dir, "*.gguf"))) logger.info( "GGUF export complete. Final files in %s:\n %s", @@ -696,7 +658,6 @@ class ExportBackend: ) output_path = str(Path(abs_save_dir).resolve()) - # Push to hub if requested if push_to_hub: if not repo_id or not hf_token: return ( @@ -750,7 +711,6 @@ class ExportBackend: output_path: Optional[str] = None try: - # Save locally if requested if save_directory: save_directory = str(resolve_export_dir(save_directory)) logger.info(f"Saving LoRA adapter locally to: {save_directory}") @@ -766,7 +726,6 @@ class ExportBackend: logger.info(f"Adapter saved successfully to {save_directory}") output_path = str(Path(save_directory).resolve()) - # Push to hub if requested if push_to_hub: if not repo_id or not hf_token: return ( diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 90a946163e..20158d1891 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -1,15 +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 -""" -Export orchestrator — subprocess-based. +"""Export orchestrator — subprocess-based. -Provides the same API as ExportBackend, but delegates all ML work -to a persistent subprocess. The subprocess is spawned on first checkpoint -load and stays alive for subsequent export operations. +Same API as ExportBackend, but delegates all ML work to a persistent +subprocess spawned on first checkpoint load and reused for later exports. -When switching between checkpoints that need different transformers versions, -the old subprocess is killed and a new one is spawned with the correct version. +When switching between checkpoints needing different transformers +versions, the old subprocess is killed and a new one spawned. Pattern follows core/inference/orchestrator.py. """ @@ -30,9 +28,7 @@ logger = get_logger(__name__) _CTX = mp.get_context("spawn") -# Maximum number of captured log lines kept in memory per export -# orchestrator. Acts as scrollback for the live export log panel in the -# UI. 4000 lines is ~1 MB worst-case at 256 chars/line. +# Max log lines kept per orchestrator (live log panel scrollback); ~1 MB worst-case. _LOG_BUFFER_MAXLEN = 4000 @@ -41,46 +37,31 @@ class ExportOrchestrator: Export backend orchestrator — subprocess-based. Exposes the same API surface as ExportBackend so routes/export.py - needs minimal changes. Internally, all heavy ML operations happen in - a persistent subprocess. + needs minimal changes. All heavy ML work happens in a persistent + subprocess. """ def __init__(self): - # Subprocess state self._proc: Optional[mp.Process] = None self._cmd_queue: Any = None self._resp_queue: Any = None - # Serializes export operations (load_checkpoint, export_*, - # cleanup) so concurrent HTTP requests can never interleave - # commands on the subprocess queue. Previously unused. + # Serializes export ops so concurrent HTTP requests can't interleave commands. self._lock = threading.Lock() - # Local state mirrors (updated from subprocess responses) + # Local state mirrors (updated from subprocess responses). self.current_checkpoint: Optional[str] = None self.is_vision: bool = False self.is_peft: bool = False - # ── Live log capture ───────────────────────────────────── - # Thread-safe ring buffer of log lines forwarded from the - # worker subprocess. Powers the GET /api/export/logs/stream - # SSE endpoint that the export dialog consumes. + # Thread-safe ring buffer of worker log lines; powers the export logs SSE endpoint. self._log_buffer: Deque[Dict[str, Any]] = deque(maxlen = _LOG_BUFFER_MAXLEN) self._log_lock = threading.Lock() - # Monotonically increasing sequence number. Never reset across - # operations, so SSE clients can use it as a stable cursor even - # if clear_logs() is called mid-session. + # Monotonic seq, never reset, so SSE clients have a stable cursor across clear_logs(). self._log_seq: int = 0 - # Snapshot of _log_seq captured at the start of the current run - # (updated by clear_logs()). The SSE endpoint defaults its - # cursor to this value so a client that connects AFTER the - # worker has already emitted its first lines still sees the - # full run. Every line appended during the current run has seq - # strictly greater than _run_start_seq, and every line from - # prior runs has seq less than or equal to it. + # _log_seq snapshot at the current run's start; SSE defaults its cursor here so a + # late-connecting client still sees the full run. Current run has seq > this. self._run_start_seq: int = 0 - # True while an export operation (load/export/cleanup) is - # running. The SSE endpoint ends the stream 1 second after - # this flips back to False to drain any trailing log lines. + # True while an export op runs; SSE ends the stream 1s after this flips False. self._export_active: bool = False atexit.register(self._cleanup) @@ -91,13 +72,7 @@ class ExportOrchestrator: # ------------------------------------------------------------------ def _append_log(self, entry: Dict[str, Any]) -> None: - """Append a log line from the worker subprocess to the buffer. - - Entries look like {"type": "log", "stream": "stdout"|"stderr", - "line": "...", "ts": ...}. Each is stamped with a monotonic - seq number before it lands in the buffer so SSE clients can - cursor through new lines. - """ + """Append a worker log line to the buffer, stamped with a monotonic seq.""" line = entry.get("line") if not line: return @@ -113,18 +88,10 @@ class ExportOrchestrator: ) def clear_logs(self) -> None: - """Drop any buffered log lines from a previous operation. + """Drop buffered log lines from a previous op so the UI shows only this run. - Called at the start of each export op so the UI shows only the - output of the current run. The seq counter is NOT reset, so an - SSE client that captured the cursor before clear_logs() will - still see new lines (with strictly greater seq numbers). - - Also snapshots the current seq into ``_run_start_seq`` so the - SSE endpoint can anchor its default cursor at the start of - this run. Anything appended after this call has seq strictly - greater than the snapshot and is reachable via - ``get_logs_since(get_run_start_seq())``. + The seq counter is NOT reset (clients keep a stable cursor); the current seq + is snapshotted into ``_run_start_seq`` to anchor the SSE default cursor. """ with self._log_lock: self._log_buffer.clear() @@ -144,12 +111,7 @@ class ExportOrchestrator: return self._log_seq def get_run_start_seq(self) -> int: - """Return the seq value captured at the start of the current run. - - The SSE endpoint uses this as the default cursor so a client - that connects AFTER the worker has already started emitting - output still sees every line from the current run. - """ + """Return the seq captured at the current run's start (SSE default cursor).""" with self._log_lock: return self._run_start_seq @@ -193,22 +155,19 @@ class ExportOrchestrator: self._proc = None return - # 1. Drain stale responses self._drain_queue() - # 2. Send shutdown command try: self._cmd_queue.put({"type": "shutdown"}) except (OSError, ValueError): pass - # 3. Wait for graceful shutdown try: self._proc.join(timeout = timeout) except Exception: pass - # 4. Force kill if still alive + # Force kill if still alive. if self._proc is not None and self._proc.is_alive(): logger.warning("Export subprocess did not exit gracefully, terminating") try: @@ -268,9 +227,8 @@ class ExportOrchestrator: ) -> dict: """Block until a response of the expected type arrives. - Export operations can take a very long time — GGUF conversion for - large models (30B+) easily takes 20-30 minutes. Default timeout - is 1 hour. + Export ops can take a long time — GGUF conversion for large + models (30B+) easily takes 20-30 minutes. Default timeout 1 hour. """ deadline = time.monotonic() + timeout @@ -279,7 +237,6 @@ class ExportOrchestrator: resp = self._read_resp(timeout = min(remaining, 2.0)) if resp is None: - # Check subprocess health if not self._ensure_subprocess_alive(): raise RuntimeError("Export subprocess crashed during wait") continue @@ -294,17 +251,14 @@ class ExportOrchestrator: raise RuntimeError(f"Subprocess error: {error_msg}") if rtype == "log": - # Forwarded stdout/stderr line from the worker process. + # Forwarded stdout/stderr line from the worker. self._append_log(resp) continue if rtype == "status": message = resp.get("message", "") logger.info("Export subprocess status: %s", message) - # Surface status messages in the live log panel too so - # users see high level progress (e.g. "Importing - # Unsloth...", "Loading checkpoint: ...") alongside - # subprocess output. + # Surface status in the live log panel for high-level progress. if message: self._append_log( { @@ -315,7 +269,7 @@ class ExportOrchestrator: ) continue - # Other response types during wait — skip + # Other response types during wait — skip. logger.debug( "Skipping response type '%s' while waiting for '%s'", rtype, @@ -362,12 +316,11 @@ class ExportOrchestrator: } with self._lock: - # Start a fresh log buffer for this operation so the UI - # sees only the current run's output. + # Fresh log buffer so the UI sees only this run's output. self.clear_logs() self._export_active = True try: - # Always kill existing subprocess and spawn fresh. + # Always kill any existing subprocess and spawn fresh. if self._ensure_subprocess_alive(): self._shutdown_subprocess() elif self._proc is not None: @@ -486,14 +439,10 @@ class ExportOrchestrator: ) def _run_export(self, export_type: str, params: dict) -> Tuple[bool, str, Optional[str]]: - """Send an export command to the subprocess and wait for result. + """Send an export command and wait for the result. - Returns ``(success, message, output_path)``. ``output_path`` is the - resolved on-disk directory the worker actually wrote to (None when - the export only pushed to Hub or failed before any file was - written). Surfaced via the export route's ``details.output_path`` - so the dialog's success screen can show the user where the model - landed. + Returns ``(success, message, output_path)``. ``output_path`` is the on-disk + dir the worker wrote to (None if it only pushed to Hub or failed pre-write). """ with self._lock: if not self._ensure_subprocess_alive(): @@ -527,7 +476,6 @@ class ExportOrchestrator: """Cleanup export-related models from memory.""" with self._lock: if not self._ensure_subprocess_alive(): - # No subprocess — just clear local state self.current_checkpoint = None self.is_vision = False self.is_peft = False @@ -542,7 +490,7 @@ class ExportOrchestrator: except RuntimeError: success = False - # Shut down subprocess after cleanup — no model loaded + # Shut down subprocess after cleanup — no model loaded. self._shutdown_subprocess() self.current_checkpoint = None @@ -553,7 +501,7 @@ class ExportOrchestrator: self._export_active = False def scan_checkpoints(self, outputs_dir: str = str(outputs_root())) -> List[Tuple[str, list]]: - """Scan for checkpoints — no ML imports needed, runs locally.""" + """Scan for checkpoints — runs locally, no ML imports.""" from utils.models.checkpoints import scan_checkpoints return scan_checkpoints(outputs_dir = outputs_dir) diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index 405956ec9d..fb2a893014 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -1,16 +1,12 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Export subprocess entry point. +"""Export subprocess entry point. -Each export session runs in a persistent subprocess (mp.get_context("spawn")). -This gives us a clean Python interpreter with no stale module state — -solving the transformers version-switching problem completely. - -The subprocess stays alive while a model is loaded, accepting commands -(load, export_merged, export_base, export_gguf, export_lora, cleanup, -shutdown) via mp.Queue. +Each export session runs in a persistent subprocess (mp spawn), giving a clean +interpreter with no stale module state, which solves transformers version +switching. The subprocess stays alive while a model is loaded, accepting commands +(load, export_*, cleanup, shutdown) via mp.Queue. Pattern follows core/inference/worker.py and core/training/worker.py. """ @@ -31,42 +27,31 @@ from typing import Any logger = get_logger(__name__) -# Gate that controls whether captured stdout/stderr lines are forwarded -# to the parent's resp_queue (and from there to the export-dialog SSE -# stream). Closed by default so the noisy bootstrap phase -- transformers -# venv activation, Unsloth/torch imports, base-model resolution, "Top -# GGUF/hub models" lists, vision detection, weight loading bars -- is -# suppressed in the UI. _handle_export() opens the gate at the start of -# the actual export work and leaves it open; the orchestrator always -# spawns a fresh subprocess for the next checkpoint load (see -# orchestrator._spawn_subprocess) which resets this state. -# -# Lines dropped while the gate is closed are still echoed to the saved -# original stdout/stderr fds so the server console / log file keeps the -# full output for debugging. +# Gate controlling whether captured stdout/stderr lines are forwarded to the +# parent's resp_queue (and on to the export-dialog SSE stream). Closed by default +# so the noisy bootstrap phase (imports, model resolution, loading bars) is +# suppressed in the UI; _handle_export() opens it when export work starts. The +# orchestrator spawns a fresh subprocess per checkpoint load, resetting this. +# Dropped lines are still echoed to the saved fds so the server log keeps them. _log_forward_gate = threading.Event() def _setup_log_capture(resp_queue: Any) -> None: - """Redirect fds 1 and 2 through pipes so every line printed by this - worker process and any child process it spawns is forwarded to the - parent process via resp_queue as {"type": "log", ...} messages. + """Redirect fds 1 and 2 through pipes so every line printed by this worker + and any child it spawns is forwarded to the parent via resp_queue as + {"type": "log", ...} messages. - Must be called BEFORE LogConfig.setup_logging and BEFORE any ML - imports, otherwise library handlers may capture the original stderr - reference and bypass the pipe. - - Lines are also echoed back to the original stdout/stderr so the - server console keeps receiving the full subprocess output, even - while ``_log_forward_gate`` is closed. + Must run BEFORE LogConfig.setup_logging and any ML imports, else library + handlers may capture the original stderr reference and bypass the pipe. + Lines are also echoed back to the original fds so the server console keeps + the full output even while ``_log_forward_gate`` is closed. """ try: saved_out_fd = os.dup(1) saved_err_fd = os.dup(2) except OSError: - # dup failed (exotic platforms) - give up quietly, export still - # works, just no live log streaming. + # dup failed; give up quietly (export still works, no live streaming). return try: @@ -88,13 +73,11 @@ def _setup_log_capture(resp_queue: Any) -> None: pass return - # Close the write ends we just dup2'd (fds 1 and 2 are the real - # write ends now). + # Close the write ends we just dup2'd (fds 1 and 2 are the real write ends). os.close(w_out) os.close(w_err) - # Replace Python's sys.stdout/sys.stderr with line-buffered writers - # bound to the (now-redirected) fds 1 and 2. + # Replace sys.stdout/sys.stderr with line-buffered writers on fds 1 and 2. try: sys.stdout = os.fdopen(1, "w", buffering = 1, encoding = "utf-8", errors = "replace") sys.stderr = os.fdopen(2, "w", buffering = 1, encoding = "utf-8", errors = "replace") @@ -112,8 +95,7 @@ def _setup_log_capture(resp_queue: Any) -> None: continue if not chunk: break - # Echo to the original fd so the server console still sees - # the full output. + # Echo to the original fd so the server console keeps the full output. try: os.write(echo_fd, chunk) except OSError: @@ -133,9 +115,8 @@ def _setup_log_capture(resp_queue: Any) -> None: if not line: continue if not _log_forward_gate.is_set(): - # Gate closed (bootstrap phase) -- already echoed to - # the saved console fd above; drop the line so the - # export dialog doesn't see import / vendoring noise. + # Gate closed (bootstrap): already echoed above; drop the + # line so the export dialog skips import noise. continue try: resp_queue.put_nowait( @@ -147,8 +128,7 @@ def _setup_log_capture(resp_queue: Any) -> None: } ) except Exception: - # Queue put failed (full, closed, etc.) - drop the - # line rather than crash the reader thread. + # Queue put failed; drop the line rather than crash the thread. pass if buf and _log_forward_gate.is_set(): try: @@ -181,7 +161,7 @@ def _setup_log_capture(resp_queue: Any) -> None: 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 + # Ensure backend is on sys.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) @@ -267,11 +247,9 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: export_type = cmd["export_type"] # "merged", "base", "gguf", "lora" response_type = f"export_{export_type}_done" - # Open the log forwarding gate so the user sees the actual export - # progress (Unsloth merge bars, file copies, GGUF conversion, etc.) - # in the live log panel. The gate stays open for the rest of this - # subprocess's life; the orchestrator spawns a fresh subprocess for - # the next checkpoint load, which resets the gate to closed. + # Open the log forwarding gate so the user sees export progress in the live + # log panel. Stays open for the rest of this subprocess's life; the + # orchestrator spawns a fresh subprocess per checkpoint load, resetting it. _log_forward_gate.set() output_path: Any = None @@ -372,23 +350,16 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None """ import queue as _queue - # Install fd-level stdout/stderr capture FIRST so every subsequent - # print and every child process inherits the redirected fds. This - # is what powers the live export log stream in the UI. + # Install fd-level stdout/stderr capture FIRST so every subsequent print and + # every child process inherits the redirected fds (powers the live log stream). _setup_log_capture(resp_queue) os.environ["TOKENIZERS_PARALLELISM"] = "false" - os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports - # Force unbuffered output from any child Python process (e.g. the - # GGUF converter) so their prints surface in the log stream as they - # happen rather than at the end. + os.environ["PYTHONWARNINGS"] = "ignore" # suppress C-level warnings before imports + # Unbuffered output from child Python (e.g. GGUF converter) so prints surface live. os.environ["PYTHONUNBUFFERED"] = "1" - # tqdm defaults to a 10-second mininterval when stdout is not a tty - # (which it isn't here -- we redirected fd 1/2 to a pipe). That makes - # multi-step progress bars look frozen in the export log panel. Force - # frequent flushes so the user sees movement during merge / GGUF - # conversion. Has no effect on single-step bars (e.g. "Copying 1 - # files") which only emit start/end events regardless. + # tqdm defaults to a 10s mininterval when stdout isn't a tty (we redirected + # fd 1/2 to a pipe), making multi-step bars look frozen; force frequent flushes. os.environ.setdefault("TQDM_MININTERVAL", "0.5") import warnings @@ -419,7 +390,7 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None ) return - # ── 1b. On Windows, check Triton availability (must be before import torch) ── + # ── 1b. Check Triton on Windows (must precede import torch) ── if sys.platform == "win32": try: import triton # noqa: F401 @@ -432,10 +403,8 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None ) # ── 1c. Stub torchao on Windows ROCm ── - # Shared with the training worker; see core/_torchao_stub.py for the full - # rationale (torchao -> torch.distributed._functional_collectives crashes on - # Windows ROCm because the RCCL backend is absent). No-op off Windows ROCm. - # Must run before any import of transformers / unsloth_zoo. + # See core/_torchao_stub.py: torchao crashes on Windows ROCm (RCCL absent). + # No-op off Windows ROCm. Must run before importing transformers / unsloth_zoo. from core._torchao_stub import install_torchao_windows_rocm_stub install_torchao_windows_rocm_stub() @@ -511,7 +480,7 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None try: if cmd_type == "load": - # Load a new checkpoint (reusing this subprocess) + # Load a new checkpoint, reusing this subprocess. backend.cleanup_memory() _handle_load(backend, cmd, resp_queue) diff --git a/studio/backend/core/inference/__init__.py b/studio/backend/core/inference/__init__.py index 35318f6357..2faf70bb79 100644 --- a/studio/backend/core/inference/__init__.py +++ b/studio/backend/core/inference/__init__.py @@ -2,17 +2,17 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """ -Inference submodule - Inference backend for model loading and generation +Inference submodule - backend for model loading and generation. The default get_inference_backend() returns an InferenceOrchestrator that -delegates to a subprocess. The original InferenceBackend runs inside -the subprocess and can be imported directly from .inference when needed. +delegates to a subprocess. The original InferenceBackend runs inside the +subprocess and can be imported directly from .inference when needed. """ from .orchestrator import InferenceOrchestrator, get_inference_backend from .llama_cpp import LlamaCppBackend -# Expose InferenceOrchestrator as InferenceBackend for backward compat +# Expose InferenceOrchestrator as InferenceBackend for backward compat. InferenceBackend = InferenceOrchestrator __all__ = [ diff --git a/studio/backend/core/inference/_html_to_md.py b/studio/backend/core/inference/_html_to_md.py index f999120ffb..e7de4a5312 100644 --- a/studio/backend/core/inference/_html_to_md.py +++ b/studio/backend/core/inference/_html_to_md.py @@ -5,7 +5,7 @@ Minimal HTML-to-Markdown converter using only the standard library. Replaces the external ``html2text`` (GPL-3.0) dependency with a ~250-line -``html.parser.HTMLParser`` subclass. Covers headings, links, bold/italic, +``html.parser.HTMLParser`` subclass. Covers headings, links, bold/italic, lists, tables, blockquotes, code blocks, and entity decoding. """ @@ -81,9 +81,7 @@ class _MarkdownRenderer(HTMLParser): self._pre_parts: list[str] = [] self._in_inline_code: bool = False - # Blockquote state -- stack of output buffers so nested - # blockquotes each collect their own content and get prefixed - # with the correct number of ">" markers on close. + # Blockquote state: stack of buffers so nested blockquotes get the right ">" depth. self._bq_stack: list[list[str]] = [] # ------------------------------------------------------------------ @@ -102,7 +100,7 @@ class _MarkdownRenderer(HTMLParser): # ------------------------------------------------------------------ def _prefix_blockquote(self, content: str) -> str: """Prefix every line of *content* with ``> ``.""" - # Strip trailing whitespace first, then collapse blank lines + # Strip trailing whitespace, then collapse blank lines. content = re.sub(r"[ \t]+$", "", content, flags = re.MULTILINE) content = re.sub(r"\n{3,}", "\n\n", content).strip() if not content: @@ -116,10 +114,7 @@ class _MarkdownRenderer(HTMLParser): prefixed.append(">") return "\n".join(prefixed) - # ------------------------------------------------------------------ - # Table helpers -- flush open cells and rows so that HTML with - # omitted optional end tags (, ) does not lose data. - # ------------------------------------------------------------------ + # Table helpers: flush open cells/rows so omitted / don't lose data. def _finish_cell(self) -> None: if not self._in_cell: return @@ -142,10 +137,7 @@ class _MarkdownRenderer(HTMLParser): self._current_row = [] self._row_has_th = False - # ------------------------------------------------------------------ - # Link text helper -- normalize whitespace so block-level content - # inside an does not produce multiline Markdown link labels. - # ------------------------------------------------------------------ + # Link text helper: normalize whitespace so block content in stays single-line. def _finish_link(self) -> None: text = re.sub(r"\s+", " ", "".join(self._link_text_parts)).strip() href = self._link_href or "" @@ -234,22 +226,19 @@ class _MarkdownRenderer(HTMLParser): self._emit("\n\n") elif tag == "tr": - # Flush any open cell/row from a previous row that may - # have omitted its optional or end tags. + # Flush open cell/row from a prior row that omitted /. self._finish_cell() self._finish_row() elif tag in ("th", "td"): - # Flush any open cell (handles omitted /) - self._finish_cell() + self._finish_cell() # handles omitted / self._cell_parts = [] self._in_cell = True if tag == "th": self._row_has_th = True elif tag == "img": - # Skip images -- keeps fetched page text focused on readable - # content and avoids data-URI amplification. + # Skip images: keeps text readable, avoids data-URI amplification. return def handle_endtag(self, tag: str) -> None: @@ -310,7 +299,7 @@ class _MarkdownRenderer(HTMLParser): self._finish_row() elif tag == "table": - # Flush any remaining row (handles omitted ) + # Flush remaining row (handles omitted ). self._finish_cell() self._finish_row() self._in_table = False @@ -325,15 +314,13 @@ class _MarkdownRenderer(HTMLParser): if self._in_pre: self._pre_parts.append(data) return - # Preserve literal whitespace inside inline spans + # Preserve literal whitespace inside inline spans. if self._in_inline_code: self._emit(data) return - # Collapse all whitespace (including newlines) per HTML rules + # Collapse all whitespace (including newlines) per HTML rules. text = re.sub(r"\s+", " ", data) - # Suppress whitespace-only text nodes between table structural - # elements (indentation from source HTML) to prevent leading - # spaces from breaking Markdown table row alignment. + # Suppress whitespace-only nodes between table elements (source indentation). if self._in_table and not self._in_cell and not text.strip(): return self._emit(text) @@ -348,18 +335,10 @@ class _MarkdownRenderer(HTMLParser): return self._emit(html.unescape(f"&#{name};")) - # ------------------------------------------------------------------ # Flush pending buffers (handles truncated HTML from capped fetches) - # ------------------------------------------------------------------ def flush_pending(self) -> None: - """Flush any open side-buffers into ``_out``. - - Called after ``close()`` to recover content from truncated HTML - where closing tags were never seen (common when ``_fetch_page_text`` - caps the download by byte count). - """ + """Flush open side-buffers into ``_out`` after close(), recovering truncated HTML.""" # Flush innermost buffers first so their content propagates outward. - if self._in_link: self._finish_link() @@ -376,7 +355,7 @@ class _MarkdownRenderer(HTMLParser): block = "```\n" + raw + "\n```" self._emit("\n\n" + block + "\n\n") - # Flatten any open blockquote buffers (innermost first) + # Flatten any open blockquote buffers (innermost first). while self._bq_stack: content = "".join(self._bq_stack.pop()) prefixed = self._prefix_blockquote(content) @@ -388,15 +367,9 @@ class _MarkdownRenderer(HTMLParser): self._out.append("\n\n" + prefixed + "\n\n") -# ------------------------------------------------------------------ # Post-processing -# ------------------------------------------------------------------ def _cleanup(text: str) -> str: - """Normalize whitespace and blank lines in the final output. - - Preserves content inside fenced code blocks verbatim so that - intentional blank lines in ``
`` content are not collapsed.
-    """
+    """Normalize whitespace and blank lines, preserving fenced code blocks verbatim."""
     lines = text.split("\n")
     out: list[str] = []
     in_fence = False
@@ -411,7 +384,6 @@ def _cleanup(text: str) -> str:
             continue
 
         if in_fence:
-            # Preserve code block content exactly as-is
             out.append(line)
             continue
 
@@ -427,17 +399,13 @@ def _cleanup(text: str) -> str:
     return "\n".join(out).strip()
 
 
-# ------------------------------------------------------------------
 # Public API
-# ------------------------------------------------------------------
 def html_to_markdown(source_html: str) -> str:
-    """Convert an HTML string to Markdown.
+    """Convert HTML to Markdown (headings, links, emphasis, lists, tables, blockquotes, code, entities).
 
-    Handles headings, links, bold/italic, lists (ordered and unordered),
-    tables, blockquotes, code blocks, and HTML entities.  ``")
@@ -98,9 +96,7 @@ def test_is_same_origin_request_data_url_origin_is_cross_origin():
 
 
 def test_is_same_origin_request_blob_url_origin_is_cross_origin():
-    """``blob:`` URLs carry the inner origin only in non-canonical form; the
-    canonical comparison rejects them.
-    """
+    """``blob:`` URLs carry the inner origin only in non-canonical form; the canonical comparison rejects them."""
     from main import _is_same_origin_request
 
     req = _build_request("127.0.0.1:8902", origin = "blob:http://127.0.0.1:8902/uuid")
@@ -108,8 +104,8 @@ def test_is_same_origin_request_blob_url_origin_is_cross_origin():
 
 
 def test_is_same_origin_request_file_url_origin_is_cross_origin():
-    """``file://`` pages usually send ``Origin: null``; historical engines
-    sent ``Origin: file://``. Neither is same-origin vs an http listener.
+    """``file://`` pages usually send ``Origin: null``; older engines sent
+    ``Origin: file://``. Neither is same-origin vs an http listener.
     """
     from main import _is_same_origin_request
 
@@ -121,8 +117,8 @@ def test_is_same_origin_request_file_url_origin_is_cross_origin():
 
 
 def test_is_same_origin_request_comma_joined_origins_cross_origin():
-    """Starlette concatenates repeated headers with ``, ``; the canonical
-    parser can't safely split this, so it falls to cross-origin.
+    """Starlette joins repeated headers with ``, ``; the canonical parser can't
+    safely split this, so it falls to cross-origin.
     """
     from main import _is_same_origin_request
 
@@ -137,8 +133,8 @@ def test_is_same_origin_request_comma_joined_origins_cross_origin():
 
 
 def test_is_same_origin_request_localhost_vs_127_is_cross_origin():
-    """Browsers treat ``localhost`` and ``127.0.0.1`` as distinct origins;
-    the canonical comparison must not DNS-collapse them.
+    """Browsers treat ``localhost`` and ``127.0.0.1`` as distinct origins; the
+    canonical comparison must not DNS-collapse them.
     """
     from main import _is_same_origin_request
 
@@ -157,7 +153,7 @@ def test_is_same_origin_request_127_vs_localhost_is_cross_origin():
 
 def test_is_same_origin_request_malformed_ipv6_bracket_is_cross_origin():
     """``urlparse`` raises ``ValueError('Invalid IPv6 URL')`` on unclosed
-    brackets (CVE-2024-11168 hardening). The gate must swallow and fall to
+    brackets (CVE-2024-11168 hardening). The gate must swallow it and fall to
     cross-origin rather than 500 the SPA handler.
     """
     from main import _is_same_origin_request
@@ -184,8 +180,8 @@ def test_is_same_origin_request_bracket_with_trailing_garbage_is_cross_origin():
 
 
 def test_is_same_origin_request_empty_origin_header_is_cross_origin():
-    """Explicit empty ``Origin:`` is not a valid serialised origin and must
-    not be conflated with a missing header; cross-origin, bootstrap withheld.
+    """Explicit empty ``Origin:`` is not a valid serialised origin and must not
+    be conflated with a missing header; cross-origin, bootstrap withheld.
     """
     from main import _is_same_origin_request
 
diff --git a/studio/backend/tests/test_inference_model_validation.py b/studio/backend/tests/test_inference_model_validation.py
index faf5a67873..2427ad35fa 100644
--- a/studio/backend/tests/test_inference_model_validation.py
+++ b/studio/backend/tests/test_inference_model_validation.py
@@ -170,16 +170,15 @@ def test_walkback_does_not_cross_user_turn():
         ]
     )
     last = req.messages[-1].tool_call_id
-    # The walkback must NOT pick old_call because a user turn intervenes;
-    # falls back to synth.
+    # Walkback must NOT pick old_call across a user turn; falls back to synth.
     assert last is not None
     assert last != "old_call"
     assert last.startswith("call_")
 
 
 def test_walkback_skips_explicitly_consumed_tool_call_id():
-    """Sibling tool result with an explicit id must reserve its assistant
-    slot so a follow-up missing-id result picks the OTHER tool call."""
+    """An explicit-id tool result reserves its assistant slot so a
+    follow-up missing-id result picks the OTHER tool call."""
     req = _req(
         [
             {
@@ -207,7 +206,7 @@ def test_walkback_skips_explicitly_consumed_tool_call_id():
 
 def test_walkback_handles_malformed_function_string():
     """A tool_call with ``function`` as a string (provider quirk) must not
-    raise; resolution falls back to fallback id selection."""
+    raise; resolution falls back to id selection."""
     req = _req(
         [
             {
diff --git a/studio/backend/tests/test_kv_cache_estimation.py b/studio/backend/tests/test_kv_cache_estimation.py
index 001b5f1bee..cd834b345b 100644
--- a/studio/backend/tests/test_kv_cache_estimation.py
+++ b/studio/backend/tests/test_kv_cache_estimation.py
@@ -7,8 +7,7 @@ Covers the GGUF metadata parser, _can_estimate_kv gate, all 5 estimation
 paths (MLA, Hybrid Mamba, Sliding Window, Standard GQA, Legacy), KV cache
 quantization, edge cases, and lifecycle (init/unload/reparse).
 
-Requires no GPU, network, or external libraries beyond pytest.
-Cross-platform: Linux, macOS, Windows, WSL.
+No GPU, network, or libraries beyond pytest. Cross-platform.
 """
 
 import io
@@ -20,10 +19,8 @@ from pathlib import Path
 
 import pytest
 
-# ---------------------------------------------------------------------------
-# Stub heavy / unavailable external dependencies before importing the
-# module under test.  Same pattern as test_native_context_length.py.
-# ---------------------------------------------------------------------------
+# Stub heavy / unavailable deps before importing the module under test.
+# Same pattern as test_native_context_length.py.
 
 _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
 if _BACKEND_DIR not in sys.path:
@@ -38,11 +35,9 @@ sys.modules.setdefault("loggers", _loggers_stub)
 _structlog_stub = _types.ModuleType("structlog")
 sys.modules.setdefault("structlog", _structlog_stub)
 
-# httpx -- only stub when the real library isn't installed.  Stubbing
-# unconditionally would shadow ``HTTPError`` / ``Response`` etc. that
-# ``huggingface_hub.errors`` imports at module load time, which causes
-# the transformers introspection tier to silently return None inside
-# the test process.
+# 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.
 try:
     import httpx as _httpx_real  # noqa: F401
 except ImportError:
@@ -78,15 +73,13 @@ except ImportError:
 
 from core.inference.llama_cpp import LlamaCppBackend
 
-# ---------------------------------------------------------------------------
 # Helpers
-# ---------------------------------------------------------------------------
 
 
 def _make_gguf_bytes(arch: str, kv_pairs: dict) -> bytes:
-    """Build a minimal GGUF v3 binary blob with the given KV metadata.
+    """Build a minimal GGUF v3 blob with the given KV metadata.
 
-    Supports the scalar and simple array metadata used by the parser.
+    Supports the scalar and simple array metadata the parser uses.
     """
     buf = io.BytesIO()
     # Header: magic, version, tensor_count, kv_count
@@ -134,9 +127,8 @@ def _backend_from_gguf(
 ) -> LlamaCppBackend:
     """Create a LlamaCppBackend with parsed GGUF metadata from given fields.
 
-    `general` lets a test inject extra `general.*` metadata (used to
-    verify the dynamic SWA resolver picks up source-repo hints from
-    GGUFs that ship them).
+    `general` injects extra `general.*` metadata, to verify the dynamic
+    SWA resolver picks up source-repo hints from GGUFs that ship them.
     """
     kv = {"general.architecture": arch}
     for k, v in (general or {}).items():
@@ -157,13 +149,11 @@ def _backend_from_gguf(
         os.unlink(path)
 
 
-# ---------------------------------------------------------------------------
 # A. GGUF Parser Tests
-# ---------------------------------------------------------------------------
 
 
 class TestGGUFParserNewFields:
-    """Verify that architecture-aware fields are correctly parsed."""
+    """Architecture-aware fields are parsed correctly."""
 
     @pytest.mark.parametrize(
         "field,gguf_key,value",
@@ -217,9 +207,9 @@ class TestGGUFParserNewFields:
         )
         # Per-layer KV head count is preserved exactly...
         assert b._n_kv_heads_by_layer == [8, 8, 8, 8, 8, 2]
-        # ...and mirrored into the scalar field as a conservative max so
-        # non-SWA estimator paths and any caller using
-        # `n_kv = self._n_kv_heads or ...` get a safe upper bound.
+        # ...and mirrored into the scalar field as a conservative max, so
+        # non-SWA paths and callers using `n_kv = self._n_kv_heads or ...`
+        # get a safe upper bound.
         assert b._n_kv_heads == 8
         assert b._sliding_window_pattern == [True, True, True, True, True, False]
 
@@ -270,7 +260,7 @@ class TestArchSwaPatternDefaults:
         assert b._sliding_window_pattern is None
 
     def test_explicit_pattern_overrides_arch_default(self):
-        # Period=6 is the gemma3 default; the explicit array must win.
+        # gemma3 default is period=6; the explicit array must win.
         b = _backend_from_gguf(
             "gemma3",
             {
@@ -310,8 +300,8 @@ class TestArchSwaPatternDefaults:
         "arch", ["llama", "qwen2", "qwen3", "mistral", "mistral3", "glm4", "llama4"]
     )
     def test_non_swa_arch_uses_full_attention_path(self, arch):
-        # Pure-GQA arches: GGUF has no sliding_window, no synthetic
-        # pattern, estimator hits Path 4.
+        # Pure-GQA arches: no sliding_window, no synthetic pattern,
+        # estimator hits Path 4.
         b = _backend_from_gguf(
             arch,
             {
@@ -340,7 +330,7 @@ class TestArchSwaPatternDefaults:
             "embedding_length": 5376,
         }
         with_default = _backend_from_gguf("gemma3", common)
-        # Arch not in the table -> legacy 1/4 path.
+        # Arch not in table -> legacy 1/4 path.
         without_default = _backend_from_gguf("totallymadeupv7", common)
 
         kv_default = with_default._estimate_kv_cache_bytes(131072, "f16")
@@ -430,7 +420,7 @@ class TestDynamicSwaResolver:
     def test_period_from_layer_types_finds_smallest_period(self):
         from core.inference.llama_cpp import _period_from_layer_types
 
-        # gemma3 (1 global per 6), gpt-oss (alternating), gemma3n (1 per 5).
+        # gemma3 (1 global/6), gpt-oss (alternating), gemma3n (1/5).
         assert _period_from_layer_types((["sliding_attention"] * 5 + ["full_attention"]) * 4) == 6
         assert _period_from_layer_types(["sliding_attention", "full_attention"] * 12) == 2
         assert _period_from_layer_types((["sliding_attention"] * 4 + ["full_attention"]) * 7) == 5
@@ -481,14 +471,14 @@ class TestDynamicSwaResolver:
 
     def test_disk_cache_takes_precedence_over_bootstrap(self, monkeypatch, tmp_path):
         self._isolate_cache(monkeypatch, tmp_path)
-        # Override bootstrap=6 with a cached period=3.
+        # Cached period=3 overrides bootstrap=6.
         with open(tmp_path / "swa_cache.json", "w") as f:
             json.dump({"gemma3": 3}, f)
         b = _backend_from_gguf("gemma3", dict(_SWA_FIELDS, block_count = 18))
         assert b._sliding_window_pattern == [(i + 1) % 3 != 0 for i in range(18)]
 
     def test_disk_cache_supports_array_entries(self, monkeypatch, tmp_path):
-        # Aperiodic mask gets tiled across n_layers.
+        # Aperiodic mask is tiled across n_layers.
         self._isolate_cache(monkeypatch, tmp_path)
         mask = [True, False, True, True, False, True, False, False]
         with open(tmp_path / "swa_cache.json", "w") as f:
@@ -556,7 +546,7 @@ class TestDynamicSwaResolver:
         from core.inference import llama_cpp as lc
 
         monkeypatch.setattr(lc, "_fetch_swa_entry_from_hf", lambda repo_id: None)
-        # Force the failure into the Tier 3 path; bypass Tier 2.5.
+        # Force failure into Tier 3; bypass Tier 2.5.
         monkeypatch.setattr(lc, "_resolve_swa_entry_from_transformers", lambda arch: None)
         b = _backend_from_gguf(
             "newmodel",
@@ -639,7 +629,7 @@ class TestTransformersIntrospection:
         assert _resolve_swa_entry_from_transformers("totally-fake-arch-xyz") is None
 
     def test_full_resolver_uses_transformers_before_hf_fetch(self, monkeypatch, tmp_path):
-        # With bootstrap empty, Tier 2.5 must answer before Tier 3 fires.
+        # Bootstrap empty: Tier 2.5 must answer before Tier 3 fires.
         self._isolate_cache(monkeypatch, tmp_path)
         from core.inference import llama_cpp as lc
 
@@ -660,10 +650,10 @@ class TestTransformersIntrospection:
 
 
 class TestGGUFParserReset:
-    """Verify that fields are properly reset between parses."""
+    """Fields are reset between parses."""
 
     def test_reset_between_parses(self):
-        # First parse with all fields
+        # First parse: all fields set
         b = _backend_from_gguf(
             "arch1",
             {
@@ -685,7 +675,7 @@ class TestGGUFParserReset:
         assert b._kv_value_length_swa == 64
         assert b._ssm_inner_size == 4096
 
-        # Second parse without those fields -- they should be None
+        # Second parse without those fields -- they must be None
         kv = {"general.architecture": "arch2", "arch2.block_count": 64}
         import tempfile, os
 
@@ -707,13 +697,11 @@ class TestGGUFParserReset:
         assert b._n_layers == 64
 
 
-# ---------------------------------------------------------------------------
 # B. _can_estimate_kv Gate Tests
-# ---------------------------------------------------------------------------
 
 
 class TestCanEstimateKV:
-    """Verify gate logic for all field combinations."""
+    """Gate logic for all field combinations."""
 
     def test_no_layers_returns_false(self):
         b = LlamaCppBackend()
@@ -729,7 +717,7 @@ class TestCanEstimateKV:
         assert b._can_estimate_kv()
 
     def test_key_length_alone_insufficient(self):
-        """key_length without value_length should NOT be enough."""
+        """key_length without value_length is NOT enough."""
         b = LlamaCppBackend()
         b._n_layers = 32
         b._kv_key_length = 128
@@ -767,9 +755,7 @@ class TestCanEstimateKV:
         assert not b._can_estimate_kv()
 
 
-# ---------------------------------------------------------------------------
 # C. Path 1: MLA Estimation
-# ---------------------------------------------------------------------------
 
 
 class TestMLAEstimation:
@@ -799,33 +785,33 @@ class TestMLAEstimation:
         assert b._estimate_kv_cache_bytes(163840, "f16") == expected
 
     def test_mla_ignores_value_length(self):
-        """MLA should NOT add value_length -- V is reconstructed from the latent."""
+        """MLA must NOT add value_length -- V is reconstructed from the latent."""
         b = self._mla_backend()
         result = b._estimate_kv_cache_bytes(1000, "f16")
-        # Should be n_layers * ctx * 1 * key_len(576) * 2
+        # n_layers * ctx * 1 * key_len(576) * 2
         expected = 61 * 1000 * 1 * 576 * 2
         assert result == expected
 
     def test_mla_fallback_when_no_key_length(self):
-        """If key_length is missing, fallback to kv_lora_rank + key_length_mla."""
+        """No key_length: fall back to kv_lora_rank + key_length_mla."""
         b = self._mla_backend(_kv_key_length = None)
-        # _key_length_mla=192 in default, so rope_dim=192
+        # default _key_length_mla=192, so rope_dim=192
         result = b._estimate_kv_cache_bytes(1000, "f16")
         expected = 61 * 1000 * 1 * (512 + 192) * 2  # 704
         assert result == expected
 
     def test_mla_fallback_no_key_length_mla(self):
-        """If both key_length and key_length_mla are missing, fallback to +64."""
+        """No key_length and no key_length_mla: fall back to +64."""
         b = self._mla_backend(_kv_key_length = None, _key_length_mla = None)
         result = b._estimate_kv_cache_bytes(1000, "f16")
         expected = 61 * 1000 * 1 * (512 + 64) * 2  # 576
         assert result == expected
 
     def test_mla_defaults_n_kv_to_1_when_heads_absent(self):
-        """MLA should use n_kv=1 even if n_kv_heads is None (not n_heads)."""
+        """MLA uses n_kv=1 even if n_kv_heads is None (not n_heads)."""
         b = self._mla_backend(_n_kv_heads = None)  # n_heads=128 still set
         result = b._estimate_kv_cache_bytes(1000, "f16")
-        # Should use n_kv_mla=1, NOT n_heads=128
+        # Uses n_kv_mla=1, NOT n_heads=128
         expected = 61 * 1000 * 1 * 576 * 2
         assert result == expected
 
@@ -838,9 +824,7 @@ class TestMLAEstimation:
         assert result_q4 == int(61 * 1000 * 1 * 576 * 0.5625)
 
 
-# ---------------------------------------------------------------------------
 # D. Path 2: Hybrid Mamba Estimation
-# ---------------------------------------------------------------------------
 
 
 class TestHybridMambaEstimation:
@@ -883,14 +867,14 @@ class TestHybridMambaEstimation:
         assert b._estimate_kv_cache_bytes(262144, "f16") == expected
 
     def test_hybrid_without_explicit_dims(self):
-        """Fallback to head_dim when key_length/value_length are missing."""
+        """Fall back to head_dim when key_length/value_length are missing."""
         b = self._hybrid_backend(_kv_key_length = None, _kv_value_length = None)
         head_dim = 5120 // 24  # 213
         expected = 16 * 4096 * 4 * 2 * head_dim * 2
         assert b._estimate_kv_cache_bytes(4096, "f16") == expected
 
     def test_fai_zero_safety(self):
-        """full_attention_interval=0 should not cause ZeroDivisionError."""
+        """full_attention_interval=0 must not ZeroDivisionError."""
         b = self._hybrid_backend(_full_attention_interval = 0)
         result = b._estimate_kv_cache_bytes(4096, "f16")
         # fai=0 -> n_attn = n_layers (all layers)
@@ -898,9 +882,7 @@ class TestHybridMambaEstimation:
         assert result == expected
 
 
-# ---------------------------------------------------------------------------
 # E. Path 3: Sliding Window Estimation
-# ---------------------------------------------------------------------------
 
 
 class TestSlidingWindowEstimation:
@@ -978,7 +960,7 @@ class TestSlidingWindowEstimation:
             assert b._estimate_kv_cache_bytes(ctx, "f16") == expected(ctx)
 
     def test_ctx_smaller_than_window(self):
-        """When context < 2 * sliding_window, SWA cache caps at ctx."""
+        """When ctx < 2 * sliding_window, SWA cache caps at ctx."""
         b = self._swa_backend(_sliding_window = 8192)
         n_global = max(1, 62 // 4)  # 15
         n_swa = 62 - n_global  # 47
@@ -996,9 +978,7 @@ class TestSlidingWindowEstimation:
         assert b._estimate_kv_cache_bytes(1000, "f16") == expected
 
 
-# ---------------------------------------------------------------------------
 # F. Path 4: Standard GQA Estimation
-# ---------------------------------------------------------------------------
 
 
 class TestStandardGQAEstimation:
@@ -1031,20 +1011,18 @@ class TestStandardGQAEstimation:
         assert b._estimate_kv_cache_bytes(4096, "f16") == expected
 
     def test_differs_from_legacy(self):
-        """GQA path should differ from legacy when key_length != embed//n_heads."""
+        """GQA path differs from legacy when key_length != embed//n_heads."""
         b = self._gqa_backend()
         head_dim = 1024 // 16  # 64
         gqa_result = b._estimate_kv_cache_bytes(4096, "f16")
-        # Legacy would use: 2 * 8 * 64 * 28 * 4096 * 2
+        # Legacy: 2 * 8 * 64 * 28 * 4096 * 2
         legacy_result = int(2 * 8 * head_dim * 28 * 4096 * 2)
         # GQA: 28 * 4096 * 8 * (128+128) * 2 -- uses actual key_length=128
         assert gqa_result != legacy_result
         assert gqa_result > legacy_result  # key_length (128) > head_dim (64)
 
 
-# ---------------------------------------------------------------------------
 # G. Path 5: Legacy Fallback Estimation
-# ---------------------------------------------------------------------------
 
 
 class TestLegacyEstimation:
@@ -1077,7 +1055,7 @@ class TestLegacyEstimation:
         assert b._estimate_kv_cache_bytes(4096, "f16") == expected
 
     def test_legacy_identical_to_old_formula(self):
-        """Confirm legacy path produces the same result as the pre-PR formula."""
+        """Legacy path matches the pre-PR formula."""
         b = self._legacy_backend()
         n_layers = 32
         n_kv_heads = 8
@@ -1088,16 +1066,14 @@ class TestLegacyEstimation:
         assert b._estimate_kv_cache_bytes(n_ctx, "f16") == old_formula
 
 
-# ---------------------------------------------------------------------------
 # H. Path Priority (selection order)
-# ---------------------------------------------------------------------------
 
 
 class TestPathPriority:
     """Confirm: MLA > Hybrid Mamba > SWA > GQA > Legacy."""
 
     def test_mla_takes_priority_over_all(self):
-        """If kv_lora_rank is set, MLA path is used even if other fields are present."""
+        """If kv_lora_rank is set, MLA path wins even with other fields present."""
         b = LlamaCppBackend()
         b._n_layers = 61
         b._n_kv_heads = 1
@@ -1132,9 +1108,9 @@ class TestPathPriority:
         assert b._estimate_kv_cache_bytes(1000, "f16") == expected_hybrid
 
     def test_all_paths_produce_different_values(self):
-        """With carefully chosen params, each path should yield a distinct value."""
-        # Use embedding_length=768 so legacy head_dim (768//16=48) differs from
-        # key_length (256), and MLA key_len (256) != legacy K+V (2*48=96).
+        """With chosen params, each path yields a distinct value."""
+        # embedding_length=768 so legacy head_dim (768//16=48) != key_length
+        # (256), and MLA key_len (256) != legacy K+V (2*48=96).
         params = {
             "_n_layers": 40,
             "_n_kv_heads": 4,
@@ -1185,13 +1161,11 @@ class TestPathPriority:
         assert len(set(values)) == 5, f"Expected 5 distinct values, got {values}"
 
 
-# ---------------------------------------------------------------------------
 # I. KV Cache Quantization
-# ---------------------------------------------------------------------------
 
 
 class TestQuantization:
-    """Verify all supported cache_type_kv values produce correct scaling."""
+    """All supported cache_type_kv values scale correctly."""
 
     @pytest.mark.parametrize(
         "cache_type,expected_bpe",
@@ -1222,9 +1196,7 @@ class TestQuantization:
         assert result == expected
 
 
-# ---------------------------------------------------------------------------
 # J. Edge Cases
-# ---------------------------------------------------------------------------
 
 
 class TestEdgeCases:
@@ -1285,10 +1257,8 @@ class TestEdgeCases:
         assert result == expected
 
 
-# ---------------------------------------------------------------------------
 # J2. Server-flag knobs (--swa-full, --kv-unified/--parallel,
 #     --ctx-checkpoints, --kv-offload)
-# ---------------------------------------------------------------------------
 
 
 class TestServerFlags:
@@ -1332,7 +1302,7 @@ class TestServerFlags:
         b = self._swa_backend()
         ctx = 32_768
         flagged = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True)
-        # With swa_full, every layer caches n_ctx -- equals path 4 sizing.
+        # swa_full: every layer caches n_ctx -- equals path 4 sizing.
         kv_per_token = 4 * (256 + 256) * 2  # n_kv_heads * (k+v) * f16
         expected = 26 * ctx * kv_per_token
         assert flagged == expected
@@ -1366,10 +1336,9 @@ class TestServerFlags:
         assert with_cp > b._estimate_kv_cache_bytes(8192, "f16")
 
     # ── --parallel + --kv-unified ──────────────────────────────────
-    # Empirically verified against llama-server: non-SWA caches partition
-    # n_ctx across slots (total memory constant); SWA layers are the only
-    # portion that scales with --parallel.  --kv-unified is currently a
-    # no-op for memory math (kept for API forward-compat).
+    # Verified against llama-server: non-SWA caches partition n_ctx across
+    # slots (total memory constant); only SWA layers scale with --parallel.
+    # --kv-unified is a no-op for memory math (kept for API forward-compat).
 
     def test_gqa_kv_constant_across_parallel(self):
         b = self._gqa_backend()
@@ -1394,7 +1363,7 @@ class TestServerFlags:
         b = self._swa_backend()
         ctx = 8192
         baseline = b._estimate_kv_cache_bytes(ctx, "f16")
-        # Decompose baseline by walking the same loop the estimator does.
+        # Decompose baseline by walking the estimator's own loop.
         swa = b._sliding_window
         per_token_global = 4 * (256 + 256) * 2  # n_kv * (k+v) * f16
         per_token_swa = 4 * (256 + 256) * 2  # k_swa/val_swa fall back
@@ -1409,10 +1378,10 @@ class TestServerFlags:
         )
         # Sanity: parallel=1 reproduces baseline exactly
         assert global_bytes + swa_bytes_per_slot == baseline
-        # Only SWA portion scales by parallel
+        # Only the SWA portion scales by parallel
         for slots in (1, 2, 3, 4):
             scaled = b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = slots, kv_unified = False)
-            # SWA cells get clamped to per_slot_ctx when ctx/slots < 2*swa
+            # SWA cells clamp to per_slot_ctx when ctx/slots < 2*swa
             per_slot_ctx = max(1, ctx // slots)
             cells = min(ctx, 2 * swa, per_slot_ctx)
             swa_bps = sum(
@@ -1452,7 +1421,7 @@ class TestServerFlags:
         ctx = 8192
         baseline = b._estimate_kv_cache_bytes(ctx, "f16")
         flagged = b._estimate_kv_cache_bytes(ctx, "f16", ctx_checkpoints = 4)
-        # 22 SWA layers * 4 checkpoints * 512 cells * 4 heads * (256+256) * 2 bytes
+        # 22 SWA layers * 4 cps * 512 cells * 4 heads * (256+256) * 2 bytes
         n_swa_layers = sum(1 for f in [True, True, True, True, True, False] * 4 + [True, True] if f)
         per_layer = 4 * 512 * 4 * (256 + 256) * 2
         assert flagged == baseline + n_swa_layers * per_layer
@@ -1470,7 +1439,7 @@ class TestServerFlags:
 
     def test_ctx_checkpoints_compose_with_n_parallel(self):
         # Only the SWA + checkpoint portion scales by n_parallel; the
-        # global-layer portion stays constant.
+        # global-layer portion is constant.
         b = self._swa_backend()
         ctx = 8192
         swa = b._sliding_window
@@ -1493,7 +1462,7 @@ class TestServerFlags:
 
     def test_fit_returns_requested_when_kv_off_gpu(self):
         b = self._gqa_backend()
-        # Tiny VRAM budget -- normally would force a reduction.
+        # Tiny VRAM budget -- would normally force a reduction.
         fitted = b._fit_context_to_vram(
             requested_ctx = 32_768,
             available_mib = 1,
@@ -1515,8 +1484,8 @@ class TestServerFlags:
         assert fitted < 32_768
 
     def test_fit_mtp_engaged_returns_smaller_or_equal_context(self):
-        # MTP-engaged budget is 0.85 of available; non-MTP is 0.90.
-        # On a tight budget the MTP path must yield <= the non-MTP path.
+        # MTP budget is 0.85 of available, non-MTP is 0.90; on a tight
+        # budget MTP must yield <= non-MTP.
         b = self._gqa_backend()
         common = dict(
             requested_ctx = 32_768,
@@ -1529,7 +1498,7 @@ class TestServerFlags:
         assert mtp <= baseline
 
     def test_fit_mtp_engaged_unchanged_when_kv_off_gpu(self):
-        # kv_on_gpu=False short-circuits the fit; mtp_engaged is irrelevant.
+        # kv_on_gpu=False short-circuits the fit; mtp_engaged irrelevant.
         b = self._gqa_backend()
         fitted = b._fit_context_to_vram(
             requested_ctx = 32_768,
@@ -1548,7 +1517,7 @@ class TestServerFlags:
         kv_default = b._estimate_kv_cache_bytes(ctx, "f16")
         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 should not fit.
+        # Budget = model + kv_default (rounded up) -- swa_full must not fit.
         budget_mib = (1024 * 1024 + kv_default) / (1024 * 1024) / 0.90 + 1
         fitted_default = b._fit_context_to_vram(
             requested_ctx = ctx,
@@ -1567,22 +1536,20 @@ class TestServerFlags:
         assert fitted_full < ctx
 
 
-# ---------------------------------------------------------------------------
 # J2.5. --parallel N memory accounting (per-layer-type scaling rule)
-# ---------------------------------------------------------------------------
 
 
 class TestParallelSWAScaling:
-    """Verifies the per-layer-type scaling rule against the closed form
-    measured from llama-server. Empirical formula on Gemma-3 270m at
-    ctx=8192: total_kv = 24 + parallel * 15 (MiB).
+    """Per-layer-type scaling rule vs the closed form measured from
+    llama-server. Empirical formula on Gemma-3 270m at ctx=8192:
+    total_kv = 24 + parallel * 15 (MiB).
 
     Rule (verified vs ``llama-server`` log on real GGUFs):
       * non-SWA layers: total cells = n_ctx, partitioned across slots,
         memory CONSTANT in n_parallel.
       * SWA layers: per-slot cells = 2 * sliding_window (clamped at
         n_ctx and at per_slot_ctx); memory LINEAR in n_parallel.
-      * --kv-unified is a no-op for memory math; both modes yield the
+      * --kv-unified is a no-op for memory math; both modes give the
         same total in measured cases.
     """
 
@@ -1703,7 +1670,7 @@ class TestParallelSWAScaling:
 
     def test_swa_per_slot_clamped_when_ctx_lt_slots_x_2window(self):
         # ctx=4096 / slots=8 -> per_slot_ctx=512, but 2*sliding=1024.
-        # SWA cells should clamp at per_slot_ctx (512), not 2*sliding.
+        # SWA cells clamp at per_slot_ctx (512), not 2*sliding.
         b = self._swa_backend()
         ctx = 4096
         per_slot_ctx_at_8 = ctx // 8
@@ -1719,8 +1686,8 @@ class TestParallelSWAScaling:
         assert b._estimate_kv_cache_bytes(ctx, "f16", n_parallel = 8) == expected
 
     def test_swa_full_does_not_scale_under_parallel(self):
-        # swa_full forces every layer to n_ctx; result is the all-global
-        # GQA-style total, which is constant in parallel.
+        # swa_full forces every layer to n_ctx -> all-global GQA-style
+        # total, constant in parallel.
         b = self._swa_backend()
         ctx = 8192
         baseline = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True)
@@ -1732,8 +1699,8 @@ class TestParallelSWAScaling:
     # ── kv_unified: no-op for memory math ──────────────────────────
 
     def test_kv_unified_is_no_op_for_memory_math(self):
-        # Both unified=True and unified=False must produce the same
-        # total bytes for every backend type and every parallel value.
+        # unified=True and unified=False must give the same total bytes
+        # for every backend type and parallel value.
         backends = [
             ("gqa", self._gqa_backend()),
             ("swa", self._swa_backend()),
@@ -1761,9 +1728,8 @@ class TestParallelSWAScaling:
         b._kv_key_length = 256
         b._kv_value_length = 256
         b._sliding_window = 512
-        # 5-period [swa,swa,swa,swa,full] * 3 + [swa,swa,swa]: mirrors the
-        # bootstrap-resolved pattern for gemma3 (period 6) on an 18-layer
-        # model (15 SWA, 3 global).
+        # Mirrors the bootstrap-resolved gemma3 pattern (period 6) on an
+        # 18-layer model: 15 SWA, 3 global.
         b._sliding_window_pattern = [(i + 1) % 6 != 0 for i in range(18)]
         n_global = 3
         n_swa = 15
@@ -1777,20 +1743,18 @@ class TestParallelSWAScaling:
             ), f"slots={slots}: got {got_mib} MiB, expected {expected_mib} MiB"
 
 
-# ---------------------------------------------------------------------------
 # J3. shared_kv_layers (Gemma 3n / Gemma 4)
-# ---------------------------------------------------------------------------
 
 
 class TestSharedKVLayers:
     """``.attention.shared_kv_layers`` reduces the layer count that
-    actually allocates KV.  The trailing ``shared_kv_layers`` blocks reuse
-    earlier caches (Gemma 3n: 35 layers, 15 shared -> 20 allocate; Gemma 4
-    same field).  Unset on every other arch -> no behavioural change."""
+    allocates KV. The trailing ``shared_kv_layers`` blocks reuse earlier
+    caches (Gemma 3n: 35 layers, 15 shared -> 20 allocate; Gemma 4 same
+    field). Unset on every other arch -> no behavioural change."""
 
     def _gemma3n_backend(self, **overrides):
-        # Mirrors google/gemma-3n-E4B-it: 35 layers, 15 shared,
-        # SWA window 1024, period 5 (4 sliding + 1 full repeating).
+        # Mirrors google/gemma-3n-E4B-it: 35 layers, 15 shared, SWA window
+        # 1024, period 5 (4 sliding + 1 full repeating).
         defaults = {
             "_n_layers": 35,
             "_n_kv_heads": 4,
@@ -1873,9 +1837,8 @@ class TestSharedKVLayers:
     def test_path3_pattern_loops_only_unshared_layers(self):
         b = self._gemma3n_backend()
         ctx = 8192
-        # First 20 layers contribute; layers 20..34 are skipped.
-        # Pattern: [s,s,s,s,F] repeated.  In layers 0..19:
-        #   sliding: 16, full: 4
+        # First 20 layers contribute; layers 20..34 skipped. Pattern
+        # [s,s,s,s,F] repeated -> in layers 0..19: sliding 16, full 4.
         sliding_in_unshared = sum(b._sliding_window_pattern[:20])
         full_in_unshared = 20 - sliding_in_unshared
         assert sliding_in_unshared == 16
@@ -1890,7 +1853,7 @@ class TestSharedKVLayers:
         with_shared = b._estimate_kv_cache_bytes(8192, "f16")
         b._shared_kv_layers = 0
         without_shared = b._estimate_kv_cache_bytes(8192, "f16")
-        # 20/35 = 0.571 of the work; expect ~43% reduction.
+        # 20/35 = 0.571 of the work; ~43% reduction.
         ratio = with_shared / without_shared
         assert 0.5 < ratio < 0.65
 
@@ -1898,8 +1861,8 @@ class TestSharedKVLayers:
         b = self._gemma3n_backend()
         ctx = 8192
         flagged = b._estimate_kv_cache_bytes(ctx, "f16", swa_full = True)
-        # Every unshared layer caches n_ctx; equals path-4-style sizing
-        # over only the 20 unshared layers.
+        # Every unshared layer caches n_ctx -> path-4-style sizing over
+        # only the 20 unshared layers.
         kv_per = 4 * (256 + 256) * 2
         assert flagged == 20 * ctx * kv_per
 
@@ -1917,15 +1880,15 @@ class TestSharedKVLayers:
         assert b._estimate_kv_cache_bytes(ctx, "f16") == expected
 
     def test_shared_floors_at_one_layer(self):
-        # Pathological: shared >= n_layers should not zero out the cache.
+        # Pathological: shared >= n_layers must not zero out the cache.
         b = self._gqa_backend(_shared_kv_layers = 99)
         ctx = 4096
         kv_per = 8 * (128 + 128) * 2
         assert b._estimate_kv_cache_bytes(ctx, "f16") == 1 * ctx * kv_per
 
     def test_composes_with_n_parallel(self):
-        # Only the SWA portion of the unshared layers scales by n_parallel;
-        # the global portion stays constant.
+        # Only the SWA portion of unshared layers scales by n_parallel;
+        # the global portion is constant.
         b = self._gemma3n_backend()
         ctx = 8192
         swa = b._sliding_window
@@ -1946,7 +1909,7 @@ class TestSharedKVLayers:
         ctx = 8192
         baseline = b._estimate_kv_cache_bytes(ctx, "f16")
         with_cp = b._estimate_kv_cache_bytes(ctx, "f16", ctx_checkpoints = 4)
-        # Checkpoints only count over UNSHARED SWA layers (16 of them).
+        # Checkpoints count only over UNSHARED SWA layers (16 of them).
         sliding_in_unshared = sum(b._sliding_window_pattern[:20])
         per_cp_layer = 4 * 1024 * 4 * (256 + 256) * 2  # cps * swa * heads * (k+v) * bpe
         assert with_cp == baseline + sliding_in_unshared * per_cp_layer
@@ -1958,9 +1921,7 @@ class TestSharedKVLayers:
         assert b._shared_kv_layers is None
 
 
-# ---------------------------------------------------------------------------
 # K. Lifecycle Tests
-# ---------------------------------------------------------------------------
 
 
 class TestLifecycle:
@@ -2017,7 +1978,7 @@ class TestLifecycle:
         assert b._n_kv_heads_by_layer is None
 
     def test_end_to_end_synthetic_mla(self):
-        """Full round-trip: write GGUF -> parse -> estimate."""
+        """Round-trip: write GGUF -> parse -> estimate."""
         b = _backend_from_gguf(
             "deepseek2",
             {
@@ -2075,8 +2036,8 @@ class TestLifecycle:
         )
         assert b._can_estimate_kv()
         result = b._estimate_kv_cache_bytes(131072, "f16")
-        # gemma3 -> period 6 from the bootstrap table, SWA cache
-        # double-buffered to 2 * sliding_window cells.
+        # gemma3 -> period 6 from bootstrap; SWA cache double-buffered to
+        # 2 * sliding_window cells.
         period = 6
         kv_per = 16 * 256 * 2
         expected = 0
@@ -2104,13 +2065,13 @@ class TestLifecycle:
         )
         assert b._can_estimate_kv()
         assert b._shared_kv_layers == 15
-        # Bootstrap table for gemma3n_text -> period 5; the resolver
-        # synthesises a 35-entry bool array.  The first 20 entries
-        # (n_layers - shared) are the only ones that allocate KV.
+        # Bootstrap for gemma3n_text -> period 5; resolver synthesises a
+        # 35-entry bool array. Only the first 20 (n_layers - shared)
+        # allocate KV.
         result = b._estimate_kv_cache_bytes(8192, "f16")
         assert result > 0
-        # Sanity: setting shared back to 0 must produce a strictly larger
-        # estimate (more layers allocate).
+        # Sanity: shared back to 0 -> strictly larger estimate (more
+        # layers allocate).
         b._shared_kv_layers = 0
         unshared = b._estimate_kv_cache_bytes(8192, "f16")
         assert unshared > result
diff --git a/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py b/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py
index f887f7747c..3443f53a12 100644
--- a/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py
+++ b/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py
@@ -1,11 +1,10 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-"""Validates that the installer correctly resolves lemonade ROCm prebuilt assets.
+"""Validates that the installer resolves lemonade ROCm prebuilt assets.
 
-Uses a faked HostInfo so no AMD GPU is needed. Network calls to the lemonade
-GitHub API are stubbed out so the suite runs without internet access and is
-not subject to rate limits.
+Uses a faked HostInfo so no AMD GPU is needed. The lemonade GitHub API calls
+are stubbed so the suite runs offline and isn't subject to rate limits.
 """
 
 from __future__ import annotations
@@ -33,7 +32,7 @@ if resolve_lemonade_rocm_choice is None or _LEMONADE_GFX_FAMILIES is None:
 @pytest.fixture(autouse = True)
 def _clear_lemonade_release_cache():
     """Prevent cross-test pollution of the lemonade release lru_cache when
-    future tests vary the fetch_json mock return value."""
+    tests vary the fetch_json mock return value."""
     _cache = getattr(_mod, "_fetch_lemonade_release_cached", None)
     if _cache is not None and hasattr(_cache, "cache_clear"):
         _cache.cache_clear()
@@ -90,9 +89,7 @@ def _lookup_family(gfx: str) -> str | None:
     return None
 
 
-# ---------------------------------------------------------------------------
 # GPU family mapping
-# ---------------------------------------------------------------------------
 
 
 @pytest.mark.parametrize(
@@ -114,9 +111,7 @@ def test_unknown_gpu_not_in_families():
     assert _lookup_family("gfx999") is None
 
 
-# ---------------------------------------------------------------------------
 # Asset resolution - hits real lemonade GitHub API
-# ---------------------------------------------------------------------------
 
 
 @pytest.mark.parametrize(
@@ -146,21 +141,18 @@ def test_unknown_gpu_falls_through_to_upstream():
     assert result is None
 
 
-# ---------------------------------------------------------------------------
 # Simple-policy dispatcher must plan a lemonade ROCm attempt for AMD-only hosts.
-# This is the path setup.sh actually invokes (via --simple-policy), so the
-# lemonade integration is useless if it isn't wired in here.
-# ---------------------------------------------------------------------------
+# This is the path setup.sh invokes (via --simple-policy), so the lemonade
+# integration is useless if it isn't wired in here.
 
 direct_linux_release_plan = getattr(_mod, "direct_linux_release_plan", None)
 direct_upstream_release_plan = getattr(_mod, "direct_upstream_release_plan", None)
 
 
 def _stub_unsloth_release(release_tag: str = "b9022") -> dict:
-    # Minimal payload that parse_direct_linux_release_bundle accepts. It
-    # requires at least one `app-{label}-linux-x64*.tar.gz` asset for the
-    # bundle to be recognised; we ship a bare CPU one so the planner has a
-    # baseline non-ROCm attempt to fall through to.
+    # Minimal payload parse_direct_linux_release_bundle accepts. It needs at
+    # least one `app-{label}-linux-x64*.tar.gz` asset to recognise the bundle;
+    # we ship a bare CPU one so the planner has a baseline non-ROCm fallback.
     asset_name = f"app-{release_tag}-linux-x64.tar.gz"
     return {
         "tag_name": release_tag,
@@ -255,9 +247,8 @@ def test_lemonade_release_api_url_pinned_tag():
 
 
 def test_lemonade_release_api_url_encodes_tag():
-    """Unexpected slashes / hashes in the tag must be URL-encoded so the URL
-    cannot be reshaped (defence in depth -- tags should already be sanitised
-    upstream)."""
+    """Slashes / hashes in the tag must be URL-encoded so the URL can't be
+    reshaped (defence in depth -- tags should already be sanitised upstream)."""
     url = _mod._lemonade_release_api_for("b1260/../latest")
     assert "/releases/tags/b1260%2F..%2Flatest" in url
     assert "//latest" not in url.split("/releases/tags/", 1)[1]
@@ -272,9 +263,9 @@ def test_lemonade_resolver_skipped_by_opt_out_env(monkeypatch):
 
 
 def test_lemonade_resolver_rejects_non_github_url(monkeypatch):
-    """If the GitHub API response somehow contained an off-host download URL,
-    the resolver must refuse to use it (lemonade assets are not in the
-    approved-hash manifest)."""
+    """If the GitHub API response contained an off-host download URL, the
+    resolver must refuse it (lemonade assets aren't in the approved-hash
+    manifest)."""
     bad_release = {
         "tag_name": _STUB_TAG,
         "assets": [
@@ -298,7 +289,7 @@ def test_lemonade_resolver_rejects_http_scheme():
 
 
 def test_lemonade_resolver_accepts_github_cdn():
-    # Real GitHub release CDN URLs carry the /github-production-release-asset- prefix.
+    # Real GitHub release CDN URLs carry the /github-production-release-asset- prefix
     assert _mod._is_trusted_github_release_url(
         "https://objects.githubusercontent.com/github-production-release-asset-abc123/456/789?token=x",
         "lemonade-sdk/llamacpp-rocm",
@@ -306,7 +297,7 @@ def test_lemonade_resolver_accepts_github_cdn():
 
 
 def test_lemonade_resolver_rejects_arbitrary_cdn_path():
-    # A CDN URL without the release-asset path prefix must be rejected.
+    # A CDN URL without the release-asset path prefix must be rejected
     assert not _mod._is_trusted_github_release_url(
         "https://objects.githubusercontent.com/abc/def",
         "lemonade-sdk/llamacpp-rocm",
@@ -348,7 +339,7 @@ def test_lemonade_runtime_patterns_include_hip_runtime():
 
     Lemonade ZIPs carry transitive deps (libamd_comgr, libLLVM, libclang-cpp,
     ...) whose names change across ROCm releases. A broad ``lib*.so*`` glob
-    avoids having to enumerate every transitive dependency by name.
+    avoids enumerating every transitive dependency by name.
     """
     from install_llama_prebuilt import runtime_patterns_for_choice, AssetChoice
 
@@ -362,7 +353,7 @@ def test_lemonade_runtime_patterns_include_hip_runtime():
     )
     pats = runtime_patterns_for_choice(choice)
     # The broad glob must be present so every .so in the lemonade bundle
-    # (including transitive deps added in future ROCm releases) gets overlaid.
+    # (including future transitive deps) gets overlaid.
     assert "lib*.so*" in pats, f"'lib*.so*' missing from linux-rocm patterns: {pats}"
 
 
@@ -374,9 +365,9 @@ _pick_rocm_gfx_target = getattr(_mod, "_pick_rocm_gfx_target", None)
     reason = "_pick_rocm_gfx_target not present on this branch",
 )
 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).
+    """AMD HIP honours CUDA_VISIBLE_DEVICES like 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 real tool output).
     probe_out = "gfx1151\ngfx1151\ngfx1100\ngfx1100"
     monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False)
     monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
@@ -405,7 +396,7 @@ 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).
+    # rocminfo output for 3 GPUs (2x gfx1100 dGPU + 1x gfx1151 APU).
     # Each GPU gets its own Agent section with a few token mentions.
     probe_out = (
         "***\nAgent 1\n***\n  gfx1100 some info\n  gfx1100\n"
diff --git a/studio/backend/tests/test_llama_cpp_cache_aware_disk_check.py b/studio/backend/tests/test_llama_cpp_cache_aware_disk_check.py
index 255c04a956..666a503dbb 100644
--- a/studio/backend/tests/test_llama_cpp_cache_aware_disk_check.py
+++ b/studio/backend/tests/test_llama_cpp_cache_aware_disk_check.py
@@ -1,21 +1,12 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-"""Tests for the cache-aware disk-space preflight in
-``LlamaCppBackend.load_model``.
+"""Tests for the cache-aware disk-space preflight in ``LlamaCppBackend.load_model``.
 
-The preflight used to compare the repo's total GGUF download size against
-free disk without accounting for bytes already present in the Hugging
-Face cache. That made re-loading a cached large model (e.g.
-``unsloth/MiniMax-M2.7-GGUF`` at 131 GB) fail cold whenever free disk was
-below the full weight footprint, even though nothing needed
-downloading.
-
-These tests exercise the preflight arithmetic in isolation by driving
-``get_paths_info`` and ``try_to_load_from_cache`` through ``mock.patch``.
-No network, GPU, or subprocess use.
-
-Cross-platform: Linux, macOS, Windows, WSL.
+The preflight used to compare the repo's total GGUF size against free disk
+without counting bytes already in the HF cache, so re-loading a cached large
+model failed cold even though nothing needed downloading. These tests exercise
+the preflight arithmetic in isolation (no network/GPU/subprocess).
 """
 
 from __future__ import annotations
@@ -28,10 +19,8 @@ from unittest.mock import patch
 
 import pytest
 
-# ---------------------------------------------------------------------------
-# Stub heavy / unavailable external dependencies before importing the
-# module under test.  Same pattern as test_kv_cache_estimation.py.
-# ---------------------------------------------------------------------------
+# Stub heavy / unavailable deps before importing the module under test.
+# Same pattern as test_kv_cache_estimation.py.
 
 _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
 if _BACKEND_DIR not in sys.path:
@@ -99,12 +88,11 @@ def _preflight(
     hf_repo = "unsloth/Example-GGUF",
     hf_token = None,
 ):
-    """Run the preflight arithmetic as written in llama_cpp.py and return
-    the decision outcome as a dict.
+    """Run the llama_cpp.py preflight arithmetic; return the decision as a dict.
 
     ``repo_files``: list of (filename, remote_bytes).
-    ``cached_files``: dict {filename: on_disk_bytes} for files already in cache.
-    ``free_bytes``: value returned by shutil.disk_usage(cache_dir).free.
+    ``cached_files``: {filename: on_disk_bytes} for files already cached.
+    ``free_bytes``: shutil.disk_usage(cache_dir).free.
     """
     import os
     import shutil
@@ -112,22 +100,20 @@ def _preflight(
     path_infos = [_FakePathInfo(name, size) for name, size in repo_files]
 
     with tempfile.TemporaryDirectory() as tmp:
-        # Create SPARSE files for the cached ones so os.path.exists /
-        # os.path.getsize pass without actually allocating bytes on disk.
-        # This is critical when simulating multi-GB models.
+        # Sparse files so exists/getsize pass without allocating bytes on disk
+        # (critical for multi-GB models).
         cache_paths = {}
         for name, sz in cached_files.items():
             p = Path(tmp) / name.replace("/", "_")
             with open(p, "wb") as fh:
                 if sz > 0:
-                    fh.truncate(sz)  # sparse allocation: no data blocks written
+                    fh.truncate(sz)  # sparse: no data blocks written
             cache_paths[name] = str(p)
 
         def fake_try_to_load_from_cache(repo_id, filename):
             return cache_paths.get(filename)
 
-        # Mirror the same variable names and control flow as the real code
-        # so behavioral drift is caught immediately.
+        # Mirror the real code's names and control flow so drift is caught.
         total_bytes = sum((p.size or 0) for p in path_infos)
         already_cached_bytes = 0
         for p in path_infos:
@@ -189,8 +175,8 @@ class TestCacheAwarePreflight:
         assert out["would_raise_disk_error"] is False
 
     def test_partial_cache_insufficient_disk_for_rest_still_raises(self):
-        """Two of four shards cached; remaining 70 GB still bigger than
-        free disk -> preflight correctly wants to raise."""
+        """Two of four shards cached; remaining 70 GB still exceeds free
+        disk -> preflight correctly wants to raise."""
         shards = [(f"UD-Q4_K_XL/shard-{i}.gguf", 35 * GIB) for i in range(4)]
         cached = {
             shards[0][0]: shards[0][1],
@@ -217,8 +203,8 @@ class TestCacheAwarePreflight:
         assert out["would_raise_disk_error"] is False
 
     def test_incomplete_cached_blob_is_not_credited(self):
-        """A partial file on disk (e.g. interrupted download) is not
-        counted as cached -- we still require bytes for it."""
+        """A partial file on disk (e.g. interrupted download) isn't counted
+        as cached -- we still require bytes for it."""
         shards = [("UD-Q4_K_XL/shard-0.gguf", 40 * GIB)]
         partial = {"UD-Q4_K_XL/shard-0.gguf": 10 * GIB}
         out = _preflight(
@@ -231,7 +217,7 @@ class TestCacheAwarePreflight:
         assert out["would_raise_disk_error"] is False
 
     def test_zero_size_path_infos_do_not_crash(self):
-        """A path_info with size=0 should not be credited or break the
+        """A path_info with size=0 must not be credited or break the
         arithmetic."""
         shards = [("mmproj.gguf", 0), ("UD-Q4_K_XL/shard-0.gguf", 40 * GIB)]
         out = _preflight(
diff --git a/studio/backend/tests/test_llama_cpp_context_fit.py b/studio/backend/tests/test_llama_cpp_context_fit.py
index ee8d54443a..58226f938c 100644
--- a/studio/backend/tests/test_llama_cpp_context_fit.py
+++ b/studio/backend/tests/test_llama_cpp_context_fit.py
@@ -5,26 +5,14 @@
 
 Guards two regressions in ``LlamaCppBackend.load_model``:
 
-1. **Auto mode on weights-exceed-VRAM** (``n_ctx == 0``): when the model
-   weights alone exceed 90% of every GPU subset's free memory, the
-   auto-pick loop used to exit without matching, leaving
-   ``effective_ctx`` at the model's native context (e.g. 196608 for
-   MiniMax-M2.7). The intended default per Studio's UI spec is 4096 so
-   the slider lands on a usable value; the user can still drag higher
-   and trigger ``--fit on`` with a warning.
+1. Auto mode (``n_ctx == 0``) when weights exceed every GPU subset's free
+   memory: auto-pick should fall back to 4096 (a usable slider value) rather
+   than leaving native ctx. User can still drag higher onto ``--fit on``.
+2. Explicit ctx must never be silently shrunk: when KV overflows fittable
+   weights, honor the explicit ctx with ``--fit on`` flexing ``-ngl``.
 
-2. **Explicit ctx silently shrunk when KV overflows**: with fittable
-   weights but a requested ctx whose KV cache pushes total memory over
-   90% of VRAM, the old code binary-searched a smaller ctx and emitted
-   ``-c  -ngl -1`` without informing the caller. The UI had
-   already surfaced its "might be slower" warning and expects the user's
-   explicit ctx to be honored with ``--fit on`` flexing ``-ngl`` instead.
-
-Tests avoid GPU probing, subprocess spawning, and GGUF I/O by driving the
-post-metadata decision block directly against a stubbed instance.
-
-Requires no GPU, network, or external libraries beyond pytest.
-Cross-platform: Linux, macOS, Windows, WSL.
+Drives the post-metadata decision block against a stubbed instance: no GPU,
+network, subprocess, or GGUF I/O. Cross-platform.
 """
 
 from __future__ import annotations
@@ -36,24 +24,20 @@ from pathlib import Path
 import pytest
 
 # ---------------------------------------------------------------------------
-# Stub heavy / unavailable external dependencies before importing the
-# module under test.  Same pattern as test_kv_cache_estimation.py.
+# 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
 _loggers_stub = _types.ModuleType("loggers")
 _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
 sys.modules.setdefault("loggers", _loggers_stub)
 
-# structlog
 _structlog_stub = _types.ModuleType("structlog")
 sys.modules.setdefault("structlog", _structlog_stub)
 
-# httpx
 _httpx_stub = _types.ModuleType("httpx")
 for _exc_name in (
     "ConnectError",
@@ -103,8 +87,7 @@ def _make_backend(
     kv_key_length = 128,
     kv_value_length = 128,
 ):
-    """Create a LlamaCppBackend instance with GGUF metadata fields set and
-    the helpers used by the decision block stubbed out."""
+    """LlamaCppBackend with GGUF metadata set and decision helpers stubbed."""
     inst = LlamaCppBackend.__new__(LlamaCppBackend)
     inst._context_length = native_ctx
     inst._n_layers = n_layers
@@ -136,8 +119,8 @@ def _drive(
 ):
     """Drive the post-metadata portion of load_model with stubbed inputs.
 
-    Mirrors the decision block at llama_cpp.py:1137-1296 so we can assert
-    the command that would be built, without subprocesses or GPU probes.
+    Mirrors llama_cpp.py:1137-1296 to assert the built command, without
+    subprocesses or GPU probes.
     """
     inst = _make_backend(native_ctx = native_ctx)
     model_size = int(model_gib * GIB)
@@ -154,9 +137,7 @@ def _drive(
     inst._can_estimate_kv = lambda: can_estimate_kv
 
     context_length = inst._context_length
-    # Use the production helper instead of reimplementing the conditional
-    # locally; reimplementing makes the test pass for the test's own logic
-    # rather than production's, and silent drift won't be caught.
+    # Use the production helper, not a reimplementation, to avoid testing our own logic.
     ctx_override = parse_ctx_override(extra_args)
     requested_ctx = resolve_requested_ctx(extra_args, n_ctx)
 
@@ -267,8 +248,8 @@ class TestAutoModeWeightsExceedVRAM:
         assert plan["c_arg"] == FALLBACK_CTX
         assert plan["use_fit"] is True
         assert plan["gpu_indices"] is None
-        # UI slider ceiling stays at native: user can still drag higher
-        # and get the "might be slower" path.
+        # UI slider ceiling stays at native: user can drag higher and get
+        # the "might be slower" path.
         assert plan["max_available_ctx"] == 196608
 
     def test_multi_gpu_all_subsets_fail(self):
@@ -304,9 +285,8 @@ class TestExplicitCtxRespectsUser:
     """``n_ctx > 0`` must never be silently shrunk."""
 
     def test_fittable_weights_oversized_kv(self):
-        # 8 GB weights + 131k ctx KV on 24 GB VRAM.
-        # Budget = 21.6 GB, KV at 131k >> 13.6 GB remaining, so
-        # _select_gpus flips use_fit=True.
+        # 8 GB weights + 131k ctx KV on 24 GB VRAM. Budget = 21.6 GB, KV
+        # at 131k >> 13.6 GB remaining, so _select_gpus flips use_fit=True.
         plan = _drive(
             n_ctx = 131072,
             model_gib = 8,
@@ -350,7 +330,7 @@ class TestExplicitCtxRespectsUser:
         assert plan["use_fit"] is True
 
     def test_explicit_below_floor_honored(self):
-        # 2048 is below --fit-ctx default; still honored since user set it.
+        # 2048 is below --fit-ctx default; honored since user set it.
         plan = _drive(
             n_ctx = 2048,
             model_gib = 8,
@@ -451,8 +431,8 @@ class TestTightFitPinsToGPU:
     """Models that fit at 91-95% of free VRAM must use the GPU."""
 
     def test_rtx_4090_qwen_24gb_class(self):
-        # noahterbest's #5106 log: 20.8 GB model on 22805 MiB free
-        # GPU, ctx=4096 -> ~94% utilization, ~1.4 GiB headroom.
+        # noahterbest's #5106 log: 20.8 GB model on 22805 MiB free GPU,
+        # ctx=4096 -> ~94% utilization, ~1.4 GiB headroom.
         plan = _drive(
             n_ctx = 0,
             model_gib = 20.8,
@@ -495,9 +475,8 @@ class TestTightFitPinsToGPU:
 
 @pytest.mark.parametrize("platform_tag", ["linux", "windows", "mac", "rocm"])
 def test_identical_decision_across_platforms(platform_tag):
-    """The decision function takes ``[(gpu_idx, free_mib), ...]`` regardless
-    of how upstream (nvidia-smi / nvidia-smi.exe / Metal / rocm-smi) produced
-    it. Identical inputs must yield identical plans."""
+    """Decision takes ``[(gpu_idx, free_mib), ...]`` regardless of source;
+    identical inputs must yield identical plans."""
     plan_a = _drive(n_ctx = 0, model_gib = 8, gpus = [(0, 24_000)])
     plan_b = _drive(n_ctx = 0, model_gib = 8, gpus = [(0, 24_000)])
     assert plan_a == plan_b, platform_tag
@@ -525,8 +504,8 @@ class TestClassifyGpuOffload:
         assert inst._classify_gpu_offload(True, [(0, 22805)]) is True
 
     def test_cpu_only_buffer_returns_false(self):
-        # llama-server printed buffer lines but only CPU buffers --
-        # this is the silent CPU fallback symptom we want to catch.
+        # Buffer lines printed but only CPU buffers -- the silent CPU
+        # fallback symptom we want to catch.
         inst = self._backend(
             [
                 "load_tensors:   CPU_Mapped model buffer size = 21000.0 MiB",
@@ -555,8 +534,7 @@ class TestClassifyGpuOffload:
         assert inst._classify_gpu_offload(False, []) is None
 
     def test_user_did_not_intend_gpu_returns_none(self):
-        # Studio called start_llama_server without expecting GPU use;
-        # don't warn.
+        # Studio called start_llama_server without expecting GPU; don't warn.
         inst = self._backend(
             [
                 "load_tensors:   CPU_Mapped model buffer size = 21000.0 MiB",
diff --git a/studio/backend/tests/test_llama_cpp_freshness.py b/studio/backend/tests/test_llama_cpp_freshness.py
index c7dd111ee3..cb17e0d5e7 100644
--- a/studio/backend/tests/test_llama_cpp_freshness.py
+++ b/studio/backend/tests/test_llama_cpp_freshness.py
@@ -3,8 +3,8 @@
 
 """Tests for the llama.cpp prebuilt freshness check.
 
-Pins the marker parser, the disk+memory cache, the stale decision
-matrix, and fail-open behaviour on missing data.
+Pins the marker parser, disk+memory cache, stale-decision matrix, and
+fail-open behaviour on missing data.
 """
 
 from __future__ import annotations
@@ -57,7 +57,7 @@ def _write_marker(install_dir: Path, **overrides) -> Path:
 
 
 def _fake_binary(install_dir: Path, *, layout: str = "cmake") -> Path:
-    """Stub llama-server under one of the supported install layouts."""
+    """Stub llama-server under a supported install layout."""
     if layout == "cmake":
         bin_dir = install_dir / "build" / "bin"
         bin_name = "llama-server"
@@ -77,7 +77,7 @@ def _fake_binary(install_dir: Path, *, layout: str = "cmake") -> Path:
 
 @pytest.fixture(autouse = True)
 def _reset(monkeypatch, tmp_path):
-    # Isolate disk cache per-test; never touch the user's real cache.
+    # Isolate disk cache per-test; never touch the real cache.
     monkeypatch.setattr(fr, "_cache_dir", lambda: tmp_path / ".freshness")
     fr.reset_caches()
     yield
@@ -107,8 +107,8 @@ def test_read_install_marker_finds_root_layout(tmp_path):
 
 
 def test_read_install_marker_finds_windows_cmake_layout(tmp_path):
-    # Windows cmake puts the .exe under build/bin/Release/, so the
-    # marker is four levels above the binary.
+    # Windows cmake puts the .exe under build/bin/Release/, so the marker
+    # is four levels above the binary.
     install_dir = tmp_path / "llama.cpp"
     _write_marker(install_dir, tag = "b8888")
     bin_path = _fake_binary(install_dir, layout = "windows")
@@ -119,10 +119,9 @@ def test_read_install_marker_finds_windows_cmake_layout(tmp_path):
 
 @pytest.mark.parametrize("repo", ["unslothai/llama.cpp", "ggml-org/llama.cpp"])
 def test_read_install_marker_carries_published_repo_dynamically(tmp_path, repo):
-    # The freshness check queries whichever release repo the marker
-    # records, so CUDA Linux (unslothai), CPU Linux x86_64 / macOS
-    # (ggml-org), and ROCm source-build (unslothai upstream label)
-    # all surface the right "latest" tag.
+    # The freshness check queries whichever release repo the marker records,
+    # so CUDA (unslothai), CPU/macOS (ggml-org), and ROCm all get the right
+    # "latest" tag.
     install_dir = tmp_path / "llama.cpp"
     _write_marker(install_dir, tag = "b9000", published_repo = repo)
     bin_path = _fake_binary(install_dir, layout = "cmake")
diff --git a/studio/backend/tests/test_llama_cpp_load_progress.py b/studio/backend/tests/test_llama_cpp_load_progress.py
index f95d8bf1a4..cc1c6256a8 100644
--- a/studio/backend/tests/test_llama_cpp_load_progress.py
+++ b/studio/backend/tests/test_llama_cpp_load_progress.py
@@ -3,34 +3,20 @@
 
 """Tests for ``LlamaCppBackend.load_progress()``.
 
-The chat settings flow and the training overlay both show a generic
-"Starting model..." spinner during the window after a GGUF download
-finishes and before llama-server reports healthy. For small models
-that window is a second or two and nobody notices. For large MoE GGUFs
-(MiniMax-M2.7, Qwen3.5-397B-A17B, etc.) the llama-server process spends
-minutes in kernel state D, paging tens or hundreds of GB of shards
-into the page cache. The UI has no way to show a real progress bar,
-rate, or ETA during that window.
+For large MoE GGUFs, llama-server spends minutes paging shards into the page
+cache after download. ``load_progress()`` samples ``/proc//status VmRSS``
+against the total shard size on disk so the UI can render a real bar plus
+rate/ETA. Contract pinned here:
 
-``load_progress()`` samples ``/proc//status VmRSS`` (what the
-kernel has actually paged in) against the total shard file size on
-disk, so the frontend can render a real bar plus rate/ETA. This
-module pins that contract:
-
-  * returns ``None`` when no load is in flight
-  * returns ``{"phase": "mmap", ...}`` while the subprocess is alive
-    but ``_healthy`` is False
-  * returns ``{"phase": "ready", ...}`` once ``_healthy`` flips
-  * ``bytes_total`` is derived from the resolved on-disk path
-    (which the paired fix assigns to ``self._gguf_path`` on both the
-    local-GGUF and HF-download code paths)
+  * ``None`` when no load is in flight
+  * ``{"phase": "mmap", ...}`` while the subprocess is alive but ``_healthy`` is False
+  * ``{"phase": "ready", ...}`` once ``_healthy`` flips
+  * ``bytes_total`` derived from the resolved on-disk path (``self._gguf_path``)
   * ``bytes_loaded`` is VmRSS in bytes, capped by total, rounded
-  * ``fraction`` is clamped to 0..1 and rounded to 4 decimal places
+  * ``fraction`` clamped to 0..1, rounded to 4 dp
 
-Linux-only via ``/proc``. On platforms without ``/proc`` the method
-returns ``None`` instead of raising.
-Cross-platform test: skips cleanly on macOS / Windows if ``/proc`` is
-not available.
+Linux-only via ``/proc``; returns ``None`` (not raises) without it, so tests
+skip cleanly on macOS / Windows.
 """
 
 from __future__ import annotations
@@ -44,10 +30,7 @@ from unittest.mock import patch
 
 import pytest
 
-# ---------------------------------------------------------------------------
-# Stub heavy / unavailable external dependencies before importing the
-# module under test. Same pattern as test_kv_cache_estimation.py.
-# ---------------------------------------------------------------------------
+# 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:
@@ -106,7 +89,7 @@ def _make_instance():
 
 
 class _FakeProc:
-    """Minimal stand-in for subprocess.Popen that just carries a pid."""
+    """Minimal stand-in for subprocess.Popen carrying just a pid."""
 
     def __init__(self, pid: int):
         self.pid = pid
@@ -188,7 +171,7 @@ class TestLoadProgressSingleShard:
 
 
 class TestLoadProgressMultiShard:
-    """Shard-aware total: for ``*-00001-of-00004.gguf`` primaries the
+    """Shard-aware total: for ``*-00001-of-00004.gguf`` primaries, the
     method sums sibling files with the same prefix."""
 
     def test_sharded_total_aggregates_siblings(self, tmp_path):
@@ -197,7 +180,7 @@ class TestLoadProgressMultiShard:
                 tmp_path / f"model-{i:05d}-of-00004.gguf",
                 size_bytes = 20 * 1024**3,
             )
-        # Drop an unrelated .gguf in the same folder -- must not be counted.
+        # An unrelated .gguf in the same folder -- must not be counted.
         _write_sparse_file(tmp_path / "mmproj-BF16.gguf", 2 * 1024**3)
 
         inst = _make_instance()
diff --git a/studio/backend/tests/test_llama_cpp_load_progress_live.py b/studio/backend/tests/test_llama_cpp_load_progress_live.py
index 44a8f00834..98e19944dd 100644
--- a/studio/backend/tests/test_llama_cpp_load_progress_live.py
+++ b/studio/backend/tests/test_llama_cpp_load_progress_live.py
@@ -3,20 +3,10 @@
 
 """Live, no-mock integration test for ``LlamaCppBackend.load_progress()``.
 
-The companion files (``test_llama_cpp_load_progress.py`` and
-``test_llama_cpp_load_progress_matrix.py``) patch ``builtins.open`` to
-feed synthetic VmRSS values. This file is the opposite: it uses **real**
-subprocesses, **real** file sizes, and the **real** ``/proc``
-interface. It is the sanity check that the contract we keep in the
-mocked tests still maps to what the kernel actually returns on a live
-Linux system.
-
-Why both: the mocked tests can be fooled by a buggy implementation that
-parses ``/proc`` output in a format the kernel no longer uses, or that
-makes assumptions about ``Path.stat()`` vs ``os.path.getsize``. This
-file hits the real APIs so any format drift gets caught.
-
-Skipped cleanly on non-Linux (no ``/proc``).
+The companion mocked tests patch ``builtins.open`` for synthetic VmRSS values;
+this one uses real subprocesses, file sizes, and ``/proc`` so format drift the
+mocks can't see (kernel ``/proc`` layout, stat vs getsize) gets caught. Skipped
+on non-Linux (no ``/proc``).
 """
 
 from __future__ import annotations
@@ -30,10 +20,7 @@ from pathlib import Path
 
 import pytest
 
-# ---------------------------------------------------------------------------
-# Same stubs as the matrix file (keep self-contained so the file can be
-# run standalone as well as via the full suite).
-# ---------------------------------------------------------------------------
+# Same stubs as the matrix file (self-contained for standalone + full-suite runs).
 
 _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
 if _BACKEND_DIR not in sys.path:
@@ -88,8 +75,8 @@ def _make_backend(
 
 
 def test_live_rss_matches_kernel_vmrss(tmp_path):
-    """Spawn a real child, let it allocate real bytes, confirm
-    ``bytes_loaded`` tracks the kernel's VmRSS within a sane tolerance."""
+    """Spawn a real child, let it allocate real bytes, confirm ``bytes_loaded``
+    tracks the kernel's VmRSS within a sane tolerance."""
     # Child that allocates ~100 MB of zero'd bytes and then idles.
     script = tmp_path / "burn.py"
     script.write_text(
@@ -112,7 +99,7 @@ def test_live_rss_matches_kernel_vmrss(tmp_path):
         ready = proc.stdout.readline()
         assert ready.strip() == b"ready"
 
-        # Create a fake 200 MB sparse gguf so bytes_total is concrete.
+        # Fake 200 MB sparse gguf so bytes_total is concrete.
         gguf = tmp_path / "model.gguf"
         with open(gguf, "wb") as f:
             f.truncate(200 * 1024 * 1024)
@@ -123,8 +110,8 @@ def test_live_rss_matches_kernel_vmrss(tmp_path):
         assert out is not None, "load_progress returned None for live pid"
         assert out["phase"] == "mmap"
         assert out["bytes_total"] == 200 * 1024 * 1024
-        # VmRSS for the Python child includes the interpreter + the 100MB
-        # buffer, so a realistic floor is 50 MB and ceiling is 200 MB.
+        # VmRSS for the Python child includes the interpreter + 100MB buffer,
+        # so a realistic floor is 50 MB and ceiling is 200 MB.
         assert (
             out["bytes_loaded"] >= 50 * 1024 * 1024
         ), f"bytes_loaded unexpectedly low: {out['bytes_loaded']}"
@@ -153,8 +140,8 @@ def test_live_ready_phase_when_healthy(tmp_path):
 
 
 def test_live_dead_pid_returns_none(tmp_path):
-    """A recently-dead pid may linger in /proc for ms; use a clearly
-    invalid id so the read reliably fails."""
+    """A recently-dead pid may linger in /proc for ms; use a clearly invalid id
+    so the read reliably fails."""
     gguf = tmp_path / "m.gguf"
     gguf.touch()
 
@@ -164,8 +151,8 @@ def test_live_dead_pid_returns_none(tmp_path):
 
 
 def test_live_shard_aggregation_counts_real_files(tmp_path):
-    """With 4 real sibling shards on disk, ``bytes_total`` equals their
-    summed size to the byte."""
+    """With 4 real sibling shards on disk, ``bytes_total`` equals their summed
+    size to the byte."""
     shard_size = 7 * 1024 * 1024  # 7 MB each
     for i in range(1, 5):
         f = tmp_path / f"model-{i:05d}-of-00004.gguf"
@@ -186,8 +173,8 @@ def test_live_shard_aggregation_counts_real_files(tmp_path):
 
 
 def test_live_repeated_polling_stays_sane(tmp_path):
-    """Sampling the same backend 20 times should not raise or produce
-    non-numeric output, even under normal kernel RSS jitter."""
+    """Sampling the same backend 20 times must not raise or produce non-numeric
+    output, even under normal kernel RSS jitter."""
     gguf = tmp_path / "m.gguf"
     with open(gguf, "wb") as f:
         f.truncate(500 * 1024 * 1024)
diff --git a/studio/backend/tests/test_llama_cpp_load_progress_matrix.py b/studio/backend/tests/test_llama_cpp_load_progress_matrix.py
index a88450ec0b..5c4d9106d8 100644
--- a/studio/backend/tests/test_llama_cpp_load_progress_matrix.py
+++ b/studio/backend/tests/test_llama_cpp_load_progress_matrix.py
@@ -3,29 +3,12 @@
 
 """Extended test matrix for ``LlamaCppBackend.load_progress()``.
 
-Companion to ``test_llama_cpp_load_progress.py`` (which pins the basic
-contract). This file widens coverage to the edge cases that bit users
-or were hypothesized to bite them on cross-platform installs:
+Companion to ``test_llama_cpp_load_progress.py`` (basic contract). Covers
+cross-platform edge cases: platform matrix (/proc absence), VmRSS parsing,
+filesystem edges (HF-cache symlinks, broken/missing/relative paths), shard
+aggregation, lifecycle races, concurrent sampling, and fraction bounds.
 
-  * Platform matrix — macOS/Windows simulation via ``/proc`` absence.
-  * ``VmRSS`` parsing — tab vs space delimiter, missing line, malformed
-    integer.
-  * Filesystem edges — HF-cache symlinks, broken symlinks, nonexistent
-    paths, relative paths.
-  * Shard aggregation — partial multi-shard downloads where some shards
-    are still ``.incomplete``, two shard series in the same dir,
-    ``mmproj-*.gguf`` sibling exclusion for non-sharded primaries,
-    single-file models.
-  * Lifecycle races — process set before ``_gguf_path`` is assigned,
-    process dead mid-sample, ``_healthy`` flipped to True.
-  * Concurrent sampling — 10 threads × 50 iterations against a single
-    backend, hitting real ``/proc`` (no mocks — see the note in
-    ``TestConcurrentSampling`` for why).
-  * Fraction bounds — capped at 1.0 when RSS exceeds total; 0.0 when
-    total is zero.
-
-All tests are Linux-only in practice (we stub ``/proc`` where needed).
-The stable subset runs in well under a second.
+Linux-only in practice (``/proc`` stubbed where needed).
 """
 
 from __future__ import annotations
@@ -40,10 +23,7 @@ from unittest.mock import patch
 
 import pytest
 
-# ---------------------------------------------------------------------------
-# Stub heavy / unavailable external dependencies before importing the
-# module under test. Same pattern as test_llama_cpp_load_progress.py.
-# ---------------------------------------------------------------------------
+# 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:
@@ -113,7 +93,7 @@ def _sparse(path, size):
 
 
 def _fake_proc_reader(rss_kb):
-    """Return an ``open()`` replacement that fakes /proc reads with a VmRSS line."""
+    """An ``open()`` replacement faking /proc reads with a VmRSS line."""
 
     def fake_open(path, *args, **kwargs):
         if str(path).startswith("/proc/"):
@@ -129,8 +109,8 @@ def _fake_proc_reader(rss_kb):
 
 
 class TestPlatformMatrix:
-    """The method is Linux-first via /proc. On macOS/Windows it must
-    degrade to None rather than crash."""
+    """Linux-first via /proc. On macOS/Windows must degrade to None
+    rather than crash."""
 
     def test_linux_live_proc_is_self_pid(self, tmp_path):
         """Self-pid /proc read uses the real kernel interface."""
@@ -144,7 +124,7 @@ class TestPlatformMatrix:
         assert out is not None
         assert out["phase"] == "mmap"
         assert out["bytes_total"] == 1 * 1024**3
-        # Our Python process has some RSS -- just sanity-check positive.
+        # Our process has some RSS -- sanity-check it's positive.
         assert out["bytes_loaded"] > 0
 
     def test_macos_no_proc_returns_none(self, tmp_path):
@@ -199,7 +179,7 @@ class TestVmRSSParsing:
         assert out["bytes_loaded"] == 2 * 1024**3
 
     def test_space_separated_fallback(self, tmp_path):
-        """Some kernels emit single-space rather than tab."""
+        """Some kernels emit a single space, not a tab."""
         gguf = tmp_path / "m.gguf"
         _sparse(gguf, 4 * 1024**3)
         inst = _make()
@@ -235,8 +215,8 @@ class TestVmRSSParsing:
         assert out["fraction"] == 0.0
 
     def test_malformed_vmrss_value(self, tmp_path):
-        """Non-integer VmRSS value should be treated as if the line were
-        absent (early ValueError caught)."""
+        """Non-integer VmRSS is treated like an absent line (ValueError
+        caught)."""
         gguf = tmp_path / "m.gguf"
         _sparse(gguf, 1 * 1024**3)
         inst = _make()
@@ -250,7 +230,7 @@ class TestVmRSSParsing:
 
         with patch("builtins.open", side_effect = fake_open):
             out = inst.load_progress()
-        # The implementation catches ValueError on int() and returns None.
+        # int() ValueError is caught and returns None.
         assert out is None
 
 
@@ -262,7 +242,7 @@ class TestVmRSSParsing:
 class TestFilesystemEdges:
     def test_symlink_primary_follows_to_blob(self, tmp_path):
         """HF cache stores blobs under blobs/ and symlinks them from
-        snapshots/. The method must follow the symlink."""
+        snapshots/. Must follow the symlink."""
         blob = tmp_path / "blob"
         _sparse(blob, 12 * 1024**3)
         snap = tmp_path / "snap"
@@ -300,7 +280,7 @@ class TestFilesystemEdges:
 
     def test_relative_gguf_path(self, tmp_path):
         """Relative paths shouldn't crash; behaviour depends on CWD but
-        the method must not raise."""
+        must not raise."""
         cwd = os.getcwd()
         try:
             os.chdir(tmp_path)
@@ -323,11 +303,11 @@ class TestFilesystemEdges:
 
 class TestShardAggregation:
     def test_partial_multi_shard_download(self, tmp_path):
-        """Primary present but shards 2..N still downloading as
-        ``.incomplete``. Sums only the fully-arrived ``.gguf`` files."""
+        """Primary present but shards 2..N still ``.incomplete``. Sums
+        only the fully-arrived ``.gguf`` files."""
         _sparse(tmp_path / "m-00001-of-00004.gguf", 30 * 1024**3)
         _sparse(tmp_path / "m-00002-of-00004.gguf", 30 * 1024**3)
-        # 3 and 4 still downloading as .incomplete
+        # 3 and 4 still downloading as .incomplete.
         _sparse(tmp_path / "m-00003-of-00004.gguf.incomplete", 5 * 1024**3)
         inst = _make()
         inst._process = _Proc(os.getpid())
@@ -337,8 +317,8 @@ class TestShardAggregation:
         assert out["bytes_total"] == 60 * 1024**3  # only the .gguf siblings
 
     def test_two_shard_series_in_same_dir(self, tmp_path):
-        """Defensive: if two quant series share a dir, prefix filter
-        only sums siblings of the chosen primary."""
+        """Defensive: when two quant series share a dir, the prefix
+        filter sums only siblings of the chosen primary."""
         for i in range(1, 3):
             _sparse(tmp_path / f"m_q4-{i:05d}-of-00002.gguf", 10 * 1024**3)
             _sparse(tmp_path / f"m_q8-{i:05d}-of-00002.gguf", 20 * 1024**3)
@@ -351,7 +331,7 @@ class TestShardAggregation:
 
     def test_mmproj_sibling_not_counted(self, tmp_path):
         """Vision models drop an ``mmproj-*.gguf`` alongside. For a
-        single-file (non-sharded) primary we only count the primary."""
+        single-file (non-sharded) primary, count only the primary."""
         _sparse(tmp_path / "m.gguf", 8 * 1024**3)
         _sparse(tmp_path / "mmproj-BF16.gguf", 2 * 1024**3)
         inst = _make()
@@ -359,7 +339,7 @@ class TestShardAggregation:
         inst._gguf_path = str(tmp_path / "m.gguf")
         with patch("builtins.open", side_effect = _fake_proc_reader(0)):
             out = inst.load_progress()
-        # Non-sharded primary: only the primary is counted.
+        # Non-sharded: only the primary is counted.
         assert out["bytes_total"] == 8 * 1024**3
 
     def test_single_file_model(self, tmp_path):
@@ -381,7 +361,7 @@ class TestShardAggregation:
 
 class TestLifecycleRaces:
     def test_process_set_but_gguf_path_not_yet(self, tmp_path):
-        """Moment between Popen and self._gguf_path=model_path."""
+        """Window between Popen and self._gguf_path=model_path."""
         inst = _make()
         inst._process = _Proc(os.getpid())
         inst._gguf_path = None
@@ -418,14 +398,10 @@ class TestLifecycleRaces:
 
 class TestConcurrentSampling:
     def test_parallel_invocations_never_raise(self, tmp_path):
-        """Many concurrent samplers hitting the same backend must not raise.
+        """Many concurrent samplers on one backend must not raise.
 
-        We intentionally do NOT patch ``builtins.open`` here because
-        ``unittest.mock.patch`` is not thread-safe: interleaved
-        enter/exit across threads can leak a Mock into ``builtins.open``
-        and poison every subsequent test in the session. Instead, we
-        let each thread hit the real ``/proc/self/status`` of the test
-        process, which is exactly the code path that matters in prod.
+        No ``builtins.open`` patch: ``mock.patch`` isn't thread-safe and could
+        leak a Mock into ``open``. Each thread hits the real ``/proc/self/status``.
         """
         _sparse(tmp_path / "m.gguf", 1 * 1024**3)
         inst = _make()
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 aa0198892f..310aaf6c0f 100644
--- a/studio/backend/tests/test_llama_cpp_max_context_threshold.py
+++ b/studio/backend/tests/test_llama_cpp_max_context_threshold.py
@@ -3,22 +3,20 @@
 
 """Tests for the ``max_context_length`` warning-threshold semantics.
 
-``/api/inference/status.max_context_length`` is what the ctx slider in
-the chat settings sheet reads to decide when to render the "Exceeds
-estimated VRAM capacity. The model may use system RAM." warning:
+The ctx slider in the chat settings sheet reads
+``/api/inference/status.max_context_length`` to decide when to render the
+"Exceeds estimated VRAM capacity. The model may use system RAM." warning:
 
     ctxDisplayValue > ggufMaxContextLength → show warning
 
-For models whose weights fit on some GPU subset, the warning threshold
-is the largest ctx that fits fully in VRAM (the binary-search cap from
-``_fit_context_to_vram``). For models whose weights exceed 90% of every
-GPU subset's free memory, the warning must fire as soon as the user
-drags above the 4096 spec default (otherwise a user loading e.g.
-MiniMax-M2.7 on a 97 GB GPU sees a slider up to 196608 with no
-indication that any value above 4096 will trigger ``--fit on`` and
-degrade performance).
+When weights fit on some GPU subset, the threshold is the largest ctx that
+fits fully in VRAM (the binary-search cap from ``_fit_context_to_vram``).
+When weights exceed 90% of every GPU subset's free memory, the warning must
+fire as soon as the user drags above the 4096 spec default (otherwise loading
+e.g. MiniMax-M2.7 on a 97 GB GPU shows a slider up to 196608 with no hint that
+any value above 4096 triggers ``--fit on`` and degrades performance).
 
-These tests pin both cases. No GPU probing, no subprocess, no GGUF I/O.
+These tests pin both cases. No GPU probing, subprocess, or GGUF I/O.
 Cross-platform: Linux, macOS, Windows, WSL.
 """
 
@@ -30,10 +28,8 @@ from pathlib import Path
 
 import pytest
 
-# ---------------------------------------------------------------------------
-# Stub heavy / unavailable external dependencies before importing the
-# module under test.  Same pattern as test_kv_cache_estimation.py.
-# ---------------------------------------------------------------------------
+# Stub heavy / unavailable deps before importing the module under test.
+# Same pattern as test_kv_cache_estimation.py.
 
 _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
 if _BACKEND_DIR not in sys.path:
@@ -81,9 +77,7 @@ sys.modules.setdefault("httpx", _httpx_stub)
 from core.inference.llama_cpp import LlamaCppBackend
 
 
-# ---------------------------------------------------------------------------
 # Helpers
-# ---------------------------------------------------------------------------
 
 GIB = 1024**3
 
@@ -115,9 +109,8 @@ def _compute_max_available_ctx(
     gpus,
     kv_per_token_bytes = 325_000,
 ):
-    """Run the ceiling-probe block from load_model and return the final
-    ``max_available_ctx`` value the backend would assign to
-    ``_max_context_length``.
+    """Run load_model's ceiling-probe block and return the final
+    ``max_available_ctx`` the backend would assign to ``_max_context_length``.
     """
     inst = _make_backend(native_ctx = native_ctx)
     model_size = int(model_gib * GIB)
@@ -157,14 +150,12 @@ def _compute_max_available_ctx(
     return max_available_ctx
 
 
-# ---------------------------------------------------------------------------
 # Weights exceed every GPU subset's VRAM  (MiniMax-M2.7-like)
-# ---------------------------------------------------------------------------
 
 
 class TestMaxContextLengthForWeightsExceedVRAM:
-    """The UI ``max_context_length`` threshold must fall back to 4096 so
-    the warning fires as soon as the user drags above the spec default.
+    """UI ``max_context_length`` must fall back to 4096 so the warning fires
+    as soon as the user drags above the spec default.
     """
 
     def test_minimax_like(self):
@@ -186,8 +177,8 @@ class TestMaxContextLengthForWeightsExceedVRAM:
         assert got == 4096
 
     def test_native_below_fallback_is_preserved(self):
-        """If the model's native ctx is itself smaller than 4096, do not
-        advertise a larger value than the model supports."""
+        """If native ctx is itself below 4096, don't advertise a larger value
+        than the model supports."""
         got = _compute_max_available_ctx(
             native_ctx = 2048,
             model_gib = 200,
@@ -196,9 +187,7 @@ class TestMaxContextLengthForWeightsExceedVRAM:
         assert got == 2048
 
 
-# ---------------------------------------------------------------------------
 # Fittable models (regression guard)
-# ---------------------------------------------------------------------------
 
 
 class TestMaxContextLengthForFittableModels:
@@ -236,9 +225,7 @@ class TestMaxContextLengthForFittableModels:
         assert got >= 131072 - 256  # rounded to 256 boundary
 
 
-# ---------------------------------------------------------------------------
 # Property plumbing
-# ---------------------------------------------------------------------------
 
 
 class TestMaxContextLengthProperty:
diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py
index 9e7944913a..784ab1b259 100644
--- a/studio/backend/tests/test_llama_cpp_mtp_detection.py
+++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py
@@ -187,9 +187,9 @@ def _mtp_backend(**overrides):
     backend._requested_n_ctx = 8192
     backend._cache_type_kv = None
     backend._speculative_type = "draft-mtp"
-    # Default fixture simulates Auto having auto-promoted to draft-mtp.
-    # Individual tests override _requested_spec_mode when they want a
-    # forced mode or the user---spec-type-extra-args path.
+    # Fixture simulates Auto having auto-promoted to draft-mtp. Tests
+    # override _requested_spec_mode for a forced mode or the
+    # user---spec-type-extra-args path.
     backend._requested_spec_mode = "auto"
     backend._chat_template_override = None
     backend._is_vision = False
@@ -239,11 +239,10 @@ def test_already_in_target_state_matches_when_request_uses_default_for_mtp_model
 
 
 def test_already_in_target_state_auto_request_matches_auto_backend_for_non_mtp_model():
-    # Under the requested-mode round-trip model, Auto requested against an
-    # Auto-recorded backend matches regardless of model name. The underlying
-    # resolved emission (--spec-default vs draft-mtp) is handled by the
-    # backend's own load path and reflected in _speculative_type; the
-    # short-circuit comparison only cares whether the *intent* changed.
+    # In the requested-mode round-trip model, Auto-vs-Auto matches regardless
+    # of model name. The resolved emission (--spec-default vs draft-mtp) is
+    # handled by the load path and reflected in _speculative_type; the
+    # short-circuit only cares whether the *intent* changed.
     backend = _mtp_backend(
         _model_identifier = "unsloth/Qwen3.6-27B-GGUF",
         _speculative_type = "default",
@@ -403,11 +402,9 @@ def test_already_in_target_state_vision_mtp_default_matches():
 
 
 def test_already_in_target_state_vision_off_matches_vision_backend():
-    # Vision loads silently drop speculative decoding at the route level
-    # (_request_matches_loaded_settings overrides req to "off"). At the
-    # llama_cpp.py level, _already_in_target_state compares canonical
-    # requested modes; a vision backend recorded with _requested_spec_mode
-    # = "off" matches a req of "off" or None+vision.
+    # Vision loads drop speculative decoding at the route level (req -> "off").
+    # _already_in_target_state compares canonical requested modes; a vision
+    # backend with _requested_spec_mode="off" matches req "off" or None+vision.
     backend = _mtp_backend(
         _model_identifier = "unsloth/Qwen3-VL-4B-Instruct-GGUF",
         _is_vision = True,
@@ -616,9 +613,8 @@ def test_probe_detects_legacy_ngram_mod_flavor(tmp_path):
 
 @_NEEDS_BASH
 def test_probe_ignores_removal_stub_descriptions(tmp_path):
-    # Post-rename binary: legacy flags are present but with
-    # "argument has been removed" descriptions; must not be detected
-    # as legacy.
+    # Post-rename binary: legacy flags present but with "argument has been
+    # removed" descriptions; must not be detected as legacy.
     fake = _make_fake_llama_server(tmp_path / "llama-server", _POST_RENAME_HELP)
     _clear_caps_cache()
     caps = LlamaCppBackend.probe_server_capabilities(str(fake))
@@ -784,16 +780,15 @@ def test_already_in_target_state_draft_n_max_ignored_when_not_mtp():
     )
 
 
-# Sub-3B MTP gate -- tiny dense models regress with the MTP draft
-# head, so load_model falls back to ngram-mod (when the binary supports
-# it) instead of draft-mtp. The reload-skip mirror must follow the
-# same fallback so a sub-3B reload-with-default does not bounce a
-# correctly-configured ngram-mod / off backend.
+# Sub-3B MTP gate -- tiny dense models regress with the MTP draft head, so
+# load_model falls back to ngram-mod (when the binary supports it) instead of
+# draft-mtp. The reload-skip mirror must follow the same fallback so a sub-3B
+# reload-with-default doesn't bounce a correctly-configured ngram-mod/off backend.
 
 
 def _patch_probe(monkeypatch, ngram_supported):
-    """Force probe_server_capabilities to a deterministic result so
-    tests don't depend on whatever llama-server happens to be on PATH."""
+    """Force probe_server_capabilities to a deterministic result so tests
+    don't depend on whatever llama-server is on PATH."""
     fake = {
         "found": True,
         "mtp_token": "draft-mtp",
@@ -815,8 +810,8 @@ def _patch_probe(monkeypatch, ngram_supported):
 
 
 def test_already_in_target_state_sub_3b_falls_back_to_ngram_mod_when_supported(monkeypatch):
-    # 0.8B MTP request -- load_model would have promoted to ngram-mod
-    # (no MTP head); reload check must match a ngram-mod backend.
+    # 0.8B MTP request -- load_model would have promoted to ngram-mod (no MTP
+    # head); reload check must match a ngram-mod backend.
     _patch_probe(monkeypatch, ngram_supported = True)
     backend = _mtp_backend(
         _model_identifier = "unsloth/Qwen3.5-0.8B-MTP-GGUF",
@@ -888,8 +883,8 @@ def test_already_in_target_state_4b_mtp_request_promotes_as_before(monkeypatch):
 
 
 def test_already_in_target_state_2b_falls_back_to_ngram_below_threshold(monkeypatch):
-    # 2.0B is below the 3B threshold -> ngram-mod fallback, not
-    # draft-mtp. Clean-bench shows 2B regresses with draft-mtp.
+    # 2.0B is below the 3B threshold -> ngram-mod fallback, not draft-mtp.
+    # Clean-bench shows 2B regresses with draft-mtp.
     _patch_probe(monkeypatch, ngram_supported = True)
     backend = _mtp_backend(
         _model_identifier = "unsloth/Qwen3.5-2B-MTP-GGUF",
@@ -1023,9 +1018,9 @@ def _resolver_backend(
 
 
 def _flags_dict(flags):
-    """Parse the spec-flag list into a small {flag: value} dict; collapses
-    repeated flags by keeping the last (only --spec-type can repeat and
-    never does in our resolver)."""
+    """Parse the spec-flag list into a {flag: value} dict; collapses repeated
+    flags by keeping the last (only --spec-type can repeat, and never does
+    in our resolver)."""
     out = {}
     i = 0
     while i < len(flags):
@@ -1125,8 +1120,8 @@ def test_build_speculative_flags_user_extra_args_owns_spec_type(monkeypatch):
         gpus = True,
         binary = "/fake/llama-server",
     )
-    # No flags emitted by the resolver -- the user's extra_args carries
-    # the --spec-type, and the resolver records requested_spec_mode = None.
+    # Resolver emits nothing -- the user's extra_args carries the --spec-type,
+    # and the resolver records requested_spec_mode = None.
     assert flags == []
     assert backend.requested_spec_mode is None
     assert backend.speculative_type is None
@@ -1179,7 +1174,7 @@ 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), but
-    # _requested_spec_mode still reflects the user's choice.
+    # _speculative_type stays None (resolved emission was none); the user's
+    # choice is still reflected in _requested_spec_mode.
     assert backend.requested_spec_mode == "mtp"
     assert backend.speculative_type is None
diff --git a/studio/backend/tests/test_llama_cpp_no_context_shift.py b/studio/backend/tests/test_llama_cpp_no_context_shift.py
index b9f25faf88..10b1dc7ff6 100644
--- a/studio/backend/tests/test_llama_cpp_no_context_shift.py
+++ b/studio/backend/tests/test_llama_cpp_no_context_shift.py
@@ -3,17 +3,15 @@
 
 """``--no-context-shift`` launch-flag contract.
 
-When llama-server runs with its default context-shift behavior, the UI
-has no way to tell the user that the KV cache has been rotated --
-earlier turns silently vanish from the conversation. The Studio
-backend always passes ``--no-context-shift`` so the server returns a
+With llama-server's default context-shift behavior, the UI cannot tell the user
+the KV cache was rotated -- earlier turns silently vanish from the conversation.
+The Studio backend always passes ``--no-context-shift`` so the server returns a
 clean error instead, and the chat adapter can point the user at the
 ``Context Length`` input in the settings panel.
 
-This file is a static read of the launch command: we ask
-``LlamaCppBackend`` to assemble its ``cmd`` list and assert the flag
-is always present. Testing via the real subprocess would require an
-actual GGUF on disk, which is out of scope for the fast test suite.
+This file statically reads the launch command: we ask ``LlamaCppBackend`` to
+assemble its ``cmd`` list and assert the flag is present. Testing via the real
+subprocess would need an actual GGUF on disk, out of scope for the fast suite.
 """
 
 from __future__ import annotations
@@ -68,11 +66,10 @@ from core.inference import llama_cpp as llama_cpp_module
 def _load_model_source() -> str:
     """Return the source of ``LlamaCppBackend.load_model``.
 
-    Using ``inspect.getsource`` instead of reading the file directly
-    scopes the assertions to the function that actually launches
-    llama-server, so neither the presence check nor the location check
-    can be fooled by a stray occurrence of ``"--no-context-shift"``
-    elsewhere in the module.
+    Using ``inspect.getsource`` instead of reading the file scopes the assertions
+    to the function that launches llama-server, so neither the presence nor the
+    location check can be fooled by a stray ``"--no-context-shift"`` elsewhere in
+    the module.
     """
     return inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model)
 
@@ -80,10 +77,9 @@ def _load_model_source() -> str:
 def test_no_context_shift_is_in_load_model():
     """The flag is part of the static launch-command template.
 
-    We check the source of ``load_model`` rather than mocking the whole
-    call chain (GPU probing, GGUF stat, etc.): the flag is written as
-    a literal in one place and any regression has to delete it, which
-    a text search will catch.
+    We check the source of ``load_model`` rather than mocking the whole call
+    chain (GPU probing, GGUF stat, etc.): the flag is a literal in one place and
+    any regression must delete it, which a text search catches.
     """
     assert '"--no-context-shift"' in _load_model_source(), (
         "llama-server must be launched with --no-context-shift so the "
@@ -93,21 +89,19 @@ def test_no_context_shift_is_in_load_model():
 
 
 def test_flag_sits_inside_the_base_cmd_list():
-    """Pin the flag's location so a future refactor can't accidentally
-    move it into a branch that only fires on some code paths.
+    """Pin the flag's location so a refactor can't move it into a branch that
+    only fires on some code paths.
 
-    We slice from ``cmd = [`` to the first ``]`` at the same indent.
-    Using ``inspect.getsource`` means the function lives in its own
-    string and there are no siblings to worry about, so a plain
-    bracket search would also work -- anchoring on the trailing indent
-    just keeps the slice from wandering into a later expression if the
-    opening literal ever grows an in-line comment trailing it.
+    We slice from ``cmd = [`` to the first ``]`` at the same indent. Since
+    ``inspect.getsource`` gives the function its own string with no siblings, a
+    plain bracket search would also work -- anchoring on the trailing indent just
+    keeps the slice from wandering into a later expression if the opening literal
+    ever grows a trailing in-line comment.
     """
     source = _load_model_source()
     start = source.find("cmd = [")
     assert start >= 0, "could not find the base cmd = [...] block"
     # Find the first line containing only ``]`` (possibly indented).
-    # Works for any indentation style the formatter picks.
     rest = source[start:]
     end_rel = -1
     for line_start, line in _iter_lines_with_offset(rest):
@@ -124,7 +118,7 @@ def test_flag_sits_inside_the_base_cmd_list():
         "conditional branch -- otherwise some code paths would still "
         "run with silent context shift enabled."
     )
-    # Also pin that it is next to -c / --ctx so the grouping makes sense.
+    # Pin that it sits next to -c / --ctx so the grouping makes sense.
     assert '"-c"' in block
     assert '"--flash-attn"' in block
 
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 202fd36c86..8d88a61ae3 100644
--- a/studio/backend/tests/test_llama_cpp_start_failure_classification.py
+++ b/studio/backend/tests/test_llama_cpp_start_failure_classification.py
@@ -4,10 +4,10 @@
 """Tests for LlamaCppBackend._classify_llama_start_failure.
 
 When llama-server exits before becoming healthy, load_model turns its
-captured stdout/stderr into a user-facing reason. A diffusion / image
-GGUF (FLUX, Qwen-Image, ...) is a valid file with plenty of memory, so
-the generic "invalid file or out of memory" message is actively
-misleading (issue #5842). These tests pin the classification.
+captured stdout/stderr into a user-facing reason. A diffusion/image GGUF
+(FLUX, Qwen-Image, ...) is a valid file with plenty of memory, so the
+generic "invalid file or out of memory" message is misleading (issue
+#5842). These tests pin the classification.
 """
 
 from __future__ import annotations
@@ -22,8 +22,8 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
 if _BACKEND_DIR not in sys.path:
     sys.path.insert(0, _BACKEND_DIR)
 
-# Match the stubbing pattern in sibling tests so the module imports in a
-# lightweight env without fastapi.
+# Match sibling tests' stubbing so the module imports in a lightweight
+# env without fastapi.
 _loggers_stub = _types.ModuleType("loggers")
 _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
 sys.modules.setdefault("loggers", _loggers_stub)
diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py
new file mode 100644
index 0000000000..fa583ef53d
--- /dev/null
+++ b/studio/backend/tests/test_llama_cpp_tool_loop.py
@@ -0,0 +1,1151 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Focused tests for the GGUF llama.cpp agentic tool loop.
+
+These tests drive ``LlamaCppBackend.generate_chat_completion_with_tools``
+with fake llama-server SSE streams. They require no model, subprocess, GPU,
+or network access.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import copy
+import json
+import sys
+from pathlib import Path
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+    sys.path.insert(0, _BACKEND_DIR)
+
+from core.inference.llama_cpp import LlamaCppBackend
+
+
+def _sse(delta: dict) -> str:
+    return "data: " + json.dumps({"choices": [{"index": 0, "delta": delta}]}) + "\n"
+
+
+def _done() -> str:
+    return "data: [DONE]\n"
+
+
+def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]):
+    backend = LlamaCppBackend.__new__(LlamaCppBackend)
+    backend._process = object()
+    backend._healthy = True
+    backend._port = 48847
+    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,
+    ):
+        payloads.append(copy.deepcopy(payload))
+        yield type("FakeResponse", (), {"status_code": 200, "chunks": streams.pop(0)})()
+
+    def fake_iter_text_cancellable(response, _cancel_event):
+        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_names(payload: dict) -> list[str]:
+    return [
+        (tool.get("function") or {}).get("name")
+        for tool in payload.get("tools", [])
+        if (tool.get("function") or {}).get("name")
+    ]
+
+
+def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch):
+    """llama-server may emit content first and then native delta.tool_calls.
+
+    Studio must not drop that tool call after it has streamed the preface.
+    """
+
+    tool_call_id = "call_render_late"
+    first_stream = [
+        _sse({"content": "Here is the artifact.\n\n"}),
+        _sse(
+            {
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": tool_call_id,
+                        "type": "function",
+                        "function": {
+                            "name": "render_html",
+                            "arguments": json.dumps(
+                                {
+                                    "code": "
red
",
+                                    "title": "Simple Red Square",
+                                }
+                            ),
+                        },
+                    }
+                ]
+            }
+        ),
+        _done(),
+    ]
+    second_stream = [
+        _sse({"content": "Done."}),
+        _done(),
+    ]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, [first_stream, second_stream], payloads)
+
+    calls: list[tuple[str, dict]] = []
+
+    def fake_execute_tool(name, arguments, **_kwargs):
+        calls.append((name, arguments))
+        return "Rendered HTML artifact: Simple Red Square."
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "render_html",
+                "description": "Render HTML.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"code": {"type": "string"}},
+                    "required": ["code"],
+                },
+            },
+        }
+    ]
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "Make a red square."}],
+            tools = tools,
+            max_tool_iterations = 1,
+        )
+    )
+
+    content_events = [e for e in events if e.get("type") == "content"]
+    assert content_events[0]["text"] == "Here is the artifact.\n\n"
+
+    first_content_index = next(
+        i for i, event in enumerate(events) if event.get("type") == "content"
+    )
+    actual_tool_start_index = next(
+        i
+        for i, event in enumerate(events)
+        if event.get("type") == "tool_start" and event.get("arguments", {}).get("code")
+    )
+    assert first_content_index < actual_tool_start_index
+
+    assert calls == [
+        (
+            "render_html",
+            {
+                "code": "
red
",
+                "title": "Simple Red Square",
+            },
+        )
+    ]
+    assert any(e.get("type") == "tool_end" and e.get("tool_name") == "render_html" for e in events)
+
+    # The second llama-server request should include the assistant preface
+    # 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]["tool_calls"][0]["id"] == tool_call_id
+    assert assistant_messages[-1]["tool_calls"][0]["function"]["name"] == "render_html"
+
+
+def test_repeat_render_html_nudge_is_not_user_visible_error(monkeypatch):
+    """A repeated render_html call is an internal no-op, not a visible card."""
+
+    first_stream = [
+        _sse(
+            {
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_first",
+                        "type": "function",
+                        "function": {
+                            "name": "render_html",
+                            "arguments": json.dumps(
+                                {
+                                    "code": "first",
+                                    "title": "First",
+                                }
+                            ),
+                        },
+                    }
+                ]
+            }
+        ),
+        _done(),
+    ]
+    repeat_stream = [
+        _sse(
+            {
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_repeat",
+                        "type": "function",
+                        "function": {
+                            "name": "render_html",
+                            "arguments": json.dumps(
+                                {
+                                    "code": "repeat",
+                                    "title": "Repeat",
+                                }
+                            ),
+                        },
+                    }
+                ]
+            }
+        ),
+        _done(),
+    ]
+    final_stream = [_sse({"content": "Short note."}), _done()]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, [first_stream, repeat_stream, final_stream], payloads)
+
+    calls: list[tuple[str, dict]] = []
+
+    def fake_execute_tool(name, arguments, **_kwargs):
+        calls.append((name, arguments))
+        return "Rendered HTML artifact: First."
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "render_html",
+                "description": "Render HTML.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"code": {"type": "string"}},
+                    "required": ["code"],
+                },
+            },
+        },
+        {"type": "function", "function": {"name": "web_search"}},
+    ]
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "Make a red square."}],
+            tools = tools,
+            max_tool_iterations = 2,
+        )
+    )
+
+    assert calls == [
+        (
+            "render_html",
+            {"code": "first", "title": "First"},
+        )
+    ]
+    assert _tool_names(payloads[1]) == ["web_search"]
+
+    actual_tool_starts = [
+        event
+        for event in events
+        if event.get("type") == "tool_start" and event.get("arguments", {}).get("code")
+    ]
+    tool_ends = [
+        event
+        for event in events
+        if event.get("type") == "tool_end" and event.get("tool_name") == "render_html"
+    ]
+    assert len(actual_tool_starts) == 1
+    assert len(tool_ends) == 1
+
+    assert len(payloads) == 3
+    render_tool_messages = [
+        message
+        for message in payloads[2]["messages"]
+        if message.get("role") == "tool" and message.get("name") == "render_html"
+    ]
+    assert len(render_tool_messages) == 1
+    internal_nudges = [
+        message
+        for message in payloads[2]["messages"]
+        if message.get("role") == "user"
+        and "Do not call render_html again" in message.get("content", "")
+    ]
+    assert len(internal_nudges) == 1
+
+
+def test_render_html_success_drops_tool_schema_before_final_pass(monkeypatch):
+    first_stream = [
+        _sse(
+            {
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_first",
+                        "type": "function",
+                        "function": {
+                            "name": "render_html",
+                            "arguments": json.dumps({"code": "ok"}),
+                        },
+                    }
+                ]
+            }
+        ),
+        _done(),
+    ]
+    final_stream = [_sse({"content": "Done."}), _done()]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads)
+
+    def fake_execute_tool(name, arguments, **_kwargs):
+        return "Rendered HTML artifact: Done."
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "Render this."}],
+            tools = [{"type": "function", "function": {"name": "render_html"}}],
+            max_tool_iterations = 3,
+        )
+    )
+
+    assert len(payloads) == 2
+    assert "tools" not in payloads[1]
+    assert any(event.get("type") == "content" and event.get("text") == "Done." for event in events)
+    final_user_messages = [
+        m.get("content", "") for m in payloads[1]["messages"] if m.get("role") == "user"
+    ]
+    assert not any("used all available tool calls" in message for message in final_user_messages)
+
+
+def test_non_consecutive_duplicate_web_search_is_internal_noop(monkeypatch):
+    first_search = [
+        _sse(
+            {
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_search_1",
+                        "type": "function",
+                        "function": {
+                            "name": "web_search",
+                            "arguments": json.dumps({"query": "gpu prices 2026"}),
+                        },
+                    }
+                ]
+            }
+        ),
+        _done(),
+    ]
+    python_call = [
+        _sse(
+            {
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_python",
+                        "type": "function",
+                        "function": {
+                            "name": "python",
+                            "arguments": json.dumps({"code": "print('ok')"}),
+                        },
+                    }
+                ]
+            }
+        ),
+        _done(),
+    ]
+    duplicate_search = [
+        _sse(
+            {
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_search_2",
+                        "type": "function",
+                        "function": {
+                            "name": "web_search",
+                            "arguments": json.dumps({"query": "gpu prices 2026"}),
+                        },
+                    }
+                ]
+            }
+        ),
+        _done(),
+    ]
+    final_stream = [_sse({"content": "Final answer from gathered data."}), _done()]
+    payloads: list[dict] = []
+    backend = _make_backend(
+        monkeypatch,
+        [first_search, python_call, duplicate_search, final_stream],
+        payloads,
+    )
+
+    calls: list[tuple[str, dict]] = []
+
+    def fake_execute_tool(name, arguments, **_kwargs):
+        calls.append((name, arguments))
+        return f"ok:{name}"
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+    tools = [
+        {"type": "function", "function": {"name": "web_search"}},
+        {"type": "function", "function": {"name": "python"}},
+    ]
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "search gpus in 2026 prices and use python"}],
+            tools = tools,
+            max_tool_iterations = 3,
+        )
+    )
+
+    assert calls == [
+        ("web_search", {"query": "gpu prices 2026"}),
+        ("python", {"code": "print('ok')"}),
+    ]
+    assert [
+        event.get("tool_name")
+        for event in events
+        if event.get("type") == "tool_start" and event.get("tool_name")
+    ] == ["web_search", "python"]
+    assert [
+        event.get("tool_name")
+        for event in events
+        if event.get("type") == "tool_end" and event.get("tool_name")
+    ] == ["web_search", "python"]
+    assert not [
+        event
+        for event in events
+        if event.get("tool_call_id") == "call_search_2"
+        and event.get("type") in {"tool_start", "tool_end"}
+    ]
+    assert len(payloads) == 4
+    assert _tool_names(payloads[3]) == ["web_search", "python"]
+    duplicate_nudges = [
+        message
+        for message in payloads[3]["messages"]
+        if message.get("role") == "user"
+        and "already completed successfully" in message.get("content", "")
+    ]
+    assert len(duplicate_nudges) == 1
+
+
+def test_duplicate_web_search_noop_allows_distinct_followup_tool(monkeypatch):
+    first_search = [
+        _sse(
+            {
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_search_1",
+                        "type": "function",
+                        "function": {
+                            "name": "web_search",
+                            "arguments": json.dumps({"query": "gpu prices 2026"}),
+                        },
+                    }
+                ]
+            }
+        ),
+        _done(),
+    ]
+    duplicate_search = [
+        _sse(
+            {
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_search_2",
+                        "type": "function",
+                        "function": {
+                            "name": "web_search",
+                            "arguments": json.dumps({"query": "gpu prices 2026"}),
+                        },
+                    }
+                ]
+            }
+        ),
+        _done(),
+    ]
+    python_call = [
+        _sse(
+            {
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_python",
+                        "type": "function",
+                        "function": {
+                            "name": "python",
+                            "arguments": json.dumps({"code": "print('ok')"}),
+                        },
+                    }
+                ]
+            }
+        ),
+        _done(),
+    ]
+    final_stream = [_sse({"content": "Final answer from gathered data."}), _done()]
+    payloads: list[dict] = []
+    backend = _make_backend(
+        monkeypatch,
+        [first_search, duplicate_search, python_call, final_stream],
+        payloads,
+    )
+
+    calls: list[tuple[str, dict]] = []
+
+    def fake_execute_tool(name, arguments, **_kwargs):
+        calls.append((name, arguments))
+        return f"ok:{name}"
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+    tools = [
+        {"type": "function", "function": {"name": "web_search"}},
+        {"type": "function", "function": {"name": "python"}},
+    ]
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "search gpus in 2026 prices and use python"}],
+            tools = tools,
+            max_tool_iterations = 4,
+        )
+    )
+
+    assert calls == [
+        ("web_search", {"query": "gpu prices 2026"}),
+        ("python", {"code": "print('ok')"}),
+    ]
+    assert [
+        event.get("tool_name")
+        for event in events
+        if event.get("type") == "tool_start" and event.get("tool_name")
+    ] == ["web_search", "python"]
+    assert [
+        event.get("tool_name")
+        for event in events
+        if event.get("type") == "tool_end" and event.get("tool_name")
+    ] == ["web_search", "python"]
+    assert not [
+        event
+        for event in events
+        if event.get("tool_call_id") == "call_search_2"
+        and event.get("type") in {"tool_start", "tool_end"}
+    ]
+    assert len(payloads) == 4
+    assert _tool_names(payloads[2]) == ["web_search", "python"]
+    duplicate_nudges = [
+        message
+        for message in payloads[2]["messages"]
+        if message.get("role") == "user"
+        and "already completed successfully" in message.get("content", "")
+    ]
+    assert len(duplicate_nudges) == 1
+
+
+def test_repeated_duplicate_noop_transitions_to_final_pass(monkeypatch):
+    first_search = [
+        _sse(
+            {
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_search_1",
+                        "type": "function",
+                        "function": {
+                            "name": "web_search",
+                            "arguments": json.dumps({"query": "gpu prices 2026"}),
+                        },
+                    }
+                ]
+            }
+        ),
+        _done(),
+    ]
+    duplicate_one = [
+        _sse(
+            {
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_search_2",
+                        "type": "function",
+                        "function": {
+                            "name": "web_search",
+                            "arguments": json.dumps({"query": "gpu prices 2026"}),
+                        },
+                    }
+                ]
+            }
+        ),
+        _done(),
+    ]
+    duplicate_two = [
+        _sse(
+            {
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_search_3",
+                        "type": "function",
+                        "function": {
+                            "name": "web_search",
+                            "arguments": json.dumps({"query": "gpu prices 2026"}),
+                        },
+                    }
+                ]
+            }
+        ),
+        _done(),
+    ]
+    final_stream = [_sse({"content": "Final answer from first search."}), _done()]
+    payloads: list[dict] = []
+    backend = _make_backend(
+        monkeypatch,
+        [first_search, duplicate_one, duplicate_two, final_stream],
+        payloads,
+    )
+    calls: list[tuple[str, dict]] = []
+
+    def fake_execute_tool(name, arguments, **_kwargs):
+        calls.append((name, arguments))
+        return "result"
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "search gpus"}],
+            tools = [{"type": "function", "function": {"name": "web_search"}}],
+            max_tool_iterations = 10,
+        )
+    )
+
+    assert calls == [("web_search", {"query": "gpu prices 2026"})]
+    assert [event.get("tool_call_id") for event in events if event.get("type") == "tool_end"] == [
+        "call_search_1"
+    ]
+    assert len(payloads) == 4
+    assert "tools" not in payloads[-1]
+    assert any(
+        event.get("type") == "content" and event.get("text") == "Final answer from first search."
+        for event in events
+    )
+
+
+def test_same_turn_duplicate_web_search_is_internal_noop(monkeypatch):
+    same_turn_duplicates = [
+        _sse(
+            {
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_search_1",
+                        "type": "function",
+                        "function": {
+                            "name": "web_search",
+                            "arguments": json.dumps({"query": "gpu prices 2026"}),
+                        },
+                    },
+                    {
+                        "index": 1,
+                        "id": "call_search_2",
+                        "type": "function",
+                        "function": {
+                            "name": "web_search",
+                            "arguments": json.dumps({"query": "gpu prices 2026"}),
+                        },
+                    },
+                ]
+            }
+        ),
+        _done(),
+    ]
+    final_stream = [_sse({"content": "Final answer."}), _done()]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, [same_turn_duplicates, final_stream], payloads)
+
+    calls: list[tuple[str, dict]] = []
+
+    def fake_execute_tool(name, arguments, **_kwargs):
+        calls.append((name, arguments))
+        return "search-result"
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "search gpus"}],
+            tools = [{"type": "function", "function": {"name": "web_search"}}],
+            max_tool_iterations = 2,
+        )
+    )
+
+    assert calls == [("web_search", {"query": "gpu prices 2026"})]
+    assert [event.get("tool_call_id") for event in events if event.get("type") == "tool_end"] == [
+        "call_search_1"
+    ]
+    assert not [
+        event
+        for event in events
+        if event.get("tool_call_id") == "call_search_2"
+        and event.get("type") in {"tool_start", "tool_end"}
+    ]
+
+
+def test_same_turn_repeated_render_html_does_not_emit_second_provisional_start(monkeypatch):
+    same_turn_render_calls = [
+        _sse(
+            {
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_html_1",
+                        "type": "function",
+                        "function": {
+                            "name": "render_html",
+                            "arguments": json.dumps({"code": "one"}),
+                        },
+                    },
+                    {
+                        "index": 1,
+                        "id": "call_html_2",
+                        "type": "function",
+                        "function": {
+                            "name": "render_html",
+                            "arguments": json.dumps({"code": "two"}),
+                        },
+                    },
+                ]
+            }
+        ),
+        _done(),
+    ]
+    final_stream = [_sse({"content": "Final answer."}), _done()]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, [same_turn_render_calls, final_stream], payloads)
+
+    calls: list[tuple[str, dict]] = []
+
+    def fake_execute_tool(name, arguments, **_kwargs):
+        calls.append((name, arguments))
+        return "Rendered HTML artifact: One."
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "render html"}],
+            tools = [{"type": "function", "function": {"name": "render_html"}}],
+            max_tool_iterations = 2,
+        )
+    )
+
+    assert calls == [("render_html", {"code": "one"})]
+    assert [
+        event.get("tool_call_id")
+        for event in events
+        if event.get("type") == "tool_start" and not event.get("arguments")
+    ] == ["call_html_1"]
+    assert not [
+        event
+        for event in events
+        if event.get("tool_call_id") == "call_html_2"
+        and event.get("type") in {"tool_start", "tool_end"}
+    ]
+    assert len(payloads) == 2
+    assert "tools" not in payloads[1]
+    render_nudges = [
+        message
+        for message in payloads[1]["messages"]
+        if message.get("role") == "user"
+        and "Do not call render_html again" in message.get("content", "")
+    ]
+    assert len(render_nudges) == 1
+
+
+def test_disabled_tool_call_is_internal_noop(monkeypatch):
+    disabled_python = [
+        _sse(
+            {
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_python_disabled",
+                        "type": "function",
+                        "function": {
+                            "name": "python",
+                            "arguments": json.dumps({"code": "print(1)"}),
+                        },
+                    }
+                ]
+            }
+        ),
+        _done(),
+    ]
+    final_stream = [_sse({"content": "I cannot run Python here."}), _done()]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, [disabled_python, final_stream], payloads)
+
+    def fake_execute_tool(name, arguments, **_kwargs):
+        raise AssertionError(f"unexpected tool execution: {name} {arguments}")
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "run python"}],
+            tools = [{"type": "function", "function": {"name": "web_search"}}],
+            max_tool_iterations = 1,
+        )
+    )
+
+    assert not [event for event in events if event.get("type") in {"tool_start", "tool_end"}]
+    assert len(payloads) == 2
+    disabled_nudges = [
+        message
+        for message in payloads[1]["messages"]
+        if message.get("role") == "user" and "not enabled" in message.get("content", "")
+    ]
+    assert len(disabled_nudges) == 1
+
+
+def test_render_html_success_does_not_reprompt_render_html_intent(monkeypatch):
+    """After render_html succeeds, do not force another render_html call.
+
+    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
+    exists.
+    """
+
+    first_stream = [
+        _sse(
+            {
+                "tool_calls": [
+                    {
+                        "index": 0,
+                        "id": "call_first",
+                        "type": "function",
+                        "function": {
+                            "name": "render_html",
+                            "arguments": json.dumps(
+                                {
+                                    "code": "first",
+                                    "title": "First",
+                                }
+                            ),
+                        },
+                    }
+                ]
+            }
+        ),
+        _done(),
+    ]
+    post_tool_stream = [
+        _sse({"content": "I will now use render_html again."}),
+        _done(),
+    ]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, [first_stream, post_tool_stream], payloads)
+
+    calls: list[tuple[str, dict]] = []
+
+    def fake_execute_tool(name, arguments, **_kwargs):
+        calls.append((name, arguments))
+        return "Rendered HTML artifact: First."
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "render_html",
+                "description": "Render HTML.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"code": {"type": "string"}},
+                    "required": ["code"],
+                },
+            },
+        }
+    ]
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "Make a red square."}],
+            tools = tools,
+            max_tool_iterations = 1,
+        )
+    )
+
+    assert len(payloads) == 2
+    assert len(calls) == 1
+    assert any(
+        event.get("type") == "content" and event.get("text") == "I will now use render_html again."
+        for event in events
+    )
+
+
+def test_internal_reprompt_attempts_do_not_duplicate_visible_text(monkeypatch):
+    """No-tool re-prompt attempts should not concatenate into the UI."""
+
+    streams = [
+        [_sse({"content": "I will use render_html now."}), _done()],
+        [_sse({"content": "Understood. I will use render_html now."}), _done()],
+    ]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, streams, payloads)
+
+    def fake_execute_tool(name, arguments, **_kwargs):
+        raise AssertionError(f"unexpected tool execution: {name} {arguments}")
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "render_html",
+                "description": "Render HTML.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"code": {"type": "string"}},
+                    "required": ["code"],
+                },
+            },
+        }
+    ]
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "Make a red square."}],
+            tools = tools,
+            max_tool_iterations = 1,
+        )
+    )
+
+    content_texts = [event.get("text", "") for event in events if event.get("type") == "content"]
+    assert content_texts == ["I will use render_html now."]
+    assert len(payloads) == 2
+
+
+def test_forced_reprompt_plain_final_answer_is_visible(monkeypatch):
+    """A hidden forced re-prompt may fall back to a plain final answer."""
+
+    streams = [
+        [_sse({"content": "I will use render_html now."}), _done()],
+        [
+            _sse({"content": "No tool is needed. Final answer: use a red square."}),
+            _done(),
+        ],
+    ]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, streams, payloads)
+
+    def fake_execute_tool(name, arguments, **_kwargs):
+        raise AssertionError(f"unexpected tool execution: {name} {arguments}")
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "Make a red square."}],
+            tools = [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "render_html",
+                        "description": "Render HTML.",
+                        "parameters": {
+                            "type": "object",
+                            "properties": {"code": {"type": "string"}},
+                            "required": ["code"],
+                        },
+                    },
+                }
+            ],
+            max_tool_iterations = 1,
+        )
+    )
+
+    content_texts = [event.get("text", "") for event in events if event.get("type") == "content"]
+    assert content_texts == [
+        "I will use render_html now.",
+        "No tool is needed. Final answer: use a red square.",
+    ]
+    assert len(payloads) == 2
+
+
+def test_internal_reprompt_disabled_when_auto_heal_disabled(monkeypatch):
+    streams = [[_sse({"content": "I will use render_html now."}), _done()]]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, streams, payloads)
+
+    def fake_execute_tool(name, arguments, **_kwargs):
+        raise AssertionError(f"unexpected tool execution: {name} {arguments}")
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "render_html",
+                "description": "Render HTML.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"code": {"type": "string"}},
+                    "required": ["code"],
+                },
+            },
+        }
+    ]
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "Make a red square."}],
+            tools = tools,
+            max_tool_iterations = 1,
+            auto_heal_tool_calls = False,
+        )
+    )
+
+    content_texts = [event.get("text", "") for event in events if event.get("type") == "content"]
+    assert content_texts == ["I will use render_html now."]
+    assert len(payloads) == 1
+
+
+def test_auto_heal_disabled_parses_well_formed_xml_when_tools_enabled(monkeypatch):
+    streams = [
+        [
+            _sse(
+                {
+                    "content": '{"name":"web_search","arguments":{"query":"x"}}'
+                }
+            ),
+            _done(),
+        ],
+        [_sse({"content": "done"}), _done()],
+    ]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, streams, payloads)
+    calls: list[tuple[str, dict]] = []
+
+    def fake_execute_tool(name, arguments, **_kwargs):
+        calls.append((name, arguments))
+        return "result"
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "search"}],
+            tools = [{"type": "function", "function": {"name": "web_search"}}],
+            auto_heal_tool_calls = False,
+            max_tool_iterations = 1,
+        )
+    )
+
+    assert calls == [("web_search", {"query": "x"})]
+    assert not any(
+        event.get("type") == "content" and "" in event.get("text", "")
+        for event in events
+    )
+
+
+def test_reprompted_tool_call_still_streams_final_answer(monkeypatch):
+    """Suppression ends once a forced re-prompt actually calls a tool."""
+
+    streams = [
+        [_sse({"content": "I will use render_html now."}), _done()],
+        [
+            _sse(
+                {
+                    "tool_calls": [
+                        {
+                            "index": 0,
+                            "id": "call_forced",
+                            "type": "function",
+                            "function": {
+                                "name": "render_html",
+                                "arguments": json.dumps(
+                                    {
+                                        "code": "forced",
+                                        "title": "Forced",
+                                    }
+                                ),
+                            },
+                        }
+                    ]
+                }
+            ),
+            _done(),
+        ],
+        [_sse({"content": "Final note after tool."}), _done()],
+    ]
+    payloads: list[dict] = []
+    backend = _make_backend(monkeypatch, streams, payloads)
+
+    calls: list[tuple[str, dict]] = []
+
+    def fake_execute_tool(name, arguments, **_kwargs):
+        calls.append((name, arguments))
+        return "Rendered HTML artifact: Forced."
+
+    monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "render_html",
+                "description": "Render HTML.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"code": {"type": "string"}},
+                    "required": ["code"],
+                },
+            },
+        }
+    ]
+
+    events = list(
+        backend.generate_chat_completion_with_tools(
+            messages = [{"role": "user", "content": "Make a red square."}],
+            tools = tools,
+            max_tool_iterations = 1,
+        )
+    )
+
+    assert len(calls) == 1
+    content_texts = [event.get("text", "") for event in events if event.get("type") == "content"]
+    assert content_texts == ["I will use render_html now.", "Final note after tool."]
+    assert len(payloads) == 3
diff --git a/studio/backend/tests/test_llama_cpp_wait_for_health.py b/studio/backend/tests/test_llama_cpp_wait_for_health.py
index bcf2eb1683..33f9e9d803 100644
--- a/studio/backend/tests/test_llama_cpp_wait_for_health.py
+++ b/studio/backend/tests/test_llama_cpp_wait_for_health.py
@@ -3,10 +3,10 @@
 
 """Tests for LlamaCppBackend._wait_for_health resilience.
 
-The probe loop must swallow transient httpx errors and fall through to
-the subprocess.poll() branch so a crashed llama-server surfaces a
-structured "exited with code X" log instead of bubbling an opaque
-exception up to the /api/inference/load route.
+The probe loop must swallow transient httpx errors and fall through to the
+subprocess.poll() branch so a crashed llama-server surfaces a structured
+"exited with code X" log instead of bubbling an opaque exception up to the
+/api/inference/load route.
 """
 
 from __future__ import annotations
@@ -22,8 +22,7 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
 if _BACKEND_DIR not in sys.path:
     sys.path.insert(0, _BACKEND_DIR)
 
-# Match the stubbing pattern in sibling tests so the module imports in
-# a lightweight env without fastapi.
+# Mirror sibling tests' stubbing so the module imports without fastapi.
 _loggers_stub = _types.ModuleType("loggers")
 _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
 sys.modules.setdefault("loggers", _loggers_stub)
@@ -33,11 +32,10 @@ import httpx  # noqa: E402
 
 from core.inference.llama_cpp import LlamaCppBackend  # noqa: E402
 
-# Sibling tests in this directory install lightweight httpx stubs via
-# sys.modules.setdefault. When collected together, our `httpx` symbol
-# may be one of those stubs, which lacks `get`. Ensure the production
-# code finds a working `httpx.get` and the standard exception types
-# regardless of collection order by adding the missing attributes.
+# Sibling tests install lightweight httpx stubs via sys.modules.setdefault.
+# When collected together, our `httpx` may be such a stub lacking `get`. Add
+# the missing attributes so production code finds a working `httpx.get` and
+# the standard exception types regardless of collection order.
 if not hasattr(httpx, "get"):
     httpx.get = None  # placeholder; every test below monkeypatches it
 for _exc_name in (
@@ -52,9 +50,7 @@ for _exc_name in (
 
 
 def _make_backend(port: int = 12345) -> LlamaCppBackend:
-    """Build a barebones LlamaCppBackend instance with only the
-    attributes _wait_for_health touches. Bypasses __init__ so we do not
-    pull in the full subprocess + logging stack."""
+    """Barebones LlamaCppBackend with only the attributes _wait_for_health touches (bypasses __init__)."""
     b = LlamaCppBackend.__new__(LlamaCppBackend)
     b._port = port
     b._stdout_thread = None
@@ -72,14 +68,9 @@ class TestWaitForHealthResilience:
         assert b._wait_for_health(timeout = 1.0, interval = 0.01) is True
 
     def test_read_error_loops_to_subprocess_poll(self, monkeypatch):
-        """WinError 10054 maps to httpx.ReadError. The loop must swallow
-        it and the next iteration must detect the dead subprocess via
-        poll() != None, returning False with a structured exit-code log
-        instead of bubbling the ReadError."""
+        """WinError 10054 (httpx.ReadError) must be swallowed; the next iteration sees the dead subprocess and returns False with a structured exit-code log."""
         b = _make_backend()
-        # First iteration: process alive (so we reach the httpx probe).
-        # Second iteration: process has exited (so we hit the structured
-        # exit-code branch and return False).
+        # Iter 1: alive (reach probe); iter 2: exited (exit-code branch -> False).
         b._process.poll.side_effect = [None, 1]
         b._process.returncode = 1
         b._stdout_lines = ["llama-server: ggml-cuda.dll failed to load"]
@@ -89,12 +80,12 @@ class TestWaitForHealthResilience:
 
         monkeypatch.setattr(httpx, "get", raise_read_error)
         assert b._wait_for_health(timeout = 5.0, interval = 0.01) is False
-        # Both iterations of the loop ran -- the ReadError did not bubble.
+        # Both loop iterations ran -- the ReadError did not bubble.
         assert b._process.poll.call_count >= 2
 
     def test_remote_protocol_error_also_swallowed(self, monkeypatch):
-        """Partial / malformed response on the probe (server crashed
-        mid-headers) raises RemoteProtocolError -- also non-fatal."""
+        """A partial/malformed probe response (server crashed mid-headers)
+        raises RemoteProtocolError -- also non-fatal."""
         b = _make_backend()
         b._process.poll.side_effect = [None, -1]
         b._process.returncode = -1
@@ -121,8 +112,8 @@ class TestWaitForHealthResilience:
         assert b._process.poll.call_count >= 2
 
     def test_connect_error_swallowed_until_success(self, monkeypatch):
-        """Sanity: existing ConnectError swallowing still works -- the
-        loop retries until llama-server eventually answers 200."""
+        """Sanity: existing ConnectError swallowing still works -- the loop
+        retries until llama-server answers 200."""
         b = _make_backend()
         b._process.poll.return_value = None
         calls = {"n": 0}
@@ -139,8 +130,8 @@ class TestWaitForHealthResilience:
         assert calls["n"] >= 3
 
     def test_dead_process_before_probe_returns_false(self, monkeypatch):
-        """If poll() != None on entry, _wait_for_health must return
-        False immediately without calling httpx at all."""
+        """poll() != None on entry: _wait_for_health returns False
+        immediately without calling httpx."""
         b = _make_backend()
         b._process.poll.return_value = 137
         b._process.returncode = 137
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 125b13782a..493bb93e8c 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
@@ -3,9 +3,9 @@
 
 """``_wait_for_vram_settle`` helper contract.
 
-Pins the bounded poll over ``_get_gpu_free_memory`` that bridges the
-kill -> spawn VRAM-reclaim window. Patches ``_get_gpu_free_memory``;
-no real llama-server or nvidia-smi involved.
+Pins the bounded poll over ``_get_gpu_free_memory`` bridging the kill -> spawn
+VRAM-reclaim window. Patches ``_get_gpu_free_memory``; no real llama-server or
+nvidia-smi involved.
 """
 
 from __future__ import annotations
@@ -19,10 +19,7 @@ from unittest.mock import patch
 import pytest
 
 
-# ---------------------------------------------------------------------------
-# Same external-dep stubs as the other llama_cpp tests so this module
-# imports cleanly without httpx / structlog / loggers installed.
-# ---------------------------------------------------------------------------
+# External-dep stubs so this module imports without httpx / structlog / loggers.
 _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
 if _BACKEND_DIR not in sys.path:
     sys.path.insert(0, _BACKEND_DIR)
@@ -34,8 +31,7 @@ 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)
-# Ensure get_logger is set even if a previous test module already
-# inserted a bare ``structlog`` stub via ``setdefault``.
+# Set get_logger even if a prior test inserted a bare ``structlog`` stub.
 if not hasattr(sys.modules["structlog"], "get_logger"):
     sys.modules["structlog"].get_logger = _structlog_stub.get_logger
 
@@ -74,8 +70,8 @@ def _patch_probe(samples):
     """Patch ``_get_gpu_free_memory`` to yield ``samples`` in order.
 
     Each entry is a list[(idx, free_mib)], a callable, or an exception
-    (instance or class). Calls past the end repeat the last entry so
-    tests can assert "stopped polling" via the call count.
+    (instance or class). Calls past the end repeat the last entry so tests
+    can assert "stopped polling" via the call count.
     """
     state = {"i": 0, "calls": 0}
 
@@ -112,8 +108,8 @@ def _kw(**extra):
 
 
 def test_cold_start_returns_immediately_without_probing():
-    """Default ``since_kill=0.0`` is cold-start: no kill recorded,
-    helper short-circuits without ever invoking the probe."""
+    """Default ``since_kill=0.0`` is cold-start: no kill recorded, so the
+    helper short-circuits without invoking the probe."""
     ctx, state = _patch_probe([[(0, 10000)], [(0, 10000)]])
     with ctx:
         start = time.monotonic()
@@ -154,8 +150,8 @@ def test_first_probe_raises_returns_without_polling():
 
 
 def test_two_consecutive_samples_within_tolerance_settles():
-    """The reclaim ramp from 10000 → 11500 → 11550: third sample within
-    256 MiB of the second so the helper returns after exactly three probes."""
+    """Reclaim ramp 10000 → 11500 → 11550: third sample within 256 MiB of
+    the second, so the helper returns after exactly three probes."""
     ctx, state = _patch_probe(
         [
             [(0, 10000)],
@@ -168,7 +164,7 @@ def test_two_consecutive_samples_within_tolerance_settles():
         LlamaCppBackend._wait_for_vram_settle(**_kw(max_wait = 2.0, interval = 0.05))
         elapsed = time.monotonic() - start
     assert state["calls"] == 3
-    # interval * 2 sleeps = 0.10; allow generous slack for scheduler jitter.
+    # interval * 2 sleeps = 0.10; allow slack for scheduler jitter.
     assert elapsed < 1.0
 
 
@@ -199,7 +195,7 @@ def test_max_wait_respected_when_never_settles():
         start = time.monotonic()
         LlamaCppBackend._wait_for_vram_settle(**_kw(max_wait = 0.5, interval = 0.1))
         elapsed = time.monotonic() - start
-    # We must stop near max_wait, not run forever. Generous upper bound for CI.
+    # Must stop near max_wait, not run forever. Generous upper bound for CI.
     assert 0.3 <= elapsed < 2.0, f"helper ignored max_wait: elapsed={elapsed:.3f}s"
 
 
@@ -217,8 +213,8 @@ def test_max_wait_respected_when_probe_is_slow():
             **_kw(max_wait = 0.4, interval = 0.25),
         )
         elapsed = time.monotonic() - start
-    # First probe (0.30 s) + at most one short clipped sleep + bail.
-    # Hard cap well below the old behaviour of 0.30 + 0.25 + 0.30 = 0.85.
+    # First probe (0.30 s) + at most one clipped sleep + bail.
+    # Hard cap well below the old 0.30 + 0.25 + 0.30 = 0.85.
     assert elapsed < 0.85, f"helper exceeded the deadline due to slow probes: {elapsed:.3f}s"
 
 
@@ -264,21 +260,19 @@ def test_tolerance_two_percent_for_large_cards():
 
 
 def test_load_model_calls_helper_outside_lock_and_uses_last_kill_timestamp():
-    """Pin the call site: outside Phase 3 lock, gated on the timestamp,
-    no ``had_live_process`` in-band flag regression. Mirrors the
-    ``inspect.getsource`` pattern from ``test_llama_cpp_no_context_shift``.
-    """
+    """Pin the call site: outside Phase 3 lock, gated on the timestamp, no
+    ``had_live_process`` in-band flag regression."""
     import inspect
 
     src = inspect.getsource(LlamaCppBackend.load_model)
     assert "_wait_for_vram_settle" in src
     assert "since_kill" in src
     assert "self._last_kill_monotonic" in src
-    # Must be invoked before Phase 3's broad lock so /unload, /cancel,
-    # /status are not blocked during the wait.
+    # Must run before Phase 3's broad lock so /unload, /cancel, /status
+    # are not blocked during the wait.
     assert src.index("_wait_for_vram_settle") < src.index("# ── Phase 3:")
-    # An in-band ``had_live_process`` flag would silently regress the
-    # frontend /unload+/load Apply path; use the timestamp instead.
+    # An in-band ``had_live_process`` flag would regress the frontend
+    # /unload+/load Apply path; use the timestamp instead.
     assert "had_live_process" not in src
 
 
diff --git a/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py b/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py
index a5c9d2255f..957de4bad6 100644
--- a/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py
+++ b/studio/backend/tests/test_llama_cpp_windows_nvidia_path.py
@@ -4,9 +4,9 @@
 """Tests for the Windows pip-nvidia DLL dir resolver.
 
 Studio installs torch with bundled CUDA wheels (nvidia-cuda-runtime-cu13,
-nvidia-cublas-cu13, etc.) and the prebuilt llama-server.exe must find
-those DLLs at runtime to load CUDA. Mirrors the Linux LD_LIBRARY_PATH
-block. See unslothai/unsloth#5106.
+nvidia-cublas-cu13, etc.) and the prebuilt llama-server.exe must find those
+DLLs at runtime to load CUDA. Mirrors the Linux LD_LIBRARY_PATH block.
+See unslothai/unsloth#5106.
 """
 
 from __future__ import annotations
@@ -60,8 +60,7 @@ from core.inference.llama_cpp import LlamaCppBackend  # noqa: E402
 
 
 def _make_nvidia_layout(prefix: Path, pkgs_with_layout: dict[str, str]):
-    """Build a fake /Lib/site-packages/nvidia//{bin|Library/bin}
-    tree with a stub DLL inside each leaf so isdir() picks them up."""
+    """Build a fake nvidia//{bin|Library/bin} tree with a stub DLL per leaf."""
     nv = prefix / "Lib" / "site-packages" / "nvidia"
     for pkg, layout in pkgs_with_layout.items():
         if layout == "bin":
@@ -124,9 +123,8 @@ class TestWindowsPipNvidiaDllDirs:
         assert len(result) == 4
 
     def test_does_not_walk_outside_known_paths(self, tmp_path):
-        # Only nvidia//{bin,Library/bin} and torch/lib are picked
-        # up. Unrelated site-packages contents (numpy, scipy, ...) must
-        # be ignored.
+        # Only nvidia//{bin,Library/bin} and torch/lib are picked up.
+        # Unrelated site-packages contents (numpy, scipy, ...) are ignored.
         site = tmp_path / "Lib" / "site-packages"
         (site / "numpy").mkdir(parents = True)
         (site / "scipy" / "linalg").mkdir(parents = True)
@@ -134,10 +132,8 @@ class TestWindowsPipNvidiaDllDirs:
         assert result == []
 
     def test_picks_up_torch_lib(self, tmp_path):
-        # PyTorch's Windows CUDA wheel bundles cudart64_X.dll /
-        # cublas64_X.dll directly under Lib/site-packages/torch/lib/
-        # instead of as separate nvidia-* wheels. Without this, users
-        # on torch-bundled-CUDA installs still hit #5106.
+        # PyTorch's Windows CUDA wheel bundles cudart64/cublas64 DLLs under
+        # torch/lib/ rather than as nvidia-* wheels; else still hits #5106.
         torch_lib = tmp_path / "Lib" / "site-packages" / "torch" / "lib"
         torch_lib.mkdir(parents = True)
         (torch_lib / "cudart64_12.dll").write_bytes(b"")
@@ -146,8 +142,7 @@ class TestWindowsPipNvidiaDllDirs:
         assert Path(result[0]) == torch_lib
 
     def test_torch_lib_combined_with_nvidia_wheels(self, tmp_path):
-        # Both modular nvidia-* wheels and torch/lib are returned when
-        # present together.
+        # Both modular nvidia-* wheels and torch/lib are returned together.
         _make_nvidia_layout(
             tmp_path,
             {
@@ -165,8 +160,7 @@ class TestWindowsPipNvidiaDllDirs:
         assert any(Path(p) == torch_lib for p in result)
 
     def test_torch_lib_must_be_a_directory(self, tmp_path):
-        # If torch/lib exists as a file (broken install), it is
-        # ignored, not returned.
+        # If torch/lib exists as a file (broken install), it is ignored.
         site = tmp_path / "Lib" / "site-packages" / "torch"
         site.mkdir(parents = True)
         (site / "lib").write_bytes(b"not a dir")
@@ -176,24 +170,19 @@ class TestWindowsPipNvidiaDllDirs:
     def test_skips_non_directories(self, tmp_path):
         nv = tmp_path / "Lib" / "site-packages" / "nvidia"
         (nv / "cuda_runtime").mkdir(parents = True)
-        # Create a regular file at the path where 'bin' would normally be a dir
+        # Regular file where 'bin' would normally be a dir
         (nv / "cuda_runtime" / "bin").write_bytes(b"not a dir")
         result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path))
         assert result == []
 
     def test_missing_prefix_does_not_raise(self):
-        # If sys.prefix points to a path that doesn't exist (unusual,
-        # but possible during test setup), the resolver must just
-        # return [] rather than raising.
+        # Nonexistent sys.prefix: resolver must return [], not raise.
         result = LlamaCppBackend._windows_pip_nvidia_dll_dirs("/this/path/does/not/exist/anywhere")
         assert result == []
 
     def test_picks_up_cu13_bin_x86_64_layout(self, tmp_path):
-        # Current ``nvidia-cuda-runtime`` 13.x and ``nvidia-cublas``
-        # 13.x Windows wheels ship DLLs under
-        # ``nvidia/cu13/bin/x86_64/`` instead of ``nvidia//bin/``.
-        # Without this, users on the new CUDA 13 wheel generation hit
-        # the original #5106 failure mode.
+        # nvidia 13.x Windows wheels ship DLLs under nvidia/cu13/bin/x86_64/
+        # not nvidia//bin/; else the new CUDA 13 wheels hit #5106.
         dll_dir = tmp_path / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x86_64"
         dll_dir.mkdir(parents = True)
         for name in ("cudart64_13.dll", "cublas64_13.dll", "cublasLt64_13.dll"):
@@ -203,7 +192,7 @@ class TestWindowsPipNvidiaDllDirs:
 
     def test_picks_up_bin_x64_layout(self, tmp_path):
         # Some repackaged wheels use ``bin/x64`` (Windows-x64 convention)
-        # instead of ``bin/x86_64`` (NVIDIA-internal convention).
+        # rather than ``bin/x86_64`` (NVIDIA-internal convention).
         dll_dir = tmp_path / "Lib" / "site-packages" / "nvidia" / "cu13" / "bin" / "x64"
         dll_dir.mkdir(parents = True)
         (dll_dir / "cudart64_13.dll").write_bytes(b"")
@@ -211,9 +200,8 @@ class TestWindowsPipNvidiaDllDirs:
         assert str(dll_dir) in result
 
     def test_mixed_cu12_and_cu13_layouts(self, tmp_path):
-        # A venv could have both the modular cu12 wheels (legacy) and
-        # the unsuffixed cu13 wheel installed side by side. Both must
-        # be reachable.
+        # A venv could have both the modular cu12 wheels (legacy) and the
+        # unsuffixed cu13 wheel side by side. Both must be reachable.
         site = tmp_path / "Lib" / "site-packages"
         cu12_bin = site / "nvidia" / "cuda_runtime" / "bin"
         cu13_arch = site / "nvidia" / "cu13" / "bin" / "x86_64"
@@ -225,10 +213,8 @@ class TestWindowsPipNvidiaDllDirs:
         assert cu13_arch in result_set
 
     def test_glob_meta_in_prefix_is_safe(self, tmp_path):
-        # Windows usernames / install paths can contain ``[`` or ``]``.
-        # A glob-based resolver would interpret these as a character
-        # class and silently return [] even when DLL dirs exist. The
-        # iterdir-based implementation must work on such paths.
+        # Windows paths can contain ``[``/``]``; a glob-based resolver would
+        # read these as a character class. The iterdir impl must handle them.
         prefix = tmp_path / "studio_[gpu]_install"
         dll_dir = prefix / "Lib" / "site-packages" / "nvidia" / "cuda_runtime" / "bin"
         dll_dir.mkdir(parents = True)
@@ -237,18 +223,16 @@ class TestWindowsPipNvidiaDllDirs:
         assert str(dll_dir) in result, f"bracket-prefixed path returned empty: {result}"
 
     def test_arch_subdir_listed_before_parent_bin(self, tmp_path):
-        # When both ``nvidia//bin/`` and
-        # ``nvidia//bin/x86_64/`` exist, the arch-specific subdir
-        # must be listed first so Windows DLL search picks up the
-        # cudart64_X.dll location even if the parent ``bin`` is empty.
+        # When both bin/ and bin/x86_64/ exist, the arch subdir must come first
+        # so the Windows DLL search finds cudart64_X.dll if parent bin is empty.
         site = tmp_path / "Lib" / "site-packages"
         outer_bin = site / "nvidia" / "cu13" / "bin"
         arch_bin = outer_bin / "x86_64"
         arch_bin.mkdir(parents = True)
         (arch_bin / "cudart64_13.dll").write_bytes(b"")
         result = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(tmp_path))
-        # outer_bin exists as a directory (it contains arch_bin); the
-        # arch-specific subdir should come first in the list.
+        # outer_bin exists as a dir (it holds arch_bin); the arch-specific
+        # subdir should come first in the list.
         result_paths = [Path(p) for p in result]
         assert arch_bin in result_paths
         assert outer_bin in result_paths
diff --git a/studio/backend/tests/test_llama_server_args.py b/studio/backend/tests/test_llama_server_args.py
index 2c7362a9ff..52e3848157 100644
--- a/studio/backend/tests/test_llama_server_args.py
+++ b/studio/backend/tests/test_llama_server_args.py
@@ -4,8 +4,8 @@
 """Unit tests for the llama-server pass-through args validator.
 
 The validator is the boundary between user CLI/HTTP input and the
-llama-server subprocess. These tests pin denylist behaviour so it
-doesn't quietly regress when new managed flags are added.
+llama-server subprocess. These tests pin denylist behaviour so it doesn't
+regress when new managed flags are added.
 """
 
 from __future__ import annotations
@@ -16,10 +16,8 @@ from pathlib import Path
 
 import pytest
 
-# Load llama_server_args.py directly so this test doesn't drag in the
-# full backend chain (fastapi / structlog / loggers / utils.hardware)
-# via core/inference/__init__.py. The validator is intentionally
-# dependency-free and unit-tests should reflect that.
+# Load llama_server_args.py directly to avoid dragging in the full backend
+# chain via core/inference/__init__.py. The validator is dependency-free.
 _LSA_PATH = Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_server_args.py"
 _spec = importlib.util.spec_from_file_location("_lsa_test_only", _LSA_PATH)
 _lsa = importlib.util.module_from_spec(_spec)
@@ -72,10 +70,9 @@ validate_extra_args = _lsa.validate_extra_args
         # Reasoning controls
         ["--reasoning-format", "deepseek"],
         ["-rea", "auto"],
-        # Soft-managed: user-supplied flags last-wins-override Studio's
-        # auto-set version. --parallel / -np / --n-parallel are NOT
-        # here -- they're hard-denied (KV-cache + slot count would
-        # desync). Use `unsloth studio run --parallel N` instead.
+        # Soft-managed: user flags last-wins over Studio's auto-set version.
+        # --parallel / -np / --n-parallel are hard-denied (KV-cache + slot
+        # count would desync); use `unsloth studio run --parallel N` instead.
         ["-c", "131072"],
         ["--ctx-size", "8192"],
         ["--flash-attn", "off"],
@@ -124,8 +121,8 @@ def test_non_flag_token_passes_through():
         "-np",
         "--parallel",
         "--n-parallel",
-        # Model identity (every alias; bumping llama.cpp must keep
-        # every form rejected, not just the long).
+        # Model identity (every alias; bumping llama.cpp must keep every
+        # form rejected, not just the long one).
         "-m",
         "--model",
         "-mu",
@@ -173,8 +170,8 @@ def test_non_flag_token_passes_through():
         "--models-max",
         "--models-autoload",
         "--no-models-autoload",
-        # Server-mode flips: --embedding / --rerank would restrict
-        # llama-server to those endpoints and break Studio's chat hop.
+        # Server-mode flips: --embedding / --rerank restrict llama-server to
+        # those endpoints and break Studio's chat hop.
         "--embedding",
         "--embeddings",
         "--rerank",
@@ -191,16 +188,16 @@ def test_denylist_rejects_all_aliases(denied):
 @pytest.mark.parametrize(
     "args,offending",
     [
-        # Pass-through --parallel would last-wins-override the real
-        # slot count while Studio's KV-cache fit + llama_parallel_slots
-        # stay at the typer value -- plan vs. process disagree.
+        # Pass-through --parallel would last-wins-override the real slot
+        # count while Studio's KV-cache fit + llama_parallel_slots stay at
+        # the typer value -- plan vs. process disagree.
         (["--parallel", "8"], "--parallel"),
         (["--parallel=8"], "--parallel"),
         (["--n-parallel", "16"], "--n-parallel"),
         (["--n-parallel=16"], "--n-parallel"),
         (["-np", "32"], "-np"),
-        # Attached short form: Click clusters it CLI-side; HTTP /load
-        # with `["-np8"]` must still resolve to managed.
+        # Attached short form: Click clusters it CLI-side; HTTP /load with
+        # `["-np8"]` must still resolve to managed.
         (["-np8"], "-np"),
         (["-np64"], "-np"),
         # Out-of-range values that would bypass the typer 1..64 guard.
@@ -227,8 +224,8 @@ def test_denylist_rejects_equals_form():
     [" --parallel", "--parallel ", "\t--parallel", "  -np", "-np \n", "-np\t"],
 )
 def test_denylist_rejects_whitespace_padded_forms(padded):
-    # `_flag_name` trims whitespace before lookup; otherwise a trailing
-    # space could slip a managed flag past the boundary.
+    # `_flag_name` trims whitespace before lookup; else a trailing space
+    # could slip a managed flag past the boundary.
     with pytest.raises(ValueError, match = "parallel|np"):
         validate_extra_args([padded, "8"])
 
@@ -238,15 +235,15 @@ def test_denylist_rejects_whitespace_padded_forms(padded):
     ["-np8x", "-np-1foo", "-np+1bar", "-np9zzz"],
 )
 def test_denylist_rejects_np_with_digit_prefix_and_junk(attached):
-    # Backend `_flag_name` must classify the same forms the CLI
-    # rewriter expands, else HTTP /load could smuggle `-np8x` through.
+    # Backend `_flag_name` must classify the same forms the CLI rewriter
+    # expands, else HTTP /load could smuggle `-np8x` through.
     with pytest.raises(ValueError, match = "np"):
         validate_extra_args([attached])
 
 
 def test_denylist_rejects_short_form_when_long_is_denied():
-    # `-m` is the short form of --model; rejecting only the long
-    # form would leave a trivial bypass.
+    # `-m` is the short form of --model; rejecting only the long form
+    # would leave a trivial bypass.
     with pytest.raises(ValueError, match = "-m"):
         validate_extra_args(["-m", "/some/other/path.gguf"])
 
diff --git a/studio/backend/tests/test_llm_assist_startup_opt_in.py b/studio/backend/tests/test_llm_assist_startup_opt_in.py
new file mode 100644
index 0000000000..e81b1d3775
--- /dev/null
+++ b/studio/backend/tests/test_llm_assist_startup_opt_in.py
@@ -0,0 +1,127 @@
+# 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 Helper LLM startup pre-cache opt-in behavior."""
+
+from __future__ import annotations
+
+import sys
+import types
+from pathlib import Path
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+    sys.path.insert(0, _BACKEND_DIR)
+
+from models.datasets import AiAssistMappingRequest
+from routes import datasets as datasets_route
+from routes import settings as settings_route
+from utils import helper_precache_settings
+
+
+def _install_fake_studio_db(monkeypatch, *, stored = None):
+    storage_pkg = types.ModuleType("storage")
+    studio_db = types.ModuleType("storage.studio_db")
+    values: dict[str, object] = {}
+    if stored is not None:
+        values[helper_precache_settings.HELPER_PRECACHE_SETTING_KEY] = stored
+
+    def get_app_setting(key, fallback = None):
+        return values.get(key, fallback)
+
+    def upsert_app_settings(settings):
+        values.update(settings)
+        return dict(values)
+
+    studio_db.get_app_setting = get_app_setting
+    studio_db.upsert_app_settings = upsert_app_settings
+    monkeypatch.setitem(sys.modules, "storage", storage_pkg)
+    monkeypatch.setitem(sys.modules, "storage.studio_db", studio_db)
+    return values
+
+
+def test_helper_precache_defaults_off_when_setting_missing(monkeypatch):
+    monkeypatch.delenv("UNSLOTH_HELPER_MODEL_DISABLE", raising = False)
+    _install_fake_studio_db(monkeypatch)
+
+    assert helper_precache_settings.get_helper_precache_enabled() is False
+    assert helper_precache_settings.should_preload_helper_on_startup() is False
+
+
+def test_helper_precache_opt_in_is_blocked_by_existing_disable_env(monkeypatch):
+    _install_fake_studio_db(monkeypatch, stored = True)
+    monkeypatch.setenv("UNSLOTH_HELPER_MODEL_DISABLE", "true")
+
+    assert helper_precache_settings.get_helper_precache_enabled() is True
+    assert helper_precache_settings.should_preload_helper_on_startup() is False
+
+
+def test_settings_route_persists_helper_precache_toggle(monkeypatch):
+    values = _install_fake_studio_db(monkeypatch)
+    monkeypatch.delenv("UNSLOTH_HELPER_MODEL_DISABLE", raising = False)
+
+    response = settings_route.update_helper_precache(
+        settings_route.HelperPrecachePayload(enabled = True),
+        current_subject = "test-user",
+    )
+
+    assert response.enabled is True
+    assert response.default_enabled is False
+    assert response.disabled_by_env is False
+    assert values[helper_precache_settings.HELPER_PRECACHE_SETTING_KEY] is True
+
+
+def test_main_startup_uses_helper_precache_gate_instead_of_unconditional_precache():
+    source = (Path(__file__).resolve().parent.parent / "main.py").read_text(encoding = "utf-8")
+    startup_section = source[
+        source.index("cleanup_orphaned_runs") : source.index("# Initialize RSA key pair")
+    ]
+
+    assert "_start_helper_precache_if_enabled()" in startup_section
+    assert "precache_helper_gguf" not in startup_section
+    assert "threading.Thread(target = _precache" not in startup_section
+
+
+def test_ai_assist_route_still_calls_on_demand_advisor(monkeypatch):
+    calls: list[dict] = []
+    llm_assist = types.ModuleType("utils.datasets.llm_assist")
+
+    def fake_llm_conversion_advisor(**kwargs):
+        calls.append(kwargs)
+        return {
+            "success": True,
+            "suggested_mapping": {"prompt": "user", "answer": "assistant"},
+            "system_prompt": "Answer carefully.",
+            "dataset_type": "question_answering",
+            "is_conversational": False,
+            "user_notification": "Columns mapped by AI Assist.",
+        }
+
+    llm_assist.llm_conversion_advisor = fake_llm_conversion_advisor
+    monkeypatch.setitem(sys.modules, "utils.datasets.llm_assist", llm_assist)
+
+    response = datasets_route.ai_assist_mapping(
+        AiAssistMappingRequest(
+            columns = ["prompt", "answer"],
+            samples = [{"prompt": "x" * 250, "answer": "ok", "extra": "ignored"}],
+            dataset_name = "owner/dataset",
+            hf_token = "hf_test",
+            model_name = "unsloth/test",
+            model_type = "text",
+        ),
+        current_subject = "test-user",
+    )
+
+    assert response.success is True
+    assert response.suggested_mapping == {"prompt": "user", "answer": "assistant"}
+    assert response.system_prompt == "Answer carefully."
+    assert calls == [
+        {
+            "column_names": ["prompt", "answer"],
+            "samples": [{"prompt": "x" * 200, "answer": "ok"}],
+            "dataset_name": "owner/dataset",
+            "hf_token": "hf_test",
+            "model_name": "unsloth/test",
+            "model_type": "text",
+        }
+    ]
diff --git a/studio/backend/tests/test_log_filter_no_truncation.py b/studio/backend/tests/test_log_filter_no_truncation.py
index d9a6e2bc4a..52ffa2ba32 100644
--- a/studio/backend/tests/test_log_filter_no_truncation.py
+++ b/studio/backend/tests/test_log_filter_no_truncation.py
@@ -4,9 +4,9 @@
 """
 Regression tests for loggers.handlers.filter_sensitive_data.
 
-Pins two properties: (1) long strings with commas/slashes pass through
-unchanged (the base64-truncation heuristic from PR #5246 was too aggressive),
-and (2) native-path lease redaction still fires for both inline and dict-key forms.
+Pins two properties: (1) long strings with commas/slashes pass through unchanged
+(the base64-truncation heuristic from PR #5246 was too aggressive), and
+(2) native-path lease redaction still fires for inline and dict-key forms.
 """
 
 from loggers.handlers import filter_sensitive_data
diff --git a/studio/backend/tests/test_login_rate_limit.py b/studio/backend/tests/test_login_rate_limit.py
index 1b084cf436..14b10576da 100644
--- a/studio/backend/tests/test_login_rate_limit.py
+++ b/studio/backend/tests/test_login_rate_limit.py
@@ -4,11 +4,11 @@
 """Tests for the per-(ip, username) login rate limiter.
 
 Covers:
-  - bucket key composition is (client-ip, username.lower())
-  - X-Forwarded-For is honoured only when UNSLOTH_STUDIO_TRUST_FORWARDED is set
+  - bucket key is (client-ip, username.lower())
+  - X-Forwarded-For honoured only when UNSLOTH_STUDIO_TRUST_FORWARDED is set
   - 429 detail body does NOT leak the client IP
-  - One username failing does not lock out a different user from the same IP
-  - One IP failing does not lock out the same user from a different IP
+  - One username failing doesn't lock out a different user from the same IP
+  - One IP failing doesn't lock out the same user from a different IP
 """
 
 import os
@@ -69,8 +69,7 @@ class TestClientIp:
             "127.0.0.1",
             {"x-forwarded-for": "198.51.100.7, 10.0.0.1"},
         )
-        # The proxy header could be spoofed; without the opt-in we
-        # only trust the direct connection.
+        # Proxy header is spoofable; without the opt-in, trust the direct connection.
         assert _client_ip(req) == "127.0.0.1"
 
     def test_honours_first_xff_when_trust_on(self, env_trust_proxy):
@@ -123,8 +122,8 @@ class TestClientIp:
     def test_forwarded_isolates_first_element(self, env_trust_proxy):
         from routes.auth import _client_ip
 
-        # Multi-element Forwarded must pick the first element only,
-        # otherwise suffix variations create attacker-controlled buckets.
+        # Pick the first Forwarded element only, else suffix variations create
+        # attacker-controlled buckets.
         req = _FakeRequest(
             "127.0.0.1",
             {"forwarded": "for=198.51.100.42, for=10.0.0.1;proto=https"},
@@ -155,7 +154,7 @@ class TestBucketKeyAndBlocking:
         for _ in range(_LOGIN_MAX_FAILS):
             _record_login_failure(_bucket_key(req, "alice"))
         assert _login_blocked(_bucket_key(req, "alice")) > 0
-        # bob's account from the same IP is unaffected by alice's typos.
+        # bob's account from the same IP is unaffected by alice's typos
         assert _login_blocked(_bucket_key(req, "bob")) == 0
 
     def test_record_per_ip_isolates_other_ips(self, env_no_proxy):
@@ -189,9 +188,7 @@ class TestBucketKeyAndBlocking:
         req = _FakeRequest("203.0.113.10")
         for idx in range(5):
             auth_routes._record_login_failure(auth_routes._unknown_user_key(req))
-            # Different "username" each attempt would not have throttled
-            # under per-(ip,username) only; the IP aggregate must.
-        # The next missing-user attempt is blocked.
+            # Per-(ip,username) alone wouldn't throttle distinct usernames; the IP aggregate must.
         assert auth_routes._login_blocked(auth_routes._unknown_user_key(req)) > 0
 
     def test_unknown_user_bucket_is_single_sentinel(self, env_no_proxy):
@@ -202,8 +199,7 @@ class TestBucketKeyAndBlocking:
         unknown_key = auth_routes._unknown_user_key(req)
         for _ in range(20):
             auth_routes._record_login_failure(unknown_key)
-        # Account bucket cardinality stays at exactly one sentinel entry
-        # for this IP regardless of how many distinct usernames sprayed.
+        # Exactly one sentinel bucket for this IP regardless of usernames sprayed.
         ip_keys = [k for k in auth_routes._LOGIN_BUCKETS if k[0] == "203.0.113.11"]
         assert len(ip_keys) == 1
         assert ip_keys[0][1].startswith("\x00")
@@ -216,7 +212,7 @@ class TestBucketKeyAndBlocking:
         req = _FakeRequest("203.0.113.12")
         for idx in range(50):
             auth_routes._record_login_failure((req.client.host, f"user-{idx}"))
-        # Hard cap respected; further keys do not allocate.
+        # Hard cap respected; further keys don't allocate.
         assert len(auth_routes._LOGIN_BUCKETS) <= 10
 
 
@@ -249,7 +245,7 @@ class TestLogin429Body:
     def test_429_detail_does_not_leak_ip(self, env_no_proxy, login_client):
         from routes.auth import _LOGIN_MAX_FAILS
 
-        # Drive 6 failures from the same client IP / username.
+        # Drive 6 failures from the same client IP / username
         for _ in range(_LOGIN_MAX_FAILS):
             r = login_client.post(
                 "/api/auth/login",
@@ -262,8 +258,8 @@ class TestLogin429Body:
         )
         assert r.status_code == 429
         detail = r.json()["detail"]
-        # The 429 body must not interpolate the source IP.
+        # The 429 body must not interpolate the source IP
         assert "127.0.0.1" not in detail
         assert "Too many" in detail
-        # Retry-After header is still set for clients.
+        # Retry-After header is still set for clients
         assert "Retry-After" in r.headers
diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py
index 0d263d7630..ede3cf15d4 100644
--- a/studio/backend/tests/test_mcp_servers.py
+++ b/studio/backend/tests/test_mcp_servers.py
@@ -151,8 +151,7 @@ def test_execute_tool_disabled_server(tmp_path, monkeypatch):
 
 
 def test_mcp_specs_skip_invalid_openai_function_names():
-    """OpenAI requires function.name ^[a-zA-Z0-9_-]{1,64}$; tools whose
-    names contain '.', '/', spaces, etc. would 400 the whole request."""
+    """OpenAI requires function.name ^[a-zA-Z0-9_-]{1,64}$; bad names 400 the request."""
     from core.inference.tools import _mcp_specs_for_server
 
     server = {"id": "srv", "display_name": "S"}
@@ -177,8 +176,7 @@ def test_mcp_specs_skip_empty_tool_name():
 
 
 def test_mcp_specs_drops_duplicate_names():
-    """Same tool name twice from one MCP server -> OpenAI rejects the
-    request as 'duplicates'. Drop the duplicate before forwarding."""
+    """Duplicate tool names from one server -> OpenAI rejects; drop before forwarding."""
     from core.inference.tools import _mcp_specs_for_server
 
     server = {"id": "srv", "display_name": "S"}
@@ -188,8 +186,7 @@ def test_mcp_specs_drops_duplicate_names():
 
 
 def test_call_tool_sync_respects_pre_set_cancel_event(monkeypatch):
-    """cancel_event already set before the call -> immediate Error: cancelled
-    without making a network round-trip."""
+    """Pre-set cancel_event -> immediate cancellation, no network round-trip."""
     import threading
     from core.inference import mcp_client
 
@@ -203,7 +200,7 @@ def test_call_tool_sync_respects_pre_set_cancel_event(monkeypatch):
 
         async def call_tool(self, name, args):
             import asyncio as _asyncio
-            await _asyncio.sleep(30)  # never finishes within the test
+            await _asyncio.sleep(30)  # never finishes during the test
 
     monkeypatch.setattr(mcp_client, "_client", lambda *a, **kw: _StubClient())
 
@@ -221,8 +218,8 @@ def test_call_tool_sync_respects_pre_set_cancel_event(monkeypatch):
 
 
 def test_clear_oauth_tokens_async_no_op_safe(tmp_path, monkeypatch):
-    """clear_oauth_tokens_async on a URL with no stored token must not raise --
-    the delete + update handlers call it best-effort regardless of prior state."""
+    """clear_oauth_tokens_async on a URL with no stored token must not raise;
+    the delete + update handlers call it best-effort regardless of state."""
     import asyncio
 
     monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
@@ -233,7 +230,7 @@ def test_clear_oauth_tokens_async_no_op_safe(tmp_path, monkeypatch):
 
 
 def test_delete_server_calls_oauth_cleanup_when_oauth_was_on(tmp_path, monkeypatch):
-    """delete_mcp_server route helper should invoke clear_oauth_tokens_async
+    """delete_mcp_server route helper must call clear_oauth_tokens_async
     when the deleted row had use_oauth=true."""
     import asyncio
 
@@ -255,7 +252,7 @@ def test_delete_server_calls_oauth_cleanup_when_oauth_was_on(tmp_path, monkeypat
         calls.append(url)
 
     monkeypatch.setattr(mcp_client, "clear_oauth_tokens_async", fake_clear)
-    # Re-import the route's binding through the module so the patch is seen.
+    # Patch the route's module binding too so it's seen.
     import routes.mcp_servers as routes_mcp
 
     monkeypatch.setattr(routes_mcp, "clear_oauth_tokens_async", fake_clear)
@@ -292,7 +289,7 @@ def test_delete_server_skips_oauth_cleanup_when_oauth_off(tmp_path, monkeypatch)
 
 def test_update_server_clears_oauth_on_url_change(tmp_path, monkeypatch):
     """Changing the URL on an OAuth server must drop the old URL's tokens
-    so the new URL doesn't silently inherit credentials."""
+    so the new URL doesn't inherit credentials."""
     import asyncio
 
     _reset_db(tmp_path, monkeypatch)
@@ -381,7 +378,7 @@ def test_changes_from_payload_rejects_null_use_oauth():
 
 def test_test_endpoint_surfaces_url_validation_as_400(tmp_path, monkeypatch):
     """POST /api/mcp/servers/test must 400 on invalid URL like create/update;
-    previously the same input returned 200 with {"ok": false}."""
+    it previously returned 200 with {"ok": false}."""
     import asyncio
 
     _reset_db(tmp_path, monkeypatch)
@@ -399,9 +396,8 @@ def test_test_endpoint_surfaces_url_validation_as_400(tmp_path, monkeypatch):
 
 
 def test_tool_xml_parser_handles_hyphenated_parameter_names():
-    """MCP tool schemas commonly use hyphenated property names like
-    `issue-number` / `repo-name`; the XML parser's `` regex
-    dropped those keys. Verify hyphenated parameter names round-trip."""
+    """Hyphenated property names like `issue-number` must round-trip through the
+    XML parser (the old `` regex dropped them)."""
     from core.inference.tool_call_parser import parse_tool_calls_from_text
     import json as _json
 
@@ -417,8 +413,8 @@ def test_tool_xml_parser_handles_hyphenated_parameter_names():
 
 
 def test_tool_healing_strip_handles_hyphenated_function_names():
-    """GGUF's core/tool_healing.py has its own copy of the XML strip
-    regex; the round-4 fix to the shared parser missed this file."""
+    """core/tool_healing.py has its own copy of the XML strip regex that the
+    shared-parser fix missed."""
     from core.tool_healing import strip_tool_call_markup
 
     out = strip_tool_call_markup(
@@ -428,9 +424,8 @@ def test_tool_healing_strip_handles_hyphenated_function_names():
 
 
 def test_gguf_allow_list_blocks_unadvertised_tool(monkeypatch):
-    """When the model emits a tool call not in the per-request tool list
-    the GGUF agentic loop must refuse to dispatch -- mirroring the
-    safetensors path. Previously execute_tool ran the call regardless."""
+    """A tool call not in the per-request list must be refused by the GGUF
+    agentic loop (mirroring the safetensors path)."""
     from core.inference import tools as tools_mod
 
     captured: list[str] = []
@@ -441,8 +436,7 @@ def test_gguf_allow_list_blocks_unadvertised_tool(monkeypatch):
 
     monkeypatch.setattr(tools_mod, "execute_tool", fake_execute)
 
-    # Re-create the allow-list check inline so we can unit-test the
-    # behavior without spinning up llama-server.
+    # Inline allow-list check to unit-test behavior without llama-server.
     def _gate(tools_advertised, called_name, args):
         allowed = {
             (t.get("function") or {}).get("name")
@@ -472,9 +466,8 @@ def test_gguf_allow_list_blocks_unadvertised_tool(monkeypatch):
 
 
 def test_call_tool_sync_short_circuits_on_pre_set_cancel(monkeypatch):
-    """cancel_event set BEFORE call_tool_sync runs -> no HTTP request
-    is made. Previously the call task was created before the cancel
-    check, opening a transport that the watcher then had to cancel."""
+    """Pre-set cancel_event -> no HTTP request (task used to open a transport
+    before the cancel check)."""
     from core.inference import mcp_client
 
     opened: list[str] = []
@@ -510,9 +503,8 @@ def test_call_tool_sync_short_circuits_on_pre_set_cancel(monkeypatch):
 
 
 def test_clear_oauth_tokens_swallows_constructor_errors(tmp_path, monkeypatch):
-    """clear_oauth_tokens_async is best-effort; an OAuth constructor
-    failure (e.g. missing fastmcp.client.auth) must not bubble out into
-    a 500 from the delete / update routes."""
+    """clear_oauth_tokens_async is best-effort; an OAuth constructor failure
+    must not bubble into a 500 from the delete/update routes."""
     import asyncio
     from core.inference import mcp_client
 
@@ -534,9 +526,8 @@ def test_clear_oauth_tokens_swallows_constructor_errors(tmp_path, monkeypatch):
 
 
 def test_tool_xml_parser_handles_hyphenated_function_names():
-    """MCP tool names are advertised as `mcp__srv__list-issues` (the regex
-    fix allows '-'); the XML tool-call parser must parse them too,
-    otherwise the model can call the tool but Studio cannot dispatch."""
+    """Hyphenated tool names like `mcp__srv__list-issues` must parse, else the
+    model can call the tool but Studio can't dispatch."""
     from core.inference.tool_call_parser import parse_tool_calls_from_text
 
     calls = parse_tool_calls_from_text(
@@ -552,7 +543,7 @@ def test_tool_xml_parser_handles_hyphenated_function_names():
 
 def test_tool_xml_strip_handles_hyphenated_function_names():
     """routes/inference.py:_TOOL_XML_RE must strip a ``
-    block; otherwise hyphenated MCP tool-call XML leaks into chat history."""
+    block; else hyphenated MCP tool-call XML leaks into chat history."""
     import re as _re
     from pathlib import Path
 
@@ -570,11 +561,9 @@ def test_tool_xml_strip_handles_hyphenated_function_names():
 
 
 def test_safetensors_agentic_empty_allowlist_still_means_allow_all():
-    """Document existing contract: at the safetensors_agentic layer,
-    tools=[] is still treated as "no constraint" (so existing callers
-    work unchanged). The real fix for the MCP-only-no-discovery case
-    lives at the route level in inference.py, which refuses to enter
-    use_tools when the resolved tool list is empty."""
+    """Contract: at the safetensors_agentic layer tools=[] means "no
+    constraint". The MCP-only-no-discovery fix lives at the route level in
+    inference.py, which refuses use_tools when the resolved list is empty."""
     import threading
     from core.inference.safetensors_agentic import run_safetensors_tool_loop
 
diff --git a/studio/backend/tests/test_mcp_stdio_improvements.py b/studio/backend/tests/test_mcp_stdio_improvements.py
index 515e39d5b6..b0bfd45135 100644
--- a/studio/backend/tests/test_mcp_stdio_improvements.py
+++ b/studio/backend/tests/test_mcp_stdio_improvements.py
@@ -1,8 +1,8 @@
 """Tests for the proposed PR #5863 improvements.
 
 Covers: _client() self-gating + keep_alive, OAuth normalised off for stdio
-(create + update), env/header dropped on a transport-type switch, and the
-backend rejecting a command whose first token is a URL scheme.
+(create + update), env/header dropped on a transport-type switch, and rejecting
+a command whose first token is a URL scheme.
 
 Run from studio/backend:  python -m pytest tests/test_mcp_stdio_improvements.py -q
 """
@@ -41,7 +41,7 @@ def test_client_refuses_stdio_when_disabled(monkeypatch):
 def test_client_builds_stdio_when_enabled_without_spawning(monkeypatch):
     _enable(monkeypatch)
     # Constructing the Client must not spawn the subprocess (spawn happens on
-    # __aenter__); we only assert it builds.
+    # __aenter__); only assert it builds.
     client = mcp_client._client("npx -y server /tmp", {"K": "v"})
     assert client is not None
 
@@ -122,7 +122,7 @@ def test_switch_stdio_to_http_drops_env(tmp_path, monkeypatch):
             "s1", McpServerUpdate(url = "https://remote/mcp"), current_subject = "u"
         )
     )
-    # the stdio env must NOT survive as HTTP headers on the remote endpoint
+    # stdio env must NOT survive as HTTP headers on the remote endpoint
     assert resp.headers == {}
     assert mcp_servers_db.get_server("s1")["headers_json"] is None
 
@@ -161,7 +161,7 @@ def test_same_transport_edit_keeps_headers(tmp_path, monkeypatch):
         url = "npx server",
         headers_json = '{"API_KEY": "secret"}',
     )
-    # editing only the display name (still stdio) must not wipe env vars
+    # editing only the display name (still stdio) must keep env vars
     resp = asyncio.run(
         routes_mcp.update_mcp_server("s1", McpServerUpdate(display_name = "B"), current_subject = "u")
     )
@@ -183,13 +183,12 @@ def test_validate_url_rejects_url_scheme_command_when_enabled(monkeypatch):
 def test_validate_url_allows_url_in_argument(monkeypatch):
     from routes.mcp_servers import _validate_url
     _enable(monkeypatch)
-    # :// inside an ARGUMENT (not the first token) is still a valid command
+    # :// inside an ARGUMENT (not the first token) is a valid command
     assert _validate_url("npx server --url https://x/mcp") == ("npx server --url https://x/mcp")
 
 
 # ── P6: Data Recipe stdio path obeys the same host gate ─────────────
-# build_mcp_providers needs the data_designer plugin, which is only installed in
-# the Studio test job; skip there rather than fail the core matrix.
+# build_mcp_providers needs the Studio-only data_designer plugin; skip if absent.
 
 _STDIO_RECIPE = {
     "mcp_providers": [
@@ -209,7 +208,7 @@ def test_data_recipe_skips_stdio_when_disabled(monkeypatch):
     _disable(monkeypatch)
     from core.data_recipe.service import build_mcp_providers
 
-    # gate off -> the stdio provider is dropped (no subprocess can be spawned)
+    # gate off -> the stdio provider is dropped (no subprocess spawned)
     assert build_mcp_providers(_STDIO_RECIPE) == []
 
 
diff --git a/studio/backend/tests/test_mcp_stdio_pr5863.py b/studio/backend/tests/test_mcp_stdio_pr5863.py
index 277306836b..c6a4898d5d 100644
--- a/studio/backend/tests/test_mcp_stdio_pr5863.py
+++ b/studio/backend/tests/test_mcp_stdio_pr5863.py
@@ -1,13 +1,10 @@
 """Verification tests for PR #5863 (stdio MCP server support).
 
-Covers the pure helpers (is_stdio / parse_stdio_command / stdio_mcp_enabled /
-probe_timeout), the route-level _validate_url gate, and - most importantly -
-that the UNSLOTH_STUDIO_ALLOW_STDIO_MCP gate blocks the stdio transport at all
-five enforcement points (create, update, test, refresh, discovery, execute)
-when disabled, and reaches it when enabled. The transport (_client) is stubbed
-so no real subprocess is spawned; a recorder asserts whether it was reached.
-
-Run from studio/backend:  python -m pytest tests/test_mcp_stdio_pr5863.py -q
+Covers the pure helpers, the route-level _validate_url gate, and that the
+UNSLOTH_STUDIO_ALLOW_STDIO_MCP gate blocks the stdio transport at every
+enforcement point (create/update/test/refresh/discovery/execute) when disabled
+and reaches it when enabled. The transport is stubbed so no subprocess spawns;
+a recorder asserts whether it was reached.
 """
 
 import sys
@@ -57,7 +54,7 @@ class _FakeResult:
 
 
 class _RecordingClient:
-    """Stands in for fastmcp.Client; records that the transport was opened."""
+    """Stand-in for fastmcp.Client; records that the transport was opened."""
 
     def __init__(self, url, headers, use_oauth, recorder):
         recorder.append({"url": url, "headers": headers, "use_oauth": use_oauth})
@@ -78,7 +75,7 @@ class _RecordingClient:
 @pytest.fixture
 def transport(monkeypatch):
     """Patch mcp_client._client with a recorder. Returns the recorder list;
-    empty == the stdio transport was never reached."""
+    empty == stdio transport never reached."""
     recorder = []
     monkeypatch.setattr(
         mcp_client,
@@ -156,8 +153,8 @@ def test_parse_unclosed_quote_raises_valueerror():
 
 
 def test_parse_windows_strips_wrapping_quotes(monkeypatch):
-    # gemini "medium": posix=False keeps backslash paths but also the wrapping
-    # quotes; the PR strips a matched pair so argv[0] reaches the OS clean.
+    # gemini "medium": posix=False keeps backslash paths but also the
+    # wrapping quotes; the PR strips a matched pair so argv[0] is clean.
     monkeypatch.setattr(sys, "platform", "win32")
     parts = mcp_client.parse_stdio_command(r'"C:\Program Files\node\node.exe" server.js')
     assert parts[0] == r"C:\Program Files\node\node.exe"
@@ -214,8 +211,8 @@ def test_validate_url_gate_off_rejects_stdio(monkeypatch):
 
 
 def test_validate_url_gate_off_message_depends_on_whitespace(monkeypatch):
-    # The message names a command only when the value has whitespace, and never
-    # says "desktop app only" (self-hosted hosts can opt in via the env var).
+    # The message names a command only when the value has whitespace, and
+    # never says "desktop app only" (self-hosted can opt in via the env var).
     _disable(monkeypatch)
     from routes.mcp_servers import _validate_url
 
@@ -242,8 +239,8 @@ def test_validate_url_gate_on_accepts_stdio(monkeypatch):
     assert _validate_url("https://x/mcp") == "https://x/mcp"
     # url-bearing argument accepted as a command
     assert _validate_url("npx server --url https://x/mcp") == ("npx server --url https://x/mcp")
-    # A lone token is ambiguous; keep the prior behaviour and accept it as a
-    # command rather than guessing it's a URL (no regression for single binaries).
+    # A lone token is ambiguous; accept it as a command rather than
+    # guessing it's a URL (no regression for single binaries).
     assert _validate_url("/usr/local/bin/my-mcp-server") == "/usr/local/bin/my-mcp-server"
     assert _validate_url("mcp-server-sqlite") == "mcp-server-sqlite"
     # empty / unparseable still rejected
@@ -284,7 +281,7 @@ def test_update_http_to_stdio_blocked_when_off(tmp_path, monkeypatch):
     _reset_db(tmp_path, monkeypatch)
     _disable(monkeypatch)
     mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://a/mcp")
-    # editing url -> stdio command must 400 (http->stdio edit bypass closed)
+    # editing url -> stdio command must 400 (http->stdio bypass closed)
     with pytest.raises(HTTPException) as exc:
         asyncio.run(
             routes_mcp.update_mcp_server(
@@ -321,7 +318,7 @@ def test_refresh_route_gate(tmp_path, monkeypatch, transport):
     import routes.mcp_servers as routes_mcp
 
     _reset_db(tmp_path, monkeypatch)
-    # a stdio row as if carried over from a desktop DB
+    # a stdio row, as if carried over from a desktop DB
     mcp_servers_db.create_server(id = "stdio1", display_name = "FS", url = "npx server")
 
     _disable(monkeypatch)
diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py
index 2d6efa9f8b..1005431926 100644
--- a/studio/backend/tests/test_middleware.py
+++ b/studio/backend/tests/test_middleware.py
@@ -27,9 +27,7 @@ def main_module():
     return _main
 
 
-# =====================================================================
 # MaxBodyMiddleware
-# =====================================================================
 
 
 def _make_protected_app(
@@ -109,8 +107,8 @@ class TestMaxBodyMiddleware:
         assert "too large" in r.json()["detail"].lower()
 
     def test_chunked_upload_over_cap_rejected(self, main_module):
-        # Regression: declared-Content-Length-only check could be bypassed
-        # by chunked transfer-encoding.
+        # Regression: declared-Content-Length-only check could be bypassed by
+        # chunked transfer-encoding.
         app = _make_protected_app(1024, main_module)
         c = TestClient(app)
 
@@ -205,9 +203,7 @@ class TestMaxBodyMiddleware:
         assert "Content-Length" in r.json()["detail"]
 
 
-# =====================================================================
 # SecurityHeadersMiddleware / CSP
-# =====================================================================
 
 
 def _make_csp_app(main_module, attach_nonce: str | None = None):
@@ -280,30 +276,25 @@ class TestSecurityHeadersMiddleware:
         nonced = main_module._build_csp("XYZ")
         assert "script-src 'self' 'nonce-XYZ';" in nonced
 
-    def test_img_src_allows_google_favicons(self, main_module):
-        # sources.tsx fetches https://www.google.com/s2/favicons?... ; without
-        # this allowlist entry citation favicons fall back to gray initials.
+    def test_img_and_media_allow_https_sources(self, main_module):
+        # Model-card READMEs and citation favicons pull images/media from many
+        # https origins (HF LFS/XET CDNs, shields/badge hosts, GitHub-hosted
+        # assets, audio/video samples). img-src/media-src allow any https source
+        # so they render; this mirrors the desktop CSP in tauri.conf.json.
         csp = main_module._build_csp()
-        img_directive = next(
-            chunk.strip() for chunk in csp.split(";") if chunk.strip().startswith("img-src ")
-        )
-        # Tokenise and compare with `==` so CodeQL's URL-substring rule does
-        # not read directive-string `in` membership as URL sanitisation.
-        img_sources = img_directive.split()
-        assert any(src == "https://www.google.com" for src in img_sources)
-        # Pre-existing favicon CDNs stay allowed.
-        for host in (
-            "https://t0.gstatic.com",
-            "https://t1.gstatic.com",
-            "https://t2.gstatic.com",
-            "https://t3.gstatic.com",
-        ):
-            assert any(src == host for src in img_sources)
+        directives = {
+            chunk.strip().split()[0]: chunk.strip().split()
+            for chunk in csp.split(";")
+            if chunk.strip()
+        }
+        for name in ("img-src", "media-src"):
+            assert name in directives, f"missing {name} directive"
+            # Tokenise and compare with `==` so CodeQL's URL-substring rule does
+            # not read directive-string `in` membership as URL sanitisation.
+            assert any(src == "https:" for src in directives[name])
 
 
-# =====================================================================
 # /api/health auth gate
-# =====================================================================
 
 
 @pytest.fixture
@@ -332,9 +323,9 @@ def health_app(tmp_path, monkeypatch):
 
 
 class TestHealthAuthGate:
-    # Launcher / frontend bootstrap fields are available unauth so the Tauri
-    # watchdog can re-adopt a sibling backend and the SPA can detect chat-only
-    # mode before any token exists. Version / device_type still require a bearer.
+    # Launcher / frontend bootstrap fields are unauth so the Tauri watchdog can
+    # re-adopt a sibling backend and the SPA can detect chat-only mode before
+    # any token exists. Version / device_type still require a bearer.
     LAUNCHER_BITS = (
         "service",
         "studio_root_id",
@@ -361,7 +352,7 @@ class TestHealthAuthGate:
             assert forbidden not in body
 
     def test_invalid_bearer_returns_launcher_bits_only(self, health_app):
-        # Regression: calling the async dep without await made any Bearer header pass.
+        # Regression: calling the async dep without await let any Bearer header pass.
         c = TestClient(health_app)
         r = c.get(
             "/api/health",
diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py
index 49afab048d..9871965ce8 100644
--- a/studio/backend/tests/test_mlx_inference_backend.py
+++ b/studio/backend/tests/test_mlx_inference_backend.py
@@ -159,10 +159,9 @@ def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewri
     assert isinstance(backend._tokenizer, _DummyTokenizer)
 
 
-# Regression: MLXInferenceBackend.generate_chat_response must accept the
-# four template kwargs (tools / enable_thinking / reasoning_effort /
-# preserve_thinking) so the route layer can forward what the user
-# toggled in the UI. The previous signature raised
+# Regression: generate_chat_response must accept the four template kwargs
+# (tools / enable_thinking / reasoning_effort / preserve_thinking) so the route
+# layer can forward UI toggles. The old signature raised
 # "got an unexpected keyword argument 'tools'" on Mac.
 
 
@@ -184,8 +183,8 @@ def test_mlx_generate_chat_response_accepts_template_kwargs():
 
 
 def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch):
-    """The Mac text path must route through apply_chat_template_for_
-    generation so reasoning / tool kwargs reach the tokenizer."""
+    """Mac text path must route through apply_chat_template_for_generation so
+    reasoning / tool kwargs reach the tokenizer."""
     _install_fake_mlx(monkeypatch)
     from core.inference.mlx_inference import MLXInferenceBackend
 
@@ -203,9 +202,8 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch):
         raising = True,
     )
 
-    # mlx_lm.stream_generate yields response objects with .token; make a
-    # one-token generator so _generate_text returns without touching the
-    # real stack.
+    # mlx_lm.stream_generate yields response objects with .token; use a
+    # one-token generator so _generate_text returns without the real stack.
     import types as _types
 
     mlx_lm_pkg = _types.ModuleType("mlx_lm")
@@ -250,7 +248,7 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch):
         )
     )
     assert out == ["hi"]
-    # The kwargs the user toggled must reach the chat-template helper.
+    # The toggled kwargs must reach the chat-template helper.
     assert captured["kwargs"]["tools"] == [{"function": {"name": "web_search"}}]
     assert captured["kwargs"]["enable_thinking"] is True
     assert captured["kwargs"]["reasoning_effort"] == "medium"
diff --git a/studio/backend/tests/test_multimodal_document.py b/studio/backend/tests/test_multimodal_document.py
index b2986f2c3a..5cd7c876cc 100644
--- a/studio/backend/tests/test_multimodal_document.py
+++ b/studio/backend/tests/test_multimodal_document.py
@@ -3,17 +3,16 @@
 
 """Tests for PDF / document attachment translation on external providers.
 
-Studio introduces a normalised `input_document` content part on
-ChatCompletionRequest so the frontend doesn't have to know the
-per-provider attachment shape:
+Studio adds a normalised `input_document` content part on
+ChatCompletionRequest so the frontend needn't know the per-provider
+attachment shape:
 
-- Anthropic: translates to `{type:"document", source:{type:"base64"|"url", ...}}`
-- OpenAI Responses: translates to `{type:"input_file", file_data|file_url, filename?}`
+- Anthropic: `{type:"document", source:{type:"base64"|"url", ...}}`
+- OpenAI Responses: `{type:"input_file", file_data|file_url, filename?}`
 
-These tests pin the translation shape on both paths for base64 data
-URIs and remote URLs, with optional filename metadata, and confirm
-unknown / empty document parts are dropped without breaking the
-request.
+Pins the translation shape on both paths for base64 data URIs and remote
+URLs (with optional filename), and confirms unknown / empty document
+parts are dropped without breaking the request.
 """
 
 import asyncio
@@ -86,10 +85,8 @@ _PDF_DATA_URI = f"data:application/pdf;base64,{_TINY_PDF_B64}"
 
 
 def _strip_cache(p: dict) -> dict:
-    # Studio's prompt-cache wiring attaches cache_control:{type:ephemeral}
-    # to the tail block of the last user message; strip it before
-    # comparing the document core fields so this test stays focused
-    # on the translation, not the caching layer.
+    # Strip the prompt-cache cache_control off the last user block so this
+    # test focuses on translation, not the caching layer.
     return {k: v for k, v in p.items() if k != "cache_control"}
 
 
@@ -117,8 +114,8 @@ def test_anthropic_base64_pdf_becomes_document_block(monkeypatch):
     types = [p.get("type") for p in parts]
     assert "document" in types, parts
     doc = _strip_cache(next(p for p in parts if p.get("type") == "document"))
-    # citations: {enabled: true} opts into Anthropic's natural-citation
-    # pipeline; without it the citations_delta handler is a no-op.
+    # citations:{enabled:true} opts into Anthropic's citation pipeline;
+    # without it the citations_delta handler is a no-op.
     assert doc == {
         "type": "document",
         "source": {
@@ -179,9 +176,8 @@ def test_anthropic_empty_document_part_is_dropped(monkeypatch):
 
 
 def test_anthropic_empty_only_document_drops_whole_message(monkeypatch):
-    # If the ONLY part in a user message is an unparseable input_document,
-    # the helper must NOT append an empty-content message to the outbound
-    # body (Anthropic 400s on "at least one block is required").
+    # If the only part is an unparseable input_document, the helper must not
+    # append an empty-content message (Anthropic 400s on "at least one block").
     captured = _capture(
         monkeypatch,
         provider = "anthropic",
@@ -192,14 +188,13 @@ def test_anthropic_empty_only_document_drops_whole_message(monkeypatch):
         ],
     )
     msgs = captured["body"]["messages"]
-    # The empty-content message must be skipped; only the second remains.
+    # Empty-content message skipped; only the second remains.
     assert len(msgs) == 1, msgs
 
 
 def test_anthropic_empty_data_uri_payload_is_dropped(monkeypatch):
-    # Codex P2: `data:application/pdf;base64,` with no payload (or
-    # whitespace-only) would create an empty `source.data` that
-    # Anthropic 400s on. Must be filtered before the wire.
+    # A `data:application/pdf;base64,` with empty/whitespace payload makes an
+    # empty `source.data` that Anthropic 400s on; filter it before the wire.
     captured = _capture(
         monkeypatch,
         provider = "anthropic",
@@ -228,12 +223,9 @@ def test_anthropic_empty_data_uri_payload_is_dropped(monkeypatch):
 
 
 def test_anthropic_empty_data_uri_falls_back_to_file_url(monkeypatch):
-    # Codex P2 follow-up: my previous fix added the empty-data-URI ->
-    # file_url fallback to the OpenAI side but missed the Anthropic
-    # side, where the empty-payload branch did `continue` and discarded
-    # an otherwise-valid file_url on the same part. Mirror the OpenAI
-    # behavior so a malformed inline payload + remote URL still
-    # attaches.
+    # The empty-data-URI -> file_url fallback existed on OpenAI but not
+    # Anthropic, which discarded a valid file_url on the same part. Mirror
+    # OpenAI so a malformed inline payload + remote URL still attaches.
     captured = _capture(
         monkeypatch,
         provider = "anthropic",
@@ -255,7 +247,7 @@ def test_anthropic_empty_data_uri_falls_back_to_file_url(monkeypatch):
     )
     parts = captured["body"]["messages"][0]["content"]
     doc = _strip_cache(next(p for p in parts if p.get("type") == "document"))
-    # base64 source MUST NOT have landed on the wire; URL source survived.
+    # base64 source MUST NOT reach the wire; URL source survives.
     assert doc == {
         "type": "document",
         "source": {"type": "url", "url": "https://example.com/doc.pdf"},
@@ -344,11 +336,9 @@ def test_openai_url_pdf_becomes_input_file(monkeypatch):
 
 
 def test_openai_empty_data_uri_falls_back_to_file_url(monkeypatch):
-    # Codex P2 follow-up: an empty `data:application/pdf;base64,`
-    # payload was being preferred over a perfectly valid `file_url`
-    # in the same part, sending `file_data=""` to OpenAI and 400ing
-    # the whole turn. The translator must treat empty data URIs as
-    # missing and recover via file_url.
+    # An empty `data:application/pdf;base64,` payload was preferred over a valid
+    # `file_url` in the same part, sending `file_data=""` and 400ing. The
+    # translator must treat empty data URIs as missing and recover via file_url.
     captured = _capture(
         monkeypatch,
         provider = "openai",
@@ -370,7 +360,7 @@ def test_openai_empty_data_uri_falls_back_to_file_url(monkeypatch):
     )
     parts = captured["body"]["input"][0]["content"]
     fileblk = next(p for p in parts if p.get("type") == "input_file")
-    # file_data MUST NOT be on the wire; file_url survives.
+    # file_data MUST NOT reach the wire; file_url survives.
     assert "file_data" not in fileblk, fileblk
     assert fileblk["file_url"] == "https://example.com/doc.pdf"
     assert fileblk["filename"] == "doc.pdf"
@@ -402,8 +392,8 @@ def test_openai_whitespace_only_data_uri_falls_back_to_file_url(monkeypatch):
 
 
 def test_openai_empty_data_uri_without_fallback_is_dropped(monkeypatch):
-    # If the only signal is an empty data URI (no file_url), the
-    # whole part is skipped rather than sent as `file_data=""`.
+    # Only signal is an empty data URI (no file_url): skip the whole part
+    # rather than send `file_data=""`.
     captured = _capture(
         monkeypatch,
         provider = "openai",
@@ -448,13 +438,9 @@ def test_openai_empty_document_part_is_dropped(monkeypatch):
 
 
 # ── Pydantic schema + builder pass-through ──────────────────────────
-#
-# The translation tests above call the external-provider client directly
-# with hand-built dicts, which bypasses BOTH ChatCompletionRequest's
-# discriminated Union AND routes/inference._build_external_messages. The
-# tests below close that gap: parse an input_document part through the
-# real request schema, run the builder, and assert the part survives to
-# the dict the client would receive.
+# The tests above call the client with hand-built dicts, bypassing the schema
+# and _build_external_messages. The tests below parse an input_document part
+# through the real schema + builder and assert it survives to the client dict.
 
 
 def test_chat_message_accepts_input_document_part():
@@ -482,10 +468,9 @@ def test_chat_message_accepts_input_document_part():
 
 
 def test_build_external_messages_passes_input_document_for_anthropic_and_openai():
-    # Both providers' stream helpers have explicit input_document
-    # translation logic (Anthropic -> {type:"document"}, OpenAI
-    # Responses -> {type:"input_file"}), so the part round-trips
-    # through the builder unchanged on those routes.
+    # Both providers' stream helpers translate input_document (Anthropic ->
+    # {type:"document"}, OpenAI Responses -> {type:"input_file"}), so the
+    # part round-trips through the builder unchanged on those routes.
     from models.inference import ChatMessage
     from routes.inference import _build_external_messages
 
@@ -518,10 +503,10 @@ def test_build_external_messages_passes_input_document_for_anthropic_and_openai(
 
 def test_build_external_messages_strips_input_document_for_unmapped_providers():
     # Codex P1 follow-up: gemini / mistral / kimi / openrouter / deepseek
-    # / custom go through generic /chat/completions passthrough that
-    # forwards `messages` verbatim. Handing them an `input_document`
-    # part fails the upstream validator. Builder must strip the part
-    # for every provider whose stream helper doesn't translate it.
+    # / custom use generic /chat/completions passthrough that forwards
+    # `messages` verbatim, so an `input_document` part fails the upstream
+    # validator. The builder must strip it for any provider whose stream
+    # helper doesn't translate it.
     from models.inference import ChatMessage
     from routes.inference import _build_external_messages
 
@@ -551,8 +536,8 @@ def test_build_external_messages_strips_input_document_for_unmapped_providers():
 
 
 def test_build_external_messages_strips_input_document_when_provider_type_unknown():
-    # Defensive: legacy callers that don't pass provider_type must
-    # not leak the part to an unknown destination.
+    # Defensive: legacy callers without provider_type must not leak the
+    # part to an unknown destination.
     from models.inference import ChatMessage
     from routes.inference import _build_external_messages
 
diff --git a/studio/backend/tests/test_native_context_length.py b/studio/backend/tests/test_native_context_length.py
index 01c05d3ec6..0290bb1308 100644
--- a/studio/backend/tests/test_native_context_length.py
+++ b/studio/backend/tests/test_native_context_length.py
@@ -3,11 +3,11 @@
 
 """Tests for the native_context_length feature (PR #4746).
 
-Verifies that the new `native_context_length` property on LlamaCppBackend
-and the corresponding Pydantic model fields work correctly.  The raw GGUF
-`_context_length` must never be overwritten by VRAM-capping logic.
+Verifies the `native_context_length` property on LlamaCppBackend and the
+matching Pydantic fields. The raw GGUF `_context_length` must never be
+overwritten by VRAM-capping logic.
 
-Requires no GPU, network, or external libraries beyond pytest and pydantic.
+Needs no GPU, network, or libraries beyond pytest and pydantic.
 """
 
 import io
@@ -21,8 +21,8 @@ from unittest.mock import patch
 import pytest
 
 # ---------------------------------------------------------------------------
-# Stub heavy / unavailable external dependencies before importing the
-# module under test.  Same pattern as test_kv_cache_estimation.py.
+# Stub heavy / unavailable deps before importing the module under test.
+# Same pattern as test_kv_cache_estimation.py.
 # ---------------------------------------------------------------------------
 
 _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
@@ -38,7 +38,7 @@ sys.modules.setdefault("loggers", _loggers_stub)
 _structlog_stub = _types.ModuleType("structlog")
 sys.modules.setdefault("structlog", _structlog_stub)
 
-# httpx -- stub only the names referenced at import / class-definition time
+# httpx -- stub only names referenced at import / class-definition time
 _httpx_stub = _types.ModuleType("httpx")
 for _exc_name in (
     "ConnectError",
@@ -227,7 +227,7 @@ class TestContextValueSeparation:
     def test_all_equal_when_uncapped(self, backend):
         """All three equal when no VRAM constraint."""
         backend._context_length = 8192
-        # No effective or max set -- properties fall back to _context_length
+        # No effective/max set -- properties fall back to _context_length.
         assert backend.native_context_length == 8192
         assert backend.max_context_length == 8192
         assert backend.context_length == 8192
@@ -241,16 +241,16 @@ class TestContextValueSeparation:
         backend._embedding_length = 4096
         original = backend._context_length
 
-        # Simulate a very small VRAM budget that forces capping
+        # Tiny VRAM budget forces capping.
         result = backend._fit_context_to_vram(
             requested_ctx = 131072,
             available_mib = 512,  # very small
             model_size_bytes = 0,
         )
-        # _fit_context_to_vram returns the capped value, not modifying _context_length
+        # Returns the capped value without modifying _context_length.
         assert backend._context_length == original
         assert backend.native_context_length == original
-        # The returned capped value should be <= requested
+        # Capped value must be <= requested.
         assert result <= 131072
 
     def test_native_gt_context_when_capped(self, backend):
@@ -370,7 +370,7 @@ class TestRouteCompleteness:
             start = self._source.find(f"{class_name}(", idx)
             if start == -1:
                 break
-            # Find matching closing paren (simple depth counter)
+            # Find the matching closing paren via a depth counter.
             depth = 0
             end = start
             for i, ch in enumerate(self._source[start:], start):
@@ -401,8 +401,8 @@ class TestRouteCompleteness:
         """Non-GGUF LoadResponse blocks do not set native_context_length (defaults to None)."""
         blocks = self._find_construction_blocks("LoadResponse")
         non_gguf = [b for b in blocks if "is_gguf = True" not in b and "is_gguf=True" not in b]
-        # Non-GGUF paths should not reference native_context_length
-        # (Pydantic defaults it to None, so not setting it is correct)
+        # Non-GGUF paths shouldn't reference native_context_length
+        # (Pydantic defaults it to None, so omitting it is correct).
         for block in non_gguf:
             assert (
                 "native_context_length" not in block
@@ -476,7 +476,7 @@ class TestNativeContextEdgeCases:
         backend._read_gguf_metadata(path)
         assert backend.native_context_length == 131072
 
-        # Simulate VRAM capping by setting effective and max
+        # Simulate VRAM capping via effective and max.
         backend._effective_context_length = 16384
         backend._max_context_length = 32768
         assert backend.native_context_length == 131072
diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py
index df7652a590..aab58adfff 100644
--- a/studio/backend/tests/test_offline_gguf_cache_fallback.py
+++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py
@@ -3,29 +3,14 @@
 
 """Regression tests for the offline GGUF cache fallback path (#5505).
 
-Three failure modes hit users when ``huggingface.co`` is unreachable
-but the requested GGUF repo is fully cached locally:
+When ``huggingface.co`` is unreachable but the repo is cached, three failures
+hit: ``list_gguf_variants`` 500'd (empty dropdown), ``detect_gguf_model_remote``
+returned None (GGUF-only repo misrouted), and ``_download_gguf`` synthesised a
+name absent from cache. Follow-ups: the cache filter matches the snapshot-relative
+path (subdir layouts findable), and DNS auto-detect scopes ``HF_HUB_OFFLINE`` to
+one load so a transient hiccup can't pin the singleton offline.
 
-* ``list_gguf_variants`` raised through ``HTTPException(500)`` so the
-  variant dropdown sat empty.
-* ``detect_gguf_model_remote`` returned ``None`` so a GGUF-only repo
-  was misrouted into the transformers/Unsloth backend (on macOS this
-  surfaced as a hardware error).
-* ``_download_gguf`` fell back to a synthetic ``{repo}-{variant}.gguf``
-  name that did not exist in cache when the in-repo filename did not
-  echo the repo name (e.g. ``unsloth/Qwen3.6-27B-MTP-GGUF`` ships
-  ``Qwen3.6-27B-UD-Q4_K_XL.gguf`` with no ``MTP`` token).
-
-Two follow-up regressions covered here:
-
-* P1 #1: the cache-side variant filter must match the snapshot-relative
-  path, not just the basename, so subdir layouts like
-  ``BF16/foo.gguf`` are findable.
-* P1 #2: the DNS auto-detect must scope ``HF_HUB_OFFLINE`` to one load
-  via try/finally so a transient resolver hiccup cannot lock the
-  long-lived ``LlamaCppBackend`` singleton offline forever.
-
-No GPU, no network, no subprocess. Linux, macOS, Windows compatible.
+No GPU, no network, no subprocess. Linux/macOS/Windows compatible.
 """
 
 from __future__ import annotations
@@ -44,8 +29,8 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
 if _BACKEND_DIR not in sys.path:
     sys.path.insert(0, _BACKEND_DIR)
 
-# Stub heavy/unavailable external deps before importing the modules
-# under test (same pattern as other studio backend tests).
+# Stub heavy/unavailable external deps before importing the modules under
+# test (same pattern as other studio backend tests).
 _loggers_stub = _types.ModuleType("loggers")
 _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
 sys.modules.setdefault("loggers", _loggers_stub)
@@ -182,7 +167,7 @@ class TestIterHfCacheSnapshots:
 
     def test_repo_id_match_is_case_insensitive(self, hf_cache):
         _build_cache(hf_cache, "unsloth/Foo-GGUF", {"Foo-Q4_K_M.gguf": 1})
-        # Lookup with a different casing of the org/name still resolves
+        # Lookup with different org/name casing still resolves
         out = list(_iter_hf_cache_snapshots("UNSLOTH/foo-gguf"))
         assert len(out) == 1
 
@@ -271,10 +256,8 @@ class TestDetectGgufFromCache:
         assert _detect_gguf_from_hf_cache("unsloth/a") == "a-UD-Q4_K_XL.gguf"
 
     def test_subdir_only_quant_resolves(self, hf_cache):
-        """P1 #1 regression: ``BF16/foo.gguf`` (quant only in directory).
-        Before the fix, the offline cache scan matched on basename and
-        missed this layout, falling through to the synthetic
-        ``{repo}-{variant}.gguf`` heuristic."""
+        """Regression: ``BF16/foo.gguf`` (quant only in directory). The pre-fix
+        cache scan matched on basename and missed this layout."""
         _build_cache(
             hf_cache,
             "unsloth/gpt-oss-20b-BF16",
@@ -316,7 +299,7 @@ class TestDetectGgufModelRemoteOffline:
         assert out == "a-Q4_K_M.gguf"
 
     def test_repository_not_found_does_not_consult_cache(self, hf_cache, clean_offline_env):
-        # Cache has a file but the API explicitly says repo is gone.
+        # Cache has a file but the API says the repo is gone.
         _build_cache(hf_cache, "unsloth/a", {"a-Q4_K_M.gguf": 1})
 
         class RepositoryNotFoundError(Exception):
@@ -442,9 +425,8 @@ class TestHfOfflineIfDnsDead:
 
 
 class TestExtractQuantLabelSubdir:
-    """``_extract_quant_label`` must consider the parent directories when
-    the basename has no quant token. Subdir layouts like ``BF16/foo.gguf``
-    are documented in this codebase and surface through the cache scan."""
+    """``_extract_quant_label`` must consider parent dirs when the basename has
+    no quant token (subdir layouts like ``BF16/foo.gguf``)."""
 
     def test_quant_in_basename_unchanged(self):
         assert _extract_quant_label("BF16/foo-BF16.gguf") == "BF16"
@@ -457,17 +439,13 @@ class TestExtractQuantLabelSubdir:
         assert _extract_quant_label("UD-Q4_K_XL/weight.gguf") == "UD-Q4_K_XL"
 
     def test_deeper_nesting_picks_nearest_quant_dir(self):
-        # When multiple parent segments could match, prefer the one closest
-        # to the file (innermost). This matches how repos like
-        # ``models/MXFP4_MOE/foo.gguf`` are laid out.
+        # Multiple matching parents: prefer the innermost (closest to the file).
         assert _extract_quant_label("models/MXFP4_MOE/foo.gguf") == "MXFP4_MOE"
 
 
 class TestDownloadMmprojOfflineCacheFallback:
-    """``LlamaCppBackend._download_mmproj`` must resolve cached mmproj
-    GGUFs offline, same shape as ``_download_gguf``. Without this the
-    offline vision GGUF load path returns ``None`` even when the mmproj
-    is present in cache."""
+    """``_download_mmproj`` must resolve cached mmproj GGUFs offline, like
+    ``_download_gguf``; else the offline vision load returns None despite a cache hit."""
 
     def test_cache_lookup_returns_cached_mmproj_when_list_repo_files_fails(self, hf_cache):
         _build_cache(
@@ -559,7 +537,7 @@ class TestDownloadMmprojOfflineCacheFallback:
 
 class TestListLocalGgufVariantsSubdir:
     """Subdir layouts like ``BF16/foo.gguf`` and ``Q4_K_M/foo.gguf`` must
-    produce distinct quant labels, not collapse on basename."""
+    yield distinct quant labels, not collapse on basename."""
 
     def test_two_subdir_variants_do_not_collapse(self, tmp_path):
         from utils.models.model_config import list_local_gguf_variants
@@ -642,8 +620,8 @@ class TestListGgufVariantsPermanentErrors:
 
 
 class TestDetectGgufFromCacheExcludesMmproj:
-    """A partial cache with only a vision projector must not route the
-    projector as the main model."""
+    """A partial cache with only a vision projector must not route it as
+    the main model."""
 
     def test_mmproj_only_returns_none(self, hf_cache):
         from utils.models.model_config import _detect_gguf_from_hf_cache
@@ -671,9 +649,8 @@ class TestDetectGgufFromCacheExcludesMmproj:
 
 
 class TestProbeDnsDeadNoGlobalTimeoutMutation:
-    """``_probe_dns_dead`` must not change ``socket.setdefaulttimeout``
-    process-wide -- concurrent sockets without explicit timeout would
-    inherit it for the probe window."""
+    """``_probe_dns_dead`` must not change ``socket.setdefaulttimeout`` process-wide;
+    concurrent sockets would inherit it during the probe window."""
 
     def test_default_timeout_unchanged_when_dns_up(self, monkeypatch):
         import socket as _socket
@@ -694,7 +671,7 @@ class TestProbeDnsDeadNoGlobalTimeoutMutation:
         try:
             _probe_dns_dead("example.invalid", timeout = 0.5)
         finally:
-            # Restore exact state regardless of any test-side mutation.
+            # Restore exact state regardless of test-side mutation.
             original_set(prev)
 
         assert set_calls == [], (
@@ -716,9 +693,8 @@ class TestProbeDnsDeadNoGlobalTimeoutMutation:
 
 
 class TestWaitForHealthRetriesOnReadError:
-    """A TCP RST mid-read while llama-server is still binding the port
-    (Windows: WinError 10054) must not abort the health-poll loop --
-    that masks a legitimate 'still warming up' state as a fatal load."""
+    """A TCP RST mid-read while llama-server is still binding (Windows: WinError
+    10054) must not abort the health-poll loop and mask warmup as a fatal load."""
 
     def test_read_error_then_success(self, monkeypatch):
         import httpx
diff --git a/studio/backend/tests/test_offline_inference_parent.py b/studio/backend/tests/test_offline_inference_parent.py
index fbb0aa8999..71331220d6 100644
--- a/studio/backend/tests/test_offline_inference_parent.py
+++ b/studio/backend/tests/test_offline_inference_parent.py
@@ -140,7 +140,7 @@ class TestLoraDetectOffline:
         monkeypatch.setenv("HF_HUB_OFFLINE", "1")
 
         # Studio catches Exception broadly; pin that the call still happens
-        # (so cached LoRAs aren't missed) and returns fast via mock.
+        # (so cached LoRAs aren't missed) and returns fast via the mock.
         class _OfflineModeIsEnabled(Exception):
             pass
 
diff --git a/studio/backend/tests/test_openai_citation_markers.py b/studio/backend/tests/test_openai_citation_markers.py
index ccc17be329..3549a5993e 100644
--- a/studio/backend/tests/test_openai_citation_markers.py
+++ b/studio/backend/tests/test_openai_citation_markers.py
@@ -4,8 +4,8 @@
 """Tests for the OpenAI Responses-API citation marker rewriter.
 
 The stream interleaves text deltas with ``\\ue200cite\\ue202SOURCE_ID\\ue201``
-markers. The rewriter resolves each to `[N](URL)` when the annotation has
-arrived and drops it otherwise; the URL list still flows to Sources via
+markers. The rewriter resolves each to `[N](URL)` once the annotation arrives
+and drops it otherwise; the URL list still flows to Sources via
 `_record_url_citation`.
 
 Reference: https://developers.openai.com/api/docs/guides/citation-formatting
@@ -58,8 +58,7 @@ def test_marker_rewritten_to_link_when_annotation_known():
 def test_unknown_source_marker_dropped_silently():
     text = f"Foo {_marker('turn9view9')} bar."
     out = _replace_openai_citation_markers(text, [])
-    # Marker stripped, no garbled "E202" glyph leaks through, and the
-    # surrounding text stays intact.
+    # Marker stripped, no garbled "E202" glyph leaks, surrounding text intact.
     assert not _has_marker_codepoints(out)
     assert "E202" not in out
     assert "turn9view9" not in out
@@ -67,8 +66,7 @@ def test_unknown_source_marker_dropped_silently():
 
 
 def test_multiple_concatenated_markers_resolved_in_order():
-    """Real-world wire shape: a string of markers butted up against each other
-    after a sentence, as in the user-reported bug."""
+    """Real-world wire shape: markers butted together after a sentence (user-reported bug)."""
     markers = "".join(_marker(f"turn{i}view{j}") for i, j in [(1, 0), (1, 1), (3, 0)])
     text = f"All animals ranked. {markers}"
     citations = [
@@ -136,8 +134,7 @@ def test_citation_without_source_id_does_not_crash(citation):
 
 
 def test_multiple_source_id_aliases_resolve_to_same_url():
-    """Every alias for the same URL must resolve, not just the first.
-    Regression for the Codex P1 on the original PR."""
+    """Every alias for the same URL must resolve, not just the first (Codex P1 regression)."""
     a = _marker("turn0view0")
     b = _marker("turn0view0_span_1")
     c = _marker("turn0view0_span_2")
@@ -150,15 +147,14 @@ def test_multiple_source_id_aliases_resolve_to_same_url():
         },
     ]
     out = _replace_openai_citation_markers(text, citations)
-    # All three aliases collapse onto citation [1] -- the URL is the
-    # same so it would be misleading to show three different numbers.
+    # All three aliases collapse onto citation [1] -- same URL, so showing
+    # three different numbers would mislead.
     assert out.count("[[1]](https://example.com/paris)") == 3
     assert not _has_marker_codepoints(out)
 
 
 def test_source_ids_list_and_legacy_source_id_both_resolve():
-    """Mixed-shape citation: legacy ``source_id`` plus newer
-    ``source_ids`` aliases both resolve."""
+    """Mixed-shape citation: legacy ``source_id`` plus newer ``source_ids`` aliases both resolve."""
     legacy = _marker("legacy_id")
     alias = _marker("alias_id")
     text = f"Both {legacy} and {alias} work."
@@ -174,11 +170,9 @@ def test_source_ids_list_and_legacy_source_id_both_resolve():
     assert not _has_marker_codepoints(out)
 
 
-# ---------------------------------------------------------------------------
 # _rewrite_citation_markers_partial: deferred-annotation tests. OpenAI emits
-# url_citation annotations on a subsequent SSE event; this helper reports
+# url_citation annotations on a later SSE event; this helper reports
 # `has_unresolved` so the stream loop defers emission. See PR #5713 audit.
-# ---------------------------------------------------------------------------
 
 
 def test_partial_known_marker_resolves_and_clears_unresolved():
@@ -202,7 +196,7 @@ def test_partial_unknown_marker_preserves_verbatim_and_flags():
 
 
 def test_partial_resolves_after_late_annotation():
-    """Two-pass: first call sees no citations, second resolves after annotation."""
+    """Two-pass: first call sees no citations; second resolves after annotation."""
     text = f"See {_marker('s1')} for details."
     out1, unresolved1 = _rewrite_citation_markers_partial(text, [])
     assert unresolved1 is True
@@ -214,17 +208,15 @@ def test_partial_resolves_after_late_annotation():
 
 
 def test_partial_multi_source_partial_resolution_keeps_marker_pending():
-    """Any unresolved token in a multi-source marker leaves the whole marker
-    verbatim with ``unresolved`` True; defer until every id resolves or
-    end-of-stream forces a flush (dropping unresolved tokens then)."""
+    """Any unresolved token in a multi-source marker leaves the whole marker verbatim with ``unresolved`` True until every id resolves or end-of-stream flushes."""
     cite = f"{CITE_START}cite{CITE_DELIM}known{CITE_DELIM}locator{CITE_STOP}"
     text = f"Pre {cite} post."
     citations = [{"source_id": "known", "url": "https://example.com/y"}]
     out, unresolved = _rewrite_citation_markers_partial(text, citations)
     assert unresolved is True
     assert cite in out
-    # End-of-stream force flush: drop the unresolved token, keep the
-    # resolved link. The streamer routes pending segments through
+    # End-of-stream force flush: drop the unresolved token, keep the resolved
+    # link. The streamer routes pending segments through
     # `_replace_openai_citation_markers` at force=True for this.
     forced = _replace_openai_citation_markers(out, citations)
     assert "[[1]](https://example.com/y)" in forced
diff --git a/studio/backend/tests/test_openai_citation_markers_edge.py b/studio/backend/tests/test_openai_citation_markers_edge.py
index e8d0be6246..f44975ce33 100644
--- a/studio/backend/tests/test_openai_citation_markers_edge.py
+++ b/studio/backend/tests/test_openai_citation_markers_edge.py
@@ -13,8 +13,8 @@ Reference: https://developers.openai.com/api/docs/guides/citation-formatting
 import importlib
 
 
-# Streaming integration is exercised by ``_simulate_delta_stream`` further
-# down, mirroring the head/buffer/flush dance from ``_stream_openai_responses``.
+# Streaming is exercised by ``_simulate_delta_stream`` below, mirroring the
+# head/buffer/flush dance from ``_stream_openai_responses``.
 _module = importlib.import_module("core.inference.external_provider")
 _replace_openai_citation_markers = _module._replace_openai_citation_markers
 _split_pending_citation_tail = _module._split_pending_citation_tail
@@ -27,7 +27,7 @@ CITE_DELIM = ""
 
 def _marker(*source_ids: str, locator: str | None = None) -> str:
     """Build a ``\\ue200cite\\ue202[\\ue202...][\\ue202]\\ue201``
-    marker. Accepts one or many ``source_ids`` plus an optional ``locator``."""
+    marker from one or more ``source_ids`` plus an optional ``locator``."""
     payload = f"{CITE_START}cite{CITE_DELIM}" + CITE_DELIM.join(source_ids)
     if locator:
         payload = f"{payload}{CITE_DELIM}{locator}"
@@ -39,7 +39,7 @@ def _no_private_use(text: str) -> bool:
 
 
 # Harness mirroring the head/pending-tail/flush dance in
-# `_stream_openai_responses`, so streaming tests skip the httpx mock.
+# `_stream_openai_responses` so streaming tests skip the httpx mock.
 def _simulate_delta_stream(
     deltas: list[str],
     citations: list[dict],
@@ -56,8 +56,8 @@ def _simulate_delta_stream(
             if head:
                 emitted.append(head)
     if flush and pending:
-        # Mirror `_flush_pending_marker_tail`: drop the tail entirely if no
-        # closing stop byte arrived; the literal ``cite`` would leak otherwise.
+        # Mirror `_flush_pending_marker_tail`: drop the tail if no closing stop
+        # byte arrived; the literal ``cite`` would leak otherwise.
         if CITE_STOP not in pending:
             rendered = ""
         else:
@@ -79,7 +79,7 @@ def _simulate_delta_stream(
 
 def test_multi_source_marker_all_resolve():
     """\\ue200cite\\ue202id1\\ue202id2\\ue202id3\\ue201 expands to three links
-    when every id is known. Earlier regex captured only id1 and dropped id2/id3."""
+    when every id is known. Earlier regex captured only id1."""
     text = f"All three: {_marker('id1', 'id2', 'id3')}"
     citations = [
         {"source_id": "id1", "url": "https://example.com/1"},
@@ -136,14 +136,14 @@ def test_marker_with_range_locator():
 
 
 def test_marker_split_in_source_id():
-    """Delta-1 ends mid-source-id (``\\ue200cite\\ue202tu``), delta-2 starts
-    with the rest (``rn0view0\\ue201``). The buffer stitches the halves
-    back together so they resolve to one link instead of leaking."""
+    """Delta-1 ends mid-source-id (``\\ue200cite\\ue202tu``), delta-2 has the
+    rest (``rn0view0\\ue201``). The buffer stitches the halves so they resolve
+    to one link instead of leaking."""
     full = f"See {_marker('turn0view0')} now."
     # Cut right after the second delim + "tu" inside the source id.
     cut = full.index("tu", full.index(CITE_START)) + len("tu")
     d1, d2 = full[:cut], full[cut:]
-    # Sanity check: delta-1 actually contains a partial marker.
+    # Sanity: delta-1 contains a partial marker.
     assert CITE_START in d1 and CITE_STOP not in d1
     assert CITE_STOP in d2
     citations = [{"source_id": "turn0view0", "url": "https://x"}]
@@ -153,8 +153,8 @@ def test_marker_split_in_source_id():
 
 
 def test_marker_split_at_start_byte():
-    """Split exactly after the opening ``\\ue200`` byte; the buffer must
-    hold the lone open byte until the rest arrives."""
+    """Split right after the opening ``\\ue200`` byte; the buffer must hold the
+    lone open byte until the rest arrives."""
     full = f"Text {_marker('sid')} done"
     cut = full.index(CITE_START) + 1  # right AFTER the open byte
     d1, d2 = full[:cut], full[cut:]
@@ -167,7 +167,7 @@ def test_marker_split_at_start_byte():
 def test_marker_split_across_three_deltas():
     """Worst case: marker chopped into three pieces across three deltas."""
     full = f"A {_marker('threesplit')} B"
-    # cut at two points inside the marker
+    # Cut at two points inside the marker.
     open_pos = full.index(CITE_START)
     stop_pos = full.index(CITE_STOP)
     cut1 = open_pos + 4
@@ -206,21 +206,20 @@ def test_split_marker_unknown_source_is_dropped_cleanly():
 
 
 def test_unterminated_marker_at_stream_end_dropped_on_flush():
-    """Stream ends mid-marker (e.g. response.incomplete); the tail is
-    flushed with private-use bytes stripped, no `E202` text leaks."""
+    """Stream ends mid-marker (e.g. response.incomplete); the flushed tail
+    strips private-use bytes, no `E202` text leaks."""
     deltas = ["Some text ", f"{CITE_START}citetu", "rn0view0"]  # no STOP ever
     out = _simulate_delta_stream(deltas, [], flush = True)
     assert _no_private_use(out)
     assert "E200" not in out and "E202" not in out
-    # Surrounding prose stays; we don't assert exact marker remainder.
+    # Surrounding prose stays; don't assert exact marker remainder.
     assert "Some text " in out
 
 
 def test_flush_resolves_marker_when_late_annotation_arrives():
-    """Marker in a delta, matching annotation arrives later (on
-    response.output_text.annotation.added after the final delta). The
-    rewriter reads ``all_url_citations`` LIVE at flush, so the buffered
-    marker still resolves."""
+    """Marker in a delta; the matching annotation arrives later (on
+    response.output_text.annotation.added after the final delta). The rewriter
+    reads ``all_url_citations`` LIVE at flush, so the buffered marker resolves."""
     deltas = ["Look ", f"{CITE_START}cite{CITE_DELIM}late_sid"]
     pending = ""
     citations: list[dict] = []
@@ -291,8 +290,8 @@ def test_rewriter_idempotent_on_marker_free_text():
 
 
 def test_only_marker_no_surrounding_text():
-    """A delta that is JUST a marker (no prose) still renders correctly;
-    used to leak without the empty-string short-circuit in the split helper."""
+    """A delta that is JUST a marker (no prose) renders correctly; it leaked
+    before the empty-string short-circuit in the split helper."""
     text = _marker("solo")
     citations = [{"source_id": "solo", "url": "https://solo.example"}]
     out = _replace_openai_citation_markers(text, citations)
@@ -311,15 +310,15 @@ def test_back_to_back_markers_with_no_separator():
 
 
 def test_split_helper_buffers_only_after_last_open_byte():
-    """A complete marker followed by an unterminated one: head includes
-    the complete marker, buffer holds only the trailing partial."""
+    """Complete marker followed by an unterminated one: head includes the
+    complete marker, buffer holds only the trailing partial."""
     complete = _marker("done")
     partial = f"{CITE_START}cite{CITE_DELIM}half"  # no STOP
     text = f"pre {complete} mid {partial}"
     head, tail = _split_pending_citation_tail(text)
     assert head == f"pre {complete} mid "
     assert tail == partial
-    # And the head, once rewritten, drops every private-use byte.
+    # Head, once rewritten, drops every private-use byte.
     rewritten = _replace_openai_citation_markers(head, [{"source_id": "done", "url": "https://d"}])
     assert rewritten == "pre [[1]](https://d) mid "
 
@@ -355,7 +354,7 @@ def test_unknown_marker_does_not_perturb_citation_indexing():
         {"source_id": "real_b", "url": "https://example.com/b"},
     ]
     out = _replace_openai_citation_markers(text, citations)
-    # real_a is index 1; unknown does not take a slot.
+    # real_a is index 1; unknown takes no slot.
     assert "[[1]](https://example.com/a)" in out
     assert "[[2]](https://example.com/b)" in out
     assert _no_private_use(out)
@@ -368,8 +367,8 @@ def test_unknown_marker_does_not_perturb_citation_indexing():
 
 
 def test_unterminated_marker_does_not_leak_cite_residue():
-    """Stream ends mid-marker: drop the whole tail rather than strip
-    codepoints and leave ``cite`` behind."""
+    """Stream ends mid-marker: drop the whole tail rather than strip codepoints
+    and leave ``cite`` behind."""
     half = f"Hi there {CITE_START}cite{CITE_DELIM}turn0view0"
     out = _simulate_delta_stream([half], [], flush = True)
     # Prose before the marker stays; no private-use bytes or cite residue.
diff --git a/studio/backend/tests/test_openai_code_execution.py b/studio/backend/tests/test_openai_code_execution.py
index d4f814856b..63e94613ed 100644
--- a/studio/backend/tests/test_openai_code_execution.py
+++ b/studio/backend/tests/test_openai_code_execution.py
@@ -1,30 +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
 
-"""
-Unit tests for OpenAI's server-side `shell` tool translation in
+"""Unit tests for OpenAI's server-side `shell` tool translation in
 `_stream_openai_responses`.
 
-Covers:
-- Request body: ``enabled_tools=["code_execution"]`` on the OpenAI
-  cloud base_url appends ``{"type": "shell", "environment": {"type":
-  "container_auto"}}`` to ``tools``.
-- Container reuse: when ``openai_code_exec_container_id`` is provided,
-  the outgoing ``environment.type`` flips to ``"container_reference"``
-  and the id propagates.
-- Cloud guard: code_execution on a non-cloud base_url (e.g. a local
-  OpenAI-compat preset / ollama / llama.cpp / vLLM) does NOT add the
-  shell tool, preventing a guaranteed 400 from those servers.
-- SSE translation: a `shell_call` + `shell_call_output` pair emits one
-  ``_toolEvent`` `tool_start` (`tool_name="code_execution"`,
-  `arguments.kind="bash"`) and one `tool_end` whose `result` contains
-  the joined stdout from the shell_call_output entries.
-- Container surfacing: container_id captured from
-  `response.completed.container_id` is emitted as a synthetic
-  `container_ready` `_toolEvent` (only when it differs from the
-  inbound id).
-- Stale-container handling: 400 with "container expired" body emits a
-  `container_invalidated` event before propagating the error.
+Covers: request body shaping (container_auto, container_reference), the cloud
+guard (no shell tool on non-cloud base_urls), SSE translation of a
+shell_call/shell_call_output pair into tool_start/tool_end events, container_id
+surfacing as container_ready, and stale-container invalidation.
 """
 
 import asyncio
@@ -192,8 +175,7 @@ def test_shell_tool_refused_for_non_cloud_base_url(monkeypatch):
     _drive(run())
 
     tools = captured["body"].get("tools") or []
-    # Shell tool must NOT leak to local OpenAI-compat servers — those
-    # 400 on the unknown tool type.
+    # Shell tool must not leak to local OpenAI-compat servers (they 400 on it).
     assert all(t.get("type") != "shell" for t in tools)
 
 
@@ -266,10 +248,8 @@ def test_shell_call_emits_tool_start_and_end(monkeypatch):
     assert len(ends) == 1
     assert starts[0]["tool_name"] == "code_execution"
     assert starts[0]["tool_call_id"] == "scall_1"
-    # `_server_tool: True` is the synthetic-builtin marker the
-    # backend stamps onto every provider-side tool_start so the
-    # frontend serializer can distinguish hosted tools from
-    # user-declared functions on history replay.
+    # `_server_tool: True` marks a synthetic builtin so the frontend can tell
+    # hosted tools from user-declared functions on history replay.
     assert starts[0]["arguments"] == {"kind": "bash", "command": "ls -la", "_server_tool": True}
     assert ends[0]["tool_call_id"] == "scall_1"
     assert "total 24" in ends[0]["result"]
@@ -393,24 +373,22 @@ def test_stale_container_emits_invalidated(monkeypatch):
 
 
 def test_expired_container_triggers_transparent_retry(monkeypatch):
-    """When OpenAI 400s with 'Container is expired' on a request that
-    carried container_reference, the streamer retries once with the
-    container field stripped. The user never sees an error line — only
-    container_invalidated, then the normal stream from the retry.
+    """On a 'Container is expired' 400 for a container_reference request, the
+    streamer retries once with the container stripped; the user sees only
+    container_invalidated then the retry stream, never an error line.
     """
     calls: list[dict] = []
 
     def handler(request: httpx.Request) -> httpx.Response:
         body = json.loads(request.content.decode("utf-8"))
         calls.append(body)
-        # Find the shell tool entry to inspect environment.type.
+        # Inspect the shell tool's environment.type.
         shell_env_type = None
         for tool in body.get("tools", []) or []:
             if tool.get("type") == "shell":
                 shell_env_type = tool.get("environment", {}).get("type")
                 break
-        # First call carries container_reference -> 400 expired.
-        # Retry omits container -> normal SSE stream.
+        # container_reference -> 400 expired; retry omits container -> normal stream.
         if shell_env_type == "container_reference":
             return httpx.Response(
                 400,
@@ -424,8 +402,8 @@ def test_expired_container_triggers_transparent_retry(monkeypatch):
                 ).encode("utf-8"),
                 headers = {"content-type": "application/json"},
             )
-        # Successful retry: minimal SSE — a completed response with a
-        # fresh container_id so container_ready latches.
+        # Successful retry: minimal SSE — completed response with a fresh
+        # container_id so container_ready latches.
         sse = _openai_sse(
             [
                 {
@@ -461,8 +439,8 @@ def test_expired_container_triggers_transparent_retry(monkeypatch):
     lines = _drive(run())
     events = _tool_events(lines)
 
-    # Two outbound HTTP calls were made: the expired-container attempt
-    # then the retry without the container field.
+    # Two outbound calls: the expired-container attempt, then the retry
+    # without the container field.
     assert len(calls) == 2
     shell_types = []
     for body in calls:
@@ -471,14 +449,14 @@ def test_expired_container_triggers_transparent_retry(monkeypatch):
                 shell_types.append(tool.get("environment", {}).get("type"))
     assert shell_types == ["container_reference", "container_auto"]
 
-    # container_invalidated emitted (frontend will null its stored id).
+    # container_invalidated emitted (frontend nulls its stored id).
     assert any(e.get("type") == "container_invalidated" for e in events)
     # container_ready emitted from the retry stream with the fresh id.
     assert any(
         e.get("type") == "container_ready" and e.get("container_id") == "cntr_fresh_111"
         for e in events
     )
-    # CRUCIALLY: no SSE error line surfaced to the chat — only completion.
+    # CRUCIALLY: no SSE error line surfaced to the chat.
     error_lines = [
         line
         for line in lines
@@ -488,8 +466,8 @@ def test_expired_container_triggers_transparent_retry(monkeypatch):
 
 
 def test_expired_container_retries_only_once(monkeypatch):
-    """If the retry ALSO fails (any 4xx, expired or otherwise), the
-    error is surfaced normally — no infinite retry loop.
+    """If the retry ALSO fails (any 4xx), the error surfaces normally —
+    no infinite retry loop.
     """
     call_count = {"n": 0}
 
@@ -528,8 +506,7 @@ def test_expired_container_retries_only_once(monkeypatch):
 
     lines = _drive(run())
 
-    # Exactly two calls (first + one retry). Third would mean an
-    # infinite loop.
+    # Exactly two calls (first + one retry); a third would be a loop.
     assert call_count["n"] == 2
     # The second failure surfaces normally as an error SSE line.
     error_lines = [line for line in lines if '"error"' in line and "_toolEvent" not in line]
diff --git a/studio/backend/tests/test_openai_compaction.py b/studio/backend/tests/test_openai_compaction.py
index f0599952c6..c7de0a9aed 100644
--- a/studio/backend/tests/test_openai_compaction.py
+++ b/studio/backend/tests/test_openai_compaction.py
@@ -4,14 +4,13 @@
 """Unit tests for OpenAI Responses API context_management wiring.
 
 OpenAI's Responses API supports server-side compaction via
-``context_management: [{type:"compaction", compact_threshold:N}]``.
-There is no beta header and no dated version pin; the threshold is
-silently accepted and the API runs the compaction step when the
-rendered prompt crosses it.
+``context_management: [{type:"compaction", compact_threshold:N}]``. No
+beta header, no dated version pin; the threshold is silently accepted and
+compaction runs when the rendered prompt crosses it.
 
-These tests pin: the body shape when threshold is set on cloud OpenAI,
-the silent no-op when the base URL is non-cloud, and the
-omitted-threshold pass-through.
+These pin: the body shape when threshold is set on cloud OpenAI, the
+silent no-op on non-cloud base URLs, and the omitted-threshold
+pass-through.
 """
 
 import asyncio
@@ -32,8 +31,7 @@ def _capture(monkeypatch, *, base_url: str, threshold) -> dict:
 
     def handler(request: httpx.Request) -> httpx.Response:
         captured["body"] = json.loads(request.content.decode("utf-8"))
-        # Send an empty Responses-shaped SSE stream so the helper exits
-        # cleanly.
+        # Empty Responses-shaped SSE stream so the helper exits cleanly.
         return httpx.Response(
             200,
             content = (
@@ -88,8 +86,8 @@ def test_cloud_openai_sets_compaction_block(monkeypatch):
 
 
 def test_cloud_openai_below_default_threshold_passes_through(monkeypatch):
-    # Studio doesn't clamp the OpenAI side -- the API accepts whatever
-    # the caller sends, so a small probe like 60k still goes through.
+    # Studio doesn't clamp the OpenAI side -- the API accepts whatever the
+    # caller sends, so a small probe like 60k still goes through.
     captured = _capture(
         monkeypatch,
         base_url = "https://api.openai.com/v1",
@@ -105,8 +103,8 @@ def test_cloud_openai_below_default_threshold_passes_through(monkeypatch):
 
 def test_non_cloud_base_silently_drops_compaction(monkeypatch):
     # ollama / llama.cpp / "custom" presets collapse to provider="openai"
-    # but don't implement context_management. Sending the field would
-    # 400 those servers, so it must NOT appear on the wire.
+    # but lack context_management. Sending the field would 400 them, so it
+    # must NOT appear on the wire.
     captured = _capture(
         monkeypatch,
         base_url = "http://127.0.0.1:11434/v1",
@@ -120,9 +118,9 @@ def test_non_cloud_base_silently_drops_compaction(monkeypatch):
 
 def test_azure_openai_base_url_carries_compaction_block(monkeypatch):
     # Azure OpenAI Foundry exposes the same /v1/responses extensions
-    # (context_management, prompt_cache_retention, container shell)
-    # under a *.openai.azure.com base URL. Treat it as cloud so the
-    # compaction field actually reaches the API.
+    # (context_management, prompt_cache_retention, container shell) under
+    # a *.openai.azure.com base URL. Treat it as cloud so the compaction
+    # field reaches the API.
     captured = _capture(
         monkeypatch,
         base_url = "https://my-resource.openai.azure.com/openai/v1",
@@ -132,14 +130,14 @@ def test_azure_openai_base_url_carries_compaction_block(monkeypatch):
         {"type": "compaction", "compact_threshold": 200_000}
     ]
     # Sibling Azure-cloud extension: prompt_cache_retention should also
-    # be set so caching works the same way on Azure deployments.
+    # be set so caching works the same on Azure deployments.
     assert captured["body"].get("prompt_cache_retention") == "24h"
 
 
 def test_azure_openai_mixed_case_base_url_matches(monkeypatch):
-    # The match is case-insensitive so URLs copy-pasted from the Azure
-    # portal (which sometimes capitalise the resource name) still get
-    # the cloud-only fields.
+    # Case-insensitive match so URLs copy-pasted from the Azure portal
+    # (which sometimes capitalise the resource name) still get the
+    # cloud-only fields.
     captured = _capture(
         monkeypatch,
         base_url = "https://My-Resource.OpenAI.Azure.Com/openai/v1",
@@ -151,12 +149,11 @@ def test_azure_openai_mixed_case_base_url_matches(monkeypatch):
 
 
 def test_cloud_gate_uses_hostname_not_substring(monkeypatch):
-    # CodeQL py/incomplete-url-substring-sanitization: an attacker who
-    # controls the configured base_url could embed `api.openai.com` or
-    # `.openai.azure.com` as part of a path or a subdomain on an
-    # arbitrary host to slip the cloud-only request body fields to a
-    # server they control. The hostname-anchored helper must reject
-    # both shapes.
+    # CodeQL py/incomplete-url-substring-sanitization: an attacker
+    # controlling base_url could embed `api.openai.com` or
+    # `.openai.azure.com` in a path or subdomain on an arbitrary host to
+    # slip cloud-only body fields to their own server. The
+    # hostname-anchored helper must reject both shapes.
     for evil in [
         "https://evil.com/api.openai.com/v1",
         "https://api.openai.com.attacker.com/v1",
@@ -188,19 +185,17 @@ def test_omitted_threshold_no_body_field(monkeypatch):
 
 
 def test_chat_completion_request_accepts_any_positive_compaction_threshold():
-    # Codex follow-up: the field is documented as a no-op for non-cloud
-    # OpenAI bases and every non-OpenAI provider, so a cross-provider
-    # schema floor would 422 perfectly valid Anthropic / ollama /
-    # llama.cpp requests that happen to carry the field. Keep schema
-    # floor at ge=1 (any positive int) and rely on per-provider
-    # helpers (_stream_openai_responses / _stream_anthropic) to
-    # enforce or clamp the real floor.
+    # Codex follow-up: the field is a no-op for non-cloud OpenAI bases and
+    # every non-OpenAI provider, so a cross-provider schema floor would
+    # 422 valid Anthropic / ollama / llama.cpp requests carrying it. Keep
+    # the schema floor at ge=1 (any positive int) and let per-provider
+    # helpers (_stream_openai_responses / _stream_anthropic) enforce or
+    # clamp the real floor.
     import pytest as _pytest
 
     from models.inference import ChatCompletionRequest
 
-    # Non-positive values still rejected so blank-string posts don't
-    # sneak through.
+    # Non-positive values rejected so blank-string posts don't sneak in.
     with _pytest.raises(Exception):
         ChatCompletionRequest.model_validate(
             {
@@ -211,10 +206,10 @@ def test_chat_completion_request_accepts_any_positive_compaction_threshold():
         )
 
     # Any positive int passes schema validation, including values that
-    # would be no-ops on the OpenAI cloud path. This is intentional --
-    # the OpenAI helper drops the field on non-cloud bases and
-    # forwards-as-is on cloud bases; if the value is below the model's
-    # effective floor, the upstream API surfaces the error.
+    # are no-ops on the OpenAI cloud path. Intentional -- the OpenAI
+    # helper drops the field on non-cloud bases and forwards as-is on
+    # cloud bases; if it's below the model's effective floor, the upstream
+    # API surfaces the error.
     for v in (1, 5_000, 9_999, 10_000, 200_000):
         req = ChatCompletionRequest.model_validate(
             {
diff --git a/studio/backend/tests/test_openai_container_crud.py b/studio/backend/tests/test_openai_container_crud.py
index 48acf97e1f..e0604527fd 100644
--- a/studio/backend/tests/test_openai_container_crud.py
+++ b/studio/backend/tests/test_openai_container_crud.py
@@ -4,11 +4,11 @@
 """Unit tests for the /v1/containers CRUD client methods.
 
 Covers:
-- All three calls (list / create / delete) send
-  ``OpenAI-Beta: containers=v1``. Without it, OpenAI silently no-ops
-  the DELETE while still returning 200 ``{"deleted": true}``.
-- ``delete_openai_container`` raises when the response body does not
-  report ``{"deleted": true}``, even on a 2xx response.
+- list / create / delete all send ``OpenAI-Beta: containers=v1``. Without
+  it, OpenAI silently no-ops the DELETE but still returns 200
+  ``{"deleted": true}``.
+- ``delete_openai_container`` raises when the body omits
+  ``{"deleted": true}``, even on a 2xx response.
 """
 
 from __future__ import annotations
@@ -28,11 +28,10 @@ def _drive(coro):
 
 
 def _mock_http_client(monkeypatch, handler):
-    """Wire `handler` for both the shared `_http_client` AND any
-    per-call `httpx.AsyncClient(...)` instances. delete_openai_container
-    intentionally creates a fresh AsyncClient (see comment in
-    external_provider.delete_openai_container) so the test must
-    also intercept that constructor."""
+    """Wire `handler` for the shared `_http_client` AND any per-call
+    `httpx.AsyncClient(...)`. delete_openai_container creates a fresh
+    AsyncClient (see external_provider.delete_openai_container), so we
+    must also intercept that constructor."""
     transport = httpx.MockTransport(handler)
     monkeypatch.setattr(ep_mod, "_http_client", httpx.AsyncClient(transport = transport))
     real_async_client = httpx.AsyncClient
@@ -110,13 +109,12 @@ def test_delete_sends_openai_beta_header_and_accepts_confirmation(monkeypatch):
 
 def test_delete_raises_when_response_lacks_deleted_true(monkeypatch):
     """OpenAI returns 200 ``{"deleted": true}`` even when the request is
-    silently rejected (e.g. before we started sending OpenAI-Beta).
-    Defensive guard: when the body omits ``deleted: true``, surface it
-    as an error so the UI can report the failure instead of falsely
-    reporting success."""
+    silently rejected (e.g. before we sent OpenAI-Beta). Guard: when the
+    body omits ``deleted: true``, surface an error so the UI reports the
+    failure instead of false success."""
 
     def handler(request: httpx.Request) -> httpx.Response:
-        # 200 but no deleted flag — simulate an unexpected payload shape.
+        # 200 but no deleted flag — unexpected payload shape.
         return httpx.Response(200, json = {"id": "cntr_x", "object": "container"})
 
     _mock_http_client(monkeypatch, handler)
@@ -159,10 +157,9 @@ def test_delete_propagates_openai_4xx(monkeypatch):
 
 
 def test_list_route_filters_expired_containers(monkeypatch):
-    """OpenAI keeps containers in /v1/containers indefinitely with
-    status="expired" after their idle TTL passes — they can't be
-    used but still show up. The list route must drop them so the
-    picker only surfaces usable containers."""
+    """OpenAI keeps containers in /v1/containers with status="expired"
+    after their idle TTL passes — unusable but still listed. The list
+    route must drop them so the picker shows only usable containers."""
     from routes import inference as inf_mod
     from models.inference import OpenAIContainerRequest
 
diff --git a/studio/backend/tests/test_openai_image_generation.py b/studio/backend/tests/test_openai_image_generation.py
index f5dee2561a..ace57588d3 100644
--- a/studio/backend/tests/test_openai_image_generation.py
+++ b/studio/backend/tests/test_openai_image_generation.py
@@ -3,18 +3,11 @@
 
 """Unit tests for OpenAI Responses API image_generation tool wiring.
 
-The image_generation tool is a server-side Responses-API tool:
-``{type: "image_generation"}`` in the request's tools array, and the
-result comes back as an ``image_generation_call`` output item carrying
-the base64 image on ``result``. Studio translates the output item
-into ``_toolEvent`` chunks (``tool_start`` with `kind:"image"`,
-``tool_end`` with ``image_b64`` + ``image_mime``) so the chat adapter
-can render the image inline.
-
-These tests pin: the tool is added to the outbound body only when the
-caller asks for it on a cloud OpenAI base; the SSE output_item.done
-for ``image_generation_call`` produces the expected _toolEvent chunks;
-non-cloud bases drop the tool silently.
+The tool is a server-side Responses-API tool (``{type: "image_generation"}``);
+the result comes back as an ``image_generation_call`` output item, which Studio
+translates into ``_toolEvent`` chunks so the chat adapter renders it inline.
+Tests pin: the tool is added to the body only on a cloud OpenAI base when asked
+for, the done event produces the expected chunks, and non-cloud bases drop it.
 """
 
 import asyncio
@@ -75,8 +68,8 @@ def _capture_body(monkeypatch, *, base_url: str, enabled_tools) -> dict:
 
 
 def _collect_tool_events(monkeypatch) -> list[dict]:
-    """Drive a Responses stream that emits one image_generation_call done
-    event and return the parsed _toolEvent chunks."""
+    """Drive a Responses stream with one image_generation_call done event and
+    return the parsed _toolEvent chunks."""
 
     sse = (
         b"event: response.output_item.done\n"
@@ -207,8 +200,8 @@ def test_image_generation_done_emits_tool_event_chunks(monkeypatch):
     ends = [e for e in image_events if e.get("type") == "tool_end"]
     assert len(starts) == 1, image_events
     assert len(ends) == 1, image_events
-    # `_server_tool: True` marks this as a provider-side synthetic
-    # tool card on the frontend's history serializer.
+    # `_server_tool: True` marks this as a provider-side synthetic tool card
+    # for the frontend's history serializer.
     assert starts[0]["arguments"] == {
         "kind": "image",
         "prompt": "A photorealistic cat sitting",
diff --git a/studio/backend/tests/test_openai_responses_translation.py b/studio/backend/tests/test_openai_responses_translation.py
index 95e9b2d63c..f7d7e83a43 100644
--- a/studio/backend/tests/test_openai_responses_translation.py
+++ b/studio/backend/tests/test_openai_responses_translation.py
@@ -5,14 +5,14 @@
 Unit tests for the OpenAI `/v1/responses` translation in external_provider.
 
 Covers:
-- Request body shape: system messages collapse into `instructions`, user/
-  assistant messages go into `input`, sampling knobs Responses does not
-  support (presence_penalty, top_k) are not forwarded.
-- SSE translation: `response.output_text.delta` events become OpenAI Chat
-  Completions chunks, `response.completed` emits a `finish_reason: stop`
-  chunk, the stream terminates with `data: [DONE]`.
-- Image parts in user content are rewritten from Chat Completions
-  `{type: image_url, image_url: {url}}` into Responses
+- Request body shape: system messages collapse into `instructions`,
+  user/assistant messages go into `input`, and unsupported sampling knobs
+  (presence_penalty, top_k) are not forwarded.
+- SSE translation: `response.output_text.delta` → Chat Completions chunks,
+  `response.completed` → a `finish_reason: stop` chunk, stream ends with
+  `data: [DONE]`.
+- Image parts rewritten from Chat Completions
+  `{type: image_url, image_url: {url}}` to Responses
   `{type: input_image, image_url: }`.
 """
 
@@ -101,9 +101,9 @@ def test_responses_request_body_uses_input_and_instructions(monkeypatch):
     assert body["input"] == [{"role": "user", "content": "Hi"}]
     assert body["max_output_tokens"] == 512
     assert body["stream"] is True
-    # Responses API on reasoning-class models (gpt-5.x / o3 / gpt-4.5 — the
-    # only OpenAI ids the registry allowlist exposes) rejects these as
-    # `Unsupported parameter`. Make sure we never silently forward them.
+    # Responses API on reasoning-class models (gpt-5.x / o3 / gpt-4.5 — the only
+    # OpenAI ids the registry allowlist exposes) rejects these as `Unsupported
+    # parameter`. Never silently forward them.
     assert "temperature" not in body
     assert "top_p" not in body
     assert "presence_penalty" not in body
@@ -193,7 +193,7 @@ def test_responses_sse_translates_to_chat_completions_chunks(monkeypatch):
 
     lines = _drive(run())
 
-    # Drop empty / non-data lines for assertion clarity.
+    # Keep only data lines for assertion clarity.
     data_lines = [line for line in lines if line.startswith("data:")]
     payloads = []
     for line in data_lines:
@@ -213,11 +213,11 @@ def test_responses_sse_translates_to_chat_completions_chunks(monkeypatch):
 
 
 def test_responses_function_call_output_translates_to_delta_tool_calls(monkeypatch):
-    """Round 12: caller-supplied function tools forwarded into /v1/responses
-    must have their `function_call` output items translated back into Chat
-    Completions delta.tool_calls, and the terminal chunk must emit
-    finish_reason="tool_calls" so the frontend's accumulator runs the
-    function instead of seeing finish_reason="stop"."""
+    """Round 12: function tools forwarded into /v1/responses must have their
+    `function_call` output items translated back into Chat Completions
+    delta.tool_calls, and the terminal chunk must emit
+    finish_reason="tool_calls" (not "stop") so the frontend's accumulator runs
+    the function."""
 
     def handler(request: httpx.Request) -> httpx.Response:
         events = [
@@ -288,7 +288,7 @@ def test_responses_function_call_output_translates_to_delta_tool_calls(monkeypat
     assert tc["id"] == "call_xyz"
     assert tc["function"]["name"] == "get_weather"
     assert tc["function"]["arguments"] == '{"city":"SF"}'
-    # Final chunk reports tool_calls instead of stop.
+    # Final chunk reports tool_calls, not stop.
     terminal = next(
         p
         for p in payloads
@@ -301,8 +301,8 @@ def test_responses_function_call_output_translates_to_delta_tool_calls(monkeypat
 
 def test_responses_parallel_function_calls_get_distinct_indices(monkeypatch):
     """Round 13: parallel function_call items must land on distinct
-    delta.tool_calls[].index slots so index-keyed clients don't
-    collapse the second call into the first."""
+    delta.tool_calls[].index slots so index-keyed clients don't collapse the
+    second call into the first."""
 
     def handler(request: httpx.Request) -> httpx.Response:
         events = [
@@ -388,10 +388,10 @@ def test_responses_parallel_function_calls_get_distinct_indices(monkeypatch):
 
 
 def test_responses_follow_up_tool_result_uses_function_call_output_items(monkeypatch):
-    """Round 13: a second turn after a Responses function call must
-    serialize the tool_calls history and tool result as Responses
-    `function_call` / `function_call_output` input items, not as
-    Chat Completions role="tool" content."""
+    """Round 13: a second turn after a Responses function call must serialize
+    the tool_calls history and tool result as Responses `function_call` /
+    `function_call_output` input items, not Chat Completions role="tool"
+    content."""
     captured: dict = {}
 
     def handler(request: httpx.Request) -> httpx.Response:
diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py
index 7d488ae4c9..db4c07cdc8 100644
--- a/studio/backend/tests/test_openai_tool_passthrough.py
+++ b/studio/backend/tests/test_openai_tool_passthrough.py
@@ -1,26 +1,18 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
 
-"""
-Tests for the OpenAI /v1/chat/completions client-side tool pass-through.
+"""Tests for the OpenAI /v1/chat/completions client-side tool pass-through.
 
-Covers:
-- ChatCompletionRequest accepts standard OpenAI `tools` / `tool_choice` / `stop`.
-- ChatMessage accepts role="tool" with `tool_call_id` and role="assistant"
-  with `content: None` + `tool_calls`.
-- ChatCompletionRequest carries unknown fields via `extra="allow"`.
-- anthropic_tool_choice_to_openai() covers all four Anthropic shapes.
-- _build_passthrough_payload() honors a caller-supplied tool_choice and
-  defaults to "auto" when unset.
-- _friendly_error() maps httpx transport errors to a "Lost connection"
-  message so passthrough failures are legible instead of bare 500s.
-
-No running server or GPU required.
+Covers ChatMessage tool/assistant roles, ChatCompletionRequest tool fields and
+extra="allow", anthropic_tool_choice_to_openai, _build_passthrough_payload
+tool_choice propagation, and _friendly_error's httpx-to-"Lost connection"
+mapping. No server or GPU required.
 """
 
 import os
 import sys
 import asyncio
+import json
 from types import SimpleNamespace
 
 _backend = os.path.join(os.path.dirname(__file__), "..")
@@ -34,13 +26,20 @@ from pydantic import ValidationError
 from models.inference import (
     ChatCompletionRequest,
     ChatMessage,
+    CompletionChoice,
+    CompletionMessage,
 )
 from core.inference.anthropic_compat import (
     anthropic_tool_choice_to_openai,
 )
 from routes.inference import (
+    _build_openai_passthrough_body,
     _build_passthrough_payload,
+    _clamp_finish_reason,
+    _effective_max_tokens,
+    _extract_content_parts,
     _friendly_error,
+    _openai_stream_usage_chunk,
     _set_or_prepend_system_message,
     openai_chat_completions,
 )
@@ -120,8 +119,8 @@ class TestChatMessageToolRoles:
             ChatMessage(role = "function", content = "x")
 
     def test_content_absent_on_assistant_tool_call_defaults_to_none(self):
-        # Assistant messages that carry only tool_calls are the one
-        # documented case where `content=None` is permitted.
+        # Assistant messages carrying only tool_calls are the one documented
+        # case where `content=None` is permitted.
         msg = ChatMessage(
             role = "assistant",
             tool_calls = [
@@ -136,9 +135,9 @@ class TestChatMessageToolRoles:
 
     def test_tool_role_missing_tool_call_id_left_for_request_validator(self):
         # Per-message: missing tool_call_id is now allowed at this layer.
-        # ChatCompletionRequest's walkback fills it in from the prior
-        # assistant tool_calls; see test_inference_model_validation.py for
-        # the resolution coverage.
+        # ChatCompletionRequest's walkback fills it from the prior assistant
+        # tool_calls; see test_inference_model_validation.py for resolution
+        # coverage.
         msg = ChatMessage(role = "tool", content = '{"temperature": 72}')
         assert msg.tool_call_id is None
         assert msg.content == '{"temperature": 72}'
@@ -173,7 +172,7 @@ class TestChatMessageToolRoles:
         assert "content" in str(exc_info.value)
 
     def test_assistant_without_content_or_tool_calls_tolerated(self):
-        # Stop-button leaves an empty assistant turn; tolerate so replay round-trips.
+        # Stop-button leaves an empty assistant turn; tolerate for replay.
         msg = ChatMessage(role = "assistant")
         assert msg.content is None
         assert msg.tool_calls is None
@@ -280,18 +279,19 @@ class TestChatCompletionRequestToolFields:
         assert req.stop is None
 
     def test_extra_fields_accepted(self):
-        # `frequency_penalty`, `seed`, `response_format` are not yet
-        # explicitly declared but must survive Pydantic parsing now that
-        # extra="allow" is set.
+        # `frequency_penalty` and `response_format` are not yet explicitly
+        # declared but must survive Pydantic parsing now that extra="allow" is
+        # set. `seed` is declared and should land on the typed field instead.
         req = self._make(
             frequency_penalty = 0.5,
             seed = 42,
             response_format = {"type": "json_object"},
         )
+        assert req.seed == 42
         # Extras land in model_extra
         assert req.model_extra is not None
         assert req.model_extra.get("frequency_penalty") == 0.5
-        assert req.model_extra.get("seed") == 42
+        assert "seed" not in req.model_extra
         assert req.model_extra.get("response_format") == {"type": "json_object"}
 
     def test_unsloth_extensions_still_work(self):
@@ -305,25 +305,16 @@ class TestChatCompletionRequestToolFields:
         assert req.session_id == "abc"
 
     def test_stream_defaults_false_matching_openai_spec(self):
-        # OpenAI's /v1/chat/completions spec defaults `stream` to false.
-        # Studio previously defaulted to true, which broke naive curl
-        # clients (and .NET / System.Text.Json SDKs per #5047) that omit
-        # `stream` -- they expect a JSON blob, got SSE.
-        # Pin the corrected default so it can't silently regress.
+        # OpenAI defaults `stream` to false. Studio used to default true,
+        # breaking naive curl/.NET clients (#5047) that omit it. Pin the fix.
         req = self._make()
         assert req.stream is False
 
     def test_post_without_stream_field_decodes_to_stream_false_over_http(self, monkeypatch):
-        # Wire-level guard for the same default: a POST body that omits
-        # `stream` entirely (the exact shape naive curl / .NET clients
-        # send) must deserialise into stream=False *and* the response
-        # must be `application/json`, never `text/event-stream`.
-        # Mounts the real `routes.inference.router` so this catches
-        # regressions in middleware/aliasing on the actual endpoint
-        # (e.g. someone adding a request layer that injects stream=True
-        # before pydantic builds the model). Backends are bypassed by
-        # routing through `provider_type` and stubbing the external
-        # provider proxy.
+        # Wire-level guard: a POST body omitting `stream` must deserialise to
+        # stream=False and return application/json, never text/event-stream.
+        # Mounts the real router to catch middleware/aliasing regressions;
+        # backends are bypassed via provider_type + a stubbed proxy.
         from fastapi import FastAPI
         from fastapi.responses import JSONResponse
         from fastapi.testclient import TestClient
@@ -356,6 +347,153 @@ class TestChatCompletionRequestToolFields:
         assert "text/event-stream" not in resp.headers["content-type"]
         assert captured["stream"] is False
 
+    def _v1_client(
+        self,
+        monkeypatch,
+        llama_backend,
+        inference_backend = None,
+    ):
+        from fastapi import FastAPI
+        from fastapi.testclient import TestClient
+
+        import routes.inference as inference_route
+        from auth.authentication import get_current_subject
+        from utils.api_errors import install_api_error_handlers
+
+        monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: llama_backend)
+        if inference_backend is not None:
+            monkeypatch.setattr(inference_route, "get_inference_backend", lambda: inference_backend)
+
+        app = FastAPI()
+        app.include_router(inference_route.router, prefix = "/v1")
+        install_api_error_handlers(app)
+        app.dependency_overrides[get_current_subject] = lambda: "test-user"
+        return TestClient(app)
+
+    def _assert_unsupported_param(self, response, param):
+        assert response.status_code == 400
+        body = response.json()
+        assert body["error"]["param"] == param
+        assert body["error"]["code"] == "unsupported_parameter"
+
+    def _assert_unsupported_n(self, response):
+        self._assert_unsupported_param(response, "n")
+
+    def test_n_allows_openai_chat_completion_range(self):
+        req = self._make(n = 128)
+        assert req.n == 128
+        with pytest.raises(ValidationError):
+            self._make(n = 129)
+
+    def test_n_rejected_for_external_provider_path(self, monkeypatch):
+        class _UnusedBackend:
+            is_loaded = False
+
+        client = self._v1_client(monkeypatch, _UnusedBackend())
+        resp = client.post(
+            "/v1/chat/completions",
+            json = {
+                "messages": [{"role": "user", "content": "hi"}],
+                "provider_type": "openai",
+                "n": 2,
+            },
+        )
+        self._assert_unsupported_n(resp)
+
+    def test_logprobs_rejected_until_supported(self, monkeypatch):
+        class _UnusedBackend:
+            is_loaded = False
+
+        client = self._v1_client(monkeypatch, _UnusedBackend())
+        resp = client.post(
+            "/v1/chat/completions",
+            json = {
+                "messages": [{"role": "user", "content": "hi"}],
+                "provider_type": "openai",
+                "logprobs": True,
+            },
+        )
+        self._assert_unsupported_param(resp, "logprobs")
+
+    def test_top_logprobs_rejected_until_supported(self, monkeypatch):
+        class _UnusedBackend:
+            is_loaded = False
+
+        client = self._v1_client(monkeypatch, _UnusedBackend())
+        resp = client.post(
+            "/v1/chat/completions",
+            json = {
+                "messages": [{"role": "user", "content": "hi"}],
+                "provider_type": "openai",
+                "top_logprobs": 3,
+            },
+        )
+        self._assert_unsupported_param(resp, "top_logprobs")
+
+    def test_n_rejected_for_gguf_streaming_path(self, monkeypatch):
+        class _GGUFBackend:
+            is_loaded = True
+            model_identifier = "test-gguf"
+            supports_tools = False
+            is_vision = False
+            _is_audio = False
+
+        client = self._v1_client(monkeypatch, _GGUFBackend())
+        resp = client.post(
+            "/v1/chat/completions",
+            json = {
+                "messages": [{"role": "user", "content": "hi"}],
+                "stream": True,
+                "n": 2,
+            },
+        )
+        self._assert_unsupported_n(resp)
+
+    def test_n_rejected_for_gguf_tools_passthrough_path(self, monkeypatch):
+        class _GGUFBackend:
+            is_loaded = True
+            model_identifier = "test-gguf"
+            supports_tools = True
+            is_vision = False
+            _is_audio = False
+
+        client = self._v1_client(monkeypatch, _GGUFBackend())
+        resp = client.post(
+            "/v1/chat/completions",
+            json = {
+                "messages": [{"role": "user", "content": "hi"}],
+                "tools": [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": "lookup",
+                            "parameters": {"type": "object"},
+                        },
+                    }
+                ],
+                "n": 2,
+            },
+        )
+        self._assert_unsupported_n(resp)
+
+    def test_n_rejected_for_non_gguf_path(self, monkeypatch):
+        class _NoGGUFBackend:
+            is_loaded = False
+
+        class _InferenceBackend:
+            active_model_name = "test-model"
+            models = {"test-model": {}}
+
+        client = self._v1_client(monkeypatch, _NoGGUFBackend(), _InferenceBackend())
+        resp = client.post(
+            "/v1/chat/completions",
+            json = {
+                "messages": [{"role": "user", "content": "hi"}],
+                "n": 2,
+            },
+        )
+        self._assert_unsupported_n(resp)
+
     def test_multiturn_tool_loop_messages(self):
         req = ChatCompletionRequest(
             messages = [
@@ -468,17 +606,132 @@ class TestBuildPassthroughPayloadToolChoice:
         body = _build_passthrough_payload(**self._args(), tool_choice = tc)
         assert body["tool_choice"] == tc
 
-    def test_stream_adds_include_usage(self):
+    def test_stream_omits_usage_options_when_client_did_not_request_them(self):
         args = self._args()
         args["stream"] = True
         body = _build_passthrough_payload(**args)
+        assert "stream_options" not in body
+
+    def test_stream_forwards_include_usage_when_client_requests_it(self):
+        args = self._args()
+        args["stream"] = True
+        body = _build_passthrough_payload(
+            **args,
+            stream_options = {"include_usage": True},
+        )
         assert body.get("stream_options") == {"include_usage": True}
 
+    def test_stream_forwards_include_usage_false_when_client_requests_it(self):
+        args = self._args()
+        args["stream"] = True
+        body = _build_passthrough_payload(
+            **args,
+            stream_options = {"include_usage": False},
+        )
+        assert body.get("stream_options") == {"include_usage": False}
+
     def test_repetition_penalty_renamed(self):
         body = _build_passthrough_payload(**self._args(), repetition_penalty = 1.1)
         assert body.get("repeat_penalty") == 1.1
         assert "repetition_penalty" not in body
 
+    def test_passthrough_body_merges_system_and_developer_messages(self):
+        payload = ChatCompletionRequest(
+            model = "default",
+            messages = [
+                {"role": "system", "content": "original system"},
+                {"role": "developer", "content": "developer rules"},
+                {"role": "user", "content": "hi"},
+            ],
+            tools = self._args()["openai_tools"],
+        )
+
+        body = _build_openai_passthrough_body(payload, backend_ctx = 4096)
+
+        assert body["messages"] == [
+            {"role": "system", "content": "original system\n\ndeveloper rules"},
+            {"role": "user", "content": "hi"},
+        ]
+
+
+# =====================================================================
+# OpenAI API compatibility helpers — verified spec edge cases
+# =====================================================================
+
+
+class TestOpenAICompatibilityHelpers:
+    def test_max_completion_tokens_wins_over_deprecated_max_tokens(self):
+        payload = SimpleNamespace(max_tokens = 128, max_completion_tokens = 64)
+        assert _effective_max_tokens(payload) == 64
+
+    @pytest.mark.parametrize(
+        "finish_reason",
+        ["stop", "length", "tool_calls", "content_filter", "function_call"],
+    )
+    def test_clamp_finish_reason_preserves_openai_finish_reasons(self, finish_reason):
+        assert _clamp_finish_reason(finish_reason) == finish_reason
+
+    def test_clamp_finish_reason_defaults_unknown_to_stop(self):
+        assert _clamp_finish_reason(None) == "stop"
+        assert _clamp_finish_reason("unexpected") == "stop"
+
+    def test_non_streaming_completion_choice_accepts_tool_calls_finish_reason(self):
+        choice = CompletionChoice(
+            index = 0,
+            message = CompletionMessage(content = ""),
+            finish_reason = "tool_calls",
+        )
+        assert choice.finish_reason == "tool_calls"
+
+    def test_stream_usage_chunk_requires_include_usage(self):
+        usage = {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}
+        payload = SimpleNamespace(stream_options = None)
+        assert (
+            _openai_stream_usage_chunk(payload, "chatcmpl-test", 123, "model", usage, None) is None
+        )
+
+        payload.stream_options = {"include_usage": True}
+        line = _openai_stream_usage_chunk(payload, "chatcmpl-test", 123, "model", usage, None)
+        assert line is not None
+        assert '"choices":[]' in line
+        assert '"usage"' in line
+
+    def test_stream_usage_chunk_coerces_nullable_counts(self):
+        payload = SimpleNamespace(stream_options = {"include_usage": True})
+        line = _openai_stream_usage_chunk(
+            payload,
+            "chatcmpl-test",
+            123,
+            "model",
+            {"prompt_tokens": None, "completion_tokens": 7, "total_tokens": None},
+            None,
+        )
+
+        assert line is not None
+        parsed = json.loads(line.removeprefix("data: "))
+        usage = parsed["usage"]
+        assert usage["prompt_tokens"] == 0
+        assert usage["completion_tokens"] == 7
+        assert usage["total_tokens"] == 7
+
+    def test_developer_message_preserves_existing_system_prompt(self):
+        payload = ChatCompletionRequest(
+            messages = [
+                {"role": "system", "content": "original system"},
+                {"role": "developer", "content": "developer rules"},
+                {"role": "user", "content": "hi"},
+            ]
+        )
+        for message in payload.messages:
+            if message.role == "developer":
+                message.role = "system"
+
+        system_prompt, chat_messages, image_b64 = _extract_content_parts(payload.messages)
+
+        assert system_prompt == "original system\n\ndeveloper rules"
+        assert chat_messages == [{"role": "user", "content": "hi"}]
+        assert image_b64 is None
+
 
 # =====================================================================
 # _friendly_error — httpx transport failures
@@ -486,13 +739,10 @@ class TestBuildPassthroughPayloadToolChoice:
 
 
 class TestFriendlyErrorHttpx:
-    """The async pass-through helpers talk to llama-server via httpx.
-    When the subprocess is down, httpx raises RequestError subclasses
-    whose string form (``"All connection attempts failed"``, ``"[Errno 111]
-    Connection refused"``, ...) does NOT contain the substring
-    ``"Lost connection to llama-server"`` the sync path uses, so the
-    previous substring-only `_friendly_error` returned a useless generic
-    message. These tests pin the new isinstance-based mapping.
+    """When llama-server is down, httpx RequestError strings lack the
+    "Lost connection to llama-server" substring the sync path keys off, so the
+    old substring-only `_friendly_error` returned a useless generic message.
+    These tests pin the new isinstance-based mapping.
     """
 
     def _req(self):
@@ -515,9 +765,8 @@ class TestFriendlyErrorHttpx:
         assert "Lost connection" in _friendly_error(exc)
 
     def test_non_httpx_unchanged(self):
-        # Non-httpx exceptions still fall through to the existing substring
-        # heuristics — a context-size message must still produce the
-        # "Message too long" path.
+        # Non-httpx exceptions still fall through to the substring heuristics
+        # — a context-size message must still produce "Message too long".
         ctx_msg = "request (4096 tokens) exceeds the available context size (2048 tokens)"
         assert "Message too long" in _friendly_error(ValueError(ctx_msg))
 
@@ -543,7 +792,7 @@ class TestDropEmptyAssistantSentinels:
         assert out == [{"role": "user", "content": "hi"}, {"role": "user", "content": "again"}]
 
     def test_drops_assistant_with_no_content_key(self):
-        # exclude_none=True strips the content key entirely; filter must catch this.
+        # exclude_none=True strips the content key entirely; filter must catch it.
         msgs = [
             {"role": "user", "content": "hi"},
             {"role": "assistant"},
@@ -658,8 +907,8 @@ class TestGgufVisionMessages:
         assert len(messages[2]["content"]) == 2
         assert isinstance(messages[1]["content"], str)
 
-        # Legacy top-level image_base64 must be ignored when any message-level
-        # image already exists; otherwise turn 2 ends up with two image parts.
+        # Legacy top-level image_base64 must be ignored when a message-level
+        # image exists; otherwise turn 2 ends up with two image parts.
         for msg in messages:
             content = msg.get("content")
             if isinstance(content, list):
@@ -823,3 +1072,131 @@ class TestGgufVisionToolRouting:
         assert tool_messages[0]["role"] == "system"
         assert tool_messages[1]["role"] == "user"
         assert tool_messages[1]["content"][1]["type"] == "image_url"
+
+    def test_parallel_tool_calls_false_reaches_gguf_tool_loop(self, monkeypatch):
+        import routes.inference as inf_mod
+
+        reset_tool_policy()
+        captured = {}
+
+        def _plain(**kwargs):
+            raise AssertionError("plain GGUF path should not be used")
+
+        def _tools(**kwargs):
+            captured["kwargs"] = kwargs
+            yield {"type": "content", "text": "done"}
+
+        backend = SimpleNamespace(
+            is_loaded = True,
+            is_vision = False,
+            supports_tools = True,
+            model_identifier = "test-gguf",
+            generate_chat_completion = _plain,
+            generate_chat_completion_with_tools = _tools,
+        )
+        monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend)
+
+        payload = ChatCompletionRequest(
+            model = "default",
+            enable_tools = True,
+            enabled_tools = ["web_search"],
+            parallel_tool_calls = False,
+            messages = [{"role": "user", "content": "search once"}],
+        )
+
+        response = self._drive(
+            openai_chat_completions(payload, request = self._Request(), current_subject = "test")
+        )
+        self._consume_response(response)
+
+        assert captured["kwargs"]["disable_parallel_tool_use"] is True
+
+    def test_standard_gguf_merges_system_and_developer_messages(self, monkeypatch):
+        import routes.inference as inf_mod
+
+        captured = {}
+
+        def _generate(**kwargs):
+            captured["messages"] = kwargs["messages"]
+            yield "done"
+            yield {
+                "type": "metadata",
+                "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4},
+                "finish_reason": "stop",
+            }
+
+        backend = SimpleNamespace(
+            is_loaded = True,
+            is_vision = False,
+            supports_tools = False,
+            model_identifier = "test-gguf",
+            generate_chat_completion = _generate,
+        )
+        monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend)
+
+        payload = ChatCompletionRequest(
+            model = "default",
+            messages = [
+                {"role": "system", "content": "original system"},
+                {"role": "developer", "content": "developer rules"},
+                {"role": "user", "content": "hi"},
+            ],
+        )
+
+        self._drive(
+            openai_chat_completions(payload, request = self._Request(), current_subject = "test")
+        )
+
+        assert captured["messages"] == [
+            {"role": "system", "content": "original system\n\ndeveloper rules"},
+            {"role": "user", "content": "hi"},
+        ]
+
+    @pytest.mark.parametrize(
+        ("seed", "expected"),
+        [
+            (41, [41, 42, 43]),
+            (-1, [-1, -1, -1]),
+        ],
+    )
+    def test_gguf_n_choices_vary_explicit_non_negative_seed(self, monkeypatch, seed, expected):
+        import routes.inference as inf_mod
+
+        seen_seeds = []
+
+        def _generate(**kwargs):
+            seen_seeds.append(kwargs.get("seed"))
+            yield f"choice-{len(seen_seeds)}"
+            yield {
+                "type": "metadata",
+                "usage": {
+                    "prompt_tokens": 5,
+                    "completion_tokens": 7,
+                    "total_tokens": 12,
+                },
+                "finish_reason": "stop",
+            }
+
+        backend = SimpleNamespace(
+            is_loaded = True,
+            is_vision = False,
+            supports_tools = False,
+            model_identifier = "test-gguf",
+            generate_chat_completion = _generate,
+        )
+        monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend)
+
+        payload = ChatCompletionRequest(
+            model = "default",
+            messages = [{"role": "user", "content": "hi"}],
+            n = 3,
+            seed = seed,
+        )
+
+        response = self._drive(
+            openai_chat_completions(payload, request = self._Request(), current_subject = "test")
+        )
+        body = json.loads(response.body)
+
+        assert seen_seeds == expected
+        assert [choice["index"] for choice in body["choices"]] == [0, 1, 2]
diff --git a/studio/backend/tests/test_openai_tool_result_fallbacks.py b/studio/backend/tests/test_openai_tool_result_fallbacks.py
index 7c033bc348..5ea812441f 100644
--- a/studio/backend/tests/test_openai_tool_result_fallbacks.py
+++ b/studio/backend/tests/test_openai_tool_result_fallbacks.py
@@ -3,8 +3,8 @@
 
 """Regression tests for OpenAI Responses tool-result rendering.
 
-Covers two bug classes: empty web_search cards (per-card result seeded
-with "Searching: ") and orphan shell_call cards (bundled-output
+Two bug classes: empty web_search cards (per-card result seeded with
+"Searching: ") and orphan shell_call cards (bundled-output
 fallback + final flush at response.completed / response.incomplete).
 """
 
@@ -137,8 +137,8 @@ def test_web_search_each_call_carries_its_own_query_as_result(monkeypatch):
 
 
 def test_web_search_last_call_overwritten_with_citations(monkeypatch):
-    """Last call still gets the aggregated citation list; earlier calls
-    keep their per-call `Searching:` text."""
+    """Last call gets the aggregated citations; earlier calls keep their
+    per-call `Searching:` text."""
     sse_events = [
         {
             "type": "response.output_item.done",
@@ -170,12 +170,12 @@ def test_web_search_last_call_overwritten_with_citations(monkeypatch):
     events = _tool_events(lines)
     ends = [e for e in events if e["type"] == "tool_end"]
     by_id: dict = {}
-    # Keep the LAST tool_end per id (the citation overwrite for ws_2).
+    # Keep the LAST tool_end per id (citation overwrite for ws_2).
     for e in ends:
         by_id[e["tool_call_id"]] = e
     # First call keeps its own query.
     assert by_id["ws_1"]["result"] == "Searching: first query"
-    # Last call gets overwritten with the citation block.
+    # Last call overwritten with the citation block.
     assert "Title: Example A" in by_id["ws_2"]["result"]
     assert "URL: https://example.com/a" in by_id["ws_2"]["result"]
 
diff --git a/studio/backend/tests/test_pricing.py b/studio/backend/tests/test_pricing.py
index d669747b1a..8cd7796f14 100644
--- a/studio/backend/tests/test_pricing.py
+++ b/studio/backend/tests/test_pricing.py
@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-"""Unit tests for the per-session cost calculator. Verifies math
-against ``core/inference/pricing.py`` and graceful degradation."""
+"""Unit tests for the per-session cost calculator: math against
+``core/inference/pricing.py`` plus graceful degradation."""
 
 import math
 
@@ -462,12 +462,12 @@ def test_snapshot_contains_provider_buckets_and_multipliers():
 
 
 # ── longest-prefix match: dated mini variant must not collide with the
-#    shorter family prefix. ──
+#    shorter family prefix ──
 
 
 def test_longest_prefix_match_wins_for_dated_mini_snapshot():
-    """`gpt-5.4-mini-2026-...` must inherit the mini rate, not the
-    shorter `gpt-5.4` rate (longest prefix wins)."""
+    """`gpt-5.4-mini-2026-...` inherits the mini rate, not the shorter
+    `gpt-5.4` rate (longest prefix wins)."""
     out = calculate_cost(
         "openai",
         "gpt-5.4-mini-2026-04-23",
@@ -493,7 +493,7 @@ def test_longest_prefix_match_wins_for_dated_pro_snapshot():
 
 
 def test_openai_chat_style_usage_keys_priced_correctly():
-    """Chat-style envelope (`prompt_tokens` / `completion_tokens`) must
+    """Chat-style envelope (`prompt_tokens`/`completion_tokens`) must
     produce a non-zero cost (previously silently zeroed)."""
     out = calculate_cost(
         "openai",
@@ -577,7 +577,7 @@ def test_openai_chat_style_prompt_tokens_keeps_cache_read_semantics():
 
 def test_openai_chat_style_envelope_reads_cache_from_prompt_tokens_details():
     """Chat-style envelope ships cached under prompt_tokens_details;
-    calculator must honour both this and input_tokens_details."""
+    calculator must honour both that and input_tokens_details."""
     base = OPENAI_PRICING["gpt-5.5"]["input_per_mtok"]
     raw = calculate_cost(
         "openai",
diff --git a/studio/backend/tests/test_pricing_edge.py b/studio/backend/tests/test_pricing_edge.py
index 5818b42e8a..1fcc428f90 100644
--- a/studio/backend/tests/test_pricing_edge.py
+++ b/studio/backend/tests/test_pricing_edge.py
@@ -30,8 +30,8 @@ def _isclose(
 
 
 def test_prefix_match_requires_dash_boundary_opus_variant():
-    # `claude-opus-4-15` must not inherit `claude-opus-4-1` pricing;
-    # next char must be `-` or end-of-string.
+    # `claude-opus-4-15` must not inherit `claude-opus-4-1` pricing; the next
+    # char must be `-` or end-of-string.
     assert _lookup("anthropic", "claude-opus-4-15") is None
     out = calculate_cost(
         "anthropic",
@@ -55,8 +55,8 @@ def test_prefix_match_requires_dash_boundary_gpt_variant():
 
 
 def test_prefix_match_requires_dash_boundary_pro_lookalike():
-    # `gpt-5.5-prod` must fall through `gpt-5.5-pro` (6x overcharge)
-    # and land on the canonical `gpt-5.5` row.
+    # `gpt-5.5-prod` must fall through `gpt-5.5-pro` (6x overcharge) and land on
+    # the canonical `gpt-5.5` row.
     prices = _lookup("openai", "gpt-5.5-prod")
     assert prices is not None
     assert (
@@ -81,7 +81,7 @@ def test_prefix_match_still_resolves_legit_dated_snapshots():
     assert out["priced"] is True
     assert _isclose(out["input_usd"], 0.75)
 
-    # And Anthropic dated snapshot still resolves to canonical row.
+    # Anthropic dated snapshot still resolves to the canonical row.
     out = calculate_cost(
         "anthropic",
         "claude-opus-4-7-20260414",
@@ -95,7 +95,6 @@ def test_prefix_match_still_resolves_legit_dated_snapshots():
 
 
 def test_explicit_zero_input_tokens_wins_over_stale_prompt_tokens():
-    # Input-side mirror of the output zero precedence test.
     out = calculate_cost(
         "openai",
         "gpt-5.5",
@@ -110,7 +109,7 @@ def test_explicit_zero_input_tokens_wins_over_stale_prompt_tokens():
 
 
 def test_none_input_tokens_falls_through_to_prompt_tokens():
-    # `None` is "key present but unset"; chat-style mirror wins.
+    # `None` means "key present but unset"; chat-style mirror wins.
     out = calculate_cost(
         "openai",
         "gpt-5.5",
@@ -176,8 +175,8 @@ def test_negative_prompt_tokens_chat_style_clamp():
 
 
 def test_anthropic_chat_cache_read_exceeds_prompt_no_negative_billable():
-    # cache_read > prompt_tokens clamps uncached_input at 0; billable
-    # still reflects cache buckets (we charge for what we got).
+    # cache_read > prompt_tokens clamps uncached_input at 0; billable still
+    # reflects cache buckets (we charge for what we got).
     out = calculate_cost(
         "anthropic",
         "claude-opus-4-7",
@@ -216,8 +215,8 @@ def test_openai_raw_cached_tokens_exceeds_input_clamp_non_cached():
 
 
 def test_openai_long_context_triggers_on_cache_creation_inflated_billable():
-    # cache_creation pushes billable past 272k -> long-context tier
-    # must fire to avoid undercounting.
+    # cache_creation pushes billable past 272k -> long-context tier must fire to
+    # avoid undercounting.
     out = calculate_cost(
         "openai",
         "gpt-5.5",
@@ -272,8 +271,8 @@ def test_openai_chat_envelope_long_context_parity_with_raw():
 
 
 def test_cache_creation_as_int_does_not_crash():
-    # Proxies sometimes fold cache_creation to an int; tolerate it
-    # and fall back to the 5m default.
+    # Proxies sometimes fold cache_creation to an int; tolerate it and fall back
+    # to the 5m default.
     base = ANTHROPIC_PRICING["claude-opus-4-7"]["input_per_mtok"]
     out = calculate_cost(
         "anthropic",
@@ -363,16 +362,16 @@ def test_empty_usage_dict_zero_bill():
 
 
 def test_anthropic_prompt_tokens_details_fallback_when_native_key_missing():
-    """Chat-style envelope without `cache_read_input_tokens` but with
-    mirrored `prompt_tokens_details.cached_tokens` should still apply
-    the cache_read discount."""
+    """Chat-style envelope without `cache_read_input_tokens` but with mirrored
+    `prompt_tokens_details.cached_tokens` should still apply the cache_read
+    discount."""
     r = calculate_cost(
         provider = "anthropic",
         model = "claude-opus-4-7",
         usage = {
             "prompt_tokens": 1_000_000,
             "completion_tokens": 0,
-            # Only the mirrored shape (no native key).
+            # Mirrored shape only (no native key).
             "prompt_tokens_details": {"cached_tokens": 1_000_000},
             "cache_creation_input_tokens": 0,
         },
@@ -383,8 +382,8 @@ def test_anthropic_prompt_tokens_details_fallback_when_native_key_missing():
 
 
 def test_anthropic_native_key_takes_precedence_over_mirrored():
-    """When both native and mirrored cache-read fields are present,
-    the native Anthropic field wins (mirror is fallback-only)."""
+    """When both native and mirrored cache-read fields are present, the native
+    Anthropic field wins (mirror is fallback-only)."""
     r = calculate_cost(
         provider = "anthropic",
         model = "claude-opus-4-7",
@@ -403,8 +402,8 @@ def test_anthropic_native_key_takes_precedence_over_mirrored():
 
 
 def test_anthropic_native_zero_takes_precedence_over_mirrored():
-    """Explicit `cache_read_input_tokens: 0` is authoritative; a stale
-    mirrored block from a proxy must not inflate cache_read past it."""
+    """Explicit `cache_read_input_tokens: 0` is authoritative; a stale mirrored
+    block from a proxy must not inflate cache_read past it."""
     r = calculate_cost(
         provider = "anthropic",
         model = "claude-opus-4-7",
@@ -429,8 +428,8 @@ def test_anthropic_native_zero_takes_precedence_over_mirrored():
 
 
 def test_build_usage_chunk_forwards_anthropic_cache_creation_breakdown():
-    """Chat-style envelope must carry the 5m/1h cache-write breakdown
-    so downstream cost calc applies the 2x 1h premium."""
+    """Chat-style envelope must carry the 5m/1h cache-write breakdown so
+    downstream cost calc applies the 2x 1h premium."""
     import json
     from core.inference.external_provider import _build_usage_chunk
 
diff --git a/studio/backend/tests/test_providers_api.py b/studio/backend/tests/test_providers_api.py
index 88df886b6c..5e24ed752d 100644
--- a/studio/backend/tests/test_providers_api.py
+++ b/studio/backend/tests/test_providers_api.py
@@ -4,13 +4,13 @@
 """
 Integration tests for the external providers API.
 
-Requires a running Unsloth Studio server. Configure via environment variables:
+Requires a running Unsloth Studio server. Configure via env vars:
 
     export STUDIO_TEST_URL="http://localhost:8888"   # default
     export STUDIO_TEST_USER="unsloth"                # default
     export STUDIO_TEST_PASSWORD="..."                # required — see .bootstrap_password
 
-    # Provider API keys — any left unset will have their tests automatically skipped
+    # Provider API keys — tests skip when their key is unset
     export OPENAI_API_KEY="sk-..."
     export MISTRAL_API_KEY="..."
     export GOOGLE_API_KEY="..."
@@ -38,15 +38,14 @@ BASE_URL = os.getenv("STUDIO_TEST_URL", "http://localhost:8000")
 USERNAME = os.getenv("STUDIO_TEST_USER", "unsloth")
 PASSWORD = os.getenv("STUDIO_TEST_PASSWORD", "")
 
-# These tests require a live Studio server reachable at BASE_URL with a known
-# bootstrap password. Skip the whole module when that environment is missing
-# (e.g. on CI runners) so pytest discovery does not error out.
+# Skip the whole module when no live Studio server / bootstrap password is
+# available (e.g. on CI) so pytest discovery does not error out.
 pytestmark = pytest.mark.skipif(
     not PASSWORD,
     reason = "Integration test requires a running Studio server; set STUDIO_TEST_PASSWORD to enable.",
 )
 
-# Map provider_type → (env var name, model to use for inference test)
+# provider_type → (env var name, model for inference test)
 _PROVIDER_CONFIGS: dict[str, tuple[str, str]] = {
     "openai": ("OPENAI_API_KEY", "gpt-4o-mini"),
     "mistral": ("MISTRAL_API_KEY", "mistral-small-2506"),
@@ -73,11 +72,9 @@ def _url(path: str) -> str:
 
 
 def _parse_sse_stream(response: requests.Response) -> tuple[str, bool]:
-    """
-    Read a streaming SSE response and return (assembled_text, saw_done).
+    """Read an SSE response, return (assembled_text, saw_done).
 
-    Each chunk is a JSON object with choices[0].delta.content.
-    The stream ends with `data: [DONE]`.
+    Each chunk is JSON with choices[0].delta.content; stream ends at `data: [DONE]`.
     """
     reply_parts: list[str] = []
     saw_done = False
@@ -93,7 +90,7 @@ def _parse_sse_stream(response: requests.Response) -> tuple[str, bool]:
             break
         try:
             chunk = json.loads(data)
-            # Handle both error payloads and normal chunks
+            # Handle error payloads and normal chunks
             if "error" in chunk:
                 raise RuntimeError(f"Provider error in stream: {chunk['error']}")
             delta = chunk.get("choices", [{}])[0].get("delta", {})
@@ -111,20 +108,12 @@ def _parse_sse_stream(response: requests.Response) -> tuple[str, bool]:
 
 @pytest.fixture(scope = "session")
 def auth_headers() -> dict[str, str]:
-    """
-    Log in once per session and return auth headers.
+    """Log in once per session and return auth headers.
 
-    On a fresh Studio install the bootstrap password triggers a forced password
-    change (must_change_password=True).  Any subsequent API call using that token
-    returns 403 "Password change required".  This fixture detects that state,
-    automatically completes the change-password flow, and re-logs in so all other
-    tests get a fully usable token.
-
-    The new password used during auto-change is:
-        STUDIO_TEST_NEW_PASSWORD  (env var, optional)
-        or PASSWORD + "-test"     (derived default)
-
-    On the second run, set STUDIO_TEST_PASSWORD to the new password.
+    On a fresh install the bootstrap password forces a change; this fixture
+    detects must_change_password, auto-completes the change (new password =
+    STUDIO_TEST_NEW_PASSWORD or PASSWORD + "-test"), and re-logs in. On the
+    second run, set STUDIO_TEST_PASSWORD to the new password.
     """
     assert PASSWORD, (
         "STUDIO_TEST_PASSWORD is not set.\n"
@@ -142,8 +131,8 @@ def auth_headers() -> dict[str, str]:
     assert token, "access_token is empty"
 
     if body.get("must_change_password"):
-        # Bootstrap token is restricted — only /api/auth/change-password works with it.
-        # Auto-complete the forced change so the rest of the tests get a full token.
+        # Bootstrap token only works with change-password; auto-complete the
+        # forced change so the rest of the tests get a full token.
         new_password = os.getenv("STUDIO_TEST_NEW_PASSWORD") or f"{PASSWORD}-test"
         change_resp = requests.post(
             _url("/api/auth/change-password"),
@@ -175,12 +164,10 @@ def public_key_pem(auth_headers: dict[str, str]) -> str:
 
 @pytest.fixture(scope = "session")
 def vision_image_data_url() -> str:
-    """
-    Download the sloth image once per session and return it as a base64 data URI.
+    """Download the sloth image once per session as a base64 data URI.
 
-    Using a data URI instead of a remote URL ensures every provider receives
-    the image inline — Gemini's OpenAI-compatible layer does not fetch external
-    HTTP URLs, so raw image_url links silently produce empty replies for Gemini.
+    A data URI sends the image inline; Gemini's OpenAI-compatible layer doesn't
+    fetch external HTTP URLs, so raw image_url links give empty Gemini replies.
     """
     resp = requests.get(_VISION_IMAGE_URL, timeout = 30)
     resp.raise_for_status()
@@ -191,11 +178,10 @@ def vision_image_data_url() -> str:
 
 @pytest.fixture(scope = "session")
 def encrypt_key(public_key_pem: str):
+    """Return encrypt_key(plaintext) -> base64 RSA-OAEP ciphertext.
+
+    Uses the backend's RSA public key; mirrors the frontend.
     """
-    Return a callable encrypt_key(plaintext: str) -> str (base64 RSA-OAEP ciphertext).
-    Uses the backend's RSA public key — mirrors what the frontend does.
-    """
-    # Decode PEM → load RSA public key
     pem_bytes = public_key_pem.encode("utf-8")
     rsa_pub = serialization.load_pem_public_key(pem_bytes)
 
@@ -298,8 +284,8 @@ class TestRegistry:
 
 class TestProviderCRUD:
     """
-    These tests run sequentially within the class and share state via class variables.
-    They create, read, update, and delete a single test provider config.
+    Run sequentially, sharing state via class variables. Create, read, update,
+    and delete a single test provider config.
     """
 
     _created_id: str = ""
@@ -356,7 +342,7 @@ class TestProviderCRUD:
         )
         assert resp.status_code == 204, f"Delete failed ({resp.status_code}): {resp.text}"
 
-        # Confirm gone from list
+        # Confirm gone
         list_resp = requests.get(_url("/api/providers/"), headers = auth_headers, timeout = 10)
         ids = [p["id"] for p in list_resp.json()]
         assert TestProviderCRUD._created_id not in ids, "Deleted provider still in list"
@@ -366,7 +352,7 @@ class TestProviderCRUD:
 # ── TestProviderInference ────────────────────────────────────────────
 
 
-# Build parametrize list: (provider_type, model, api_key) for configured providers only
+# Parametrize (provider_type, model, api_key) for configured providers
 _INFERENCE_PARAMS = [
     pytest.param(
         ptype,
@@ -384,8 +370,8 @@ _INFERENCE_PARAMS = [
 
 class TestProviderInference:
     """
-    Live inference tests — one parametrized set per provider.
-    Each test is automatically skipped when the provider's API key env var is not set.
+    Live inference tests, one parametrized set per provider. Each is skipped
+    when the provider's API key env var is unset.
     """
 
     @pytest.mark.parametrize("provider_type,model,api_key", _INFERENCE_PARAMS)
@@ -475,7 +461,7 @@ class TestProviderInference:
 
 # ── TestVisionInference ─────────────────────────────────────────────
 
-# Sloth photo — used to test vision routing across providers
+# Sloth photo for testing vision routing across providers
 _VISION_IMAGE_URL = "https://www.travelexcellence.com/images/where-to-see-sloths-in-costa-rica.jpg"
 
 _VISION_PARAMS = [
@@ -496,8 +482,8 @@ _VISION_PARAMS = [
 
 class TestVisionInference:
     """
-    Send a 1×1 white PNG alongside a text question to each vision-capable provider.
-    Verifies that image content parts survive the proxy and the provider replies.
+    Send a 1×1 white PNG plus a text question to each vision-capable provider.
+    Verifies image content parts survive the proxy and the provider replies.
     """
 
     @pytest.mark.parametrize("provider_type,model,api_key", _VISION_PARAMS)
@@ -556,12 +542,10 @@ class TestVisionInference:
 
 class TestLocalInferenceUnaffected:
     def test_chat_without_provider(self, auth_headers: dict[str, str]):
-        """
-        POST /v1/chat/completions without provider fields must not return 422 or 500.
+        """POST /v1/chat/completions without provider fields must not 422/500.
 
-        200 = a local model is loaded and responded.
-        503 = no model loaded (expected in test environment — that's fine).
-        Any other 4xx/5xx (except 503) = regression in request handling.
+        200 = local model responded; 503 = no model loaded (fine in tests);
+        any other 4xx/5xx = request-handling regression.
         """
         resp = requests.post(
             _url("/v1/chat/completions"),
diff --git a/studio/backend/tests/test_pytorch_mirror.py b/studio/backend/tests/test_pytorch_mirror.py
index 5844f209b6..59842214ef 100644
--- a/studio/backend/tests/test_pytorch_mirror.py
+++ b/studio/backend/tests/test_pytorch_mirror.py
@@ -19,8 +19,8 @@ OFFICIAL_URL = "https://download.pytorch.org/whl"
 
 
 def _reload_whl_base(monkeypatch, mirror_value = None):
-    """(Re-)import install_python_stack with a controlled env and return _PYTORCH_WHL_BASE."""
-    # Remove cached module so the module-level assignment re-executes
+    """(Re-)import install_python_stack with a controlled env, return _PYTORCH_WHL_BASE."""
+    # Drop cached module so the module-level assignment re-executes.
     sys.modules.pop("install_python_stack", None)
 
     if mirror_value is None:
@@ -28,7 +28,7 @@ def _reload_whl_base(monkeypatch, mirror_value = None):
     else:
         monkeypatch.setenv("UNSLOTH_PYTORCH_MIRROR", mirror_value)
 
-    # Temporarily add the script's directory to sys.path for import
+    # Add the script's directory to sys.path for import.
     script_dir = str(_INSTALL_SCRIPT.parent)
     monkeypatch.syspath_prepend(script_dir)
 
diff --git a/studio/backend/tests/test_recommended_folders_permission.py b/studio/backend/tests/test_recommended_folders_permission.py
index 659c3b547d..33a457755e 100644
--- a/studio/backend/tests/test_recommended_folders_permission.py
+++ b/studio/backend/tests/test_recommended_folders_permission.py
@@ -6,17 +6,16 @@ Regression test for the /recommended-folders (and /browse-folders) 500
 caused by an unreadable model directory, e.g. a stock root-owned
 ``ollama`` install at ``/usr/share/ollama/.ollama/models``.
 
-Root cause: the folder-scan helpers in ``routes.models`` probed candidate
-paths with a bare ``Path(p).is_dir()``. On Python <= 3.11 that returned
-``False`` for an unreadable path; on Python >= 3.12 ``is_dir()`` propagates
+Root cause: ``routes.models`` folder-scan helpers probed candidates with a
+bare ``Path(p).is_dir()``. On Python <= 3.11 that returned ``False`` for an
+unreadable path; on Python >= 3.12 ``is_dir()`` propagates
 ``PermissionError`` (EACCES), so the endpoint 500-ed through the whole
-middleware stack instead of just skipping the directory. The probes now go
-through the module-level ``_safe_is_dir`` helper.
+middleware stack instead of skipping the directory. Probes now go through
+the module-level ``_safe_is_dir`` helper.
 
-``routes.models`` pulls the full backend dependency tree (fastapi,
-structlog, the models package, ...), so rather than stand up the app we
-extract the real ``_safe_is_dir`` definition from the source file and
-exercise that exact function in isolation. The test therefore stays
+``routes.models`` pulls the full backend dep tree (fastapi, structlog, the
+models package, ...), so rather than stand up the app we extract the real
+``_safe_is_dir`` from the source file and exercise it in isolation —
 dependency-free while still running the shipped code.
 
 Run:
@@ -36,7 +35,7 @@ _models_src = _backend_root / "routes" / "models.py"
 
 def _load_safe_is_dir():
     """Return the real ``_safe_is_dir`` from routes/models.py without
-    importing the (heavily dependency-laden) module."""
+    importing the dependency-laden module."""
     tree = ast.parse(_models_src.read_text())
     fn = next(
         node
@@ -51,8 +50,8 @@ def _load_safe_is_dir():
 
 safe_is_dir = _load_safe_is_dir()
 
-# Permission bits are bypassed for the superuser, so the chmod-000 setup
-# below would not actually deny access when running as root.
+# The superuser bypasses permission bits, so the chmod-000 setup below
+# would not deny access when running as root.
 _skip_as_root = pytest.mark.skipif(
     hasattr(os, "geteuid") and os.geteuid() == 0,
     reason = "root bypasses filesystem permission bits",
@@ -60,8 +59,7 @@ _skip_as_root = pytest.mark.skipif(
 
 
 def test_helper_exists_in_source():
-    # Guards against a refactor silently dropping the helper the fix
-    # depends on (the extractor would then raise StopIteration).
+    # Guard against a refactor silently dropping the helper the fix needs.
     assert callable(safe_is_dir)
 
 
@@ -83,8 +81,8 @@ def test_file_is_false(tmp_path):
 def test_mode000_dir_itself_is_still_a_dir(tmp_path):
     """A mode-000 directory is still stat-able via its (traversable)
     parent, so _safe_is_dir reports True without raising. Filtering out
-    dirs we cannot actually *read* is the caller's separate
-    os.access(R_OK|X_OK) check, not this helper's job."""
+    dirs we can't *read* is the caller's separate os.access(R_OK|X_OK)
+    check, not this helper's job."""
     locked = tmp_path / "locked"
     locked.mkdir()
     os.chmod(locked, 0o000)
@@ -96,8 +94,8 @@ def test_mode000_dir_itself_is_still_a_dir(tmp_path):
 
 @_skip_as_root
 def test_path_under_unreadable_parent_returns_false_not_raises(tmp_path):
-    """The exact production scenario: stat()-ing a child of a mode-700
-    system directory, e.g. ``/usr/share/ollama/.ollama/models``."""
+    """The production scenario: stat()-ing a child of a mode-700 system
+    directory, e.g. ``/usr/share/ollama/.ollama/models``."""
     parent = tmp_path / "ollama"
     parent.mkdir()
     os.chmod(parent, 0o000)
@@ -113,8 +111,8 @@ def test_path_under_unreadable_parent_returns_false_not_raises(tmp_path):
     reason = "is_dir() only propagates PermissionError on Python >= 3.12",
 )
 def test_demonstrates_the_underlying_stdlib_regression(tmp_path):
-    """Documents *why* _safe_is_dir exists: the old bare pattern raises
-    on the interpreters Studio ships on (3.12+)."""
+    """Documents *why* _safe_is_dir exists: the old bare pattern raises on
+    the interpreters Studio ships on (3.12+)."""
     parent = tmp_path / "ollama"
     parent.mkdir()
     os.chmod(parent, 0o000)
diff --git a/studio/backend/tests/test_responses_api.py b/studio/backend/tests/test_responses_api.py
index c0ec876ad0..693e832113 100644
--- a/studio/backend/tests/test_responses_api.py
+++ b/studio/backend/tests/test_responses_api.py
@@ -1,18 +1,15 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
 
-"""
-Tests for the OpenAI Responses API schemas and input normalisation.
-These tests do NOT require a running server or GPU -- they validate
-the Pydantic models and the _normalise_responses_input helper.
-"""
+"""Tests for OpenAI Responses API Pydantic schemas and the
+_normalise_responses_input helper. No server or GPU required."""
 
 import sys
 import os
 import json
 import re
 
-# Ensure backend is on path
+# Ensure backend is on path.
 _backend = os.path.join(os.path.dirname(__file__), "..")
 sys.path.insert(0, _backend)
 
@@ -33,34 +30,32 @@ from models.inference import (
 )
 
 
-# ── _normalise_responses_input: copied from routes/inference.py ──
-# We cannot import routes.inference directly because routes/__init__.py
-# pulls in heavy dependencies (structlog/twisted/torch). This is a
-# direct copy of the function for testing purposes.
+# Copied from routes/inference.py: can't import it directly because
+# routes/__init__.py pulls in heavy deps (structlog/twisted/torch).
 
 
 def _normalise_responses_input(payload: ResponsesRequest) -> list:
-    """Convert a ResponsesRequest into a list of ChatMessage for the completions backend."""
+    """Convert a ResponsesRequest into ChatMessages for the completions backend."""
     messages = []
 
-    # System / developer instructions
+    # System / developer instructions.
     if payload.instructions:
         messages.append(ChatMessage(role = "system", content = payload.instructions))
 
-    # Simple string input
+    # Simple string input.
     if isinstance(payload.input, str):
         if payload.input:
             messages.append(ChatMessage(role = "user", content = payload.input))
         return messages
 
-    # List of ResponsesInputMessage
+    # List of ResponsesInputMessage.
     for msg in payload.input:
         role = "system" if msg.role == "developer" else msg.role
 
         if isinstance(msg.content, str):
             messages.append(ChatMessage(role = role, content = msg.content))
         else:
-            # Convert Responses content parts -> Chat content parts
+            # Convert Responses content parts -> Chat content parts.
             parts = []
             for part in msg.content:
                 if isinstance(part, ResponsesInputTextPart):
@@ -130,7 +125,7 @@ class TestResponsesRequest:
         assert req.instructions == "You are a helpful assistant."
 
     def test_extra_fields_accepted(self):
-        """OpenAI SDK may send fields we don't model -- extra='allow' should pass."""
+        """OpenAI SDK may send unmodeled fields -- extra='allow' must pass."""
         req = ResponsesRequest(
             input = "test",
             tools = [{"type": "web_search_preview"}],
@@ -167,7 +162,7 @@ class TestResponsesRequest:
 
 
 class TestResponsesResponse:
-    """Validate response models serialise correctly."""
+    """Response models serialise correctly."""
 
     def test_basic_response(self):
         resp = ResponsesResponse(
@@ -224,7 +219,7 @@ class TestResponsesResponse:
 
 
 class TestNormaliseResponsesInput:
-    """Test _normalise_responses_input converts Responses input to ChatMessages."""
+    """_normalise_responses_input converts Responses input to ChatMessages."""
 
     def test_string_input(self):
         payload = ResponsesRequest(input = "Hello world")
diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py
index 146d6017d0..f0f2714214 100644
--- a/studio/backend/tests/test_responses_tool_passthrough.py
+++ b/studio/backend/tests/test_responses_tool_passthrough.py
@@ -6,19 +6,19 @@ Tests for the OpenAI /v1/responses client-side function-calling pass-through.
 
 Covers:
 - ResponsesRequest accepts Responses-shape `tools`, `tool_choice`,
-  `parallel_tool_calls`, and the `function_call` / `function_call_output`
-  input items used for multi-turn tool loops.
-- _translate_responses_tools_to_chat() converts the flat Responses tool
-  shape to the nested Chat Completions shape, drops non-function built-in
-  tools, and returns None for empty lists.
-- _translate_responses_tool_choice_to_chat() passes string choices through
-  and converts {type:function,name:X} to Chat Completions' nested shape.
-- _normalise_responses_input() maps function_call_output items to
+  `parallel_tool_calls`, and `function_call` / `function_call_output`
+  input items for multi-turn tool loops.
+- _translate_responses_tools_to_chat(): flat Responses tool shape ->
+  nested Chat Completions shape, drops non-function built-in tools,
+  returns None for empty lists.
+- _translate_responses_tool_choice_to_chat(): passes string choices
+  through, converts {type:function,name:X} to the nested shape.
+- _normalise_responses_input(): maps function_call_output items to
   role="tool" ChatMessages with tool_call_id, and function_call items to
   assistant messages with tool_calls.
-- _chat_tool_calls_to_responses_output() preserves call_id and drops
+- _chat_tool_calls_to_responses_output(): keeps call_id, drops
   non-function tool calls.
-- ResponsesOutputFunctionCall and ResponsesResponse round-trip tool-call
+- ResponsesOutputFunctionCall / ResponsesResponse round-trip tool-call
   outputs without losing fields.
 
 No running server or GPU required.
@@ -26,12 +26,15 @@ No running server or GPU required.
 
 import os
 import sys
+import asyncio
+from types import SimpleNamespace
 
 _backend = os.path.join(os.path.dirname(__file__), "..")
 sys.path.insert(0, _backend)
 
 import json
 
+import httpx
 import pytest
 from pydantic import ValidationError
 
@@ -52,8 +55,10 @@ from models.inference import (
     ResponsesUsage,
 )
 from routes.inference import (
+    _build_chat_request,
     _chat_tool_calls_to_responses_output,
     _normalise_responses_input,
+    _responses_stream,
     _translate_responses_tool_choice_to_chat,
     _translate_responses_tools_to_chat,
 )
@@ -104,9 +109,9 @@ class TestResponsesRequestTools:
         assert req.parallel_tool_calls is True
 
     def test_builtin_tool_type_passes_validation(self):
-        """Non-function built-in tools (web_search, file_search, mcp, ...) must
-        not raise at request validation so SDKs that default to them don't
-        fail on Studio; they are filtered out during translation."""
+        """Non-function built-in tools (web_search, file_search, mcp, ...)
+        must not raise at validation so SDKs that default to them don't
+        fail on Studio; they're filtered out during translation."""
         req = ResponsesRequest(
             input = "hi",
             tools = [{"type": "web_search_preview"}],
@@ -249,8 +254,8 @@ class TestToolChoiceTranslation:
         ) == {"type": "function", "function": {"name": "get_weather"}}
 
     def test_already_chat_nested_shape_passes_through(self):
-        """If a client happens to send the Chat Completions nested shape,
-        we don't double-wrap it."""
+        """A client sending the Chat Completions nested shape isn't
+        double-wrapped."""
         already_nested = {"type": "function", "function": {"name": "get_weather"}}
         assert _translate_responses_tool_choice_to_chat(already_nested) == already_nested
 
@@ -259,6 +264,26 @@ class TestToolChoiceTranslation:
         assert _translate_responses_tool_choice_to_chat(obj) == obj
 
 
+class TestBuildChatRequest:
+    def test_parallel_tool_calls_false_is_preserved_for_passthrough_caps(self):
+        payload = ResponsesRequest(
+            input = "hi",
+            tools = [
+                {
+                    "type": "function",
+                    "name": "lookup",
+                    "parameters": {"type": "object"},
+                }
+            ],
+            parallel_tool_calls = False,
+        )
+        messages = [ChatMessage(role = "user", content = "hi")]
+
+        chat_req = _build_chat_request(payload, messages, stream = True)
+
+        assert chat_req.parallel_tool_calls is False
+
+
 # =====================================================================
 # _normalise_responses_input — multi-turn tool mapping
 # =====================================================================
@@ -297,10 +322,10 @@ class TestNormaliseResponsesInputWithTools:
 
     def test_instructions_plus_developer_message_are_merged(self):
         """Codex CLI sends `instructions` (system prompt) AND a developer
-        message in `input`. Strict chat templates (harmony / gpt-oss, Qwen3,
-        ...) raise "System message must be at the beginning" when two
-        separate system-role messages appear, so we must emit exactly one
-        merged system message at the top.
+        message in `input`. Strict chat templates (harmony / gpt-oss,
+        Qwen3, ...) raise "System message must be at the beginning" on two
+        separate system-role messages, so we emit exactly one merged
+        system message at the top.
         """
         payload = ResponsesRequest(
             instructions = "Base instructions.",
@@ -314,14 +339,14 @@ class TestNormaliseResponsesInputWithTools:
         assert len(system_roles) == 1
         assert "Base instructions." in system_roles[0].content
         assert "Developer override." in system_roles[0].content
-        # System must be the very first message for strict templates.
+        # System must be the first message for strict templates.
         assert msgs[0].role == "system"
         assert msgs[1].role == "user"
 
     def test_developer_message_after_user_is_still_hoisted(self):
-        """Multi-turn conversations where a developer message appears after
-        user turns must still produce a single leading system message, not
-        a mid-conversation system that strict templates reject."""
+        """A developer message appearing after user turns must still
+        produce a single leading system message, not a mid-conversation
+        system that strict templates reject."""
         payload = ResponsesRequest(
             input = [
                 {"role": "user", "content": "Hello"},
@@ -424,6 +449,130 @@ class TestChatToolCallsToResponsesOutput:
         assert items[0]["arguments"] == ""
 
 
+# =====================================================================
+# Streaming Responses adapter
+# =====================================================================
+
+
+class TestResponsesStreamAdapter:
+    class _Request:
+        async def is_disconnected(self):
+            return False
+
+    @staticmethod
+    async def _collect(response):
+        chunks = []
+        async for chunk in response.body_iterator:
+            chunks.append(chunk.decode() if isinstance(chunk, bytes) else chunk)
+        return chunks
+
+    @staticmethod
+    def _payloads(lines, event_name):
+        prefix = f"event: {event_name}\n"
+        return [
+            json.loads(line.split("data: ", 1)[1].strip())
+            for line in lines
+            if line.startswith(prefix)
+        ]
+
+    def test_requests_usage_and_caps_parallel_tool_calls(self, monkeypatch):
+        import routes.inference as inf_mod
+
+        captured = {}
+
+        def handler(request: httpx.Request) -> httpx.Response:
+            captured["body"] = json.loads(request.content.decode())
+            chunks = [
+                {
+                    "choices": [
+                        {
+                            "delta": {
+                                "tool_calls": [
+                                    {
+                                        "index": 0,
+                                        "id": "call_0",
+                                        "type": "function",
+                                        "function": {"name": "first", "arguments": "{}"},
+                                    },
+                                    {
+                                        "index": 1,
+                                        "id": "call_1",
+                                        "type": "function",
+                                        "function": {"name": "second", "arguments": "{}"},
+                                    },
+                                ]
+                            }
+                        }
+                    ]
+                },
+                {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}},
+            ]
+            content = "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks)
+            content += "data: [DONE]\n\n"
+            return httpx.Response(
+                200,
+                content = content.encode(),
+                headers = {"content-type": "text/event-stream"},
+            )
+
+        transport = httpx.MockTransport(handler)
+        real_async_client = httpx.AsyncClient
+
+        def _client(*args, **kwargs):
+            return real_async_client(
+                transport = transport,
+                timeout = kwargs.get("timeout", 600),
+            )
+
+        monkeypatch.setattr(inf_mod.httpx, "AsyncClient", _client)
+        monkeypatch.setattr(
+            inf_mod,
+            "get_llama_cpp_backend",
+            lambda: SimpleNamespace(
+                is_loaded = True,
+                is_vision = False,
+                context_length = 4096,
+                base_url = "http://llama.test",
+            ),
+        )
+
+        payload = ResponsesRequest(
+            input = "hi",
+            stream = True,
+            parallel_tool_calls = False,
+            tools = [
+                {
+                    "type": "function",
+                    "name": "first",
+                    "parameters": {"type": "object"},
+                },
+                {
+                    "type": "function",
+                    "name": "second",
+                    "parameters": {"type": "object"},
+                },
+            ],
+        )
+        messages = [ChatMessage(role = "user", content = "hi")]
+
+        async def run():
+            response = await _responses_stream(payload, messages, self._Request())
+            return await self._collect(response)
+
+        lines = asyncio.run(run())
+
+        assert captured["body"]["stream_options"] == {"include_usage": True}
+        joined = "".join(lines)
+        assert "call_0" in joined
+        assert "call_1" not in joined
+        completed = self._payloads(lines, "response.completed")[0]
+        assert completed["response"]["usage"] == {
+            "input_tokens": 2,
+            "output_tokens": 3,
+            "total_tokens": 5,
+        }
+
+
 # =====================================================================
 # Response model — ResponsesOutputFunctionCall / mixed output
 # =====================================================================
@@ -486,8 +635,8 @@ class TestCodexStyleRequestShapes:
     """Regression tests for the request shapes OpenAI Codex CLI sends."""
 
     def test_assistant_replay_output_text_accepted(self):
-        """Codex replays prior assistant turns with `output_text` content.
-        Before, this triggered a 422 on every turn after the first."""
+        """Codex replays prior assistant turns with `output_text` content;
+        this used to 422 on every turn after the first."""
         req = ResponsesRequest(
             input = [
                 {"role": "user", "content": "Hi"},
@@ -514,7 +663,7 @@ class TestCodexStyleRequestShapes:
 
     def test_reasoning_item_accepted_as_unknown(self):
         """`reasoning` items replayed from prior o-series turns must not
-        fail validation — Codex preserves them in multi-turn."""
+        fail validation — Codex keeps them in multi-turn."""
         req = ResponsesRequest(
             input = [
                 {"role": "user", "content": "Hi"},
@@ -532,7 +681,7 @@ class TestCodexStyleRequestShapes:
 
     def test_unknown_content_part_type_accepted(self):
         """Unknown content-part types (e.g. future input_audio) validate as
-        ResponsesUnknownContentPart so the whole request doesn't 422."""
+        ResponsesUnknownContentPart so the request doesn't 422."""
         req = ResponsesRequest(
             input = [
                 {
@@ -596,15 +745,15 @@ class TestCodexStyleRequestShapes:
             ],
         )
         msgs = _normalise_responses_input(payload)
-        # Single leading merged system; no mid-conversation system.
+        # One leading merged system; no mid-conversation system.
         assert msgs[0].role == "system"
         assert sum(1 for m in msgs if m.role == "system") == 1
         assert "Base instructions." in msgs[0].content
         assert "Dev override." in msgs[0].content
 
         roles = [m.role for m in msgs[1:]]
-        # Reasoning item is dropped. Order: user, assistant(tool_calls),
-        # tool, assistant(text), user.
+        # Reasoning dropped. Order: user, assistant(tool_calls), tool,
+        # assistant(text), user.
         assert roles == ["user", "assistant", "tool", "assistant", "user"]
         assert msgs[2].tool_calls is not None
         assert msgs[3].role == "tool"
@@ -612,9 +761,9 @@ class TestCodexStyleRequestShapes:
         assert msgs[4].content == "It's 20°C."
 
     def test_single_output_text_part_flattens_to_string(self):
-        """ChatMessage assistant role prefers plain string content — tests
-        confirm we don't forward a single-part array that would otherwise
-        force legacy chat templates into multimodal handling."""
+        """ChatMessage assistant role prefers plain string content — we
+        don't forward a single-part array that would force legacy chat
+        templates into multimodal handling."""
         payload = ResponsesRequest(
             input = [
                 {
@@ -630,9 +779,9 @@ class TestCodexStyleRequestShapes:
 
 
 class TestTranslatedMessagesValidate:
-    """Verify that the messages produced by _normalise_responses_input
-    satisfy ChatMessage's role-shape validator so the downstream /v1/chat/
-    completions pass-through does not reject them."""
+    """Messages from _normalise_responses_input satisfy ChatMessage's
+    role-shape validator so the downstream /v1/chat/completions
+    pass-through doesn't reject them."""
 
     def test_round_trip_multi_turn(self):
         payload = ResponsesRequest(
@@ -654,6 +803,6 @@ class TestTranslatedMessagesValidate:
         )
         msgs = _normalise_responses_input(payload)
         for m in msgs:
-            # Constructing a fresh ChatMessage from the dump round-trips the
-            # role-shape validator — the key invariant for the passthrough.
+            # Building a fresh ChatMessage from the dump round-trips the
+            # role-shape validator — the passthrough's key invariant.
             ChatMessage(**m.model_dump(exclude_none = True))
diff --git a/studio/backend/tests/test_rocm_oom_guard.py b/studio/backend/tests/test_rocm_oom_guard.py
index d27ee83f1d..21edabd142 100644
--- a/studio/backend/tests/test_rocm_oom_guard.py
+++ b/studio/backend/tests/test_rocm_oom_guard.py
@@ -3,16 +3,12 @@
 
 """Unit tests for _rocm_classify_unified_memory (ROCm OOM-guard classifier).
 
-Covers the three classification paths:
-  Path 1 – canonical gcnArchName attribute present.
-  Path 2 – gcnArchName absent, alternate-spelling attribute present.
-  Path 3 – ALL arch attrs absent; falls back to device-name substring match.
+Three paths: (1) canonical gcnArchName, (2) alternate-spelling attr, (3) all
+arch attrs absent -> device-name substring match.
 
-Regression for: Strix Halo (gfx1151) misclassified as discrete on AMD SDK /
-Radeon wheels that populate props.name = "Radeon 8060S Graphics" but do NOT
-set any gcnArchName attribute.  Without the 8060s/8050s name patterns the
-fallback returned is_unified=False, applying the 0.90 fraction instead of
-0.80 and leaving only ~12.8 GiB OS headroom on a 128 GiB unified-memory pool.
+Regression: Strix Halo (gfx1151) was misclassified as discrete on Radeon wheels
+that set props.name="Radeon 8060S Graphics" but no gcnArchName, applying the
+wrong headroom factor on a 128 GiB unified-memory pool.
 """
 
 from __future__ import annotations
@@ -62,7 +58,7 @@ class TestCanonicalGcnArchName:
         assert is_unified is True
 
     def test_canonical_attr_wins_over_name(self) -> None:
-        """Arch attr takes priority; device name should be ignored."""
+        """Arch attr takes priority; device name is ignored."""
         # Discrete arch, but name looks like a unified SKU — arch must win.
         props = _props(gcnArchName = "gfx1100", name = "Radeon 890M")
         gcn, is_unified = _rocm_classify_unified_memory(props)
@@ -97,7 +93,7 @@ class TestAlternateSpellingFallback:
         assert is_unified is False
 
     def test_first_non_empty_attr_wins(self) -> None:
-        """When multiple alternate attrs are present the first non-empty one wins."""
+        """With multiple alternate attrs, the first non-empty one wins."""
         props = _props(gcn_arch_name = "gfx1151", arch_name = "gfx1100", name = "irrelevant")
         gcn, is_unified = _rocm_classify_unified_memory(props)
         assert gcn == "gfx1151"
@@ -147,7 +143,7 @@ class TestDeviceNameFallback:
             "Radeon RX 6900 XT",
             "Radeon Pro W7900",
             "AMD Instinct MI300X",
-            # Names that contain superficially similar substrings but are discrete
+            # Superficially similar substrings but discrete
             "Radeon RX 580",
             "Radeon VII",
         ],
@@ -161,7 +157,7 @@ class TestDeviceNameFallback:
         ), f"discrete device {device_name!r} should NOT be classified as unified-memory"
 
     def test_empty_name_returns_false(self) -> None:
-        """Completely absent name must not crash and must default to discrete."""
+        """Absent name must not crash and must default to discrete."""
         props = _props()  # no 'name' attr at all
         gcn, is_unified = _rocm_classify_unified_memory(props)
         assert gcn == ""
diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py
index 5e3d2cc9d1..1e8fb9e2b2 100644
--- a/studio/backend/tests/test_safetensors_capability_advertise.py
+++ b/studio/backend/tests/test_safetensors_capability_advertise.py
@@ -1,11 +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
 
-"""
-Capability advertisement contract: classifier honesty, worker→
-orchestrator IPC hop, and route-layer end-to-end. Pure helpers + fakes;
-no torch / transformers import.
-"""
+"""Capability advertisement contract: classifier honesty, worker->orchestrator
+IPC hop, route-layer end-to-end. Pure helpers + fakes; no torch/transformers."""
 
 from __future__ import annotations
 
@@ -116,7 +113,7 @@ def test_detect_safetensors_features_none_template_returns_all_false():
 
 
 def test_detect_safetensors_features_gptoss_disables_tools():
-    """gpt-oss Harmony: tools intentionally off even if template marks it."""
+    """gpt-oss Harmony: tools off even if template marks it."""
     from routes.inference import _detect_safetensors_features
 
     backend = MagicMock()
@@ -129,11 +126,9 @@ def test_detect_safetensors_features_gptoss_disables_tools():
     assert flags["supports_tools"] is False
 
 
-# Llama-3 / Mistral templates advertise tool handling but the model emits
-# tool calls in <|python_tag|> / [TOOL_CALLS] format -- not the
-#  /  / [TOOL_CALLS],
+# which our parser can't read. The route helper must not flip supports_tools=True
+# for them, else the UI enables a pill the agentic loop can't honour.
 
 LLAMA3_TEMPLATE = """
 {%- if tools %}
@@ -207,9 +202,8 @@ def test_detect_safetensors_features_function_xml_format_keeps_tools_on():
     assert flags["supports_tools"] is True
 
 
-# Qwen3.5 family pins -- the live GGUF + safetensors templates fetched
-# from the unsloth/Qwen3.5-0.8B(-GGUF) repos both wrap tool calls as
-# ``\n...``. Capture a faithful slice so the
+# Qwen3.5 family pin: the live GGUF + safetensors templates both wrap tool
+# calls as ``\n...``. Faithful slice so the
 # classifier never silently regresses for this family.
 
 QWEN35_TOOL_INSTRUCTION = (
@@ -234,7 +228,7 @@ QWEN35_TOOL_INSTRUCTION = (
 
 
 def test_detect_safetensors_features_qwen35_keeps_tools_on():
-    """unsloth/Qwen3.5-0.8B family must surface tools+reasoning enabled."""
+    """unsloth/Qwen3.5-0.8B family must surface tools+reasoning on."""
     from routes.inference import _detect_safetensors_features
 
     backend = SimpleNamespace(active_model_name = "unsloth/Qwen3.5-0.8B")
@@ -248,7 +242,7 @@ def test_detect_safetensors_features_qwen35_keeps_tools_on():
 
 
 def test_orchestrator_mirrors_chat_template_info_into_models_dict():
-    """Worker → orchestrator must copy chat_template_info verbatim."""
+    """Worker → orchestrator copies chat_template_info verbatim."""
     from core.inference.orchestrator import InferenceOrchestrator
 
     orch = InferenceOrchestrator.__new__(InferenceOrchestrator)
@@ -274,7 +268,7 @@ def test_orchestrator_mirrors_chat_template_info_into_models_dict():
         },
     }
 
-    # Replay orchestrator.load_model's mirror block verbatim.
+    # Replay orchestrator.load_model's mirror block.
     orch.active_model_name = model_info["identifier"]
     orch.models[orch.active_model_name] = {
         "is_vision": model_info.get("is_vision", False),
@@ -386,7 +380,7 @@ def test_worker_load_reply_payload_includes_chat_template_info():
 
 
 def test_worker_load_reply_payload_survives_missing_template():
-    """Tokenizer with no chat_template still produces a valid reply."""
+    """Tokenizer with no chat_template still yields a valid reply."""
 
     class _StubBackend:
         def __init__(self):
@@ -421,7 +415,7 @@ def test_worker_load_reply_payload_survives_missing_template():
 
 
 def test_route_layer_emits_supports_tools_true_for_qwen3_safetensors():
-    """End-to-end: Qwen3 safetensors flips supports_tools=True."""
+    """E2E: Qwen3 safetensors flips supports_tools=True."""
     from routes.inference import _detect_safetensors_features
 
     backend = SimpleNamespace(
diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py
index ff5653c742..16ed1dc182 100644
--- a/studio/backend/tests/test_safetensors_tool_loop.py
+++ b/studio/backend/tests/test_safetensors_tool_loop.py
@@ -1,30 +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
 
-"""
-Tests for the safetensors agentic tool loop.
+"""Tests for the safetensors agentic tool loop.
 
-Covers the shared ``tool_call_parser`` helpers and the cumulative-text
-state machine inside ``safetensors_agentic.run_safetensors_tool_loop``.
-The loop is exercised with hand-crafted fake single-turn generators so
-no model load is needed; the tests run in CI under a few seconds.
-
-Edge cases under coverage:
-* Plain answers (no tool calls) flush full content.
-* Single ``{json}`` triggers the tool and re-enters.
-* Single ``...`` XML form triggers the same path.
-* Truncated unclosed ```` is still parsed.
-* Tool result is fed back as ``role=tool`` for the next iteration.
-* Bad JSON inside ```` does not raise and (when healed) is
-  routed as a ``{"query": ...}`` web search call.
-* Duplicate tool calls produce a synthetic "do not repeat" result the
-  second time.
-* ``__IMAGES__`` sentinel is stripped before the model sees the result.
-* Tool execution errors are tagged so the model gets a nudge but the
-  loop keeps streaming.
-* Cancel is honoured between iterations.
-* ``max_tool_iterations`` cap is respected and a final-answer attempt
-  closes the stream cleanly.
+Covers the ``tool_call_parser`` helpers and the cumulative-text state machine in
+``run_safetensors_tool_loop``, run against fake single-turn generators (no model
+load). Edge cases: plain answers, JSON and XML tool-call forms, truncated/unclosed
+calls, tool-result feedback, bad-JSON heal, duplicate-call short-circuit,
+``__IMAGES__`` sentinel stripping, executor errors, cancel, and the iteration cap.
 """
 
 import threading
@@ -37,6 +20,7 @@ from core.inference.safetensors_agentic import (
     _coerce_arguments,
     _detect_render_html_tool_start,
     run_safetensors_tool_loop,
+    strip_tool_markup_streaming,
 )
 from core.inference.tool_call_parser import (
     has_tool_signal,
@@ -64,12 +48,17 @@ class TestParser:
         assert "hello" in tc["function"]["arguments"]
 
     def test_json_tool_call_unclosed(self):
-        # No ; balanced-brace extractor must still close.
+        # No ; balanced-brace extractor must still close it.
         text = '{"name":"python","arguments":{"code":"print(1)"}}'
         result = parse_tool_calls_from_text(text)
         assert len(result) == 1
         assert result[0]["function"]["name"] == "python"
 
+    def test_json_tool_call_unclosed_requires_healing(self):
+        text = '{"name":"python","arguments":{"code":"print(1)"}}'
+        assert parse_tool_calls_from_text(text)[0]["function"]["name"] == "python"
+        assert parse_tool_calls_from_text(text, allow_incomplete = False) == []
+
     def test_xml_function_call(self):
         text = "print('hi')"
         result = parse_tool_calls_from_text(text)
@@ -85,10 +74,14 @@ class TestParser:
         assert result[0]["function"]["name"] == "terminal"
         assert "ls -la" in result[0]["function"]["arguments"]
 
+    def test_xml_unclosed_requires_healing(self):
+        text = "ls -la"
+        assert parse_tool_calls_from_text(text)[0]["function"]["name"] == "terminal"
+        assert parse_tool_calls_from_text(text, allow_incomplete = False) == []
+
     def test_code_with_embedded_xml(self):
-        # A code parameter contains the literal . Must not
-        # truncate the value because the parser uses end-of-body as the
-        # only boundary for single-parameter calls.
+        # A code parameter with a literal  must not truncate: the
+        # parser uses end-of-body as the only boundary for single-param calls.
         text = (
             "html = ''\n"
             "print('hi')"
@@ -121,7 +114,7 @@ class TestParser:
     def test_bad_json_does_not_raise(self):
         text = "{not valid json}"
         result = parse_tool_calls_from_text(text)
-        # Bad JSON is silently dropped; caller can fall back to text.
+        # Bad JSON is dropped silently; caller can fall back to text.
         assert result == []
 
     def test_has_tool_signal(self):
@@ -147,11 +140,28 @@ class TestParser:
 
     def test_strip_markup_unclosed_final(self):
         text = "before {partial"
-        # With final=True the trailing run is dropped.
+        # final=True drops the trailing run.
         assert strip_tool_markup(text, final = True) == "before"
         # Without final=True the unclosed run is preserved.
         assert "partial" in strip_tool_markup(text)
 
+    def test_streaming_strip_respects_disabled_healing(self):
+        raw = 'before {"name":"web_search"'
+        assert strip_tool_markup_streaming(raw, auto_heal_tool_calls = False) == raw
+        assert strip_tool_markup_streaming(raw) == "before "
+
+    def test_streaming_strip_respects_disabled_healing_without_tool_protocol(self):
+        raw = 'before {"name":"web_search"'
+        assert strip_tool_markup_streaming(raw, auto_heal_tool_calls = False) == raw
+        assert (
+            strip_tool_markup_streaming(
+                raw,
+                auto_heal_tool_calls = False,
+                tool_protocol_active = True,
+            )
+            == "before "
+        )
+
 
 # ────────────────────────────────────────────────────────────────────
 # run_safetensors_tool_loop
@@ -220,8 +230,7 @@ def _make_loop(
 ):
     """Build a configured loop with a multi-turn fake generator.
 
-    ``turns`` is a list of chunk-lists; iteration N yields chunks from
-    ``turns[N]``.
+    ``turns`` is a list of chunk-lists; iteration N yields chunks from ``turns[N]``.
     """
     turn_iter = iter(turns)
 
@@ -249,6 +258,41 @@ def _make_loop(
     ), exec_fn
 
 
+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."])
+
+    def fake_single_turn(_messages, *, active_tools = None):
+        captured_tool_names.append(
+            [
+                (tool.get("function") or {}).get("name")
+                for tool in (active_tools or [])
+                if (tool.get("function") or {}).get("name")
+            ]
+        )
+        if len(captured_tool_names) == 1:
+            yield '{"name":"render_html","arguments":{"code":"one"}}'
+        else:
+            yield "Done."
+
+    events = _collect_events(
+        run_safetensors_tool_loop(
+            single_turn = fake_single_turn,
+            messages = [{"role": "user", "content": "make html"}],
+            tools = [
+                {"type": "function", "function": {"name": "render_html"}},
+                {"type": "function", "function": {"name": "web_search"}},
+            ],
+            execute_tool = exec_fn,
+            max_tool_iterations = 3,
+        )
+    )
+
+    assert exec_fn.calls == [("render_html", {"code": "one"})]
+    assert captured_tool_names == [["render_html", "web_search"], ["web_search"]]
+    assert any(event.get("type") == "content" and event.get("text") == "Done." for event in events)
+
+
 class TestLoopBasic:
     def test_plain_answer(self):
         # No tool XML; loop should yield content then status="".
@@ -260,7 +304,7 @@ class TestLoopBasic:
         contents = [e for e in events if e["type"] == "content"]
         statuses = [e for e in events if e["type"] == "status"]
         assert contents, "expected at least one content event"
-        # Final cumulative content should contain the answer.
+        # Final cumulative content must contain the answer.
         final_text = contents[-1]["text"]
         assert "Hello world!" in final_text
         assert statuses and statuses[-1]["text"] == ""
@@ -268,13 +312,13 @@ class TestLoopBasic:
     def test_single_tool_then_answer(self):
         loop, exec_fn = _make_loop(
             turns = [
-                # : tool call only.
+                # Tool call only.
                 [
                     '{"name":"web_search",',
                     '"arguments":{"query":"weather"}}',
                     "",
                 ],
-                # : final answer.
+                # Final answer.
                 ["The ", "weather is ", "sunny."],
             ],
             exec_results = ["Sunny and 22C"],
@@ -284,7 +328,7 @@ class TestLoopBasic:
 
         assert "tool_start" in kinds
         assert "tool_end" in kinds
-        # Tool was actually called with the parsed arguments.
+        # Tool was called with the parsed arguments.
         assert exec_fn.calls == [("web_search", {"query": "weather"})]
 
         tool_start = next(e for e in events if e["type"] == "tool_start")
@@ -402,8 +446,8 @@ class TestLoopBasic:
     def test_truncated_unclosed_tool_call(self):
         loop, exec_fn = _make_loop(
             turns = [
-                # No ; balanced-brace parser must still
-                # succeed because the JSON itself is balanced.
+                # No ; balanced-brace parser still succeeds because
+                # the JSON itself is balanced.
                 ['{"name":"web_search","arguments":{"query":"x"}}'],
                 ["done"],
             ],
@@ -413,14 +457,10 @@ class TestLoopBasic:
         assert exec_fn.calls == [("web_search", {"query": "x"})]
 
     def test_bad_json_healed_to_query(self):
-        # Tool call with non-JSON string arguments. With auto_heal_tool_calls
-        # the string is routed as {"query": ...}.
+        # Non-JSON string arguments heal to {"query": ...} under auto_heal_tool_calls.
         loop, exec_fn = _make_loop(
             turns = [
-                # JSON inside the tool call is well-formed; the
-                # ``arguments`` is a string that is not itself valid
-                # JSON for ``_coerce_arguments`` to parse, so the
-                # heal path runs.
+                # ``arguments`` is a string _coerce_arguments can't parse, so heal runs.
                 ['{"name":"web_search","arguments":"hello world"}'],
                 ["ok"],
             ],
@@ -432,29 +472,162 @@ class TestLoopBasic:
 
 
 class TestLoopBehaviour:
-    def test_duplicate_tool_call_synthetic_result(self):
-        # Two identical successful calls in a row: the second is short-
-        # circuited with a "do not repeat" message and execute_tool is
-        # called only once.
-        loop, exec_fn = _make_loop(
-            turns = [
+    def test_duplicate_tool_call_internal_noop(self):
+        captured_messages: list[list[dict]] = []
+        turns = iter(
+            [
                 ['{"name":"web_search","arguments":{"query":"x"}}'],
                 ['{"name":"web_search","arguments":{"query":"x"}}'],
                 ["final"],
-            ],
-            exec_results = ["search-result-1"],
+            ]
+        )
+
+        def fake_single_turn(messages):
+            captured_messages.append([dict(message) for message in messages])
+            chunks = next(turns)
+            acc = ""
+            for chunk in chunks:
+                acc += chunk
+                yield acc
+
+        exec_fn = FakeExecuteTool(["search-result-1"])
+        events = _collect_events(
+            run_safetensors_tool_loop(
+                single_turn = fake_single_turn,
+                messages = [{"role": "user", "content": "hi"}],
+                tools = [{"type": "function", "function": {"name": "web_search"}}],
+                execute_tool = exec_fn,
+                max_tool_iterations = 3,
+            )
+        )
+
+        assert exec_fn.calls == [("web_search", {"query": "x"})]
+        assert [e["tool_call_id"] for e in events if e["type"] == "tool_end"] == ["call_0"]
+        assert not [
+            e
+            for e in events
+            if e.get("tool_call_id") == "call_1" and e.get("type") in {"tool_start", "tool_end"}
+        ]
+        duplicate_nudges = [
+            message
+            for message in captured_messages[-1]
+            if message.get("role") == "user"
+            and "already completed successfully" in message.get("content", "")
+        ]
+        assert len(duplicate_nudges) == 1
+
+    def test_duplicate_tool_call_internal_noop_allows_distinct_followup_tool(self):
+        captured_messages: list[list[dict]] = []
+        captured_tool_names: list[list[str]] = []
+        turns = iter(
+            [
+                ['{"name":"web_search","arguments":{"query":"x"}}'],
+                ['{"name":"web_search","arguments":{"query":"x"}}'],
+                ['{"name":"python","arguments":{"code":"print(1)"}}'],
+                ["final"],
+            ]
+        )
+
+        def fake_single_turn(messages, active_tools = None):
+            captured_messages.append([dict(message) for message in messages])
+            captured_tool_names.append(
+                [
+                    tool["function"]["name"]
+                    for tool in (active_tools or [])
+                    if tool.get("function", {}).get("name")
+                ]
+            )
+            chunks = next(turns)
+            acc = ""
+            for chunk in chunks:
+                acc += chunk
+                yield acc
+
+        exec_fn = FakeExecuteTool(["search-result-1", "python-result"])
+        events = _collect_events(
+            run_safetensors_tool_loop(
+                single_turn = fake_single_turn,
+                messages = [{"role": "user", "content": "hi"}],
+                tools = [
+                    {"type": "function", "function": {"name": "web_search"}},
+                    {"type": "function", "function": {"name": "python"}},
+                ],
+                execute_tool = exec_fn,
+                max_tool_iterations = 4,
+            )
+        )
+
+        assert exec_fn.calls == [
+            ("web_search", {"query": "x"}),
+            ("python", {"code": "print(1)"}),
+        ]
+        assert [e["tool_call_id"] for e in events if e["type"] == "tool_end"] == [
+            "call_0",
+            "call_2",
+        ]
+        assert not [
+            e
+            for e in events
+            if e.get("tool_call_id") == "call_1" and e.get("type") in {"tool_start", "tool_end"}
+        ]
+        duplicate_nudges = [
+            message
+            for message in captured_messages[2]
+            if message.get("role") == "user"
+            and "already completed successfully" in message.get("content", "")
+        ]
+        assert len(duplicate_nudges) == 1
+        assert captured_tool_names[2] == ["web_search", "python"]
+
+    def test_repeated_duplicate_noop_transitions_to_final_attempt(self):
+        captured_tool_names: list[list[str]] = []
+        turns = iter(
+            [
+                ['{"name":"web_search","arguments":{"query":"x"}}'],
+                ['{"name":"web_search","arguments":{"query":"x"}}'],
+                ['{"name":"web_search","arguments":{"query":"x"}}'],
+                ["final from first result"],
+            ]
+        )
+
+        def fake_single_turn(messages, active_tools = None):
+            captured_tool_names.append(
+                [
+                    (tool.get("function") or {}).get("name")
+                    for tool in (active_tools or [])
+                    if (tool.get("function") or {}).get("name")
+                ]
+            )
+            chunks = next(turns)
+            acc = ""
+            for chunk in chunks:
+                acc += chunk
+                yield acc
+
+        exec_fn = FakeExecuteTool(["search-result"])
+        events = _collect_events(
+            run_safetensors_tool_loop(
+                single_turn = fake_single_turn,
+                messages = [{"role": "user", "content": "hi"}],
+                tools = [{"type": "function", "function": {"name": "web_search"}}],
+                execute_tool = exec_fn,
+                max_tool_iterations = 10,
+            )
+        )
+
+        assert exec_fn.calls == [("web_search", {"query": "x"})]
+        assert [
+            event.get("tool_call_id") for event in events if event.get("type") == "tool_end"
+        ] == ["call_0"]
+        assert captured_tool_names[-1] == []
+        assert any(
+            event.get("type") == "content" and "final from first result" in event.get("text", "")
+            for event in events
         )
-        events = _collect_events(loop)
-        # Only one real call.
-        assert len(exec_fn.calls) == 1
-        tool_end_events = [e for e in events if e["type"] == "tool_end"]
-        assert len(tool_end_events) == 2
-        assert "do not repeat" in tool_end_events[1]["result"].lower()
 
     def test_image_sentinel_stripped_from_model_feed(self):
-        # The tool result has a frontend image sentinel that should be
-        # stripped before being fed back into the next turn, BUT the
-        # tool_end event still carries the raw result for the UI.
+        # The image sentinel is stripped before the next turn, but tool_end still
+        # carries the raw result for the UI.
         loop, exec_fn = _make_loop(
             turns = [
                 ['{"name":"python","arguments":{"code":"plot()"}}'],
@@ -490,7 +663,7 @@ class TestLoopBehaviour:
                 auto_heal_tool_calls = True,
             )
         )
-        # Model's second turn must not see "__IMAGES__".
+        # The model's second turn must not see "__IMAGES__".
         assert len(captured) >= 2
         tool_msgs = [m for m in captured[1] if m.get("role") == "tool"]
         assert tool_msgs, "no tool message reached the model"
@@ -539,7 +712,7 @@ class TestLoopBehaviour:
         events = _collect_events(loop)
         tool_end = next(e for e in events if e["type"] == "tool_end")
         assert tool_end["result"].startswith("Error")
-        # The loop must still produce a content event after the failure.
+        # The loop must still emit a content event after the failure.
         contents = [e for e in events if e["type"] == "content"]
         assert contents
 
@@ -560,8 +733,7 @@ class TestLoopControl:
     def test_cancel_event_breaks_loop(self):
         cancel = threading.Event()
         cancel.set()
-        # Even with a fake stream that emits tool calls, the loop must
-        # bail before invoking execute_tool when cancel is set.
+        # With cancel set, the loop bails before invoking execute_tool.
         exec_fn = FakeExecuteTool([])
         events = list(
             run_safetensors_tool_loop(
@@ -578,13 +750,13 @@ class TestLoopControl:
         assert exec_fn.calls == []
 
     def test_max_iterations_caps_loop(self):
-        # The loop should stop after max_tool_iterations even if the
-        # model keeps asking for tools, then emit a final-attempt round.
+        # The loop stops after max_tool_iterations even if the model keeps
+        # asking for tools, then emits a final-attempt round.
         loop, exec_fn = _make_loop(
             turns = [
-                # : tool call (executes once)
+                # Tool call (executes once).
                 ['{"name":"web_search","arguments":{"query":"a"}}'],
-                # : model gives a final answer when nudged.
+                # Model gives a final answer when nudged.
                 ["here is the final answer"],
             ],
             exec_results = ["result"],
@@ -592,13 +764,13 @@ class TestLoopControl:
         )
         events = _collect_events(loop)
         contents = [e for e in events if e["type"] == "content"]
-        # Final content must include the final answer.
+        # Final content must contain the final answer.
         assert contents and "final answer" in contents[-1]["text"]
 
 
 class TestStatusFormatting:
     def test_status_for_known_tools(self):
-        # Use the private helper directly to verify status formatting.
+        # Call the private helper directly to verify status formatting.
         assert (
             safetensors_agentic._status_for_tool("web_search", {"query": "abc"}) == "Searching: abc"
         )
@@ -617,16 +789,13 @@ class TestStatusFormatting:
 
 class TestProseMentioningToolCall:
     def test_assistant_prose_with_literal_tool_call_text_survives(self):
-        # Regression: if the assistant text legitimately mentions
-        # ```` as a literal string and the parser finds no
-        # actual call, the loop must surface the full content instead
-        # of silently stripping everything past the literal marker.
+        # Regression: prose that mentions a literal ```` (no real call)
+        # must surface in full, not be stripped past the marker.
         loop, exec_fn = _make_loop(
             turns = [
-                # : a real tool call so the loop moves to
-                # .
+                # A real tool call so the loop advances a turn.
                 ['{"name":"web_search","arguments":{"query":"x"}}'],
-                # : prose that mentions the literal text.
+                # Prose that mentions the literal text.
                 ["the docs say  means an LLM tool call wrapper"],
             ],
             exec_results = ["result"],
@@ -640,9 +809,8 @@ class TestProseMentioningToolCall:
         ), f"prose mentioning  should not be truncated; got {final!r}"
 
     def test_tool_result_with_tool_call_text_does_not_retrigger(self):
-        # Tool result text contains the literal ```` string.
-        # The loop must only parse the MODEL output, not the tool
-        # result, so we should see exactly one call.
+        # A literal ```` in the tool result must not re-trigger: the
+        # loop parses only model output, so exactly one call.
         loop, exec_fn = _make_loop(
             turns = [
                 ['{"name":"web_search","arguments":{"query":"x"}}'],
@@ -725,34 +893,62 @@ class TestChatTemplateHelper:
 
 class TestGuardrails:
     def test_disabled_tool_is_not_executed(self):
-        exec_fn = FakeExecuteTool([])
-        loop = run_safetensors_tool_loop(
-            single_turn = _fake_stream(
-                ['{"name":"terminal","arguments":{"command":"echo bypass"}}']
-            ),
-            messages = [{"role": "user", "content": "hi"}],
-            tools = [{"type": "function", "function": {"name": "web_search"}}],
-            execute_tool = exec_fn,
-            max_tool_iterations = 2,
-        )
-        events = _collect_events(loop)
-        assert exec_fn.calls == []
-        tool_ends = [e for e in events if e["type"] == "tool_end"]
-        assert tool_ends and "not enabled" in tool_ends[0]["result"].lower()
+        captured_messages: list[list[dict]] = []
 
-    def test_empty_tools_list_does_not_enforce_allowlist(self):
-        exec_fn = FakeExecuteTool(["OK"])
-        loop = run_safetensors_tool_loop(
-            single_turn = _fake_stream(
-                ['{"name":"python","arguments":{"code":"print(1)"}}']
-            ),
-            messages = [{"role": "user", "content": "hi"}],
-            tools = [],
-            execute_tool = exec_fn,
-            max_tool_iterations = 2,
+        def fake_single_turn(messages):
+            captured_messages.append([dict(message) for message in messages])
+            if len(captured_messages) == 1:
+                yield '{"name":"terminal","arguments":{"command":"echo bypass"}}'
+            else:
+                yield "final"
+
+        exec_fn = FakeExecuteTool([])
+        events = _collect_events(
+            run_safetensors_tool_loop(
+                single_turn = fake_single_turn,
+                messages = [{"role": "user", "content": "hi"}],
+                tools = [{"type": "function", "function": {"name": "web_search"}}],
+                execute_tool = exec_fn,
+                max_tool_iterations = 2,
+            )
+        )
+
+        assert exec_fn.calls == []
+        assert not [event for event in events if event.get("type") in {"tool_start", "tool_end"}]
+        disabled_nudges = [
+            message
+            for message in captured_messages[-1]
+            if message.get("role") == "user" and "not enabled" in message.get("content", "")
+        ]
+        assert len(disabled_nudges) == 1
+
+    def test_empty_tools_list_means_allow_all_in_core_loop(self):
+        turns = iter(
+            [
+                ['{"name":"python","arguments":{"code":"print(1)"}}'],
+                ["done"],
+            ]
+        )
+
+        def fake_single_turn(_messages, active_tools = None):
+            assert active_tools == []
+            acc = ""
+            for chunk in next(turns):
+                acc += chunk
+                yield acc
+
+        exec_fn = FakeExecuteTool(["OK"])
+        events = _collect_events(
+            run_safetensors_tool_loop(
+                single_turn = fake_single_turn,
+                messages = [{"role": "user", "content": "hi"}],
+                tools = [],
+                execute_tool = exec_fn,
+                max_tool_iterations = 2,
+            )
         )
-        _collect_events(loop)
         assert exec_fn.calls == [("python", {"code": "print(1)"})]
+        assert any(event.get("type") == "tool_end" for event in events)
 
     def test_max_iterations_zero_executes_no_tools(self):
         loop, exec_fn = _make_loop(
@@ -797,6 +993,67 @@ class TestGuardrails:
         _collect_events(loop)
         assert exec_fn.calls == [("web_search", {"query": "x"})]
 
+    def test_auto_heal_disabled_preserves_xml_on_final_no_tools_pass(self):
+        turns = iter(
+            [
+                ['{"name":"web_search","arguments":{"query":"x"}}'],
+                ['{"name":"web_search","arguments":{"query":"literal"}}'],
+            ]
+        )
+
+        def fake_single_turn(_messages, active_tools = None):
+            acc = ""
+            for chunk in next(turns):
+                acc += chunk
+                yield acc
+
+        exec_fn = FakeExecuteTool(["OK"])
+        events = _collect_events(
+            run_safetensors_tool_loop(
+                single_turn = fake_single_turn,
+                messages = [{"role": "user", "content": "show literal"}],
+                tools = [{"type": "function", "function": {"name": "web_search"}}],
+                execute_tool = exec_fn,
+                max_tool_iterations = 1,
+                auto_heal_tool_calls = False,
+            )
+        )
+        assert exec_fn.calls == [("web_search", {"query": "x"})]
+        assert any(
+            event.get("type") == "content" and "" in event.get("text", "")
+            for event in events
+        )
+
+    def test_auto_heal_disabled_does_not_repair_unclosed_tool_call(self):
+        loop, exec_fn = _make_loop(
+            turns = [
+                ['{"name":"web_search","arguments":{"query":"x"}}'],
+            ],
+            exec_results = ["OK"],
+            auto_heal_tool_calls = False,
+            max_tool_iterations = 1,
+        )
+        events = _collect_events(loop)
+        assert exec_fn.calls == []
+        assert any(
+            event.get("type") == "content" and "" in event.get("text", "")
+            for event in events
+        )
+
+    def test_auto_heal_enabled_strips_unparseable_xml_tool_call(self):
+        loop, exec_fn = _make_loop(
+            turns = [["{not valid json}"]],
+            exec_results = ["OK"],
+            auto_heal_tool_calls = True,
+            max_tool_iterations = 1,
+        )
+        events = _collect_events(loop)
+        assert exec_fn.calls == []
+        assert not any(
+            event.get("type") == "content" and "" in event.get("text", "")
+            for event in events
+        )
+
     def test_non_consecutive_duplicate_is_short_circuited(self):
         loop, exec_fn = _make_loop(
             turns = [
@@ -810,8 +1067,39 @@ class TestGuardrails:
         )
         events = _collect_events(loop)
         assert exec_fn.calls == [("web_search", {"query": "A"}), ("web_search", {"query": "B"})]
-        tool_ends = [e for e in events if e["type"] == "tool_end"]
-        assert "already made this exact call" in tool_ends[-1]["result"]
+        assert [
+            event.get("tool_call_id") for event in events if event.get("type") == "tool_end"
+        ] == ["call_0", "call_1"]
+        assert not [
+            event
+            for event in events
+            if event.get("tool_call_id") == "call_2"
+            and event.get("type") in {"tool_start", "tool_end"}
+        ]
+
+    def test_same_turn_duplicate_is_short_circuited(self):
+        loop, exec_fn = _make_loop(
+            turns = [
+                [
+                    '{"name":"web_search","arguments":{"query":"A"}}'
+                    '{"name":"web_search","arguments":{"query":"A"}}'
+                ],
+                ["final"],
+            ],
+            exec_results = ["res-A"],
+            max_tool_iterations = 2,
+        )
+        events = _collect_events(loop)
+        assert exec_fn.calls == [("web_search", {"query": "A"})]
+        assert [
+            event.get("tool_call_id") for event in events if event.get("type") == "tool_end"
+        ] == ["call_0"]
+        assert not [
+            event
+            for event in events
+            if event.get("tool_call_id") == "call_1"
+            and event.get("type") in {"tool_start", "tool_end"}
+        ]
 
     def test_coerce_string_args_python_uses_code_key(self):
         assert _coerce_arguments("print(1)", heal = True, tool_name = "python") == {"code": "print(1)"}
diff --git a/studio/backend/tests/test_sandbox_tools.py b/studio/backend/tests/test_sandbox_tools.py
index 92fee2e8e5..24b1da1772 100644
--- a/studio/backend/tests/test_sandbox_tools.py
+++ b/studio/backend/tests/test_sandbox_tools.py
@@ -117,7 +117,7 @@ class TestUntrustedHostBlock:
         )
 
     def test_dynamic_url_not_statically_blocked(self):
-        # Static AST cannot resolve runtime URLs; bash blocklist is the fallback.
+        # Static AST can't resolve runtime URLs; bash blocklist is the fallback.
         _ok('import requests; url = "https://example.com/"; requests.get(url)')
 
 
@@ -223,11 +223,8 @@ class TestUploadDenylist:
 
 
 class TestSandboxEnvIsolation:
-    """The sandbox subprocess env is built from a whitelist, not by stripping.
-
-    Confirm every credential-shaped parent var is absent regardless of how the
-    operator's process is configured. Covers Linux/macOS/WSL/Windows shapes.
-    """
+    """Sandbox env is built from a whitelist, so credential-shaped parent
+    vars stay absent regardless of operator config (Linux/macOS/WSL/Windows)."""
 
     _SECRET_KEYS = (
         # HF + ML tooling
@@ -313,8 +310,8 @@ class TestSandboxEnvIsolation:
     def test_term_is_dumb(self, tmp_path):
         from core.inference.tools import _build_safe_env
 
-        # Prevents the sandbox from re-using the operator's TERM (e.g. xterm-256color)
-        # which could trigger color-escape parsing in downstream tools.
+        # Avoid re-using the operator's TERM (e.g. xterm-256color) that
+        # could trigger color-escape parsing in downstream tools.
         env = _build_safe_env(str(tmp_path))
         assert env["TERM"] == "dumb"
 
@@ -346,12 +343,8 @@ class TestMaxBodyDefault:
 
 
 class TestBashBlocklistPosition:
-    """The blocklist must fire at command position only.
-
-    Pre-fix the per-token loop fired on any token, so `grep -r curl .`
-    and `echo source` were rejected. The position-anchored regex plus a
-    shlex-aware command-position-only token check is sufficient.
-    """
+    """The blocklist must fire at command position only, so args like
+    `grep -r curl .` and `echo source` are not falsely rejected."""
 
     @staticmethod
     def _find():
@@ -366,8 +359,7 @@ class TestBashBlocklistPosition:
         assert self._find()("echo source the data") == set()
 
     def test_cat_with_word_source_allowed(self):
-        # The 'source' word is an argument to echo; not blocked.
-        # `echo` itself isn't blocked. Only legit allowed tokens here.
+        # 'source' is an argument to echo, and echo isn't blocked either.
         assert self._find()("cat README.md && echo source") == set()
         assert "source" not in self._find()("cat README.md && echo source")
         assert "echo" not in self._find()("cat README.md && echo source")
@@ -397,14 +389,14 @@ class TestBashBlocklistPosition:
         assert "wget" in self._find()("cd /tmp && wget https://bad")
 
     def test_split_quotes_obfuscation_blocked(self):
-        # shlex collapses 'r''m' -> 'rm' as a single token at command position.
+        # shlex collapses 'r''m' -> 'rm' at command position.
         assert "rm" in self._find()("r''m -rf /")
 
     def test_path_prefixed_command_blocked(self):
         assert "sudo" in self._find()("/usr/bin/sudo whoami")
 
     def test_nested_bash_c_blocked(self):
-        # Recursion into the nested command string still catches command-position curl.
+        # Recursion into the nested command string catches command-position curl.
         assert "curl" in self._find()("bash -c 'curl https://x'")
 
     def test_subshell_command_blocked(self):
@@ -465,9 +457,8 @@ class TestBashBlocklistPosition:
 
 
 class TestHfUploadImportGate:
-    """HfApi-style upload-method blocking should require an HF import in
-    scope; otherwise paramiko / boto3 / internal SDKs with the same
-    method names hit a false positive."""
+    """Upload-method blocking requires an HF import in scope, so paramiko /
+    boto3 / internal SDKs with the same method names don't false-positive."""
 
     def test_paramiko_upload_file_allowed_without_hf_import(self):
         _ok("import paramiko; sftp=None; sftp.upload_file('a','b')")
@@ -476,14 +467,14 @@ class TestHfUploadImportGate:
         _ok("client=None; client.create_commit(Repo='x')")
 
     def test_hf_api_upload_safe_path_allowed(self):
-        # Sandbox-local relative path -- the call shape we want to permit.
+        # Sandbox-local relative path -- the permitted call shape.
         _ok("from huggingface_hub import HfApi; HfApi().upload_file('a','b','c')")
 
     def test_hf_upload_file_fq_safe_path_allowed(self):
         _ok("import huggingface_hub; huggingface_hub.upload_file('a','b','c')")
 
     def test_dynamic_builtin_import_safe_path_allowed(self):
-        # `__import__('huggingface_hub')` puts HF in scope; relative-literal path is safe.
+        # `__import__('huggingface_hub')` puts HF in scope; relative literal is safe.
         _ok("hf=__import__('huggingface_hub'); hf.HfApi().upload_file('a','b','c')")
 
     def test_dynamic_importlib_safe_path_allowed(self):
@@ -499,8 +490,8 @@ class TestHfUploadImportGate:
         )
 
     def test_hf_bare_name_upload_safe_path_allowed(self):
-        # `from huggingface_hub import upload_file` then bare `upload_file(...)`
-        # with a sandbox-local relative-path literal is allowed.
+        # Bare `upload_file(...)` (imported from huggingface_hub) with a
+        # sandbox-local relative-path literal is allowed.
         _ok(
             "from huggingface_hub import upload_file;"
             " upload_file(path_or_fileobj='x', path_in_repo='x', repo_id='r')"
@@ -519,15 +510,14 @@ class TestHfUploadImportGate:
         )
 
     def test_bare_name_upload_file_without_hf_import_allowed(self):
-        # No HF import -- local helper named upload_file should pass.
+        # No HF import -- local helper named upload_file passes.
         _ok("def upload_file(*a, **k):\n    pass\nupload_file('x', 'y', 'z')")
 
 
 class TestHfUploadSandboxLocalPaths:
-    """The HF upload gate must only allow uploads of files that already live in
-    the sandbox workdir. Absolute paths, `..` traversal, home expansion, and
-    Windows drive letters are rejected because the LLM can use them to lift
-    secrets from outside the sandbox."""
+    """HF upload gate allows only files in the sandbox workdir. Absolute paths,
+    `..` traversal, home expansion, and Windows drives are rejected (they could
+    lift secrets from outside the sandbox)."""
 
     def test_relative_literal_allowed(self):
         _ok(
@@ -621,8 +611,8 @@ class TestHfUploadSandboxLocalPaths:
         )
 
     def test_dynamic_variable_path_blocked(self):
-        # A non-literal expression could resolve to any path at runtime;
-        # the static checker cannot prove safety, so block.
+        # A non-literal expr could resolve to any path at runtime; the
+        # static checker can't prove safety, so block.
         _blocked(
             "import huggingface_hub, os\n"
             "p = os.path.join('outputs', 'x.bin')\n"
@@ -674,11 +664,9 @@ class TestHfUploadSandboxLocalPaths:
 
 
 class TestHfUploadEnvAndSecretLeakBlock:
-    """The HF upload gate must reject any positional / keyword arg sourced from
-    `os.environ` / `os.getenv` / subprocess env reads. Even though
-    `_build_safe_env` strips HF_TOKEN/WANDB/AWS upfront for the sandbox shell,
-    a Python script can still reach the parent process env if it bypasses the
-    safe-env wrapper at the source -- so block statically."""
+    """HF upload gate rejects any arg sourced from os.environ / os.getenv /
+    subprocess env reads, since a script can reach the parent env directly
+    despite the safe-env shell wrapper."""
 
     def test_path_from_os_environ_subscript_blocked(self):
         _blocked(
@@ -756,7 +744,7 @@ class TestHfUploadEnvAndSecretLeakBlock:
         )
 
     def test_env_dict_unpacked_via_environ_attr_blocked(self):
-        # `os.environ` as a bare reference (passed somewhere it gets serialized).
+        # Bare `os.environ` reference (passed somewhere it gets serialized).
         _blocked(
             "import huggingface_hub, os\n"
             "huggingface_hub.upload_file(path_or_fileobj=str(os.environ),"
@@ -765,8 +753,8 @@ class TestHfUploadEnvAndSecretLeakBlock:
         )
 
     def test_repo_id_from_env_also_blocked(self):
-        # Even non-path args must not source env vars -- an attacker could
-        # encode secrets in repo_id or path_in_repo.
+        # Non-path args must not source env vars either -- an attacker
+        # could encode secrets in repo_id or path_in_repo.
         _blocked(
             "import huggingface_hub, os\n"
             'huggingface_hub.upload_file(path_or_fileobj="x.bin",'
diff --git a/studio/backend/tests/test_studio_api.py b/studio/backend/tests/test_studio_api.py
index 70a103a2f8..27f695b744 100644
--- a/studio/backend/tests/test_studio_api.py
+++ b/studio/backend/tests/test_studio_api.py
@@ -4,9 +4,9 @@
 """
 End-to-end tests for Unsloth Studio's HTTP API surface.
 
-Covers the OpenAI-compatible and Anthropic-compatible endpoints exposed
-by the server that ``unsloth studio run`` boots, plus API key
-authentication and the CLI's ``--help`` output:
+Covers the OpenAI- and Anthropic-compatible endpoints exposed by the
+server that ``unsloth studio run`` boots, plus API key authentication and
+the CLI's ``--help`` output:
 
     1. curl -- basic chat completions (non-streaming)
     2. curl -- streaming chat completions
@@ -38,14 +38,14 @@ Usage:
     export UNSLOTH_E2E_API_KEY=sk-unsloth-...   # from the server banner
     pytest tests/test_studio_api.py -v
 
-    # Pytest mode, fixture-managed server — pytest launches and tears
-    # down the server itself. One-shot verification, CI-friendly.
+    # Pytest mode, fixture-managed server — pytest launches and tears down
+    # the server itself. One-shot verification, CI-friendly.
     pytest tests/test_studio_api.py -v \\
         --unsloth-model unsloth/Qwen3-1.7B-GGUF \\
         --unsloth-gguf-variant UD-Q4_K_XL
 
-The ``base_url`` / ``api_key`` parameters on the test functions resolve
-via the ``studio_server`` session fixture in ``conftest.py``.
+The ``base_url`` / ``api_key`` parameters on the test functions resolve via
+the ``studio_server`` session fixture in ``conftest.py``.
 
 Requires a GPU and ~2 GB of disk for the GGUF download.
 """
@@ -65,17 +65,17 @@ import urllib.request
 from pathlib import Path
 
 
-# ── Configuration ────────────────────────────────────────────────────
+# Configuration
 
 DEFAULT_MODEL = "unsloth/Qwen3-1.7B-GGUF"
 DEFAULT_VARIANT = "UD-Q4_K_XL"
 PORT = 18222  # high port unlikely to collide
 HOST = "127.0.0.1"
-STARTUP_TIMEOUT = 120  # seconds to wait for banner
+STARTUP_TIMEOUT = 120  # seconds
 LOG_FILE = Path(__file__).resolve().parent.parent.parent.parent / "temp" / "test_studio_api.log"
 
 
-# ── Helpers ──────────────────────────────────────────────────────────
+# Helpers
 
 
 def _http(
@@ -125,7 +125,7 @@ def _stream_http(
         return exc.code, []
 
 
-# ── Test functions ───────────────────────────────────────────────────
+# Test functions
 
 
 def test_help_output():
@@ -233,10 +233,10 @@ def test_openai_sdk(base_url: str, api_key: str):
 def test_curl_with_tools(base_url: str, api_key: str):
     """Example 4: chat completion with tool calling enabled.
 
-    Note: when ``enable_tools`` is set the server always returns SSE
-    streaming regardless of the ``stream`` flag, so we parse SSE chunks.
-    The model may or may not produce visible content -- tool orchestration
-    can intercept the response -- so we only assert the endpoint succeeds.
+    When ``enable_tools`` is set the server always returns SSE streaming
+    regardless of the ``stream`` flag, so we parse SSE chunks. The model may
+    not produce visible content (tool orchestration can intercept the
+    response), so we only assert the endpoint succeeds.
     """
     status, chunks = _stream_http(
         f"{base_url}/v1/chat/completions",
@@ -265,19 +265,13 @@ def test_curl_with_tools(base_url: str, api_key: str):
     print(f"  PASS  curl with tools: {len(chunks)} chunks, {len(full)} chars content")
 
 
-# ── Standard OpenAI function-calling pass-through tests ─────────────
+# Standard OpenAI function-calling pass-through tests.
 #
-# Regression coverage for unslothai/unsloth#4999: Studio's
-# /v1/chat/completions used to silently strip standard OpenAI `tools`
-# and `tool_choice` fields, so clients (opencode, Claude Code, Cursor,
-# Continue, ...) could never get structured tool_calls back. These
-# tests exercise the client-side pass-through path that forwards those
-# fields to llama-server verbatim.
-#
-# They require a tool-capable GGUF (``supports_tools=True`` — e.g.
-# Qwen3, Qwen2.5-Coder, Llama-3.1-Instruct). The default test model
-# ``unsloth/Qwen3-1.7B-GGUF`` advertises tool support via its chat
-# template metadata.
+# Regression coverage for unslothai/unsloth#4999: /v1/chat/completions used
+# to strip standard OpenAI `tools`/`tool_choice`, so clients never got
+# structured tool_calls back. These exercise the pass-through that forwards
+# those fields to llama-server verbatim. Require a tool-capable GGUF
+# (supports_tools=True); the default unsloth/Qwen3-1.7B-GGUF qualifies.
 
 _WEATHER_TOOL = {
     "type": "function",
@@ -301,9 +295,9 @@ _WEATHER_TOOL = {
 def _collect_streamed_tool_calls(chunks: list[dict]) -> list[dict]:
     """Reassemble OpenAI streaming delta.tool_calls into full tool calls.
 
-    OpenAI streams partial tool calls across chunks — the first chunk for
-    a given index carries ``id`` + ``function.name``, and subsequent
-    chunks append fragments to ``function.arguments``.
+    OpenAI streams partial tool calls across chunks — the first chunk for a
+    given index carries ``id`` + ``function.name``, and later chunks append
+    fragments to ``function.arguments``.
     """
     by_index: dict[int, dict] = {}
     for c in chunks:
@@ -346,8 +340,8 @@ def _final_finish_reason(chunks: list[dict]) -> str | None:
 def test_openai_tools_nonstream(base_url: str, api_key: str):
     """Standard OpenAI function calling, non-streaming, tool_choice='required'.
 
-    Regression: before the fix, Studio silently stripped `tools` and the
-    model returned plain text with finish_reason='stop'. After the fix,
+    Regression: before the fix, Studio stripped `tools` and the model
+    returned plain text with finish_reason='stop'. After the fix,
     llama-server's response is forwarded verbatim so the client sees
     finish_reason='tool_calls' with a structured tool_calls array and
     non-zero usage.prompt_tokens.
@@ -428,9 +422,9 @@ def test_openai_tools_multiturn(base_url: str, api_key: str):
     messages and assistant messages carrying tool_calls are accepted.
 
     Regression: before the fix, ChatMessage.role was restricted to
-    {system,user,assistant} and rejected role='tool' at the Pydantic
-    validation stage. This test sends a full round trip so the model
-    receives the simulated tool result and responds with final text.
+    {system,user,assistant} and rejected role='tool' at Pydantic
+    validation. This test sends a full round trip so the model receives the
+    simulated tool result and responds with final text.
     """
     status, text = _http(
         "POST",
@@ -467,7 +461,7 @@ def test_openai_tools_multiturn(base_url: str, api_key: str):
     assert status == 200, f"Expected 200, got {status}: {text[:500]}"
     data = json.loads(text)
     msg = data["choices"][0]["message"]
-    # The model should respond with text now that it has the tool result
+    # The model should respond with text now it has the tool result
     content = msg.get("content") or ""
     assert len(content) > 0 or msg.get(
         "tool_calls"
@@ -532,7 +526,7 @@ def test_no_key_rejected(base_url: str):
     print(f"  PASS  no API key rejected ({status})")
 
 
-# ── Anthropic SSE helper ─────────────────────────────────────────────
+# Anthropic SSE helper
 
 
 def _stream_anthropic_http(
@@ -580,7 +574,7 @@ def _collect_anthropic_text(events: list[tuple[str, dict]]) -> str:
     return "".join(parts)
 
 
-# ── Anthropic /v1/messages test functions ────────────────────────────
+# Anthropic /v1/messages test functions
 
 
 def test_anthropic_basic(base_url: str, api_key: str):
@@ -701,9 +695,9 @@ def test_anthropic_tool_choice_any(base_url: str, api_key: str):
     """Anthropic Messages API: ``tool_choice: {"type": "any"}`` must be
     honored (forwarded as OpenAI ``tool_choice: "required"`` to
     llama-server). Regression for the secondary fix bundled with #4999 —
-    previously this field was accepted on the request model but silently
-    dropped with a warning log, so the model was free to answer from
-    memory instead of using the tool.
+    previously this field was accepted on the request model but dropped with
+    a warning log, so the model could answer from memory instead of using
+    the tool.
     """
     status, events = _stream_anthropic_http(
         f"{base_url}/v1/messages",
@@ -711,7 +705,7 @@ def test_anthropic_tool_choice_any(base_url: str, api_key: str):
             "model": "default",
             "max_tokens": 256,
             "messages": [
-                # A question the model could easily answer from memory if
+                # A question the model could answer from memory if
                 # tool_choice were not enforced.
                 {
                     "role": "user",
@@ -740,7 +734,7 @@ def test_anthropic_tool_choice_any(base_url: str, api_key: str):
     assert status == 200, f"Expected 200, got {status}"
     assert len(events) > 0, "No SSE events received"
 
-    # With tool_choice=any, stop_reason must be tool_use (not end_turn)
+    # With tool_choice=any, stop_reason must be tool_use, not end_turn
     stop_reason = None
     for etype, data in events:
         if etype == "message_delta":
@@ -763,7 +757,7 @@ def test_anthropic_tool_choice_any(base_url: str, api_key: str):
     )
 
 
-# ── Server lifecycle ─────────────────────────────────────────────────
+# Server lifecycle
 
 
 def _start_server(model: str, variant: str | None) -> tuple[subprocess.Popen, str]:
@@ -837,7 +831,7 @@ def _kill_server(proc: subprocess.Popen):
         proc.wait(timeout = 5)
 
 
-# ── Main ─────────────────────────────────────────────────────────────
+# Main
 
 
 def main():
@@ -870,11 +864,11 @@ def main():
             failed += 1
             print(f"  ERROR {fn.__name__}: {type(exc).__name__}: {exc}")
 
-    # ── 1. Test --help (no server needed) ────────────────────────────
+    # 1. --help (no server needed)
     print("\n[1/16] Testing --help output")
     run_test(test_help_output)
 
-    # ── 2-16. Start server and run API tests ─────────────────────────
+    # 2-16. Start server and run API tests
     print(f"\nStarting server: {args.model} (variant={args.gguf_variant}) on port {PORT}...")
     proc = None
     try:
@@ -929,14 +923,14 @@ def main():
 
     except RuntimeError as exc:
         print(f"\nFATAL: Server failed to start: {exc}")
-        failed += 16  # count remaining tests as failed
+        failed += 16  # remaining tests count as failed
     finally:
         if proc:
             print("\nStopping server...")
             _kill_server(proc)
             print("Server stopped.")
 
-    # ── Summary ──────────────────────────────────────────────────────
+    # Summary
     total = passed + failed
     print(f"\n{'=' * 40}")
     print(f"Results: {passed}/{total} passed, {failed} failed")
diff --git a/studio/backend/tests/test_studio_train_validation.py b/studio/backend/tests/test_studio_train_validation.py
index 6df491610b..0ffecb3ce4 100644
--- a/studio/backend/tests/test_studio_train_validation.py
+++ b/studio/backend/tests/test_studio_train_validation.py
@@ -24,7 +24,7 @@ from models.training import (
 
 
 def _check_field(field_name: str, value):
-    """Run the field validator without constructing a full TrainingStartRequest."""
+    """Run the field validator without building a full TrainingStartRequest."""
     from models.training import TrainingStartRequest
 
     schema_field = TrainingStartRequest.model_fields[field_name]
@@ -87,15 +87,14 @@ class TestVisionImageSizeCap:
 
     @pytest.mark.parametrize("value", [True, False])
     def test_bool_error_says_integer_not_range(self, value):
-        # Regression guard: bools must say "integer or null", not "in [256, 2048]".
+        # Regression guard: bools say "integer or null", not "in [256, 2048]".
         with pytest.raises(ValidationError) as exc:
             _check_field("vision_image_size", value)
         assert "integer or null" in str(exc.value)
 
     @pytest.mark.parametrize("value", ["++512", "--256", "+-+512", "+", "-"])
     def test_multi_sign_string_says_integer_not_raw(self, value):
-        # Regression guard: multi-sign strings must not leak int()'s raw
-        # "invalid literal" message; precise contract is "integer or null".
+        # Regression guard: multi-sign strings say "integer or null", not int()'s raw message.
         with pytest.raises(ValidationError) as exc:
             _check_field("vision_image_size", value)
         assert "integer or null" in str(exc.value)
@@ -103,8 +102,7 @@ class TestVisionImageSizeCap:
 
     @pytest.mark.parametrize("value", ["５１２", "٥١٢", "१०२४"])
     def test_unicode_digit_string_rejected(self, value):
-        # Full-width / Arabic-Indic / Devanagari digits must be rejected so the
-        # value reaching the backend equals the ASCII the user typed.
+        # Reject non-ASCII (full-width/Arabic-Indic/Devanagari) digits.
         with pytest.raises(ValidationError) as exc:
             _check_field("vision_image_size", value)
         assert "integer or null" in str(exc.value)
diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py
new file mode 100644
index 0000000000..8ff41342d7
--- /dev/null
+++ b/studio/backend/tests/test_tool_call_parser_strict.py
@@ -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
+
+"""Strict-mode (Auto-Heal disabled) tool-call parsing.
+
+With ``allow_incomplete=False`` the parser must accept a well-formed
+``...`` call even when the model appends prose
+after the closing tag -- matching the JSON-style ``...`` path,
+which already tolerates trailing text -- while still rejecting genuinely
+truncated calls that never close.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+    sys.path.insert(0, _BACKEND_DIR)
+
+from core.inference.tool_call_parser import parse_tool_calls_from_text
+
+
+def _only(text: str) -> dict:
+    calls = parse_tool_calls_from_text(text, allow_incomplete = False)
+    assert len(calls) == 1, f"expected exactly one call, got {len(calls)}: {calls!r}"
+    fn = calls[0]["function"]
+    return {"name": fn["name"], "arguments": json.loads(fn["arguments"])}
+
+
+class TestFunctionStyleTrailingText:
+    def test_closed_function_with_trailing_prose_is_accepted(self):
+        text = (
+            "weather london"
+            " Let me check that for you."
+        )
+        call = _only(text)
+        assert call == {"name": "web_search", "arguments": {"query": "weather london"}}
+
+    def test_closed_function_with_trailing_whitespace_is_accepted(self):
+        text = "cats   \n\n"
+        call = _only(text)
+        assert call == {"name": "web_search", "arguments": {"query": "cats"}}
+
+    def test_closed_function_without_trailing_text_still_parses(self):
+        text = "cats"
+        call = _only(text)
+        assert call == {"name": "web_search", "arguments": {"query": "cats"}}
+
+    def test_multi_param_with_trailing_prose(self):
+        text = (
+            "ls -la"
+            "home running it now"
+        )
+        call = _only(text)
+        assert call == {
+            "name": "terminal",
+            "arguments": {"command": "ls -la", "workdir": "home"},
+        }
+
+    def test_code_value_containing_literal_close_tag_is_preserved(self):
+        # The real closing  is the last one; the literal inside
+        # the code argument must survive (rfind, not the first match).
+        text = (
+            ""
+            'print("")'
+            " all done"
+        )
+        call = _only(text)
+        assert call == {"name": "python", "arguments": {"code": 'print("")'}}
+
+    def test_incomplete_function_without_close_is_still_rejected(self):
+        text = "weather london"
+        assert parse_tool_calls_from_text(text, allow_incomplete = False) == []
+
+    def test_param_without_close_tag_is_rejected_in_strict_mode(self):
+        # Closing  present, but the single parameter never closes.
+        text = "weather london"
+        assert parse_tool_calls_from_text(text, allow_incomplete = False) == []
+
+
+class TestParityWithJsonStyle:
+    def test_json_tool_call_with_trailing_prose_is_accepted(self):
+        text = (
+            '{"name":"web_search","arguments":{"query":"weather london"}}'
+            " Let me check that for you."
+        )
+        calls = parse_tool_calls_from_text(text, allow_incomplete = False)
+        assert len(calls) == 1
+        assert calls[0]["function"]["name"] == "web_search"
+
+    def test_function_and_json_styles_agree_on_trailing_text(self):
+        q = "weather london"
+        func = parse_tool_calls_from_text(
+            f"{q} trailing",
+            allow_incomplete = False,
+        )
+        js = parse_tool_calls_from_text(
+            f'{{"name":"web_search","arguments":{{"query":"{q}"}}}} trailing',
+            allow_incomplete = False,
+        )
+        assert len(func) == len(js) == 1
+        assert json.loads(func[0]["function"]["arguments"]) == {"query": q}
+        assert json.loads(js[0]["function"]["arguments"]) == {"query": q}
+
+
+class TestHealingPathUnaffected:
+    def test_auto_heal_still_repairs_unclosed_function(self):
+        text = "cats"
+        calls = parse_tool_calls_from_text(text, allow_incomplete = True)
+        assert len(calls) == 1
+        assert calls[0]["function"]["name"] == "web_search"
diff --git a/studio/backend/tests/test_tool_loop_controller.py b/studio/backend/tests/test_tool_loop_controller.py
new file mode 100644
index 0000000000..dea5de6d6e
--- /dev/null
+++ b/studio/backend/tests/test_tool_loop_controller.py
@@ -0,0 +1,212 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+    sys.path.insert(0, _BACKEND_DIR)
+
+from core.inference.tool_loop_controller import (
+    ToolLoopController,
+    canonical_tool_call_key,
+    coerce_tool_arguments,
+    status_for_tool,
+    strip_result_for_model,
+    tool_event_provenance,
+)
+
+
+def _tool(name: str) -> dict:
+    return {"type": "function", "function": {"name": name}}
+
+
+def _call(
+    name: str,
+    args,
+    call_id: str = "call_0",
+) -> dict:
+    return {
+        "id": call_id,
+        "type": "function",
+        "function": {
+            "name": name,
+            "arguments": json.dumps(args) if isinstance(args, dict) else args,
+        },
+    }
+
+
+def test_canonical_tool_call_key_sorts_arguments():
+    a = canonical_tool_call_key("web_search", {"query": "gpu", "limit": 5})
+    b = canonical_tool_call_key("web_search", {"limit": 5, "query": "gpu"})
+    c = canonical_tool_call_key("python", {"limit": 5, "query": "gpu"})
+
+    assert a == b
+    assert a != c
+    assert a == 'web_search:{"limit":5,"query":"gpu"}'
+
+
+def test_coerce_tool_arguments_parses_json_and_heals_raw_strings():
+    parsed = coerce_tool_arguments('{"query":"gpu prices"}', heal = True)
+    healed = coerce_tool_arguments("print(1)", heal = True, tool_name = "python")
+    raw = coerce_tool_arguments("not-json", heal = False, tool_name = "python")
+
+    assert parsed.arguments == {"query": "gpu prices"}
+    assert not parsed.healed
+    assert healed.arguments == {"code": "print(1)"}
+    assert healed.healed
+    assert raw.arguments == {"raw": "not-json"}
+    assert not raw.healed
+
+
+def test_status_and_provenance_match_local_event_conventions():
+    assert status_for_tool("web_search", {"query": "gpus"}) == "Searching: gpus"
+    assert (
+        status_for_tool("web_search", {"url": "https://www.example.com/a"})
+        == "Reading: example.com"
+    )
+    assert status_for_tool("python", {"code": "print(1)\nprint(2)"}) == "Running Python: print(1)"
+    assert tool_event_provenance(healed = True, forced = False, provisional = None) == {
+        "source": "local",
+        "healed": True,
+    }
+
+
+def test_prepare_execute_builds_visible_events_and_model_tool_message():
+    controller = ToolLoopController(tools = [_tool("web_search")])
+    decision = controller.prepare_call(_call("web_search", {"query": "gpu prices"}))
+
+    assert decision.should_execute
+    assert decision.emit_visible_events
+    assert decision.status_text == "Searching: gpu prices"
+    assert decision.tool_start_payload()["arguments"] == {"query": "gpu prices"}
+    assert decision.tool_start_event()["type"] == "tool_start"
+    assert decision.as_assistant_tool_call()["function"]["arguments"] == '{"query":"gpu prices"}'
+
+    completion = controller.record_result(decision, "Search result\n__IMAGES__:{...}")
+
+    assert completion.tool_end_payload()["result"] == "Search result\n__IMAGES__:{...}"
+    assert completion.tool_end_event()["type"] == "tool_end"
+    assert completion.tool_message() == {
+        "role": "tool",
+        "name": "web_search",
+        "content": "Search result",
+        "tool_call_id": "call_0",
+    }
+
+
+def test_successful_duplicate_is_internal_noop_and_keeps_remaining_tools():
+    controller = ToolLoopController(tools = [_tool("web_search"), _tool("python")])
+    first = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_a"))
+    controller.record_result(first, "ok")
+
+    duplicate = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_b"))
+    completion = controller.record_noop(duplicate)
+
+    assert duplicate.action == "duplicate"
+    assert not duplicate.should_execute
+    assert not duplicate.emit_visible_events
+    duplicate_nudge = completion.model_message()["content"]
+    assert "already completed successfully" in duplicate_nudge
+    assert "different enabled tool" in duplicate_nudge
+    assert completion.model_message()["role"] == "user"
+    assert not controller.force_final_answer
+    assert [tool["function"]["name"] for tool in controller.active_tools()] == [
+        "web_search",
+        "python",
+    ]
+
+
+def test_repeated_successful_duplicate_becomes_terminal_after_one_recovery_nudge():
+    controller = ToolLoopController(tools = [_tool("web_search"), _tool("python")])
+    first = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_a"))
+    controller.record_result(first, "ok")
+
+    duplicate_one = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_b"))
+    completion_one = controller.record_noop(duplicate_one)
+
+    assert duplicate_one.action == "duplicate"
+    assert "already completed successfully" in completion_one.model_message()["content"]
+    assert not controller.force_final_answer
+    assert [tool["function"]["name"] for tool in controller.active_tools()] == [
+        "web_search",
+        "python",
+    ]
+
+    duplicate_two = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_c"))
+    completion_two = controller.record_noop(duplicate_two)
+
+    assert duplicate_two.action == "duplicate"
+    assert "already completed successfully" in completion_two.model_message()["content"]
+    assert controller.force_final_answer
+    assert controller.active_tools() == []
+
+
+def test_failed_call_does_not_block_retry():
+    controller = ToolLoopController(tools = [_tool("web_search")])
+    first = controller.prepare_call(_call("web_search", {"query": "gpu prices"}))
+    controller.record_result(first, "Error: temporary failure")
+
+    retry = controller.prepare_call(_call("web_search", {"query": "gpu prices"}))
+
+    assert retry.should_execute
+    assert retry.action == "execute"
+
+
+def test_empty_enabled_tool_list_blocks_all_tool_calls():
+    controller = ToolLoopController(tools = [])
+    decision = controller.prepare_call(_call("web_search", {"query": "gpu prices"}))
+    completion = controller.record_noop(decision)
+
+    assert decision.action == "disabled"
+    assert not decision.emit_visible_events
+    assert completion.model_message()["role"] == "user"
+    assert "not enabled" in completion.model_message()["content"]
+    assert controller.force_final_answer
+    assert controller.active_tools() == []
+
+
+def test_disabled_tool_is_internal_noop_not_visible_tool_error():
+    controller = ToolLoopController(tools = [_tool("web_search")])
+    decision = controller.prepare_call(_call("python", {"code": "print(1)"}))
+    completion = controller.record_noop(decision)
+
+    assert decision.action == "disabled"
+    assert not decision.emit_visible_events
+    assert completion.model_message()["role"] == "user"
+    assert "not enabled" in completion.model_message()["content"]
+    assert controller.force_final_answer
+    assert controller.active_tools() == []
+
+
+def test_render_html_success_filters_active_tools_and_repeat_is_internal():
+    controller = ToolLoopController(tools = [_tool("render_html"), _tool("web_search")])
+    assert [t["function"]["name"] for t in controller.active_tools()] == [
+        "render_html",
+        "web_search",
+    ]
+
+    first = controller.prepare_call(_call("render_html", {"code": ""}, "call_html_1"))
+    controller.record_result(first, "Rendered HTML artifact: Demo")
+
+    assert [t["function"]["name"] for t in controller.active_tools()] == ["web_search"]
+
+    repeat = controller.prepare_call(_call("render_html", {"code": ""}, "call_html_2"))
+    completion = controller.record_noop(repeat)
+
+    assert repeat.action == "render_html_repeat"
+    assert not repeat.emit_visible_events
+    assert completion.model_message()["role"] == "user"
+    assert "Do not call render_html again" in completion.model_message()["content"]
+    assert controller.force_final_answer
+    assert controller.active_tools() == []
+
+
+def test_strip_result_for_model_removes_frontend_image_sentinel():
+    assert strip_result_for_model('text\n__IMAGES__:{"paths":[]}') == "text"
+    assert strip_result_for_model("text __IMAGES__:payload") == "text"
+    assert strip_result_for_model("plain text") == "plain text"
diff --git a/studio/backend/tests/test_tool_policy_gates.py b/studio/backend/tests/test_tool_policy_gates.py
index 01f6bbbc3f..fad121a4a1 100644
--- a/studio/backend/tests/test_tool_policy_gates.py
+++ b/studio/backend/tests/test_tool_policy_gates.py
@@ -2,8 +2,8 @@
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
 
 """
-Tests for `_effective_enable_tools` -- the helper that folds the
-process-level `tool_policy` over a request's `enable_tools` field.
+Tests for `_effective_enable_tools` -- folds the process-level `tool_policy`
+over a request's `enable_tools` field.
 
 Truth table (policy x payload.enable_tools -> effective):
   policy=None  + payload=None  -> None
diff --git a/studio/backend/tests/test_tool_xml_strip.py b/studio/backend/tests/test_tool_xml_strip.py
index 857302d543..2ba3310fbe 100644
--- a/studio/backend/tests/test_tool_xml_strip.py
+++ b/studio/backend/tests/test_tool_xml_strip.py
@@ -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
 
-"""Tests for `_TOOL_XML_RE` (routes/inference.py) -- strips tool-call
-XML that leaks past the speculative buffer in core/inference/llama_cpp.py
-when the open/close pair is split across the visible/DRAIN boundary.
+"""Tests for `_TOOL_XML_RE` (routes/inference.py) -- strips tool-call XML that
+leaks past the speculative buffer in core/inference/llama_cpp.py when the
+open/close pair is split across the visible/DRAIN boundary.
 """
 
 from __future__ import annotations
@@ -27,11 +27,25 @@ assert _m, "could not extract _TOOL_XML_RE source"
 _ns = {"_re": _re}
 exec(f"_TOOL_XML_RE = _re.compile({_m.group(1)})", _ns)
 _TOOL_XML_RE = _ns["_TOOL_XML_RE"]
+_helper = _re.search(
+    r"def _strip_tool_xml_for_display\(text: str, \*, auto_heal_tool_calls: bool\) -> str:\n"
+    r"(?:    .+\n)+",
+    _src,
+)
+assert _helper, "could not extract _strip_tool_xml_for_display source"
+exec(_helper.group(0), _ns)
+_strip_tool_xml_for_display = _ns["_strip_tool_xml_for_display"]
 
 
 # ── Well-formed pairs ─────────────────────────────────────────────
 
 
+def test_route_display_strip_respects_disabled_auto_heal_contract():
+    text = 'literal {"name":"web_search"} survives'
+    assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text
+    assert "" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
+
+
 def test_strips_well_formed_tool_call():
     text = (
         "Let me search.\n"
@@ -134,7 +148,7 @@ def test_strips_tail_only_parameter_orphan_no_trailing_ws():
 
 
 def test_preserves_mid_string_parameter_in_code_sample():
-    # Tail-anchor on `` is required so doc/example prose survives.
+    # Tail-anchor on `` so doc/example prose survives.
     text = (
         "Here is the Qwen tool-call format:\n"
         "```xml\n"
@@ -207,8 +221,8 @@ def test_real_world_sweep_leaks_get_stripped(leak):
 # ── Real-world tail-only  from gdpval sweep ──────────
 
 
-# All end-anchored: outer  truncated by EOS,
-# inner  open DRAINED, leaving bare  tail.
+# All end-anchored: outer  truncated by EOS, inner
+#  open DRAINED, leaving bare  tail.
 GDPVAL_PARAMETER_LEAKS = [
     # Qwen3.5-27B Q8_0 / worldbank s00
     "the page contains image data and the text is not readable.\n\n\n",
diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py
index edc47c705e..3c5d6cd094 100644
--- a/studio/backend/tests/test_training_worker_flash_attn.py
+++ b/studio/backend/tests/test_training_worker_flash_attn.py
@@ -221,7 +221,7 @@ def test_mamba_ssm_path_preserves_wheel_first_install_args(monkeypatch):
 
 
 def _force_missing_fla_imports(monkeypatch):
-    """Make fla.modules / fla.ops.gated_delta_rule imports raise ImportError."""
+    """Force fla.modules / fla.ops imports to raise ImportError."""
     real_import = builtins.__import__
 
     def fake_import(name, *a, **kw):
@@ -267,8 +267,8 @@ def test_flash_linear_attention_skips_for_unrelated_models(monkeypatch):
 
 
 def test_flash_linear_attention_skips_for_ssm_only_models(monkeypatch):
-    # Nemotron-H / Falcon-H1 / Granite-H / LFM2 take the mamba_ssm path
-    # and never call FLA's gated_delta_rule kernels.
+    # Nemotron-H / Falcon-H1 / Granite-H / LFM2 take the mamba_ssm path,
+    # never FLA's gated_delta_rule kernels.
     run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
     monkeypatch.setattr(worker._sp, "run", run_mock)
 
@@ -289,7 +289,7 @@ def test_flash_linear_attention_matches_full_qwen3_family(monkeypatch):
     monkeypatch.setattr(worker._sp, "run", run_mock)
     _force_missing_fla_imports(monkeypatch)
     monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None)
-    # Hermetic discovery: pretend installed transformers ships all the Qwen GDN families.
+    # Hermetic discovery: pretend transformers ships all Qwen GDN families.
     monkeypatch.setattr(
         worker,
         "_discover_fla_model_types",
@@ -370,9 +370,8 @@ def test_flash_linear_attention_install_includes_einops(monkeypatch):
 
     args = run_mock.call_args[0][0]
     assert "--no-deps" in args
-    # einops is declared by fla-core; packaging and triton are pulled in
-    # because fla/utils.py imports them at module load but neither is
-    # declared in fla-core's METADATA (an upstream FLA gap).
+    # packaging and triton are added because fla/utils.py imports them at load
+    # but neither is in fla-core's METADATA (an upstream FLA gap).
     assert "einops" in args
     assert "packaging" in args
     assert "triton" in args
@@ -389,8 +388,8 @@ def test_flash_linear_attention_logs_post_install_import_failure(monkeypatch):
 
     def fake_importable():
         import_calls["count"] += 1
-        # First call (pre-install probe) -> False so we attempt install.
-        # Second call (post-install verify) -> still False.
+        # Pre-install probe -> False (attempt install); post-install
+        # verify -> still False.
         return False
 
     monkeypatch.setattr(worker, "_flash_linear_attention_importable", fake_importable)
@@ -433,13 +432,13 @@ def test_tilelang_backend_pins_only_binary(monkeypatch):
     run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
     monkeypatch.setattr(worker._sp, "run", run_mock)
     monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None)
-    # Need to bypass the post-install probe too.
+    # Bypass the post-install probe too.
     probe_calls = {"count": 0}
 
     def fake_probe():
         probe_calls["count"] += 1
-        # First probe (pre-install): False so install runs.
-        # Second probe (post-install): True so success branch taken.
+        # Pre-install probe: False (install runs); post-install: True
+        # (success branch taken).
         return probe_calls["count"] > 1
 
     monkeypatch.setattr(worker, "_tilelang_importable", fake_probe)
@@ -491,14 +490,10 @@ def test_tilelang_backend_installs_pinned_pair_for_qwen3_5(monkeypatch):
 def test_tilelang_backend_reinstalls_when_tvm_ffi_is_broken(monkeypatch):
     """Repair path issues TWO pip calls:
 
-    Call 1 (repair): `--force-reinstall --no-deps apache-tvm-ffi==0.1.9`
-      — surgically downgrades the broken package only. `--no-deps` here
-      is REQUIRED to prevent --force-reinstall from cascading through
-      apache-tvm-ffi's dep graph and replacing torch / the CUDA stack.
-
-    Call 2 (install): plain `apache-tvm-ffi==0.1.9 tilelang==0.1.8`
-      — resolves missing transitive deps (z3-solver, ml-dtypes) without
-      --force-reinstall, so it never replaces already-correct packages.
+    1 (repair): --force-reinstall --no-deps apache-tvm-ffi -- downgrades only
+      the broken package; --no-deps stops the cascade through its deps to torch.
+    2 (install): plain apache-tvm-ffi + tilelang -- resolves missing transitive
+      deps without --force-reinstall, so it never replaces correct packages.
     """
     monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False)
     monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
@@ -515,14 +510,14 @@ def test_tilelang_backend_reinstalls_when_tvm_ffi_is_broken(monkeypatch):
     assert run_mock.call_count == 2
     repair_args, install_args = (call[0][0] for call in run_mock.call_args_list)
 
-    # Repair: --force-reinstall --no-deps, apache-tvm-ffi ONLY (no tilelang).
+    # Repair: --force-reinstall --no-deps, apache-tvm-ffi ONLY.
     assert "--force-reinstall" in repair_args
     assert "--no-deps" in repair_args, "Repair MUST use --no-deps to avoid replacing torch / CUDA"
     assert "--only-binary=:all:" in repair_args
     assert f"apache-tvm-ffi=={worker._APACHE_TVM_FFI_PACKAGE_VERSION}" in repair_args
     assert all("tilelang" not in a for a in repair_args), "Repair MUST only touch apache-tvm-ffi"
 
-    # Install: regular dep-resolving install, NO --force-reinstall.
+    # Install: regular dep-resolving install, no --force-reinstall.
     assert "--force-reinstall" not in install_args
     assert "--no-deps" not in install_args
     assert "--only-binary=:all:" in install_args
@@ -573,7 +568,7 @@ def test_tilelang_backend_swallows_install_timeout(monkeypatch):
     statuses: list[str] = []
     monkeypatch.setattr(worker, "_send_status", lambda queue, msg: statuses.append(msg))
 
-    # Should not raise.
+    # Must not raise.
     worker._ensure_tilelang_backend(
         event_queue = [],
         model_name = "unsloth/Qwen3.5-2B",
@@ -588,7 +583,7 @@ def test_tilelang_backend_skipped_for_ssm_models(monkeypatch):
     monkeypatch.setattr(worker._sp, "run", run_mock)
 
     # Nemotron-H / Falcon-H1 / Granite-H take the mamba_ssm path, not FLA's
-    # gated_delta_rule -> tilelang has no effect on them.
+    # gated_delta_rule -> tilelang doesn't affect them.
     for name in (
         "tiiuae/Falcon-H1-0.5B-Instruct",
         "nvidia/Nemotron-H-8B-Base",
@@ -633,27 +628,23 @@ def test_tilelang_backend_swallows_install_failure(monkeypatch):
     assert any("failed" in s.lower() for s in statuses)
 
 
-# ───────────────────────────────────────────────────────────────────
-# Runtime hook on `is_flash_linear_attention_available` /
-# `is_causal_conv1d_available`. These are the primary gate in
-# normal operation; the substring tests above cover the
-# UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS=1 fallback.
-# ───────────────────────────────────────────────────────────────────
+# Runtime hook on is_flash_linear_attention_available /
+# is_causal_conv1d_available -- the primary gate in normal operation. The
+# substring tests above cover the SKIP_FAST_PATH_HOOKS=1 fallback.
 
 
 class _FakeQueue(list):
-    """List with `.put` so worker._send_status can send into it during tests."""
+    """List with `.put` so worker._send_status can send into it in tests."""
 
     def put(self, item):
         self.append(item)
 
 
 def _make_fake_gate(initial_return: bool):
-    """Build a callable that mimics transformers' lru_cache-decorated gates.
+    """Callable mimicking transformers' lru_cache-decorated gates.
 
-    Tracks call count and exposes a `cache_clear` attribute. The return
-    value can be flipped to mimic install-then-True behaviour by setting
-    `.next_return`.
+    Tracks call count and exposes `cache_clear`. Flip `.next_return` to
+    mimic install-then-True behaviour.
     """
 
     class Gate:
@@ -707,7 +698,7 @@ def test_hook_installs_when_gate_returns_false(monkeypatch):
 
     from transformers.utils import import_utils as _iu
 
-    # Both gates are now wrapped. Call them — the hook should drive the install.
+    # Both gates wrapped; calling them should drive the install.
     assert _iu.is_flash_linear_attention_available() is True
     fla_install.assert_called_once()
     tile_install.assert_called_once()
@@ -716,9 +707,9 @@ def test_hook_installs_when_gate_returns_false(monkeypatch):
 
 
 def test_hook_skips_install_when_gate_already_true(monkeypatch):
-    """When both gates are already True AND tilelang is healthy, the hook
-    must do zero install work. (Tilelang repair on the already-True path
-    is covered by test_hook_runs_tilelang_repair_when_fla_already_true.)
+    """Both gates already True AND tilelang healthy -> zero install work.
+    (Tilelang repair on the already-True path is covered by
+    test_hook_runs_tilelang_repair_when_fla_already_true.)
     """
     fla_gate = _make_fake_gate(initial_return = True)
     conv_gate = _make_fake_gate(initial_return = True)
@@ -730,9 +721,8 @@ def test_hook_skips_install_when_gate_already_true(monkeypatch):
     monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install)
     monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install)
     monkeypatch.setattr(worker, "_install_package_wheel_first", conv_install)
-    # Tilelang healthy so the post_available path is a no-op (otherwise
-    # it would call tile_install, which is correct behaviour but
-    # outside the scope of this test).
+    # Tilelang healthy -> post_available path is a no-op (otherwise it
+    # would call tile_install, correct but out of scope here).
     monkeypatch.setattr(worker, "_tilelang_importable", lambda: True)
     monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: "0.1.9")
     monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False)
@@ -776,7 +766,7 @@ def test_hook_idempotent_on_repeat_call(monkeypatch):
 
     # First call: hook fires.
     _iu.is_flash_linear_attention_available()
-    # Subsequent calls: must not re-trigger the installer.
+    # Later calls: must not re-trigger the installer.
     _iu.is_flash_linear_attention_available()
     _iu.is_flash_linear_attention_available()
     assert fla_install.call_count == 1
@@ -800,7 +790,7 @@ def test_hook_handles_install_failure_gracefully(monkeypatch):
 
     from transformers.utils import import_utils as _iu
 
-    # Must not raise; returns False so transformers falls back to torch loop.
+    # Must not raise; returns False so transformers uses the torch loop.
     assert _iu.is_flash_linear_attention_available() is False
 
 
@@ -817,7 +807,7 @@ def test_hook_can_be_disabled_via_env(monkeypatch):
 
     from transformers.utils import import_utils as _iu
 
-    # Hook should NOT have been installed; gates remain the fakes.
+    # Hook not installed; gates remain the fakes.
     assert _iu.is_flash_linear_attention_available is fla_gate
     assert _iu.is_causal_conv1d_available is conv_gate
     fla_install.assert_not_called()
@@ -837,21 +827,20 @@ def test_hook_clears_lru_cache_before_first_check(monkeypatch):
     from transformers.utils import import_utils as _iu
 
     _iu.is_flash_linear_attention_available()
-    # The wrapper called cache_clear at least once before delegating.
+    # Wrapper called cache_clear at least once before delegating.
     assert fla_gate.cache_clear_count >= 1
 
 
 def test_hook_rewrites_previously_imported_module_bindings(monkeypatch):
-    """Modeling files bind `is_flash_linear_attention_available` locally
-    via `from ... import is_X`. Reassigning the attribute on
-    transformers.utils.import_utils alone does NOT reach those local
-    bindings. The hook installer sweeps sys.modules and rebinds them.
+    """Modeling files bind is_flash_linear_attention_available locally via
+    `from ... import is_X`. Reassigning the attribute on import_utils alone
+    misses those; the hook installer sweeps sys.modules and rebinds them.
     """
     fla_gate = _make_fake_gate(initial_return = False)
     conv_gate = _make_fake_gate(initial_return = True)
     _patch_iu_gates(monkeypatch, fla_gate, conv_gate)
 
-    # Create a fake modeling module that did `from ... import is_flash_linear_attention_available`.
+    # Fake modeling module that did `from ... import is_flash_linear_attention_available`.
     fake_mod = sys.modules.setdefault(
         "_test_fake_modeling_qwen35", type(sys)("_test_fake_modeling_qwen35")
     )
@@ -868,9 +857,9 @@ def test_hook_rewrites_previously_imported_module_bindings(monkeypatch):
 
     worker._install_fast_path_hooks(event_queue = _FakeQueue(), model_name = "unsloth/Qwen3.5-2B")
 
-    # The fake module's local binding has been rewritten to the wrapper.
+    # The fake module's local binding is rewritten to the wrapper.
     assert fake_mod.is_flash_linear_attention_available is not fla_gate
-    # Calling through the fake module's reference triggers the install.
+    # Calling through the fake module's reference triggers install.
     assert fake_mod.is_flash_linear_attention_available() is True
 
     del sys.modules["_test_fake_modeling_qwen35"]
@@ -894,7 +883,7 @@ def test_hook_skips_when_import_utils_unavailable(monkeypatch):
 
 
 def test_substring_fallback_unchanged_when_hook_skipped(monkeypatch):
-    """Hook disabled -> legacy gate falls back to auto-discovered model types."""
+    """Hook disabled -> legacy gate falls back to auto-discovered types."""
     install_mock = mock.Mock()
     monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", install_mock)
     monkeypatch.setattr(worker, "_discover_fla_model_types", lambda: frozenset({"qwen3_5"}))
@@ -907,8 +896,7 @@ def test_substring_fallback_unchanged_when_hook_skipped(monkeypatch):
     assert install_mock.call_count == 1
 
 
-# ───────────────────────────────────────────────────────────────────
-# Regression tests for the 10-reviewer findings:
+# Regression tests for the reviewer findings:
 #   1. tilelang Qwen-guard on hook path (non-Qwen FLA models)
 #   2. tilelang repair must not replace torch / CUDA stack
 #   3. hook must trust installer's bool, not transformers metadata
@@ -917,12 +905,11 @@ def test_substring_fallback_unchanged_when_hook_skipped(monkeypatch):
 #   6. tilelang skipped when FLA was skipped / failed
 #   7. tilelang repair runs when FLA is already True
 #   8. older FLA detected as stale and reinstalled
-# ───────────────────────────────────────────────────────────────────
 
 
 def test_hook_does_not_install_tilelang_for_model_outside_allowlist(monkeypatch):
-    """A model whose name is not in the auto-discovered FLA allowlist calls
-    is_flash_linear_attention_available but should NOT get tilelang."""
+    """A model not in the auto-discovered FLA allowlist calls
+    is_flash_linear_attention_available but must NOT get tilelang."""
     fla_gate = _make_fake_gate(initial_return = False)
     conv_gate = _make_fake_gate(initial_return = True)
     _patch_iu_gates(monkeypatch, fla_gate, conv_gate)
@@ -939,7 +926,7 @@ def test_hook_does_not_install_tilelang_for_model_outside_allowlist(monkeypatch)
     monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False)
     # Hermetize the auto-discovered set so the test stays valid as new
     # transformers releases add FLA-using model_types (eg olmo_hybrid in
-    # 5.4.0). The semantic under test is "outside-allowlist -> no tilelang".
+    # 5.4.0). Test semantic: "outside-allowlist -> no tilelang".
     monkeypatch.setattr(
         worker,
         "_discover_fla_model_types",
@@ -986,7 +973,7 @@ def test_hook_does_install_tilelang_for_qwen35(monkeypatch):
 
 def test_tilelang_repair_does_not_touch_torch_cuda_stack(monkeypatch):
     """Finding #2: the broken-tvm-ffi repair must use --no-deps on the
-    forced step so --force-reinstall does not cascade through
+    forced step so --force-reinstall doesn't cascade through
     apache-tvm-ffi's dep graph and pull a different torch wheel.
     """
     monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising = False)
@@ -1000,35 +987,25 @@ def test_tilelang_repair_does_not_touch_torch_cuda_stack(monkeypatch):
 
     assert run_mock.call_count == 2
     repair_args = run_mock.call_args_list[0][0][0]
-    # The forced step MUST be --no-deps so torch / CUDA stack is untouched.
+    # Forced step MUST be --no-deps so torch / CUDA stack is untouched.
     assert "--force-reinstall" in repair_args and "--no-deps" in repair_args
-    # And it touches ONLY apache-tvm-ffi, not tilelang / torch.
+    # Touches ONLY apache-tvm-ffi, not tilelang / torch.
     assert all("tilelang" not in a for a in repair_args)
     assert all("torch" not in a for a in repair_args)
 
 
 def test_hook_trusts_installer_bool_not_metadata(monkeypatch):
-    """Finding #3: if pip exits 0 but deep imports fail, the installer
-    returns False; the hook must propagate False even if the underlying
-    `original()` gate (which only checks metadata) returns True after
-    pip succeeds.
-
-    Setup mirrors the real bug:
-      1. Pre-install: gate=False (FLA not present) → wrapper triggers install.
-      2. Installer's `_flash_linear_attention_importable` post-probe fails,
-         so the installer returns False. (pip exited 0 but `import fla.modules`
-         raised because of a missing transitive dep.)
-      3. Post-install: gate would return True (metadata check sees fla-core
-         version) — but the wrapper must IGNORE that and use the installer's
-         False so transformers takes the torch fallback.
+    """Finding #3: if pip exits 0 but deep imports fail, the installer returns
+    False; the hook must propagate that False even if the metadata-only gate
+    returns True after pip succeeds, so transformers takes the torch fallback.
     """
     # Gate flips True after install (simulating "metadata sees fla").
     fla_gate = _make_fake_gate(initial_return = False)
     conv_gate = _make_fake_gate(initial_return = True)
     _patch_iu_gates(monkeypatch, fla_gate, conv_gate)
 
-    # Installer "succeeds" at pip, AND flips the gate to True (metadata
-    # sees fla post-install), BUT returns False (deep import broken).
+    # Installer "succeeds" at pip and flips the gate to True (metadata
+    # sees fla post-install), but returns False (deep import broken).
     def _bad_install(eq):
         fla_gate.next_return = True  # metadata says yes after pip
         return False  # but deep import is broken
@@ -1051,9 +1028,9 @@ def test_hook_trusts_installer_bool_not_metadata(monkeypatch):
 
 
 def test_rebind_does_not_trigger_module_getattr(monkeypatch):
-    """Finding #5: the rebind sweep must use __dict__, not getattr(),
-    to avoid invoking transformers' lazy module __getattr__ which spits
-    out hundreds of "Accessing X from .models..." warnings.
+    """Finding #5: the rebind sweep must use __dict__, not getattr(), to
+    avoid invoking transformers' lazy module __getattr__ which spits out
+    hundreds of "Accessing X from .models..." warnings.
     """
     original = object()
     replacement = object()
@@ -1068,8 +1045,8 @@ def test_rebind_does_not_trigger_module_getattr(monkeypatch):
     lazy = _GetattrTripwire("_lazy_test_module")
     sys.modules["_lazy_test_module"] = lazy
     try:
-        # No module-level binding to `is_flash_linear_attention_available`
-        # in __dict__, so the sweep must NOT trip the tripwire.
+        # No `is_flash_linear_attention_available` in __dict__, so the
+        # sweep must NOT trip the tripwire.
         worker._rebind_in_already_imported_modules(
             attr_name = "is_flash_linear_attention_available",
             old_obj = original,
@@ -1085,7 +1062,7 @@ def test_rebind_does_not_trigger_module_getattr(monkeypatch):
 def test_hook_skips_tilelang_when_fla_install_is_skipped(monkeypatch):
     """Finding #6: env-skipped FLA returns False from
     _ensure_flash_linear_attention_unconditional; tilelang must NOT
-    install in that case.
+    install then.
     """
     fla_gate = _make_fake_gate(initial_return = False)
     conv_gate = _make_fake_gate(initial_return = True)
@@ -1107,9 +1084,9 @@ def test_hook_skips_tilelang_when_fla_install_is_skipped(monkeypatch):
 
 
 def test_hook_runs_tilelang_repair_when_fla_already_true(monkeypatch):
-    """Finding #7: when FLA is already importable (gate returns True at
-    first probe) but tilelang is missing or apache-tvm-ffi is on the
-    broken list, the post-available action must still run tilelang.
+    """Finding #7: when FLA is already importable (gate True at first
+    probe) but tilelang is missing or apache-tvm-ffi is on the broken
+    list, the post-available action must still run tilelang.
     """
     fla_gate = _make_fake_gate(initial_return = True)
     conv_gate = _make_fake_gate(initial_return = True)
@@ -1120,7 +1097,7 @@ def test_hook_runs_tilelang_repair_when_fla_already_true(monkeypatch):
     monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", fla_install)
     monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", tile_install)
     monkeypatch.setattr(worker, "_install_package_wheel_first", mock.Mock(return_value = True))
-    # tilelang missing AND tvm-ffi is on broken list — both trigger repair.
+    # tilelang missing AND tvm-ffi on broken list — both trigger repair.
     monkeypatch.setattr(worker, "_tilelang_importable", lambda: False)
     monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: "0.1.11")
     monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising = False)
@@ -1130,19 +1107,19 @@ def test_hook_runs_tilelang_repair_when_fla_already_true(monkeypatch):
     from transformers.utils import import_utils as _iu
 
     _iu.is_flash_linear_attention_available()
-    # FLA install was NOT needed; tilelang repair WAS still triggered.
+    # FLA install NOT needed; tilelang repair still triggered.
     fla_install.assert_not_called()
     tile_install.assert_called_once()
 
 
 def test_fla_installer_force_reinstalls_when_older_version_present(monkeypatch):
-    """Finding #8: when an older `flash-linear-attention` is importable
-    but below the pin, the installer must force a reinstall (not no-op).
+    """Finding #8: an older `flash-linear-attention` that is importable
+    but below the pin must force a reinstall (not no-op).
     """
     monkeypatch.delenv(worker._FLA_SKIP_ENV, raising = False)
     monkeypatch.setattr(worker.shutil, "which", lambda name: "/usr/bin/uv")
     monkeypatch.setattr(worker, "_installed_torch_version_tuple", lambda: (2, 9))
-    # Importable but stale (current() reports False even though importable() is True).
+    # Importable but stale (current()=False though importable()=True).
     monkeypatch.setattr(worker, "_flash_linear_attention_importable", lambda: True)
     monkeypatch.setattr(worker, "_flash_linear_attention_current", lambda **kw: False)
     run_mock = mock.Mock(return_value = mock.Mock(returncode = 0, stdout = ""))
@@ -1161,23 +1138,19 @@ def test_fla_installer_force_reinstalls_when_older_version_present(monkeypatch):
 
 
 def test_run_training_process_eagerly_installs_causal_conv1d_in_normal_mode():
-    """Finding #4: SSM modeling files use `lazy_load_kernel("causal-conv1d")`
-    and never call `is_causal_conv1d_available()`, so the hook would not
-    fire for them. The orchestrator must always run the eager
-    substring installer regardless of hook mode.
-
-    This test reads the worker source rather than running the full
-    orchestrator (which requires a configured training config). It
-    asserts the eager install is OUTSIDE the if/else hook branch.
+    """Finding #4: SSM modeling files use lazy_load_kernel and never call
+    is_causal_conv1d_available(), so the hook won't fire; the orchestrator must
+    always run the eager installer regardless of hook mode. Reads the worker
+    source and asserts the eager install is OUTSIDE the if/else hook branch.
     """
     import inspect
 
     src = inspect.getsource(worker.run_training_process)
-    # Find the orchestration block.
+    # Orchestration block.
     assert "_ensure_causal_conv1d_fast_path(event_queue, model_name)" in src
     assert "_install_fast_path_hooks(event_queue, model_name)" in src
-    # The eager causal_conv1d call must appear BEFORE the hook-mode if/else,
-    # not nested inside the `if _FAST_PATH_HOOKS_SKIP_ENV` branch.
+    # Eager causal_conv1d call must come BEFORE the hook-mode if/else, not
+    # nested inside the `if _FAST_PATH_HOOKS_SKIP_ENV` branch.
     eager_pos = src.find("_ensure_causal_conv1d_fast_path(event_queue, model_name)")
     skip_check_pos = src.find('os.getenv(_FAST_PATH_HOOKS_SKIP_ENV) == "1"')
     assert eager_pos < skip_check_pos, (
@@ -1187,19 +1160,16 @@ def test_run_training_process_eagerly_installs_causal_conv1d_in_normal_mode():
     )
 
 
-# ───────────────────────────────────────────────────────────────────
-# HIP / ROCm regression coverage (h34v3nzc0dex Strix Halo report).
-# tilelang 0.1.8 has no HIP GEMM backend; FLA's TileLang dispatch
-# crashes mid-backward on AMD with "Unsupported target for gemm: hip".
-# The fix: skip the install on HIP-built torch AND setdefault
-# FLA_TILELANG=0 so already-installed tilelang doesn't get used either.
-# ───────────────────────────────────────────────────────────────────
+# HIP / ROCm regression coverage (Strix Halo report).
+# tilelang 0.1.8 has no HIP GEMM backend; FLA's TileLang dispatch crashes
+# mid-backward on AMD ("Unsupported target for gemm: hip"). Fix: skip install on
+# HIP torch AND setdefault FLA_TILELANG=0 so an existing tilelang isn't used.
 
 
 def test_tilelang_platform_unsupported_on_hip_torch(monkeypatch):
-    """Strix Halo / MI300 with ROCm torch: linux + x86_64 looks
-    identical to a CUDA box at the OS level, so the platform check
-    must consult torch.version.hip explicitly.
+    """Strix Halo / MI300 with ROCm torch: linux + x86_64 looks identical
+    to a CUDA box at the OS level, so the platform check must consult
+    torch.version.hip explicitly.
     """
     monkeypatch.setattr(worker, "_torch_has_hip", lambda: True)
     assert worker._tilelang_platform_supported() is False
@@ -1220,9 +1190,9 @@ def test_tilelang_install_skipped_on_hip_torch(monkeypatch):
 
 
 def test_install_fast_path_hooks_sets_fla_tilelang_zero_on_hip(monkeypatch):
-    """When HIP torch is detected, hook installer must set
-    FLA_TILELANG=0 (via setdefault — respects user override) so any
-    PRE-EXISTING tilelang install isn't used by FLA's dispatcher.
+    """On HIP torch, the hook installer must setdefault FLA_TILELANG=0
+    (respecting user override) so a PRE-EXISTING tilelang install isn't
+    used by FLA's dispatcher.
     """
     import os as _os
 
@@ -1239,8 +1209,8 @@ def test_install_fast_path_hooks_sets_fla_tilelang_zero_on_hip(monkeypatch):
 
 
 def test_install_fast_path_hooks_respects_user_fla_tilelang_override(monkeypatch):
-    """If the user explicitly set FLA_TILELANG (even on HIP), don't
-    overwrite — they may know they have a HIP-aware tilelang fork.
+    """If the user set FLA_TILELANG (even on HIP), don't overwrite — they
+    may have a HIP-aware tilelang fork.
     """
     import os as _os
 
@@ -1278,7 +1248,7 @@ def test_install_fast_path_hooks_does_not_set_fla_tilelang_on_cuda(monkeypatch):
 
 
 def _make_fake_transformers_tree(tmp_path, fla_types: list[str], non_fla_types: list[str]):
-    """Lay out a tmp dir as `transformers/models/{type}/modeling_{type}.py`."""
+    """Lay out tmp dir as `transformers/models/{type}/modeling_{type}.py`."""
     pkg = tmp_path / "transformers"
     models = pkg / "models"
     models.mkdir(parents = True)
@@ -1341,7 +1311,7 @@ def test_discover_fla_model_types_caches_across_calls(tmp_path, monkeypatch):
     second = worker._discover_fla_model_types()
 
     assert first == second
-    assert read_calls[0] == after_first  # cache hit: no extra disk reads
+    assert read_calls[0] == after_first  # cache hit: no extra reads
 
 
 def test_discover_fla_model_types_handles_missing_transformers(monkeypatch):
@@ -1382,7 +1352,7 @@ def test_discover_fla_model_types_handles_unreadable_file(tmp_path, monkeypatch)
 
     monkeypatch.setattr(_Path, "read_text", boom_read)
     result = worker._discover_fla_model_types()
-    assert result == frozenset()  # unreadable file simply doesn't contribute
+    assert result == frozenset()  # unreadable file doesn't contribute
 
 
 def test_model_wants_tilelang_handles_real_repo_names(monkeypatch):
@@ -1423,22 +1393,17 @@ def test_model_wants_tilelang_normalizes_separators(monkeypatch):
         assert worker._model_wants_tilelang(variant) is True, variant
 
 
-# ────────────────────────────────────────────────────────────────────
-# HIP source-build gcc-install-dir coverage (h34v3nzc0dex Strix Halo).
-# Ubuntu 24.04 ships gcc-14's runtime dir without /usr/include/c++/14,
-# so ROCm clang-20 picks it and fails with 'cstdlib' file not found
-# when building causal-conv1d (or any other HIP source fallback).
-# _hipcc_gcc_install_dir() finds a gcc dir that has both halves; the
-# _install_package_wheel_first HIP branch passes it to clang via
-# HIPCC_COMPILE_FLAGS_APPEND. Parallel to bbf004c's setup.sh fix for
-# the llama.cpp HIP build (PR #5301).
-# ────────────────────────────────────────────────────────────────────
+# HIP source-build gcc-install-dir coverage (Strix Halo).
+# Ubuntu 24.04 ships gcc-14's runtime dir without /usr/include/c++/14, so ROCm
+# clang-20 picks it and fails ('cstdlib' not found) building causal-conv1d.
+# _hipcc_gcc_install_dir() finds a gcc dir with both halves; the HIP branch of
+# _install_package_wheel_first passes it via HIPCC_COMPILE_FLAGS_APPEND.
+# Parallels bbf004c's setup.sh fix for the llama.cpp HIP build (PR #5301).
 
 
 def _isdir_for_layout(*existing: str):
-    """Return an os.path.isdir replacement that only treats the given
-    absolute paths as directories. Lets a test simulate exactly which
-    gcc runtime dirs and C++ header dirs exist on the host."""
+    """os.path.isdir replacement treating only the given absolute paths as
+    directories, to simulate which gcc runtime / C++ header dirs exist."""
     valid = set(existing)
 
     def fake_isdir(path: str) -> bool:
@@ -1449,7 +1414,7 @@ def _isdir_for_layout(*existing: str):
 
 def test_hipcc_gcc_install_dir_picks_highest_with_headers(monkeypatch):
     """gcc-14 has runtime but no /usr/include/c++/14; loop falls through
-    to gcc-13 which has both. This is the exact Ubuntu 24.04 layout."""
+    to gcc-13 which has both. The exact Ubuntu 24.04 layout."""
     monkeypatch.setattr(sys, "platform", "linux")
     import platform as _platform
 
@@ -1485,8 +1450,8 @@ def test_hipcc_gcc_install_dir_picks_14_when_headers_exist(monkeypatch):
 
 
 def test_hipcc_gcc_install_dir_returns_none_when_no_match(monkeypatch):
-    """No gcc dir has both halves → return None and skip the env injection
-    rather than guessing wrong and surfacing a confusing build failure."""
+    """No gcc dir has both halves → return None and skip env injection
+    rather than guessing wrong and causing a confusing build failure."""
     monkeypatch.setattr(sys, "platform", "linux")
     import platform as _platform
 
@@ -1516,10 +1481,9 @@ def test_hipcc_gcc_install_dir_returns_none_on_non_x86_64(monkeypatch):
 
 
 def _make_hip_install_env(monkeypatch, *, gcc_dir: str | None):
-    """Common scaffolding for tests that exercise the HIP source-build
-    branch of _install_package_wheel_first end-to-end. The package isn't
-    installed yet, no prebuilt wheel exists, hipcc is on PATH, and the
-    fake env reports an HIP torch."""
+    """Scaffolding for end-to-end tests of the HIP source-build branch of
+    _install_package_wheel_first: package not installed, no prebuilt
+    wheel, hipcc on PATH, fake env reports HIP torch."""
     monkeypatch.setattr(builtins, "__import__", _missing_module_import("causal_conv1d"))
     monkeypatch.setattr(
         worker,
@@ -1574,8 +1538,8 @@ def test_install_injects_gcc_install_dir_on_hip_source_build(monkeypatch):
 
 
 def test_install_appends_to_existing_hipcc_compile_flags(monkeypatch):
-    """User has HIPCC_COMPILE_FLAGS_APPEND='-O3 -DFOO' set → final value
-    keeps the user's flags AND adds --gcc-install-dir at the end."""
+    """User has HIPCC_COMPILE_FLAGS_APPEND='-O3 -DFOO' → final value keeps
+    the user's flags AND appends --gcc-install-dir."""
     monkeypatch.setenv("HIPCC_COMPILE_FLAGS_APPEND", "-O3 -DFOO")
     _make_hip_install_env(monkeypatch, gcc_dir = "/usr/lib/gcc/x86_64-linux-gnu/13")
 
@@ -1636,9 +1600,9 @@ def test_install_respects_user_gcc_install_dir(monkeypatch):
         release_base_url = "https://example.com",
     )
 
-    # subprocess.run was invoked without env override (the user already
-    # set HIPCC_COMPILE_FLAGS_APPEND with --gcc-install-dir, so we left
-    # the env alone — the existing value is inherited normally).
+    # subprocess.run invoked without env override (user already set
+    # HIPCC_COMPILE_FLAGS_APPEND with --gcc-install-dir, so we left the
+    # env alone — the existing value is inherited).
     assert captured == {"_called": "yes_no_env"}
 
 
@@ -1660,7 +1624,7 @@ def test_install_does_not_inject_env_on_cuda(monkeypatch):
     monkeypatch.setattr(worker, "direct_wheel_url", lambda **kw: None)
     monkeypatch.setattr(worker.shutil, "which", lambda name: None)
     monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None)
-    # If _hipcc_gcc_install_dir were called on CUDA we'd want to know.
+    # _hipcc_gcc_install_dir must not be called on CUDA.
     monkeypatch.setattr(
         worker,
         "_hipcc_gcc_install_dir",
diff --git a/studio/backend/tests/test_transformers_version.py b/studio/backend/tests/test_transformers_version.py
index ff5e1f1381..915bc4b13b 100644
--- a/studio/backend/tests/test_transformers_version.py
+++ b/studio/backend/tests/test_transformers_version.py
@@ -10,9 +10,8 @@ from unittest.mock import patch
 
 
 # ---------------------------------------------------------------------------
-# We need to be able to import the module under test.  The studio backend
-# uses relative-style imports (``from utils.…``), so we add the backend
-# directory to *sys.path* if it is not already there.
+# The studio backend uses relative-style imports (``from utils.…``), so
+# add the backend directory to *sys.path* if not already present.
 # ---------------------------------------------------------------------------
 import sys
 
@@ -20,8 +19,8 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
 if _BACKEND_DIR not in sys.path:
     sys.path.insert(0, _BACKEND_DIR)
 
-# Stub the custom logger before importing the module under test so it
-# doesn't fail on the ``from loggers import get_logger`` line.
+# Stub the custom logger before import so ``from loggers import
+# get_logger`` doesn't fail.
 import types as _types
 
 _loggers_stub = _types.ModuleType("loggers")
@@ -90,7 +89,7 @@ class TestResolveBaseModel:
         (tmp_path / "config.json").write_text(json.dumps(config_cfg))
 
         result = _resolve_base_model(str(tmp_path))
-        # Should fall through, not return the self-referencing path
+        # Falls through; does not return the self-referencing path.
         assert result == str(tmp_path)
 
     def test_no_config_files(self, tmp_path: Path):
@@ -174,7 +173,7 @@ class TestNeedsTransformers5:
 
     def test_llama_does_not_need_v5(self):
         """Standard models should not trigger v5."""
-        # Patch network call to avoid real fetch
+        # Patch network call to avoid a real fetch.
         with patch(
             "utils.transformers_version._check_tokenizer_config_needs_v5",
             return_value = False,
@@ -182,13 +181,12 @@ class TestNeedsTransformers5:
             assert needs_transformers_5("meta-llama/Llama-3-8B") is False
 
     def test_local_checkpoint_resolved_via_config(self, tmp_path: Path):
-        """A local checkpoint with config.json pointing to Qwen3.5 should need v5."""
+        """Local checkpoint with config.json pointing to Qwen3.5 needs v5."""
         config_cfg = {"model_name": "Qwen/Qwen3.5-9B"}
         (tmp_path / "config.json").write_text(json.dumps(config_cfg))
 
-        # _resolve_base_model is called by ensure_transformers_version,
-        # but needs_transformers_5 just does substring matching.
-        # We test the full resolution chain here:
+        # needs_transformers_5 only does substring matching, so test the
+        # full resolution chain via _resolve_base_model here.
         resolved = _resolve_base_model(str(tmp_path))
         assert needs_transformers_5(resolved) is True
 
@@ -230,7 +228,7 @@ class TestCheckConfigNeeds550:
 
     def test_no_config_json(self, tmp_path: Path):
         """Missing config.json should return False (fail-open)."""
-        # Patch network call to avoid real fetch
+        # Patch network call to avoid a real fetch.
         with patch("urllib.request.urlopen") as mock_urlopen:
             mock_urlopen.side_effect = Exception("no network")
             assert _check_config_needs_550(str(tmp_path)) is False
@@ -311,8 +309,7 @@ class TestGetTransformersTier:
             assert get_transformers_tier("meta-llama/Llama-3-8B") == "default"
 
     def test_550_checked_before_530(self):
-        """Ensure 5.5.0 is checked first — a model matching both should get 550."""
-        # This shouldn't happen in practice, but verifies priority
+        """5.5.0 is checked first — a model matching both gets 550."""
         assert get_transformers_tier("gemma-4-model") == "550"
 
     def test_needs_transformers_5_compat(self):
diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py
index bdb1cd2ce8..c66d56528a 100644
--- a/studio/backend/tests/test_utils.py
+++ b/studio/backend/tests/test_utils.py
@@ -1,20 +1,10 @@
 # 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 utils/hardware and utils/utils — device detection, GPU memory, error formatting.
+"""Tests for utils/hardware and utils/utils: device detection, GPU memory, error formatting.
 
-These tests are designed to pass on ANY platform:
-  • NVIDIA GPU  (CUDA backend, requires torch)
-  • Apple Silicon (MLX backend, requires mlx)
-  • CPU-only     (no GPU at all)
-
-No ML framework is imported at the top level.
-Tests that need torch/mlx internals for mocking are skipped when unavailable.
-
-Run with:
-    cd studio/backend
-    python -m pytest tests/test_utils.py -v
+Passes on any platform (NVIDIA/CUDA, Apple Silicon/MLX, CPU-only). No ML framework
+is imported at top level; tests needing torch/mlx internals skip when unavailable.
 """
 
 import platform
@@ -189,10 +179,8 @@ class TestGetGpuMemoryInfo:
         assert "backend" in get_gpu_memory_info()
 
     def test_backend_matches_device(self):
-        # The backend field uses _backend_label, which swaps "cuda" for
-        # "rocm" when running on an AMD host (IS_ROCM=True) so the UI
-        # can render the correct label. On CUDA / XPU / MLX / CPU hosts
-        # it is equivalent to `get_device().value`.
+        # _backend_label swaps "cuda" for "rocm" on AMD hosts; elsewhere it
+        # equals get_device().value.
         from utils.hardware.hardware import _backend_label
         result = get_gpu_memory_info()
         assert result["backend"] == _backend_label(get_device())
diff --git a/studio/backend/tests/test_vision_cache.py b/studio/backend/tests/test_vision_cache.py
index 2af64dac91..2fee50d842 100644
--- a/studio/backend/tests/test_vision_cache.py
+++ b/studio/backend/tests/test_vision_cache.py
@@ -3,14 +3,13 @@
 
 """Tests for is_vision_model() caching behaviour.
 
-The vision detection cache (``_vision_detection_cache``) mirrors the existing
-``_audio_detection_cache`` pattern used by ``detect_audio_type()``.  These
-tests verify that:
+``_vision_detection_cache`` mirrors the ``_audio_detection_cache``
+pattern used by ``detect_audio_type()``. These tests verify:
 
-* Repeated calls for the same model hit the cache (no redundant work).
+* Repeated calls for the same model hit the cache.
 * Different models each trigger their own detection.
 * Both True and False results are cached.
-* The subprocess path (transformers 5.x models) is also cached.
+* The subprocess path (transformers 5.x models) is cached.
 * Exceptions that fall back to False are cached.
 """
 
@@ -21,9 +20,7 @@ from unittest.mock import patch, MagicMock
 
 import pytest
 
-# ---------------------------------------------------------------------------
 # sys.path + logger stub — same pattern as the rest of the test suite
-# ---------------------------------------------------------------------------
 _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
 if _BACKEND_DIR not in sys.path:
     sys.path.insert(0, _BACKEND_DIR)
@@ -39,9 +36,7 @@ from utils.models.model_config import (
 )
 
 
-# ---------------------------------------------------------------------------
 # Helpers
-# ---------------------------------------------------------------------------
 
 
 @pytest.fixture(autouse = True)
@@ -52,9 +47,7 @@ def _clear_vision_cache():
     _vision_detection_cache.clear()
 
 
-# ---------------------------------------------------------------------------
 # Cache hit / miss tests
-# ---------------------------------------------------------------------------
 
 
 class TestVisionCacheHitMiss:
@@ -62,8 +55,7 @@ class TestVisionCacheHitMiss:
 
     @patch("utils.models.model_config._is_vision_model_uncached", return_value = True)
     def test_second_call_uses_cache(self, mock_uncached):
-        """Calling is_vision_model() twice for the same model should invoke
-        the uncached function only once."""
+        """Two calls for the same model invoke the uncached fn once."""
         assert is_vision_model("org/my-vlm") is True
         assert is_vision_model("org/my-vlm") is True
         mock_uncached.assert_called_once_with("org/my-vlm", None)
@@ -95,21 +87,19 @@ class TestVisionCacheStoresFalse:
         assert _vision_detection_cache[("org/text-only", None)] is False
 
 
-# ---------------------------------------------------------------------------
 # Subprocess path (transformers 5.x) caching
-# ---------------------------------------------------------------------------
 
 
 class TestVisionCacheSubprocessPath:
-    """Models needing transformers 5.x go through _is_vision_model_subprocess.
-    The cache should prevent the subprocess from being spawned more than once
-    per model per process."""
+    """transformers 5.x models go through _is_vision_model_subprocess.
+    The cache should spawn the subprocess at most once per model per
+    process."""
 
     @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 should only fire on the first call; second is cached."""
-        # First call: goes through uncached → subprocess
+        """Subprocess fires only on the first call; second is cached."""
+        # First call: uncached → 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
@@ -117,17 +107,26 @@ class TestVisionCacheSubprocessPath:
         mock_subprocess.assert_called_once()
         assert _vision_detection_cache[("unsloth/Qwen3.5-2B", None)] is True
 
+    @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(
+        self, mock_needs_t5, mock_subprocess, mock_raw_config
+    ):
+        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)
+
 
-# ---------------------------------------------------------------------------
 # Exception handling — cache the False fallback
-# ---------------------------------------------------------------------------
 
 
 class TestVisionCacheOnException:
-    """When detection raises an exception, _is_vision_model_uncached
-    distinguishes permanent failures (cached as False) from transient
-    failures (returned as None, not cached so the next call can retry).
-    Verify both contracts."""
+    """On exception, _is_vision_model_uncached distinguishes permanent
+    failures (cached as False) from transient ones (returned as None,
+    not cached, so the next call retries). Verify both contracts."""
 
     @patch(
         "utils.models.model_config.load_model_config",
@@ -136,16 +135,11 @@ class TestVisionCacheOnException:
     @patch("utils.transformers_version.needs_transformers_5", return_value = False)
     def test_permanent_exception_result_cached(self, mock_needs_t5, mock_load_config):
         """A permanent failure (ValueError / RepositoryNotFoundError /
-        GatedRepoError / JSONDecodeError) should be caught, return False,
-        and that False should be cached so subsequent calls don't retry.
-
-        ValueError is used here because it's the simplest of the
-        code-path's cacheable exception types and does not require an
-        import of huggingface_hub errors (whose module path varies
-        across versions)."""
-        # First call: load_model_config raises -> except branch -> False.
+        GatedRepoError / JSONDecodeError) is caught, returns False, and
+        that False is cached so subsequent calls don't retry. ValueError
+        stands in as the simplest cacheable exception type."""
+        # First call raises -> False; second is a cache hit.
         assert is_vision_model("broken/model") is False
-        # Second call: cache hit, load_model_config not called again.
         assert is_vision_model("broken/model") is False
         mock_load_config.assert_called_once()
 
@@ -155,28 +149,20 @@ class TestVisionCacheOnException:
     )
     @patch("utils.transformers_version.needs_transformers_5", return_value = False)
     def test_transient_exception_not_cached(self, mock_needs_t5, mock_load_config):
-        """A transient failure (OSError, timeouts) should return None from
-        _is_vision_model_uncached, surface as False to the caller, and
-        NOT be cached, so the next call retries detection.  This matches
-        the documented behaviour on _vision_detection_cache:
-        'transient failures (network errors, timeouts) are NOT cached so
-        they can be retried.'"""
-        # First call: load_model_config raises OSError -> uncached None
-        # -> caller returns False without caching.
+        """A transient failure (OSError, timeouts) returns None from
+        _is_vision_model_uncached, surfaces as False, and is NOT cached
+        so the next call retries."""
+        # First call: OSError -> False, not cached; second call retries.
         assert is_vision_model("broken/model") is False
-        # Second call: cache miss again, load_model_config called a
-        # second time.
         assert is_vision_model("broken/model") is False
         assert mock_load_config.call_count == 2
 
 
-# ---------------------------------------------------------------------------
 # Direct detection path (non-transformers-5 models) caching
-# ---------------------------------------------------------------------------
 
 
 class TestVisionCacheDirectPath:
-    """For models that do NOT need transformers 5.x, the detection goes through
+    """Models that do NOT need transformers 5.x detect via
     load_model_config directly. The cache must work the same way."""
 
     @patch("utils.transformers_version.needs_transformers_5", return_value = False)
@@ -202,7 +188,7 @@ class TestVisionCacheDirectPath:
         cfg.architectures = ["LlamaForCausalLM"]
         mock_load_config.return_value = cfg
 
-        # LlamaForCausalLM doesn't end with VLM suffixes, no vision_config, etc.
+        # No VLM suffix, no vision_config, etc.
         assert is_vision_model("meta-llama/Llama-3-8B") is False
         assert is_vision_model("meta-llama/Llama-3-8B") is False
         mock_load_config.assert_called_once()
@@ -221,6 +207,42 @@ class TestVisionCacheDirectPath:
         assert is_vision_model("Qwen/Qwen2-VL-7B") is True
         mock_load_config.assert_called_once()
 
+    @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):
+        cfg = MagicMock(spec = [])
+        cfg.model_type = "gemma4"
+        cfg.architectures = ["Gemma4ForConditionalGeneration"]
+        mock_load_config.return_value = cfg
+
+        assert is_vision_model("google/gemma-4-E4B-it") is True
+        assert is_vision_model("google/gemma-4-E4B-it") is True
+        mock_load_config.assert_called_once()
+
+    @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):
+        cfg = MagicMock(spec = [])
+        cfg.model_type = "gemma4_audio"
+        cfg.architectures = ["Gemma4AudioModel"]
+        mock_load_config.return_value = cfg
+
+        assert is_vision_model("local/gemma4-audio-encoder") is False
+        assert is_vision_model("local/gemma4-audio-encoder") is False
+        mock_load_config.assert_called_once()
+
+    @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):
+        cfg = MagicMock(spec = [])
+        cfg.model_type = "gemma4_text"
+        cfg.architectures = ["Gemma4ForCausalLM"]
+        mock_load_config.return_value = cfg
+
+        assert is_vision_model("local/gemma-4-text") is False
+        assert is_vision_model("local/gemma-4-text") is False
+        mock_load_config.assert_called_once()
+
     @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):
@@ -236,21 +258,18 @@ class TestVisionCacheDirectPath:
         mock_load_config.assert_called_once()
 
 
-# ---------------------------------------------------------------------------
 # hf_token handling
-# ---------------------------------------------------------------------------
 
 
 class TestVisionCacheTokenHandling:
-    """The cache is keyed on (model_name, hf_token).
-    Different tokens for the same model should trigger separate detections
-    to handle gated models correctly."""
+    """The cache is keyed on (model_name, hf_token). Different tokens
+    for the same model trigger separate detections for gated models."""
 
     @patch("utils.models.model_config._is_vision_model_uncached", return_value = True)
     def test_different_tokens_trigger_new_detection(self, mock_uncached):
-        """Calls with different tokens should trigger separate detections to
-        handle gated models correctly (e.g. unauthenticated probe → False,
-        then authenticated call should re-check)."""
+        """Different tokens trigger separate detections for gated models
+        (e.g. unauthenticated probe → False, then authenticated
+        re-check)."""
         assert is_vision_model("gated/model", hf_token = "token-a") is True
         assert is_vision_model("gated/model", hf_token = "token-b") is True
         assert mock_uncached.call_count == 2
@@ -261,3 +280,157 @@ class TestVisionCacheTokenHandling:
         assert is_vision_model("gated/model", hf_token = "token-a") is True
         assert is_vision_model("gated/model", hf_token = "token-a") is True
         mock_uncached.assert_called_once()
+
+
+# ---------------------------------------------------------------------------
+# Direct unit tests for _raw_config_has_vision_config
+# ---------------------------------------------------------------------------
+
+
+import json as _json
+
+from utils.models.model_config import (
+    _AUDIO_ONLY_MODEL_TYPES,
+    _VISION_CHECK_INLINE_HELPERS,
+    _VISION_CHECK_SCRIPT,
+    _is_vlm,
+    _raw_config_has_vision_config,
+)
+
+
+def _write_config(tmp_path, config):
+    (tmp_path / "config.json").write_text(_json.dumps(config))
+    return tmp_path
+
+
+class TestRawConfigVlmDetection:
+    """Direct coverage of _raw_config_has_vision_config across the same
+    indicator set used by _is_vlm. The cache integration tests above mock
+    this function; these exercise its real implementation."""
+
+    def test_truthy_vision_config(self, tmp_path):
+        p = _write_config(tmp_path, {"vision_config": {"hidden_size": 1024}})
+        assert _raw_config_has_vision_config(str(p)) is True
+
+    def test_empty_vision_config_key(self, tmp_path):
+        p = _write_config(tmp_path, {"vision_config": {}})
+        assert _raw_config_has_vision_config(str(p)) is True
+
+    def test_arch_suffix_detection(self, tmp_path):
+        p = _write_config(
+            tmp_path,
+            {
+                "architectures": ["Gemma4ForConditionalGeneration"],
+                "model_type": "gemma4",
+            },
+        )
+        assert _raw_config_has_vision_config(str(p)) is True
+
+    def test_img_processor_key(self, tmp_path):
+        p = _write_config(tmp_path, {"img_processor": {"image_size": 336}})
+        assert _raw_config_has_vision_config(str(p)) is True
+
+    def test_image_token_index_key(self, tmp_path):
+        p = _write_config(tmp_path, {"image_token_index": 32000})
+        assert _raw_config_has_vision_config(str(p)) is True
+
+    def test_known_vlm_model_type(self, tmp_path):
+        p = _write_config(tmp_path, {"model_type": "gemma4"})
+        assert _raw_config_has_vision_config(str(p)) is True
+
+    def test_plain_text_model_returns_false(self, tmp_path):
+        p = _write_config(
+            tmp_path,
+            {"model_type": "llama", "architectures": ["LlamaForCausalLM"]},
+        )
+        assert _raw_config_has_vision_config(str(p)) is False
+
+    def test_missing_config_returns_none(self, tmp_path):
+        assert _raw_config_has_vision_config(str(tmp_path)) is None
+
+
+# ---------------------------------------------------------------------------
+# Self-contained subprocess script (no parent backend imports)
+# ---------------------------------------------------------------------------
+
+
+class TestSubprocessScript:
+    def test_does_not_import_parent_module(self):
+        assert "from utils.models.model_config" not in _VISION_CHECK_SCRIPT
+
+    def test_inline_is_vlm_executes_correctly(self):
+        ns: dict = {}
+        exec(_VISION_CHECK_INLINE_HELPERS, ns)
+        inline_is_vlm = ns["_is_vlm"]
+
+        class _C:
+            def __init__(self, **kw):
+                for k, v in kw.items():
+                    setattr(self, k, v)
+
+        assert (
+            inline_is_vlm(
+                _C(
+                    model_type = "gemma4",
+                    architectures = ["Gemma4ForConditionalGeneration"],
+                )
+            )
+            is True
+        )
+        assert (
+            inline_is_vlm(_C(model_type = "gemma4_text", architectures = ["Gemma4ForCausalLM"]))
+            is False
+        )
+        assert inline_is_vlm(_C(model_type = "llama", architectures = ["LlamaForCausalLM"])) is False
+
+
+# ---------------------------------------------------------------------------
+# Audio-only model exclusion must apply across every detection path
+# ---------------------------------------------------------------------------
+
+
+class TestVlmAudioExclusion:
+    """The {csm, whisper} guard previously lived only in the direct caller
+    branch. These tests assert it now applies inside _is_vlm, the raw
+    fallback, and the inlined subprocess helper too."""
+
+    def test_audio_only_set_canonical(self):
+        assert _AUDIO_ONLY_MODEL_TYPES == {"csm", "whisper"}
+
+    def test_is_vlm_excludes_whisper(self):
+        cfg = MagicMock(spec = [])
+        cfg.model_type = "whisper"
+        cfg.architectures = ["WhisperForConditionalGeneration"]
+        assert _is_vlm(cfg) is False
+
+    def test_raw_fallback_excludes_whisper(self, tmp_path):
+        p = _write_config(
+            tmp_path,
+            {
+                "architectures": ["WhisperForConditionalGeneration"],
+                "model_type": "whisper",
+            },
+        )
+        assert _raw_config_has_vision_config(str(p)) is False
+
+    def test_inline_subprocess_helper_excludes_whisper(self):
+        ns: dict = {}
+        exec(_VISION_CHECK_INLINE_HELPERS, ns)
+        cfg = MagicMock(spec = [])
+        cfg.model_type = "whisper"
+        cfg.architectures = ["WhisperForConditionalGeneration"]
+        assert ns["_is_vlm"](cfg) is False
+
+    @patch("utils.models.model_config._is_vision_model_subprocess", return_value = None)
+    @patch("utils.transformers_version.needs_transformers_5", return_value = True)
+    def test_t5_subprocess_none_falls_back_through_raw_for_whisper(
+        self, mock_needs_t5, mock_subprocess, tmp_path
+    ):
+        _write_config(
+            tmp_path,
+            {
+                "architectures": ["WhisperForConditionalGeneration"],
+                "model_type": "whisper",
+            },
+        )
+        assert is_vision_model(str(tmp_path)) is False
diff --git a/studio/backend/tests/test_vram_estimation.py b/studio/backend/tests/test_vram_estimation.py
index 65964908d7..2def8738e2 100644
--- a/studio/backend/tests/test_vram_estimation.py
+++ b/studio/backend/tests/test_vram_estimation.py
@@ -531,7 +531,7 @@ class TestQuantizationSkips(unittest.TestCase):
         )
 
     def test_vlm_prefix_skip_module_does_not_match_text_alias(self):
-        # vision_tower-prefixed skips must not shadow text aliases sharing the
+        # vision_tower-prefixed skips must not shadow text aliases with the
         # same suffix.
         baseline = replace(QUANT_SKIP_STRUCTURED, quantization_skip_modules = [])
         vlm_skip = replace(
@@ -995,9 +995,9 @@ class TestParallelDenseMoE(unittest.TestCase):
             + with_parallel.num_experts * with_parallel.hidden_size
         )
         dense_only = with_parallel.hidden_size * with_parallel.intermediate_size * 3
-        # why: under gemma4 enable_moe_block, the layer's `self.experts` is a
-        # sibling of `self.mlp`; the `text.layers..mlp` aggregate must
-        # cover the dense path only, with experts in their own aggregate.
+        # why: under gemma4 enable_moe_block, `self.experts` is a sibling of
+        # `self.mlp`; the `text.layers..mlp` aggregate covers the dense path
+        # only, with experts in their own aggregate.
         self.assertEqual(elements["text.layers.0.mlp"], dense_only)
         self.assertEqual(elements["text.layers.0.experts"], moe_only)
 
@@ -1148,9 +1148,9 @@ class TestPerLayerInputAccounting(unittest.TestCase):
     def test_per_layer_input_modules_count_quantizable_block(self):
         with_ple = self._arch()
         without_ple = replace(with_ple, hidden_size_per_layer_input = 0)
-        # The PLE block adds: model_projection (hd*nl*pli), per_layer_input_gate
-        # (hd*pli per layer) + per_layer_projection (pli*hd per layer) as
-        # quantizable text linears.
+        # PLE block adds these quantizable text linears: model_projection
+        # (hd*nl*pli), per_layer_input_gate (hd*pli per layer),
+        # per_layer_projection (pli*hd per layer).
         n_layers = with_ple.num_hidden_layers
         hd = with_ple.hidden_size
         pli = with_ple.hidden_size_per_layer_input
@@ -1161,10 +1161,10 @@ class TestPerLayerInputAccounting(unittest.TestCase):
         self.assertGreaterEqual(delta, expected_quantizable_extra)
 
     def test_all_linear_lora_excludes_per_layer_input_modules(self):
-        # why: Unsloth's get_peft_regex requires module names to contain a
-        # component tag (mlp/attn/...); PLE module names (per_layer_input_gate,
-        # per_layer_projection, per_layer_model_projection) lack any tag, so
-        # all-linear training does NOT attach LoRA to them.
+        # why: Unsloth's get_peft_regex requires a component tag (mlp/attn/...)
+        # in module names; PLE names (per_layer_input_gate, per_layer_projection,
+        # per_layer_model_projection) lack one, so all-linear does NOT attach
+        # LoRA to them.
         arch = self._arch()
         without_ple = replace(arch, hidden_size_per_layer_input = 0)
         self.assertEqual(
@@ -1234,11 +1234,10 @@ class TestExpertsSkipGranularity(unittest.TestCase):
         bytes_skip_experts = compute_model_weights_bytes(skip_experts, "qlora", True)
         bytes_skip_mlp = compute_model_weights_bytes(skip_full_mlp, "qlora", True)
         # why: under gemma4 enable_moe_block, `self.experts` is a sibling of
-        # `self.mlp`; skipping `model.layers.0.mlp` should cover only the
-        # dense MLP, while `model.layers.0.mlp.experts` covers the routed
-        # experts. Routed experts have far more params than the dense MLP,
-        # so skipping experts must add more bytes than skipping the dense
-        # path.
+        # `self.mlp`; skipping `model.layers.0.mlp` covers only the dense MLP,
+        # while `model.layers.0.mlp.experts` covers the routed experts. Routed
+        # experts have far more params than the dense MLP, so skipping experts
+        # must add more bytes than skipping the dense path.
         self.assertGreater(bytes_skip_experts, bytes_no_skip)
         self.assertGreater(bytes_skip_mlp, bytes_no_skip)
         self.assertGreater(bytes_skip_experts, bytes_skip_mlp)
@@ -1485,8 +1484,8 @@ class TestPerLayerInputSkipAlias(unittest.TestCase):
         )
 
         arch_with = extract_arch_config(self._hf(["model.layers.0"]))
-        # The text.layers.0 aggregate must include the PLE per-layer modules,
-        # so the same skip on a config without PLE produces a smaller value.
+        # text.layers.0 aggregate includes the PLE per-layer modules, so the
+        # same skip on a no-PLE config produces a smaller value.
         arch_without = extract_arch_config(
             SimpleNamespace(
                 text_config = SimpleNamespace(
@@ -1558,7 +1557,7 @@ class TestSharedExpertVariants(unittest.TestCase):
         arch_separate = extract_arch_config(self._hf(shared_expert_intermediate_size = 64))
         arch_implicit = extract_arch_config(self._hf(n_shared_experts = 1))
         # Different shared sizes (64 vs default moe_intermediate_size=128) must
-        # produce different MoE element counts.
+        # give different MoE element counts.
         self.assertNotEqual(
             _compute_moe_mlp_elements(arch_separate),
             _compute_moe_mlp_elements(arch_implicit),
@@ -1567,7 +1566,7 @@ class TestSharedExpertVariants(unittest.TestCase):
     def test_shared_expert_gate_counted_only_for_qwen_style(self):
         from utils.hardware.vram_estimation import _compute_moe_mlp_elements
 
-        # Qwen-style: shared_expert_intermediate_size set -> shared_expert_gate counted.
+        # Qwen-style: shared_expert_intermediate_size set -> gate counted.
         qwen_arch = extract_arch_config(self._hf(shared_expert_intermediate_size = 64))
         hd = qwen_arch.hidden_size
         ms = qwen_arch.moe_intermediate_size
@@ -1624,8 +1623,8 @@ class TestSharedExpertActivation(unittest.TestCase):
         )
 
     def test_shared_expert_plus_dense_block_compose(self):
-        # gemma4 enable_moe_block with hypothetical shared expert: dense + routed
-        # + shared all live per layer; mlp_size should sum all three terms.
+        # gemma4 enable_moe_block with a hypothetical shared expert: dense +
+        # routed + shared all live per layer; mlp_size sums all three.
         from utils.hardware.vram_estimation import _layer_qkv_mlp_sizes
 
         arch = self._make(
@@ -1773,7 +1772,7 @@ class TestSparseMoeSkipAliases(unittest.TestCase):
                 shared_expert_intermediate_size = 32,
             )
         )
-        # shared_expert delta only -- routed mlp.experts is NOT skipped.
+        # shared_expert delta only -- routed mlp.experts NOT skipped.
         delta = _compute_skipped_quantizable_elements(arch)
         self.assertGreater(delta, 0)
         full_layer = extract_arch_config(
@@ -1944,10 +1943,10 @@ class TestErnieMoEListConfig(unittest.TestCase):
                 moe_intermediate_size = [1536, 512],
             )
         )
-        # why: ERNIE 4.5 VL MoE encodes [text_routed, vision_routed]; the
-        # second element is the vision-routed expert width, not the shared
-        # expert width. Shared experts are sized from the text-routed width
-        # (= moe_intermediate_size[0]) when moe_num_shared_experts is set.
+        # why: ERNIE 4.5 VL MoE encodes [text_routed, vision_routed]; element 1
+        # is the vision-routed width, not the shared-expert width. Shared
+        # experts size from the text-routed width (moe_intermediate_size[0])
+        # when moe_num_shared_experts is set.
         self.assertEqual(arch.moe_intermediate_size, 1536)
         self.assertIsNone(arch.shared_expert_intermediate_size)
         self.assertEqual(arch.n_shared_experts, 0)
@@ -2070,7 +2069,7 @@ class TestMultimodalFullModelBytes(unittest.TestCase):
                 load_in_4bit = True,
             )
         self.assertEqual(metadata.get("estimation_mode"), "detailed")
-        # model_weights_gb must reflect the extra non-text bytes (>5 GB
+        # model_weights_gb must reflect the extra non-text bytes (>5 GB,
         # since text-only arch_fp16 is small for these dims).
         self.assertGreater(metadata["vram_breakdown"]["model_weights_gb"], 5.0)
 
diff --git a/studio/backend/tests/test_windows_gpu_detection_mock.py b/studio/backend/tests/test_windows_gpu_detection_mock.py
index bc887f8cc5..88a1a28d14 100644
--- a/studio/backend/tests/test_windows_gpu_detection_mock.py
+++ b/studio/backend/tests/test_windows_gpu_detection_mock.py
@@ -3,12 +3,12 @@
 
 """Windows GPU-detection regression test on a synthetic layout.
 
-The bug (#5106): on Windows without a system CUDA toolkit, the prebuilt
-llama-server.exe could not LoadLibrary cudart64_X / cublas64_X /
+Bug (#5106): on Windows without a system CUDA toolkit, the prebuilt
+llama-server.exe couldn't LoadLibrary cudart64_X / cublas64_X /
 cublasLt64_X, so ggml-cuda.dll's static import on cublas64_X.dll failed
 and the model fell back to CPU even when nvidia-smi reported the GPU.
 
-The fix:
+Fix:
   * #5322 overlays upstream's paired cudart bundle into
     install_dir/build/bin/Release/ next to llama-server.exe.
   * #5324 prepends pip-installed nvidia//{bin,bin/x86_64,Library/
@@ -34,13 +34,9 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
 if _BACKEND_DIR not in sys.path:
     sys.path.insert(0, _BACKEND_DIR)
 
-# Stub heavy deps only if they actually fail to import -- unconditional
-# stubs would shadow the real module for sibling tests in this dir.
-# Use try-import rather than find_spec: loggers/__init__.py re-exports
-# handlers.get_logger, which does `from fastapi import Request,
-# Response` at module load. find_spec("loggers") returns a spec even
-# without fastapi, but the import then raises. CI has fastapi, so this
-# is dev-machine ergonomics only.
+# Stub heavy deps only if they fail to import (unconditional stubs would shadow
+# the real module for sibling tests). Use try-import, not find_spec: loggers
+# imports fastapi at load, so find_spec succeeds but the import then raises.
 import importlib as _importlib  # noqa: E402
 
 
@@ -100,7 +96,7 @@ from core.inference.llama_cpp import LlamaCppBackend  # noqa: E402
 
 
 # Upstream b9103 cudart bundle: exactly these three DLLs per CUDA major,
-# no executables, no subdirectories. Verified by direct unzip.
+# no executables or subdirectories. Verified by direct unzip.
 REAL_UPSTREAM_CUDART_BUNDLE = {
     "12.4": ("cudart64_12.dll", "cublas64_12.dll", "cublasLt64_12.dll"),
     "13.1": ("cudart64_13.dll", "cublas64_13.dll", "cublasLt64_13.dll"),
@@ -132,24 +128,21 @@ REAL_PIP_NVIDIA_WHEEL_LAYOUTS = {
 
 
 def _populate_studio_venv(prefix: Path) -> None:
-    """Lay out fake nvidia + torch wheels in /Lib/site-packages
-    matching the real win_amd64 wheel layouts. Contents are stub bytes;
-    only directory structure matters."""
+    """Lay out fake nvidia + torch wheels matching real win_amd64 layouts (stub bytes)."""
     site = prefix / "Lib" / "site-packages"
     for rel, dlls in REAL_PIP_NVIDIA_WHEEL_LAYOUTS.items():
         d = site / Path(rel)
         d.mkdir(parents = True, exist_ok = True)
         for name in dlls:
             (d / name).write_bytes(b"PE-stub")
-    # install_python_stack always installs torch alongside nvidia.
+    # install_python_stack always installs torch beside nvidia.
     (site / "torch" / "lib").mkdir(parents = True, exist_ok = True)
     for fn in ("c10.dll", "torch.dll", "torch_cpu.dll", "torch_python.dll"):
         (site / "torch" / "lib" / fn).write_bytes(b"PE-stub")
 
 
 def _populate_studio_install(install_dir: Path, runtime: str = "13.1") -> None:
-    """Lay out install_dir/build/bin/Release/ as #5322 leaves it: main
-    archive payload + paired cudart bundle overlay."""
+    """Lay out install_dir/build/bin/Release/ as #5322 leaves it: payload + cudart overlay."""
     rel = install_dir / "build" / "bin" / "Release"
     rel.mkdir(parents = True, exist_ok = True)
     for fn in (
@@ -163,7 +156,7 @@ def _populate_studio_install(install_dir: Path, runtime: str = "13.1") -> None:
         "mtmd.dll",
     ):
         (rel / fn).write_bytes(b"PE-stub")
-    # The cudart overlay #5322 contributes.
+    # The cudart overlay from #5322.
     for fn in REAL_UPSTREAM_CUDART_BUNDLE[runtime]:
         (rel / fn).write_bytes(b"PE-stub")
 
@@ -173,15 +166,13 @@ def _build_path_dirs_like_start_llama_server(
     prefix: Path,
     cuda_path: str = "",
 ) -> list[str]:
-    """Path-friendly wrapper around LlamaCppBackend._build_windows_path_dirs.
-    Asserting against the staticmethod (not a hand-copy) is the point:
-    if the win32 PATH order drops _windows_pip_nvidia_dll_dirs, tests fail."""
+    """Wrapper around the real _build_windows_path_dirs staticmethod."""
     return LlamaCppBackend._build_windows_path_dirs(str(binary_dir), str(prefix), cuda_path)
 
 
 def _mock_nvidia_smi_run(fake_output: str, returncode: int = 0) -> "mock._patch":
     """Patch subprocess.run so the nvidia-smi probe returns fake_output;
-    other subprocess.run calls pass through."""
+    other calls pass through."""
     real_run = subprocess.run
 
     def fake_run(cmd, *args, **kwargs):
@@ -199,11 +190,11 @@ def _mock_nvidia_smi_run(fake_output: str, returncode: int = 0) -> "mock._patch"
 # --------------------------------------------------------------------- #
 class TestWindowsGpuDetectionAfter5106Fix:
     """End-to-end #5106 fix on a synthetic Windows layout. nvidia-smi
-    mocked; resolver, PATH builder and install layout exercised live."""
+    mocked; resolver, PATH builder, and install layout run live."""
 
     def test_nvidia_smi_probe_reports_synthetic_gpu(self, monkeypatch):
         """Probe parses CSV output and returns (index, free_mib)."""
-        # Clear inherited masks so the synthetic CSV is not filtered.
+        # Clear inherited masks so the synthetic CSV isn't filtered.
         monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False)
         monkeypatch.delenv("NVIDIA_VISIBLE_DEVICES", raising = False)
         # The #5106 reporter's exact reproducer: RTX 4090, 22805 MiB.
@@ -222,7 +213,7 @@ class TestWindowsGpuDetectionAfter5106Fix:
 
     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; missing any one breaks ggml-cuda.dll's PE import chain."""
+        Release; any missing one breaks ggml-cuda.dll's PE import chain."""
         install = tmp_path / "studio_install"
         _populate_studio_install(install, runtime = "13.1")
         rel = install / "build" / "bin" / "Release"
@@ -232,8 +223,8 @@ class TestWindowsGpuDetectionAfter5106Fix:
         assert (rel / "ggml-cuda.dll").exists()
 
     def test_resolver_finds_real_pypi_wheel_layouts(self, tmp_path):
-        """Resolver must pick up every real-world wheel layout:
-        nvidia//bin, nvidia//bin/x86_64, torch/lib."""
+        """Resolver must pick up every wheel layout: nvidia//bin,
+        nvidia//bin/x86_64, torch/lib."""
         prefix = tmp_path / "studio_venv"
         _populate_studio_venv(prefix)
         out = LlamaCppBackend._windows_pip_nvidia_dll_dirs(str(prefix))
@@ -249,8 +240,8 @@ class TestWindowsGpuDetectionAfter5106Fix:
 
     def test_path_assembly_makes_cudart_reachable_without_toolkit(self, tmp_path):
         """The #5106 scenario: GPU detected, pip nvidia wheels present,
-        no system CUDA toolkit. cudart must be reachable from PATH, and
-        from BOTH binary_dir (#5322) and a pip nvidia dir (#5324)."""
+        no system CUDA toolkit. cudart must be reachable from PATH via
+        BOTH binary_dir (#5322) and a pip nvidia dir (#5324)."""
         prefix = tmp_path / "studio_venv"
         install = tmp_path / "studio_install"
         _populate_studio_venv(prefix)
@@ -278,7 +269,7 @@ class TestWindowsGpuDetectionAfter5106Fix:
         ), f"#5324's pip nvidia dir not contributing cudart: {cudart_locations}"
 
     def test_cublas_and_cublasLt_also_reachable(self, tmp_path):
-        """ggml-cuda imports cublas64; cublas64 imports cublasLt64. All
+        """ggml-cuda imports cublas64, which imports cublasLt64. All
         three must resolve or LoadLibrary returns NULL."""
         prefix = tmp_path / "studio_venv"
         install = tmp_path / "studio_install"
@@ -293,7 +284,7 @@ class TestWindowsGpuDetectionAfter5106Fix:
             )
 
     def test_no_pip_nvidia_wheels_still_works_via_install_dir(self, tmp_path):
-        """No pip nvidia wheels (CPU-only torch / unsloth run standalone):
+        """No pip nvidia wheels (CPU-only torch / standalone unsloth):
         cudart still resolves via #5322's binary_dir drop."""
         prefix = tmp_path / "bare_venv"
         prefix.mkdir()
@@ -309,7 +300,7 @@ class TestWindowsGpuDetectionAfter5106Fix:
 
     def test_no_install_dir_still_works_via_pip_wheels(self, tmp_path):
         """Pre-#5322 install (binary_dir lacks cudart): #5324's pip
-        wheel directories on PATH still resolve cudart."""
+        wheel dirs on PATH still resolve cudart."""
         prefix = tmp_path / "studio_venv"
         _populate_studio_venv(prefix)
         install = tmp_path / "studio_install_pre5322"
@@ -339,9 +330,9 @@ class TestWindowsGpuDetectionAfter5106Fix:
         assert cublas_reachable, "cublas unreachable on cudart-less install"
 
     def test_pre_pr_scenario_would_have_failed(self, tmp_path):
-        """Negative control: pre-#5322 + pre-#5324 world leaves cudart
+        """Negative control: pre-#5322 + pre-#5324 leaves cudart
         unreachable -- the original failure mode. Confirms the test
-        actually catches a regression."""
+        catches a regression."""
         prefix = tmp_path / "studio_venv"
         _populate_studio_venv(prefix)
         install = tmp_path / "pre_pr_install"
@@ -349,7 +340,7 @@ class TestWindowsGpuDetectionAfter5106Fix:
         rel.mkdir(parents = True)
         for fn in ("llama-server.exe", "llama.dll", "ggml-cuda.dll"):
             (rel / fn).write_bytes(b"PE-stub")
-        # Pre-PR PATH: binary_dir only. No pip nvidia dirs, no toolkit.
+        # Pre-PR PATH: binary_dir only, no pip nvidia dirs, no toolkit.
         pre_pr_path_dirs = [str(rel)]
         cudart_reachable_pre = any(
             (Path(d) / "cudart64_12.dll").exists() or (Path(d) / "cudart64_13.dll").exists()
@@ -362,9 +353,9 @@ class TestWindowsGpuDetectionAfter5106Fix:
 
 
 class TestWindowsSysPlatformMocked:
-    """Confirm the win32 branch in start_llama_server is what we test
-    (not the linux fallback). Patches sys.platform and re-runs the
-    branch-selecting helper."""
+    """Confirm we test the win32 branch in start_llama_server, not the
+    linux fallback. Patches sys.platform and re-runs the branch-selecting
+    helper."""
 
     def test_sys_platform_win32_uses_pip_nvidia_resolver(self, monkeypatch, tmp_path):
         monkeypatch.setattr(sys, "platform", "win32")
diff --git a/studio/backend/utils/api_errors.py b/studio/backend/utils/api_errors.py
new file mode 100644
index 0000000000..b1c55b61b9
--- /dev/null
+++ b/studio/backend/utils/api_errors.py
@@ -0,0 +1,252 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Error-envelope helpers for the OpenAI/Anthropic-compatible ``/v1/*`` API surface.
+
+FastAPI's defaults emit ``{"detail": ...}`` bodies (status 422 for validation,
+``exc.status_code`` for ``HTTPException``). Real OpenAI/Anthropic clients expect
+provider-specific error envelopes instead, so this module re-wraps Unsloth's own
+client-error responses on the ``/v1/*`` surface:
+
+- OpenAI surface (``/v1/chat/completions``, ``/v1/completions``, ``/v1/models``,
+  ``/v1/responses``, ``/v1/embeddings``, ...)::
+
+      {"error": {"message": str, "type": str, "param": None|str, "code": None|str}}
+
+- Anthropic surface (any path starting with ``/v1/messages``)::
+
+      {"type": "error", "error": {"type": str, "message": str}}
+
+CRITICAL: the exception handlers installed by :func:`install_api_error_handlers`
+are global, but they ONLY transform responses for paths that start with ``/v1/``.
+For every other path (``/api/...``, frontend routes) they reproduce FastAPI's
+default behavior byte-for-byte, because the Studio frontend depends on the
+``{"detail": ...}`` shape for ``/api/*``.
+
+Public contract (other modules depend on these):
+
+- ``OPENAI_TYPE_BY_STATUS`` / ``ANTHROPIC_TYPE_BY_STATUS``: status -> type maps.
+- ``openai_error_body(message, *, status=400, err_type=None, code=None, param=None)``
+- ``anthropic_error_body(message, *, status=400, err_type=None)``
+- ``is_anthropic_path(path)``
+- ``error_body_for_path(path, message, *, status, err_type=None, code=None, param=None)``
+- ``install_api_error_handlers(app)``
+"""
+
+from fastapi.encoders import jsonable_encoder
+from fastapi.responses import JSONResponse, Response
+from fastapi.exceptions import RequestValidationError
+from fastapi.utils import is_body_allowed_for_status_code
+from starlette.exceptions import HTTPException as StarletteHTTPException
+
+
+# Status-code -> error ``type`` string for the OpenAI error envelope.
+OPENAI_TYPE_BY_STATUS = {
+    400: "invalid_request_error",
+    401: "authentication_error",
+    403: "permission_error",
+    404: "not_found_error",
+    409: "conflict_error",
+    413: "invalid_request_error",
+    422: "invalid_request_error",
+    429: "rate_limit_error",
+    500: "api_error",
+    502: "api_error",
+    503: "api_error",
+}
+
+# Status-code -> error ``type`` string for the Anthropic error envelope.
+ANTHROPIC_TYPE_BY_STATUS = {
+    400: "invalid_request_error",
+    401: "authentication_error",
+    403: "permission_error",
+    404: "not_found_error",
+    409: "conflict_error",
+    413: "request_too_large",
+    422: "invalid_request_error",
+    429: "rate_limit_error",
+    500: "api_error",
+    502: "api_error",
+    503: "api_error",
+    529: "overloaded_error",
+}
+
+
+def openai_error_body(
+    message,
+    *,
+    status = 400,
+    err_type = None,
+    code = None,
+    param = None,
+) -> dict:
+    """Build an OpenAI-style error envelope.
+
+    Returns ``{"error": {"message", "type", "param", "code"}}``. The ``param``
+    and ``code`` keys are always present (value may be ``None``). ``err_type``
+    defaults to :data:`OPENAI_TYPE_BY_STATUS` for ``status`` (``"api_error"``
+    fallback).
+    """
+    return {
+        "error": {
+            "message": str(message),
+            "type": err_type or OPENAI_TYPE_BY_STATUS.get(status, "api_error"),
+            "param": param,
+            "code": code,
+        }
+    }
+
+
+def anthropic_error_body(
+    message,
+    *,
+    status = 400,
+    err_type = None,
+) -> dict:
+    """Build an Anthropic-style error envelope.
+
+    Returns ``{"type": "error", "request_id": None, "error": {"type", "message"}}``.
+    ``request_id`` is a required (nullable) field on the spec's ErrorResponse;
+    Studio has no request-id system, so it is null. ``err_type`` defaults to
+    :data:`ANTHROPIC_TYPE_BY_STATUS` for ``status`` (``"api_error"`` fallback).
+    """
+    return {
+        "type": "error",
+        "request_id": None,
+        "error": {
+            "type": err_type or ANTHROPIC_TYPE_BY_STATUS.get(status, "api_error"),
+            "message": str(message),
+        },
+    }
+
+
+def is_anthropic_path(path: str) -> bool:
+    """True iff ``path`` belongs to the Anthropic surface (``/v1/messages*``)."""
+    return path.startswith("/v1/messages")
+
+
+def error_body_for_path(
+    path,
+    message,
+    *,
+    status,
+    err_type = None,
+    code = None,
+    param = None,
+) -> dict:
+    """Dispatch to the correct envelope builder based on ``path``.
+
+    Anthropic surface paths use :func:`anthropic_error_body` (``code``/``param``
+    are not part of that envelope and are ignored); all other ``/v1/*`` paths use
+    :func:`openai_error_body`.
+    """
+    if is_anthropic_path(path):
+        return anthropic_error_body(message, status = status, err_type = err_type)
+    return openai_error_body(message, status = status, err_type = err_type, code = code, param = param)
+
+
+def _summarize_validation_errors(errors) -> tuple:
+    """Derive a readable one-line message and (optional) body param from ``exc.errors()``.
+
+    Returns ``(summary, param)``. ``summary`` is a human-readable string like
+    ``"messages: Field required"``. ``param`` is the offending body field name when
+    one can be extracted (used as the OpenAI envelope ``param``), else ``None``.
+
+    Malformed-JSON bodies surface here as ``type == "json_invalid"`` and get a
+    dedicated message.
+    """
+    if not errors:
+        return "Invalid request", None
+
+    first = errors[0]
+    if first.get("type") == "json_invalid":
+        return "Invalid JSON in request body", None
+
+    loc = first.get("loc", ()) or ()
+    msg = first.get("msg", "Invalid request")
+
+    # Extract the body field name (the loc element after a leading "body").
+    param = None
+    loc_parts = [p for p in loc if p not in ("body",)]
+    if loc and loc[0] == "body" and loc_parts:
+        # First non-"body" element that is a field name (string).
+        for part in loc_parts:
+            if isinstance(part, str):
+                param = part
+                break
+
+    label = ".".join(str(p) for p in loc_parts) if loc_parts else ".".join(str(p) for p in loc)
+    summary = f"{label}: {msg}" if label else str(msg)
+    return summary, param
+
+
+def install_api_error_handlers(app) -> None:
+    """Register validation + HTTPException handlers that emit ``/v1/*`` envelopes.
+
+    Both handlers are global but only transform responses for paths starting with
+    ``/v1/``. Non-``/v1/`` paths reproduce FastAPI's default ``{"detail": ...}``
+    behavior exactly so the Studio frontend keeps working.
+    """
+
+    @app.exception_handler(RequestValidationError)
+    async def _handle_validation_error(request, exc):
+        path = request.url.path
+        if path.startswith("/v1/"):
+            summary, param = _summarize_validation_errors(exc.errors())
+            return JSONResponse(
+                status_code = 400,
+                content = error_body_for_path(path, summary, status = 400, param = param),
+            )
+        # Default FastAPI behavior for every other path.
+        return JSONResponse(
+            status_code = 422,
+            content = {"detail": jsonable_encoder(exc.errors())},
+        )
+
+    @app.exception_handler(StarletteHTTPException)
+    async def _handle_http_exception(request, exc):
+        path = request.url.path
+        headers = getattr(exc, "headers", None)
+        # Statuses like 204/304/1xx must not carry a body — mirror FastAPI's
+        # default http_exception_handler, which returns a bodiless Response.
+        if not is_body_allowed_for_status_code(exc.status_code):
+            return Response(status_code = exc.status_code, headers = headers)
+        if path.startswith("/v1/"):
+            detail = exc.detail
+            # Already a fully-formed envelope: pass through untouched.
+            if isinstance(detail, dict) and ("error" in detail or detail.get("type") == "error"):
+                return JSONResponse(
+                    status_code = exc.status_code,
+                    content = detail,
+                    headers = headers,
+                )
+            # A dict carrying our individual fields.
+            if isinstance(detail, dict):
+                message = detail.get("message", detail)
+                err_type = detail.get("type")
+                code = detail.get("code")
+                param = detail.get("param")
+            else:
+                # Plain message string (the common HTTPException case).
+                message = detail
+                err_type = None
+                code = None
+                param = None
+            return JSONResponse(
+                status_code = exc.status_code,
+                content = error_body_for_path(
+                    path,
+                    message,
+                    status = exc.status_code,
+                    err_type = err_type,
+                    code = code,
+                    param = param,
+                ),
+                headers = headers,
+            )
+        # Default FastAPI behavior for every other path.
+        return JSONResponse(
+            status_code = exc.status_code,
+            content = {"detail": exc.detail},
+            headers = headers,
+        )
diff --git a/studio/backend/utils/cache_cleanup.py b/studio/backend/utils/cache_cleanup.py
index 9d01b40add..210735973d 100644
--- a/studio/backend/utils/cache_cleanup.py
+++ b/studio/backend/utils/cache_cleanup.py
@@ -1,14 +1,12 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-"""
-Utility for cleaning up the Unsloth compiled cache directory.
+"""Clean up the Unsloth compiled cache directory.
 
-The unsloth_compiled_cache is created by unsloth_zoo/compiler.py during
-FastModel.from_pretrained() and contains model-type-specific compiled Python
-files. It should be selectively cleared between model loads to avoid stale
-artefacts, while preserving model-agnostic components (like Trainers) needed
-by spawned subprocesses.
+unsloth_compiled_cache (created by unsloth_zoo/compiler.py during
+FastModel.from_pretrained) holds model-type-specific compiled files. Clear it
+selectively between model loads, preserving model-agnostic components (Trainers)
+that spawned subprocesses need.
 """
 
 import shutil
@@ -38,9 +36,8 @@ def get_existing_cache_dirs() -> List[Path]:
 def register_compiled_cache_on_path() -> None:
     """Add all existing compiled-cache directories to sys.path and PYTHONPATH.
 
-    This ensures spawned workers (on platforms using the 'spawn' start method,
-    i.e. Windows and macOS) can import dynamically compiled modules such as
-    UnslothSFTTrainer.
+    Ensures spawned workers (on 'spawn'-start platforms, i.e. Windows and macOS)
+    can import dynamically compiled modules such as UnslothSFTTrainer.
     """
     import os
     import sys
@@ -48,8 +45,8 @@ def register_compiled_cache_on_path() -> None:
     pypath = os.environ.get("PYTHONPATH", "")
     pypath_entries = [p for p in pypath.split(os.pathsep) if p]
 
-    # Iterate in reverse so that earlier _CACHE_DIRS entries (higher priority)
-    # are inserted last and therefore end up first in sys.path / PYTHONPATH.
+    # Iterate in reverse so earlier _CACHE_DIRS entries (higher priority) are
+    # inserted last and thus end up first in sys.path / PYTHONPATH.
     for cache_dir in reversed(get_existing_cache_dirs()):
         resolved = str(cache_dir.resolve())
         if resolved not in sys.path:
@@ -65,7 +62,7 @@ def clear_unsloth_compiled_cache(preserve_patterns: Optional[List[str]] = None)
     Remove compiled files from the cache directory (idempotent).
 
     Args:
-        preserve_patterns: A list of glob patterns for files to keep
+        preserve_patterns: glob patterns for files to keep
                            (e.g., ["Unsloth*Trainer.py"]). If None or empty,
                            the entire cache directory is deleted (legacy behavior).
     """
@@ -80,7 +77,6 @@ def clear_unsloth_compiled_cache(preserve_patterns: Optional[List[str]] = None)
 
             for item in cache_dir.iterdir():
                 if item.is_file():
-                    # Check if the file matches any of the patterns we want to keep
                     preserve = any(item.match(pattern) for pattern in preserve_patterns)
                     if not preserve:
                         try:
@@ -92,6 +88,6 @@ def clear_unsloth_compiled_cache(preserve_patterns: Optional[List[str]] = None)
                     # Always clear __pycache__ and other subdirectories
                     shutil.rmtree(item, ignore_errors = True)
         else:
-            # Legacy behavior: nuke the entire directory
+            # Legacy: remove the entire directory
             logger.info(f"Removing unsloth compiled cache: {cache_dir}")
             shutil.rmtree(cache_dir, ignore_errors = True)
diff --git a/studio/backend/utils/cpu_threads.py b/studio/backend/utils/cpu_threads.py
index 922f95ce55..4ed0021054 100644
--- a/studio/backend/utils/cpu_threads.py
+++ b/studio/backend/utils/cpu_threads.py
@@ -18,9 +18,9 @@ _THREAD_POOL_ENV_VARS = (
 def configure_cpu_threads(env: Optional[MutableMapping[str, str]] = None) -> None:
     """Apply ``UNSLOTH_CPU_THREADS`` to native CPU pools when configured.
 
-    This must run before importing libraries that initialize an OpenMP or
-    BLAS thread pool. Library-specific variables are left untouched so users
-    can override a single runtime independently.
+    Must run before importing libraries that initialize an OpenMP or BLAS
+    pool. Library-specific vars are left untouched so users can override a
+    single runtime independently.
     """
     environ = os.environ if env is None else env
     configured = environ.get("UNSLOTH_CPU_THREADS", "").strip()
diff --git a/studio/backend/utils/datasets/__init__.py b/studio/backend/utils/datasets/__init__.py
index 7988b09972..caa471bde5 100644
--- a/studio/backend/utils/datasets/__init__.py
+++ b/studio/backend/utils/datasets/__init__.py
@@ -1,22 +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
 
-"""
-Dataset utilities package.
+"""Dataset utilities for LLM/VLM fine-tuning: detection, conversion, templating, collators, mappings."""
 
-This package provides utilities for dataset format detection, conversion,
-and processing for LLM and VLM fine-tuning workflows.
-
-Modules:
-- format_detection: Detect dataset formats (Alpaca, ShareGPT, ChatML)
-- format_conversion: Convert between dataset formats
-- chat_templates: Apply chat templates to datasets
-- vlm_processing: Vision-Language Model processing utilities
-- data_collators: Custom data collators for training
-- model_mappings: Model-to-template mapping constants
-"""
-
-# Format detection
 from .format_detection import (
     detect_dataset_format,
     detect_custom_format_heuristic,
@@ -24,7 +10,6 @@ from .format_detection import (
     detect_vlm_dataset_structure,
 )
 
-# Format conversion
 from .format_conversion import (
     standardize_chat_format,
     convert_chatml_to_alpaca,
@@ -34,7 +19,6 @@ from .format_conversion import (
     convert_sharegpt_with_images_to_vlm_format,
 )
 
-# Chat templates
 from .chat_templates import (
     apply_chat_template_to_dataset,
     get_dataset_info_summary,
@@ -42,19 +26,16 @@ from .chat_templates import (
     DEFAULT_ALPACA_TEMPLATE,
 )
 
-# VLM processing
 from .vlm_processing import (
     generate_smart_vlm_instruction,
 )
 
-# Data collators
 from .data_collators import (
     DataCollatorSpeechSeq2SeqWithPadding,
     DeepSeekOCRDataCollator,
     VLMDataCollator,
 )
 
-# Model mappings (constants)
 from .model_mappings import (
     TEMPLATE_TO_MODEL_MAPPER,
     MODEL_TO_TEMPLATE_MAPPER,
@@ -62,15 +43,13 @@ from .model_mappings import (
     is_gpt_oss_model_name,
 )
 
-# Legacy imports from the original dataset_utils.py for backward compatibility
-# These functions have not yet been refactored into separate modules
+# Legacy dataset_utils.py imports kept for backward compat
 from .dataset_utils import (
     check_dataset_format,
     format_and_template_dataset,
     format_dataset,
 )
 
-# Public API
 __all__ = [
     # Detection
     "detect_dataset_format",
diff --git a/studio/backend/utils/datasets/chat_templates.py b/studio/backend/utils/datasets/chat_templates.py
index cfdd811853..82a30fd55b 100644
--- a/studio/backend/utils/datasets/chat_templates.py
+++ b/studio/backend/utils/datasets/chat_templates.py
@@ -1,11 +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
 
-"""
-Chat template application utilities for dataset processing.
+"""Chat template utilities for dataset processing.
 
-This module contains functions for applying chat templates to datasets
-and generating dataset info summaries.
+Apply chat templates to datasets and generate dataset info summaries.
 """
 
 from .format_detection import detect_dataset_format, detect_multimodal_dataset, detect_custom_format_heuristic
@@ -46,30 +44,25 @@ def _chat_template_kwargs() -> dict:
 
 
 def get_tokenizer_chat_template(tokenizer, model_name):
-    """
-    Gets appropriate chat template for tokenizer based on model.
-    Uses Unsloth's get_chat_template if model is in the mapper.
+    """Apply a chat template to the tokenizer, using Unsloth's
+    get_chat_template when the model class name is in the mapper.
 
     Args:
         tokenizer: HuggingFace tokenizer
         model_name: Model class name (e.g., "Gemma3ForCausalLM")
 
     Returns:
-        tokenizer: Tokenizer with appropriate chat template applied
+        tokenizer with the chat template applied
     """
     try:
         from unsloth.chat_templates import get_chat_template
     except ImportError:
-        # Unsloth not available, return tokenizer as-is
         return tokenizer
 
-    # Normalize model_name to lowercase for matching
     model_name_lower = model_name.lower()
 
-    # Check if model matches any template in mapper
     matched_template = None
 
-    # Direct match in MODEL_TO_TEMPLATE_MAPPER
     if model_name_lower in MODEL_TO_TEMPLATE_MAPPER:
         matched_template = MODEL_TO_TEMPLATE_MAPPER[model_name_lower]
         logger.info(f"📝 Applying Unsloth chat template: {matched_template}")
@@ -83,7 +76,6 @@ def get_tokenizer_chat_template(tokenizer, model_name):
             logger.info(f"⚠️ Failed to apply Unsloth template '{matched_template}': {e}")
             logger.info(f"   Falling back to tokenizer's default chat template")
     else:
-        # Check if tokenizer actually has a chat_template set
         has_chat_template = (
             hasattr(tokenizer, 'chat_template')
             and tokenizer.chat_template is not None
@@ -91,7 +83,7 @@ def get_tokenizer_chat_template(tokenizer, model_name):
         if has_chat_template:
             logger.info(f"📝 Using tokenizer's own chat template (no Unsloth template match)")
         else:
-            # Base model with no chat template — apply default ChatML
+            # Base model with no chat template: apply default ChatML.
             logger.info(f"📝 No chat template found — applying default ChatML template (base model)")
             try:
                 tokenizer = get_chat_template(
@@ -107,9 +99,7 @@ def get_tokenizer_chat_template(tokenizer, model_name):
 
 
 def get_dataset_info_summary(dataset_info):
-    """
-    Returns a human-readable summary for UI display.
-    """
+    """Return a human-readable summary for UI display."""
     detected_format = dataset_info["detected_format"]
     final_format = dataset_info["final_format"]
 
@@ -146,15 +136,14 @@ def apply_chat_template_to_dataset(
     num_proc = None,
     progress_callback = None,
 ):
-    """
-    Applies chat template to dataset based on its format.
+    """Apply the chat template to a dataset based on its format.
 
     Args:
         dataset_info: Output from format_dataset() with metadata
         tokenizer: Tokenizer with chat template
         custom_prompt_template: Optional string template for custom formatting
-        add_eos_token: If True, appends tokenizer.eos_token to each text
-        remove_bos_prefix: If True, removes '' prefix (for Gemma, etc.)
+        add_eos_token: If True, append tokenizer.eos_token to each text
+        remove_bos_prefix: If True, remove '' prefix (Gemma, etc.)
         custom_format_mapping: Dict mapping custom columns to standard format
         batch_size: Batch size for processing
         num_proc: Number of processes
@@ -170,7 +159,6 @@ def apply_chat_template_to_dataset(
     warnings = list(dataset_info.get("warnings", []))
     errors = []
 
-    # Get EOS token if needed
     eos_token = ""
     if add_eos_token:
         if hasattr(tokenizer, 'eos_token') and tokenizer.eos_token:
@@ -180,9 +168,8 @@ def apply_chat_template_to_dataset(
 
     # CUSTOM FORMAT MAPPING (for non-standard datasets)
     if final_format == "unknown":
-        # Try auto-detection if no custom mapping provided
         if custom_format_mapping is None and auto_detect_mapping:
-            # Check if format_dataset already tried and failed
+            # Skip if format_dataset already tried and failed.
             if not dataset_info.get("auto_detection_attempted", False):
                 custom_format_mapping = detect_custom_format_heuristic(dataset)
                 if custom_format_mapping:
@@ -196,7 +183,7 @@ def apply_chat_template_to_dataset(
                         "errors": errors
                     }
             else:
-                # Already failed once in format_dataset, don't retry
+                # Already failed once in format_dataset; don't retry.
                 errors.append(
                     "Format remains unknown after detection attempts. "
                     "Please provide custom_format_mapping to specify column roles manually."
@@ -216,7 +203,7 @@ def apply_chat_template_to_dataset(
                 conversations = []
                 num_examples = len(examples[list(examples.keys())[0]])
 
-                # Only preserve unmapped columns if auto-detected
+                # Preserve unmapped columns only if auto-detected.
                 preserved_columns = {}
                 if not is_user_provided:
                     all_columns = set(examples.keys())
@@ -236,10 +223,10 @@ def apply_chat_template_to_dataset(
                                 content = examples[col_name][i]
 
                                 if is_user_provided:
-                                    # User explicitly mapped - include even if empty
+                                    # User-mapped: include even if empty.
                                     convo.append({"role": role, "content": str(content) if content else ""})
                                 else:
-                                    # Auto-detected - skip empty
+                                    # Auto-detected: skip empty.
                                     if content and str(content).strip():
                                         convo.append({"role": role, "content": str(content)})
 
@@ -252,7 +239,6 @@ def apply_chat_template_to_dataset(
 
             try:
                 dataset = dataset.map(_apply_custom_mapping, batched = True, batch_size = batch_size)
-                # Update to use conversations format
                 final_format = "chatml_conversations"
                 chat_column = "conversations"
                 is_standardized = True
@@ -269,8 +255,7 @@ def apply_chat_template_to_dataset(
     # ALPACA FORMAT
     if final_format == "alpaca":
 
-        # Set alpaca chat template on tokenizer for saving (if not already set)
-        # This ensures the template is saved with the model for inference
+        # Set alpaca chat template (if unset) so it's saved for inference.
         if not (hasattr(tokenizer, 'chat_template') and tokenizer.chat_template):
             try:
                 from unsloth.chat_templates import get_chat_template
@@ -283,7 +268,6 @@ def apply_chat_template_to_dataset(
             except Exception as e:
                 logger.info(f"⚠️ Could not set alpaca template on tokenizer: {e}")
 
-        # Use custom template if provided
         def _format_alpaca_custom(examples):
             texts = []
             for i in range(len(examples["instruction"])):
@@ -349,7 +333,7 @@ def apply_chat_template_to_dataset(
         if not is_standardized:
             warnings.append("Dataset may not be fully standardized")
 
-        # Apply Unsloth chat template if model matches
+        # Apply Unsloth chat template if the model matches.
         if model_name:
             tokenizer = get_tokenizer_chat_template(tokenizer, model_name)
 
@@ -398,7 +382,7 @@ def apply_chat_template_to_dataset(
                 dataset_map_kwargs['num_proc'] = num_proc
                 dataset_map_kwargs['desc'] = f"Applying chat template to {final_format}"
 
-            # Monitor tqdm progress from dataset.map() and relay to callback
+            # Monitor dataset.map() tqdm progress and relay it.
             _tqdm_monitor_stop = None
             if progress_callback and not _is_torch_iterable:
                 import threading
diff --git a/studio/backend/utils/datasets/data_collators.py b/studio/backend/utils/datasets/data_collators.py
index 2955ec9023..9bfb60ba17 100644
--- a/studio/backend/utils/datasets/data_collators.py
+++ b/studio/backend/utils/datasets/data_collators.py
@@ -1,12 +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
 
-"""
-Data collators for dataset processing.
-
-This module contains custom data collators for training,
-particularly for VLM/OCR processing.
-"""
+"""Custom training data collators, particularly for VLM/OCR processing."""
 
 from dataclasses import dataclass
 from typing import Any, List, Optional, Union
@@ -20,9 +15,9 @@ class DataCollatorSpeechSeq2SeqWithPadding:
     """
     Data collator for Whisper speech-to-text training.
 
-    Pads input features (audio) and label sequences (text) separately,
-    masks padding in labels with -100, and strips leading BOS token.
-    Mirrors the collator from the Whisper.ipynb notebook.
+    Pads audio input features and text labels separately, masks label padding
+    with -100, and strips the leading BOS token. Mirrors the Whisper.ipynb
+    notebook collator.
     """
 
     processor: Any
@@ -45,13 +40,10 @@ class DataCollatorSpeechSeq2SeqWithPadding:
 
 @dataclass
 class DeepSeekOCRDataCollator:
-    """
-    Data collator for DeepSeek OCR VLM training.
+    """Data collator for DeepSeek OCR VLM training.
 
-    Handles:
-    - Image processing via processor
-    - Text tokenization
-    - Proper label masking for instruction fine-tuning
+    Handles image processing, text tokenization, and label masking for
+    instruction fine-tuning.
     """
 
     processor: Any  # Qwen2VLProcessor or similar
@@ -71,7 +63,6 @@ class DeepSeekOCRDataCollator:
         """
         from PIL import Image
 
-        # Extract messages and images
         all_messages = []
         all_images = []
 
@@ -79,7 +70,6 @@ class DeepSeekOCRDataCollator:
             messages = sample["messages"]
             all_messages.append(messages)
 
-            # Extract PIL images from content
             for msg in messages:
                 content = msg.get("content", [])
                 if isinstance(content, list):
@@ -89,9 +79,7 @@ class DeepSeekOCRDataCollator:
                             if img is not None and hasattr(img, "size"):  # PIL Image
                                 all_images.append(img)
 
-        # Process with the VL processor
         try:
-            # Qwen2VL style processing
             texts = [
                 self.processor.apply_chat_template(
                     msgs, tokenize = False, add_generation_prompt = False
@@ -99,7 +87,6 @@ class DeepSeekOCRDataCollator:
                 for msgs in all_messages
             ]
 
-            # Process with images
             inputs = self.processor(
                 text = texts,
                 images = all_images if all_images else None,
@@ -109,10 +96,7 @@ class DeepSeekOCRDataCollator:
                 max_length = self.max_length,
             )
 
-            # Create labels (mask input, keep output)
             labels = inputs["input_ids"].clone()
-
-            # Simple masking: mask padding tokens
             labels[labels == self.processor.tokenizer.pad_token_id] = self.ignore_index
 
             inputs["labels"] = labels
@@ -126,24 +110,15 @@ class DeepSeekOCRDataCollator:
 
 @dataclass
 class VLMDataCollator:
-    """
-    Generic VLM data collator that works with various processors.
-
-    Supports:
-    - Qwen2VL
-    - LLaVA
-    - Other VL models with compatible processors
-    """
+    """Generic VLM data collator for various processors (Qwen2VL, LLaVA, etc.)."""
 
     processor: Any
     max_length: int = 2048
     ignore_index: int = -100
-    mask_input_tokens: bool = True  # Whether to mask user tokens in labels
+    mask_input_tokens: bool = True  # Mask user tokens in labels
 
     def __call__(self, batch: List[dict]) -> dict:
-        """
-        Collate a batch of VLM samples.
-        """
+        """Collate a batch of VLM samples."""
         all_messages = []
         all_images = []
 
@@ -151,7 +126,6 @@ class VLMDataCollator:
             messages = sample.get("messages", [])
             all_messages.append(messages)
 
-            # Extract images
             for msg in messages:
                 content = msg.get("content", [])
                 if isinstance(content, list):
@@ -161,13 +135,11 @@ class VLMDataCollator:
                             if img is not None:
                                 all_images.append(img)
 
-        # Apply chat template
         texts = [
             self.processor.apply_chat_template(msgs, tokenize = False, add_generation_prompt = False)
             for msgs in all_messages
         ]
 
-        # Process inputs
         inputs = self.processor(
             text = texts,
             images = all_images if all_images else None,
@@ -177,10 +149,9 @@ class VLMDataCollator:
             max_length = self.max_length,
         )
 
-        # Create labels
         labels = inputs["input_ids"].clone()
 
-        # Mask padding
+        # Mask padding.
         if hasattr(self.processor, "tokenizer"):
             pad_token_id = self.processor.tokenizer.pad_token_id
         else:
diff --git a/studio/backend/utils/datasets/dataset_none_detect.py b/studio/backend/utils/datasets/dataset_none_detect.py
index e48e153fb8..a2fd8ef667 100644
--- a/studio/backend/utils/datasets/dataset_none_detect.py
+++ b/studio/backend/utils/datasets/dataset_none_detect.py
@@ -1,8 +1,6 @@
 """
-dataset_none_detect.py
-
-Detect None/empty content turns in conversation datasets.
-Reports findings without modifying data.
+Detect None/empty content turns in conversation datasets. Reports findings
+without modifying data.
 
 Usage:
     from .dataset_none_detect import scan_dataset, print_report
@@ -18,8 +16,8 @@ Supported formats (via FORMAT_REGISTRY):
     sharegpt   conversations                 from/value per turn
     gptoss     messages (alias: gpt-oss)     role/content; has a developer turn
 
-Any role/content chat template matches the chatml entry, so new templates need
-no change; add a FORMAT_REGISTRY entry only for a genuinely new column/turn shape.
+Any role/content chat template matches chatml, so new templates need no change;
+add a FORMAT_REGISTRY entry only for a genuinely new column/turn shape.
 """
 
 from datasets import Dataset
@@ -31,7 +29,7 @@ from datasets import Dataset
 # Candidate column names for conversational datasets, checked in priority order.
 CONVERSATION_COLUMNS = ("messages", "conversations", "texts")
 
-# Minimum turn key sets that identify a column as conversational (not e.g. messages=[{"id":1}]).
+# Minimum turn key sets identifying a column as conversational (not e.g. messages=[{"id":1}]).
 _CHAT_KEY_SETS = (frozenset({"role", "content"}), frozenset({"from", "value"}))
 
 
@@ -39,13 +37,13 @@ def _probe_conversation(dataset: Dataset, candidates = None):
     """
     Probe a dataset for its conversation column and turn structure.
 
-    candidates - iterable of column names to try, in priority order.
+    candidates - column names to try, in priority order.
                  Defaults to CONVERSATION_COLUMNS when None.
 
     Returns a dict with:
-        column    - name of the conversation column found
-        turn_keys - set of keys present in the first turn dict
-        roles     - set of all role values seen across the first few samples
+        column    - conversation column found
+        turn_keys - keys present in the first turn dict
+        roles     - all role values seen across the first few samples
 
     Returns None if no conversation column is found.
     """
@@ -53,12 +51,12 @@ def _probe_conversation(dataset: Dataset, candidates = None):
         candidates = CONVERSATION_COLUMNS
     columns = set(dataset.column_names)
     # Remember the first all-corrupt candidate, but keep probing: a later column
-    # may be healthy and should win (e.g. bad messages, good conversations).
+    # may be healthy and win (e.g. bad messages, good conversations).
     all_corrupt_fallback = None
     for col in candidates:
         if col not in columns:
             continue
-        # Scan up to 100 rows - row 0 alone may be empty or malformed.
+        # Scan up to 100 rows - row 0 alone may be empty/malformed.
         first = None
         for i in range(min(len(dataset), 100)):
             sample = dataset[i][col]
@@ -71,10 +69,8 @@ def _probe_conversation(dataset: Dataset, candidates = None):
                 break
         if first is None:
             # No usable dict turn in 100 rows. Record an all_corrupt fallback,
-            # marking it plausible only if we saw turn-shaped data (a None cell or
-            # a list holding a dict/None turn); scalars and list-of-strings must
-            # not look like chatml. Upgrade a non-plausible fallback when a later
-            # candidate is plausible, so probe order keeps the best match.
+            # plausible only with turn-shaped data (None cell or list of dict/None
+            # turns); a later plausible candidate upgrades a non-plausible one.
             if all_corrupt_fallback is None or not all_corrupt_fallback.get("has_plausible_turns"):
                 has_plausible_turns = False
                 for i in range(min(len(dataset), 100)):
@@ -86,8 +82,8 @@ def _probe_conversation(dataset: Dataset, candidates = None):
                     # not chat: leave it for "unknown format", matching
                     # format_detection.py.
                     if isinstance(cell, list):
-                        # Plausible only if the list holds a dict/None turn; empty
-                        # lists and list-of-strings are not chat data.
+                        # Plausible only if the list holds a dict/None turn;
+                        # empty lists and list-of-strings are not chat data.
                         if any(t is None or isinstance(t, dict) for t in cell):
                             has_plausible_turns = True
                             break
@@ -112,11 +108,11 @@ def _probe_conversation(dataset: Dataset, candidates = None):
                         r = t.get("role") or t.get("from")
                         if r:
                             roles.add(str(r))
-        # Column lacks a full chat key pair. If it still has a conversational key
-        # (role/from/content/value) it is a corrupt-but-real chat column, so save
-        # a plausible fallback for find_none_chatml to flag. Pure metadata (e.g.
-        # [{"id":1}]) is not plausible, so a later real-but-corrupt column (e.g.
-        # conversations=None) can still win.
+        # Column lacks a full chat key pair. If it has a conversational key
+        # (role/from/content/value) it is a corrupt-but-real chat column, so
+        # save a plausible fallback for find_none_chatml to flag. Pure metadata
+        # (e.g. [{"id":1}]) is not plausible, so a later real-but-corrupt column
+        # (e.g. conversations=None) can still win.
         _CONV_KEYS = {"role", "from", "content", "value"}
         if not any(keys <= turn_keys for keys in _CHAT_KEY_SETS):
             schema_less_plausible = bool(turn_keys & _CONV_KEYS)
@@ -143,7 +139,7 @@ def is_none_or_empty(value) -> bool:
     if value is None:
         return True
     if isinstance(value, str):
-        # Treat zero-width/BOM chars (U+FEFF/200B/200C/200D/2060) as empty too;
+        # Treat zero-width/BOM chars (U+FEFF/200B/200C/200D/2060) as empty;
         # they render invisibly. Two-pass strip (ws, invisibles, ws) catches
         # mixed cases like "\u200b \u200b".
         stripped = value.strip().strip("\ufeff\u200b\u200c\u200d\u2060").strip()
@@ -156,7 +152,7 @@ def is_none_or_empty(value) -> bool:
         # exists (an image-only turn is valid).
         if len(value) == 0:
             return True
-        # No dict blocks at all (e.g. [None], ['  ']) -> malformed/empty.
+        # No dict blocks (e.g. [None], ['  ']) -> malformed/empty.
         dict_blocks = [item for item in value if isinstance(item, dict)]
         if not dict_blocks:
             return True
@@ -190,7 +186,7 @@ def _classify_empty(value) -> str:
         if len(value) == 0:
             return "empty_list"
         return "empty_vlm_content"
-    return "valid"  # should not reach here if is_none_or_empty was True
+    return "valid"  # unreachable if is_none_or_empty was True
 
 
 # ---------------------------------------------------------------------------
@@ -201,7 +197,7 @@ def _classify_empty(value) -> str:
 def find_none_alpaca(dataset: Dataset) -> dict:
     """
     Scan alpaca dataset for None/empty instruction or output fields.
-    Returns stats dict with a detailed 'findings' list.
+    Returns a stats dict with a detailed 'findings' list.
     """
     stats = {
         "total_rows": len(dataset),
@@ -242,8 +238,8 @@ def find_none_chatml(dataset: Dataset, col: str = None) -> dict:
     Scan chatml/sharegpt/gptoss dataset for turns with None/empty content.
     Auto-detects the conversation column if col=None.
 
-    Returns a stats dict that includes a complete 'findings' list - one entry
-    per bad turn with row_index, turn_index, role, value_type, and raw_value.
+    Returns a stats dict with a complete 'findings' list - one entry per bad
+    turn with row_index, turn_index, role, value_type, and raw_value.
     """
     if col is None:
         # Reuse _probe_conversation so the all_corrupt path is handled here too.
@@ -292,7 +288,7 @@ def find_none_chatml(dataset: Dataset, col: str = None) -> dict:
             continue
 
         if len(conversation) == 0:
-            # Zero-turn conversation: flag so it does not scan as clean.
+            # Zero-turn conversation: flag so it doesn't scan as clean.
             stats["bad_row_indices"].append(i)
             stats["rows_with_none_turns"] += 1
             stats["total_none_turns"] += 1
@@ -340,8 +336,8 @@ def find_none_chatml(dataset: Dataset, col: str = None) -> dict:
                 role = r
             else:
                 role = str(r)
-            # Pick the content key: from+value -> value (ShareGPT, even if role is
-            # also set); role -> content (or value); from only -> value (None when
+            # Pick the content key: from+value -> value (ShareGPT, even if role
+            # is set); role -> content (or value); from only -> value (None when
             # missing, so it is flagged); neither -> content then value.
             if "from" in turn and "value" in turn:
                 content = turn.get("value")
@@ -388,7 +384,7 @@ def find_none_sharegpt(dataset: Dataset, col: str = None) -> dict:
     """ShareGPT uses 'from'/'value' keys - same scan logic handles both."""
     if col is None:
         # ShareGPT lives in 'conversations'; probe only that column so a corrupt
-        # one is still scanned, not replaced by a healthy 'messages' (P1 fix).
+        # one is still scanned, not replaced by healthy 'messages' (P1 fix).
         conv_info = _probe_conversation(dataset, candidates = ("conversations",))
         if conv_info is None:
             raise ValueError(
@@ -403,7 +399,7 @@ def find_none_gptoss(dataset: Dataset, col: str = None) -> dict:
     """gptoss: role/content plus optional thinking/tool_calls. Only content checked."""
     if col is None:
         # gptoss lives in 'messages': target it whenever present (even if
-        # corrupt), and fall back to 'conversations' only if 'messages' is absent.
+        # corrupt); fall back to 'conversations' only if 'messages' is absent.
         if "messages" in dataset.column_names:
             conv_info = _probe_conversation(dataset, candidates = ("messages",))
         else:
@@ -420,7 +416,7 @@ def find_none_gptoss(dataset: Dataset, col: str = None) -> dict:
 # ---------------------------------------------------------------------------
 # Format registry - first match wins; detect_format() auto-scales.
 # Each entry: name (label/--format value), match(dataset, conv_info) -> bool,
-# scan (find_none_* function). Put specific formats before generalisations
+# scan (find_none_* function). Put specific formats before general ones
 # (gptoss before chatml, since gptoss is chatml with a 'developer' role).
 # To add a format: write find_none_() (or reuse find_none_chatml) and
 # append an entry; detect_format(), --format, and scan_dataset() pick it up.
@@ -460,7 +456,7 @@ FORMAT_REGISTRY = [
             and (
                 {"role", "content"} <= conv["turn_keys"]
                 # all_corrupt: column found but every row malformed; require
-                # has_plausible_turns so scalar/string columns are not chatml.
+                # has_plausible_turns so scalar/string columns aren't chatml.
                 or (conv.get("all_corrupt") and conv.get("has_plausible_turns"))
             )
         ),
@@ -479,7 +475,7 @@ def detect_format(dataset: Dataset) -> str:
     """
     Auto-detect dataset format by probing columns and turn structure.
 
-    Returns one of the format names in FORMAT_REGISTRY, or 'unknown'.
+    Returns a format name from FORMAT_REGISTRY, or 'unknown'.
     Walks the registry in order; first match wins.
     """
     conv_info = _probe_conversation(dataset)
@@ -507,7 +503,7 @@ def scan_dataset(dataset: Dataset, fmt: str = "auto") -> dict:
     # Reject a DatasetDict / IterableDatasetDict (load_dataset without split):
     # its column_names is a split map and would yield a confusing "unknown
     # format". Check both (IterableDatasetDict is not a DatasetDict subclass);
-    # import locally so this module never hard-requires those symbols.
+    # import locally so this module never hard-requires them.
     _dict_types = []
     try:
         from datasets import DatasetDict as _DatasetDict
@@ -526,7 +522,7 @@ def scan_dataset(dataset: Dataset, fmt: str = "auto") -> dict:
             "Pass dataset[] or use load_dataset(..., split='train')."
         )
     # Streaming IterableDataset has no len()/column_names; give a clear error
-    # instead of a confusing TypeError downstream.
+    # instead of a confusing downstream TypeError.
     try:
         from datasets import IterableDataset as _IterableDataset
         if isinstance(dataset, _IterableDataset):
@@ -555,8 +551,8 @@ def scan_dataset(dataset: Dataset, fmt: str = "auto") -> dict:
             if entry["match"](dataset, conv_info):
                 fmt = entry["name"]
                 break
-        # No format matched: return clean stats (format="unknown") rather than
-        # raise, so callers can branch on stats["format"].
+        # No format matched: return clean stats (format="unknown") instead of
+        # raising, so callers can branch on stats["format"].
         if fmt == "unknown":
             return {
                 "format": "unknown",
@@ -567,7 +563,7 @@ def scan_dataset(dataset: Dataset, fmt: str = "auto") -> dict:
     scanner = get_scanner(fmt)
     if scanner is None:
         raise ValueError(f"Unknown or unsupported format: '{fmt}'")
-    # Column forwarding: on auto-detect pass the probed column (already the best
+    # Column forwarding: on auto-detect pass the probed column (the best
     # choice). On an explicit format let that scanner pick its own column, so
     # e.g. fmt='sharegpt' always scans 'conversations', not 'messages' (P1 fix);
     # gptoss has its own messages-first rule. alpaca never takes a column.
@@ -620,7 +616,7 @@ def _print_summary_header(stats: dict, fmt: str) -> bool:
         if rows_all:
             print(f"  Rows ALL bad: {rows_all} (every turn is None/empty)")
 
-    # Rows with no Nones - compute the count directly instead of allocating a
+    # Rows with no Nones - compute the count directly rather than allocating a
     # full set of row indices, which OOMs on large (10M+ row) datasets.
     bad_indices = set(stats.get("bad_row_indices", []))
     clean_count = total - len(bad_indices)
@@ -696,8 +692,8 @@ def show_row(
         print(f"  Row {ri}")
         print(f"{'=' * 64}")
 
-        # Print non-conversation columns. For alpaca, skip the fields the
-        # alpaca block below prints with status markers (avoid double render).
+        # Print non-conversation columns. For alpaca, skip fields the alpaca
+        # block below prints with status markers (avoid double render).
         _ALPACA_FIELDS = {"instruction", "input", "output"}
         for key in dataset.column_names:
             if key == col:
@@ -732,7 +728,7 @@ def show_row(
                         c = t.get("value")
                     else:
                         c = t.get("content") if "content" in t else t.get("value")
-                    # Mirror scanner logic: tool_calls exemption is assistant-only;
+                    # Mirror scanner: tool_calls exemption is assistant-only;
                     # other roles with empty content + tool_calls are still bad.
                     r = t.get("role") if t.get("role") is not None else t.get("from")
                     if is_none_or_empty(c) and not (str(r) == "assistant" and t.get("tool_calls")):
@@ -743,7 +739,7 @@ def show_row(
                 print(f"  {col}: {len(conversation)} turns ({none_count} None)")
                 print(f"  {'-' * 60}")
                 for i, turn in enumerate(conversation):
-                    # Non-dict turn - can't extract role or content normally.
+                    # Non-dict turn - can't extract role/content normally.
                     if not isinstance(turn, dict):
                         label = "None" if turn is None else "invalid_type"
                         print(f"  [{i:>3d}] {'unknown':<12s} [{label}]  << NONE")
diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py
index 792e0ea9bc..faa3deac70 100644
--- a/studio/backend/utils/datasets/dataset_utils.py
+++ b/studio/backend/utils/datasets/dataset_utils.py
@@ -4,12 +4,12 @@
 """
 Dataset utilities for format detection, conversion, and template application.
 
-This module provides the main entry points for dataset processing:
-- check_dataset_format: Lightweight check if manual mapping is needed (for frontend)
-- format_dataset: Detects and normalizes dataset formats
-- format_and_template_dataset: End-to-end processing with chat template application
+Main entry points for dataset processing:
+- check_dataset_format: lightweight check if manual mapping is needed (frontend)
+- format_dataset: detects and normalizes dataset formats
+- format_and_template_dataset: end-to-end processing with chat template
 
-All internal utilities have been moved to separate modules:
+Internal utilities live in separate modules:
 - format_detection: detect_dataset_format, detect_multimodal_dataset, etc.
 - format_conversion: standardize_chat_format, convert_chatml_to_alpaca, etc.
 - chat_templates: apply_chat_template_to_dataset, get_tokenizer_chat_template, etc.
@@ -20,7 +20,6 @@ All internal utilities have been moved to separate modules:
 
 import json
 
-# Import from modular files
 from .format_detection import (
     detect_dataset_format,
     detect_multimodal_dataset,
@@ -54,8 +53,8 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
     """
     Lightweight format check without processing - for frontend validation.
 
-    Use this to quickly determine if user needs to manually map columns
-    before calling the full format_and_template_dataset().
+    Quickly determines if the user must manually map columns before the full
+    format_and_template_dataset().
 
     Args:
         dataset: HuggingFace dataset
@@ -121,7 +120,7 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
         }
 
     if is_audio:
-        # Audio dataset — require manual mapping only when columns can't be auto-detected
+        # Audio dataset — require manual mapping only when columns aren't auto-detected
         detected_audio = multimodal_info.get("detected_audio_column")
         detected_text = multimodal_info.get("detected_text_column")
         needs_mapping = not detected_audio or not detected_text
@@ -181,6 +180,7 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
         "suggested_mapping": None,
         "detected_image_column": None,
         "detected_text_column": None,
+        "chat_column": detected.get("chat_column"),
         "is_image": multimodal_info["is_image"],
         "multimodal_columns": multimodal_info.get("multimodal_columns"),
         **audio_fields,
@@ -200,6 +200,17 @@ _TO_CHATML = {
 }
 _CHATML_ROLE_ORDER = ("system", "user", "assistant")
 _CHATML_TO_ALPACA = {"user": "instruction", "system": "input", "assistant": "output"}
+_KNOWN_CHAT_COLUMNS = {"messages", "conversations", "texts"}
+
+
+def _chatml_final_format(chat_column: str | None) -> str:
+    return "chatml_messages" if chat_column == "messages" else "chatml_conversations"
+
+
+def _chatml_detected_format_label(chat_column: str | None) -> str:
+    if chat_column in _KNOWN_CHAT_COLUMNS:
+        return f"chatml_{chat_column}"
+    return "chatml_conversations"
 
 
 def _apply_user_mapping(
@@ -211,9 +222,9 @@ def _apply_user_mapping(
     Apply user-provided column mapping to convert dataset to conversations format.
 
     Accepts chatml (user/assistant/system), sharegpt (human/gpt/system), and
-    alpaca (instruction/input/output) role names — all normalised to chatml output.
+    alpaca (instruction/input/output) role names — all normalised to chatml.
 
-    If the mapping contains ``__``-prefixed metadata keys (from the conversion
+    If the mapping has ``__``-prefixed metadata keys (from the conversion
     advisor), routes to template-based conversion instead of simple role mapping.
 
     Returns:
@@ -262,7 +273,7 @@ def _apply_user_mapping(
 
 def _extract_column_value(val, col: str, label_mapping: dict) -> str:
     """Extract a string value from a column, handling complex types and label mapping."""
-    # Handle complex types (dicts, lists) — extract useful text instead of raw repr
+    # Complex types (dicts, lists): extract useful text instead of raw repr
     if isinstance(val, dict):
         # Common pattern: {"text": [...]} in QA datasets
         if "text" in val:
@@ -291,10 +302,9 @@ def _apply_template_mapping(
     """
     Apply advisor-driven mapping for non-conversational datasets.
 
-    Groups columns by their assigned role (user/assistant), concatenates
-    values within each role into a single message, and injects an optional
-    system prompt.  Label mapping is applied to convert integer labels
-    to human-readable strings.
+    Groups columns by assigned role (user/assistant), concatenates values
+    within each role into one message, and injects an optional system prompt.
+    Label mapping converts integer labels to human-readable strings.
 
     Returns:
         Dataset with single 'conversations' column
@@ -327,7 +337,7 @@ def _apply_template_mapping(
             if system_prompt:
                 convo.append({"role": "system", "content": system_prompt})
 
-            # User message: concatenate all user-role column values
+            # User message: concatenate user-role column values
             user_parts = []
             for col in role_groups["user"]:
                 if col in examples:
@@ -335,7 +345,7 @@ def _apply_template_mapping(
             if user_parts:
                 convo.append({"role": "user", "content": "\n".join(user_parts)})
 
-            # Assistant message: concatenate all assistant-role column values
+            # Assistant message: concatenate assistant-role column values
             asst_parts = []
             for col in role_groups["assistant"]:
                 if col in examples:
@@ -433,7 +443,7 @@ def format_dataset(
             "final_format": final format after processing,
             "chat_column": column name with chat data,
             "is_standardized": whether role names are standardized,
-            "requires_manual_mapping": True if format detection failed and user must map columns,
+            "requires_manual_mapping": True if detection failed and user must map columns,
             "warnings": list of warning messages
         }
     """
@@ -455,7 +465,7 @@ def format_dataset(
             "warnings": [notice.message for notice in raw_result.notices],
         }
 
-    # If user provided explicit mapping, skip detection and apply in the requested format
+    # If user provided explicit mapping, skip detection and apply it
     if custom_format_mapping:
         try:
             if format_type == "alpaca":
@@ -465,8 +475,8 @@ def format_dataset(
                 final_format = "alpaca"
                 chat_column = None
             else:
-                # auto / chatml / sharegpt / conversational — all produce chatml conversations
-                # (sharegpt is always standardized to role/content internally)
+                # auto / chatml / sharegpt / conversational all produce chatml
+                # conversations (sharegpt standardized to role/content internally)
                 mapped_dataset = _apply_user_mapping(dataset, custom_format_mapping, batch_size)
                 final_format = "chatml_conversations"
                 chat_column = "conversations"
@@ -524,7 +534,7 @@ def format_dataset(
             }
 
         # ShareGPT - needs standardization
-        elif detected["format"] == "sharegpt":
+        elif detected["format"] == "sharegpt" and detected.get("chat_column"):
             try:
                 standardized = standardize_chat_format(
                     dataset,
@@ -534,11 +544,12 @@ def format_dataset(
                     aliases_for_assistant,
                     batch_size,
                     num_proc,
+                    chat_column = detected["chat_column"],
                 )
                 return {
                     "dataset": standardized,
                     "detected_format": "sharegpt",
-                    "final_format": f"chatml_{detected['chat_column']}",
+                    "final_format": _chatml_final_format(detected["chat_column"]),
                     "chat_column": detected["chat_column"],
                     "is_standardized": True,
                     "requires_manual_mapping": False,
@@ -560,15 +571,11 @@ def format_dataset(
                     "warnings": warnings,
                 }
 
-        elif detected["format"] == "chatml" and detected["chat_column"] in [
-            "conversations",
-            "messages",
-            "texts",
-        ]:
+        elif detected["format"] == "chatml" and detected.get("chat_column"):
             return {
                 "dataset": dataset,
-                "detected_format": f"chatml_{detected['chat_column']}",
-                "final_format": f"chatml_{detected['chat_column']}",
+                "detected_format": _chatml_detected_format_label(detected["chat_column"]),
+                "final_format": _chatml_final_format(detected["chat_column"]),
                 "chat_column": detected["chat_column"],
                 "is_standardized": True,
                 "requires_manual_mapping": False,
@@ -577,11 +584,11 @@ def format_dataset(
                 "warnings": warnings,
             }
 
-        # Unknown - try standardization, if fails pass as is
+        # Unknown - try standardization, pass as-is on failure
         else:
             warnings.append(f"Unknown format detected. Keys found: {detected['sample_keys']}")
 
-            # NEW: Try heuristic detection
+            # Try heuristic detection
             if auto_detect_custom:
                 custom_mapping = detect_custom_format_heuristic(dataset)
                 if custom_mapping:
@@ -639,12 +646,13 @@ def format_dataset(
                         aliases_for_assistant,
                         batch_size,
                         num_proc,
+                        chat_column = detected["chat_column"],
                     )
                     warnings.append("Successfully standardized unknown format")
                     return {
                         "dataset": standardized,
                         "detected_format": "unknown",
-                        "final_format": f"chatml_{detected['chat_column']}",
+                        "final_format": _chatml_final_format(detected["chat_column"]),
                         "chat_column": detected["chat_column"],
                         "is_standardized": True,
                         "requires_manual_mapping": False,
@@ -683,32 +691,52 @@ def format_dataset(
                 "warnings": [],
             }
 
-        elif detected["format"] in ["sharegpt", "chatml"]:
-            # First standardize if ShareGPT
-            if detected["format"] == "sharegpt":
-                dataset = standardize_chat_format(
+        elif detected["format"] in ["sharegpt", "chatml"] and detected.get("chat_column"):
+            try:
+                # First standardize if ShareGPT
+                if detected["format"] == "sharegpt":
+                    dataset = standardize_chat_format(
+                        dataset,
+                        tokenizer,
+                        aliases_for_system,
+                        aliases_for_user,
+                        aliases_for_assistant,
+                        batch_size,
+                        num_proc,
+                        chat_column = detected["chat_column"],
+                    )
+
+                # Then convert to Alpaca
+                converted = convert_chatml_to_alpaca(
                     dataset,
-                    tokenizer,
-                    aliases_for_system,
-                    aliases_for_user,
-                    aliases_for_assistant,
                     batch_size,
                     num_proc,
+                    chat_column = detected["chat_column"],
                 )
-
-            # Then convert to Alpaca
-            converted = convert_chatml_to_alpaca(dataset, batch_size, num_proc)
-            return {
-                "dataset": converted,
-                "detected_format": detected["format"],
-                "final_format": "alpaca",
-                "chat_column": None,
-                "is_standardized": True,
-                "requires_manual_mapping": False,
-                "is_image": multimodal_info["is_image"],
-                "multimodal_info": multimodal_info,
-                "warnings": [],
-            }
+                return {
+                    "dataset": converted,
+                    "detected_format": detected["format"],
+                    "final_format": "alpaca",
+                    "chat_column": None,
+                    "is_standardized": True,
+                    "requires_manual_mapping": False,
+                    "is_image": multimodal_info["is_image"],
+                    "multimodal_info": multimodal_info,
+                    "warnings": [],
+                }
+            except Exception as e:
+                warnings.append(f"Failed to convert chat dataset to Alpaca: {e}")
+                return {
+                    "dataset": dataset,
+                    "detected_format": detected["format"],
+                    "final_format": "unknown",
+                    "chat_column": detected["chat_column"],
+                    "is_standardized": False,
+                    "requires_manual_mapping": True,
+                    "is_image": multimodal_info["is_image"],
+                    "multimodal_info": multimodal_info,
+                    "warnings": warnings,
+                }
 
         else:
             warnings.append(f"Cannot convert unknown format to Alpaca")
@@ -740,33 +768,48 @@ def format_dataset(
                 "warnings": [],
             }
 
-        elif detected["format"] == "sharegpt":
-            standardized = standardize_chat_format(
-                dataset,
-                tokenizer,
-                aliases_for_system,
-                aliases_for_user,
-                aliases_for_assistant,
-                batch_size,
-                num_proc,
-            )
-            return {
-                "dataset": standardized,
-                "detected_format": "sharegpt",
-                "final_format": f"chatml_{detected['chat_column']}",
-                "chat_column": detected["chat_column"],
-                "is_standardized": True,
-                "requires_manual_mapping": False,
-                "is_image": multimodal_info["is_image"],
-                "multimodal_info": multimodal_info,
-                "warnings": [],
-            }
+        elif detected["format"] == "sharegpt" and detected.get("chat_column"):
+            try:
+                standardized = standardize_chat_format(
+                    dataset,
+                    tokenizer,
+                    aliases_for_system,
+                    aliases_for_user,
+                    aliases_for_assistant,
+                    batch_size,
+                    num_proc,
+                    chat_column = detected["chat_column"],
+                )
+                return {
+                    "dataset": standardized,
+                    "detected_format": "sharegpt",
+                    "final_format": _chatml_final_format(detected["chat_column"]),
+                    "chat_column": detected["chat_column"],
+                    "is_standardized": True,
+                    "requires_manual_mapping": False,
+                    "is_image": multimodal_info["is_image"],
+                    "multimodal_info": multimodal_info,
+                    "warnings": [],
+                }
+            except Exception as e:
+                warnings.append(f"Failed to standardize ShareGPT format: {e}")
+                return {
+                    "dataset": dataset,
+                    "detected_format": "sharegpt",
+                    "final_format": "sharegpt",
+                    "chat_column": detected["chat_column"],
+                    "is_standardized": False,
+                    "requires_manual_mapping": True,
+                    "is_image": multimodal_info["is_image"],
+                    "multimodal_info": multimodal_info,
+                    "warnings": warnings,
+                }
 
-        elif detected["format"] == "chatml":
+        elif detected["format"] == "chatml" and detected.get("chat_column"):
             return {
                 "dataset": dataset,
-                "detected_format": f"chatml_{detected['chat_column']}",
-                "final_format": f"chatml_{detected['chat_column']}",
+                "detected_format": _chatml_detected_format_label(detected["chat_column"]),
+                "final_format": _chatml_final_format(detected["chat_column"]),
                 "chat_column": detected["chat_column"],
                 "is_standardized": True,
                 "requires_manual_mapping": False,
@@ -787,11 +830,12 @@ def format_dataset(
                         aliases_for_assistant,
                         batch_size,
                         num_proc,
+                        chat_column = detected["chat_column"],
                     )
                     return {
                         "dataset": standardized,
                         "detected_format": "unknown",
-                        "final_format": f"chatml_{detected['chat_column']}",
+                        "final_format": _chatml_final_format(detected["chat_column"]),
                         "chat_column": detected["chat_column"],
                         "is_standardized": True,
                         "requires_manual_mapping": False,
@@ -853,8 +897,8 @@ def format_and_template_dataset(
     progress_callback = None,
 ):
     """
-    Convenience function that combines format_dataset and apply_chat_template_to_dataset.
-    Perfect for UI workflows - one function does everything!
+    Combines format_dataset and apply_chat_template_to_dataset. Convenient for
+    UI workflows: one function does everything.
 
     Returns:
         dict: {
@@ -862,7 +906,7 @@ def format_and_template_dataset(
             "detected_format": Original format,
             "final_format": Format after processing,
             "success": Whether template application succeeded,
-            "requires_manual_mapping": True if format detection failed and user must map columns,
+            "requires_manual_mapping": True if detection failed and user must map columns,
             "warnings": List of warnings,
             "errors": List of errors,
             "summary": Human-readable summary
@@ -876,7 +920,7 @@ def format_and_template_dataset(
 
         multimodal_info = detect_multimodal_dataset(dataset)
 
-        # NEW: If user provided explicit mapping for VLM, use it directly
+        # If user provided explicit mapping for VLM, use it directly
         if custom_format_mapping:
             # Expect mapping like: {"image_col": "image", "caption_col": "text"}
             user_vlm_image_column = None
@@ -916,15 +960,14 @@ def format_and_template_dataset(
                         "errors": [],
                     }
                 except Exception as e:
-                    # User mapping failed — fall back to auto-detection instead
-                    # of giving up (handles stale cached mappings gracefully)
+                    # User mapping failed; fall back to auto-detection (handles stale cached mappings).
                     warnings.append(
                         f"User VLM mapping (image='{user_vlm_image_column}', "
                         f"text='{user_vlm_text_column}') failed: {e} — "
                         f"falling back to auto-detection"
                     )
                     logger.info(f"⚠️ User VLM mapping failed, falling back to auto-detection...")
-                    custom_format_mapping = None  # clear so auto-detection runs below
+                    custom_format_mapping = None  # so auto-detection runs below
             else:
                 errors.append(
                     f"Invalid VLM mapping: need 'image' and 'text' roles. Got: {custom_format_mapping}"
@@ -967,7 +1010,7 @@ def format_and_template_dataset(
                     "errors": errors,
                 }
 
-        # Handle ShareGPT/ChatML + image column (e.g. ShareGPT4V, LLaVA-style)
+        # ShareGPT/ChatML + image column (e.g. ShareGPT4V, LLaVA-style)
         elif vlm_structure["format"] == "sharegpt_with_images":
             try:
                 dataset = convert_sharegpt_with_images_to_vlm_format(
@@ -1087,7 +1130,7 @@ def format_and_template_dataset(
             "errors": errors,
         }
 
-    # LLM FLOW (Existing code)
+    # LLM FLOW
     else:
         # Step 1: Format the dataset
         n_rows = len(dataset) if hasattr(dataset, "__len__") else None
@@ -1127,7 +1170,7 @@ def format_and_template_dataset(
             progress_callback(
                 status_message = f"Applying chat template to {detected} ({n_rows:,} rows)..."
             )
-        # Gemma emits a leading  that must be stripped for text-only chatml/sharegpt.
+        # Gemma emits a leading , stripped for text-only chatml/sharegpt.
         is_alpaca = format_type == "alpaca" or (
             format_type == "auto" and dataset_info["detected_format"] == "alpaca"
         )
@@ -1155,8 +1198,7 @@ def format_and_template_dataset(
         all_warnings = dataset_info.get("warnings", []) + template_result.get("warnings", [])
         all_errors = template_result.get("errors", [])
 
-        # If format_dataset returned "unknown" but apply_chat_template rescued
-        # it via heuristic detection, update final_format to reflect reality.
+        # If apply_chat_template rescued an "unknown" format, update final_format.
         final_format = dataset_info["final_format"]
         requires_manual = dataset_info.get("requires_manual_mapping", False)
         if final_format == "unknown" and template_result["success"]:
@@ -1170,7 +1212,7 @@ def format_and_template_dataset(
             "detected_format": dataset_info["detected_format"],
             "final_format": final_format,
             "chat_column": dataset_info.get("chat_column"),
-            "is_vlm": False,  # This is LLM flow
+            "is_vlm": False,  # LLM flow
             "success": template_result["success"],
             "requires_manual_mapping": requires_manual,
             "warnings": all_warnings,
diff --git a/studio/backend/utils/datasets/format_conversion.py b/studio/backend/utils/datasets/format_conversion.py
index 5433d3115c..cb24bd96ba 100644
--- a/studio/backend/utils/datasets/format_conversion.py
+++ b/studio/backend/utils/datasets/format_conversion.py
@@ -1,12 +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
 
-"""
-Format conversion utilities for dataset processing.
-
-This module contains functions for converting between dataset formats
-(Alpaca, ShareGPT, ChatML) and standardizing chat formats.
-"""
+"""Dataset format conversion between Alpaca, ShareGPT, and ChatML."""
 
 import os
 
@@ -34,16 +29,17 @@ def standardize_chat_format(
     ],
     batch_size = 1000,
     num_proc = None,
+    chat_column: str | None = None,
 ):
     """
-    Our own standardization function that handles BOTH messages and conversations.
-    Converts non-standard role names and keys to standard format.
+    Standardize BOTH messages and conversations: map non-standard role
+    names and keys to the standard format.
     """
     import collections
     import itertools
     from datasets import IterableDataset
 
-    # Check if vision tokenizer is used
+    # Detect a vision tokenizer
     is_vlm = False
     if tokenizer is not None:
         if hasattr(tokenizer, "image_processor") or hasattr(tokenizer, "tokenizer"):
@@ -51,9 +47,10 @@ def standardize_chat_format(
 
     column_names = set(next(iter(dataset)).keys())
 
-    #   Check for both 'conversations' and 'messages'
-    chat_column = None
-    if "conversations" in column_names:
+    if chat_column:
+        if chat_column not in column_names:
+            return dataset
+    elif "conversations" in column_names:
         chat_column = "conversations"
     elif "messages" in column_names:
         chat_column = "messages"
@@ -62,30 +59,48 @@ def standardize_chat_format(
     else:
         return dataset  # No chat column found
 
-    # Inspect structure
-    examples = itertools.islice(dataset, 10)
+    def _iter_probe_rows():
+        try:
+            total = min(len(dataset), 100)
+            for index in range(total):
+                yield dataset[index]
+            return
+        except Exception:
+            pass
+        for example in itertools.islice(dataset, 100):
+            yield example
+
     uniques = collections.defaultdict(list)
-    for example in examples:
-        for message in example[chat_column]:
+    for example in _iter_probe_rows():
+        chat_data = example.get(chat_column)
+        if not isinstance(chat_data, list) or len(chat_data) == 0:
+            continue
+        for message in chat_data:
+            if not isinstance(message, dict):
+                continue
             for key, value in message.items():
                 if type(value) is not str:
-                    continue  # Skip non-string values
+                    continue  # Skip non-strings
                 uniques[key].append(value)
 
-    if len(uniques.keys()) != 2:
-        return dataset  # Unexpected structure
-
-    keys = list(uniques.keys())
-    length_first = len(set(uniques[keys[0]]))
-    length_second = len(set(uniques[keys[1]]))
-
-    # Determine which is role and which is content
-    if length_first < length_second:
-        role_key = keys[0]
-        content_key = keys[1]
+    if "from" in uniques and "value" in uniques:
+        role_key = "from"
+        content_key = "value"
+    elif "role" in uniques and "content" in uniques:
+        role_key = "role"
+        content_key = "content"
+    elif len(uniques.keys()) == 2:
+        keys = list(uniques.keys())
+        length_first = len(set(uniques[keys[0]]))
+        length_second = len(set(uniques[keys[1]]))
+        if length_first < length_second:
+            role_key = keys[0]
+            content_key = keys[1]
+        else:
+            role_key = keys[1]
+            content_key = keys[0]
     else:
-        role_key = keys[1]
-        content_key = keys[0]
+        raise ValueError(f"Could not infer role/content keys for chat column '{chat_column}'")
 
     # Mapping for aliases
     aliases_mapping = {}
@@ -100,20 +115,30 @@ def standardize_chat_format(
         convos = examples[chat_column]
         all_convos = []
         for convo in convos:
+            if not isinstance(convo, list):
+                all_convos.append([])
+                continue
+
             new_convo = []
             for message in convo:
-                # Get original role and content
-                original_role = message.get(role_key, "")
-                original_content = message.get(content_key, "")
+                if not isinstance(message, dict):
+                    continue
+
+                # Use the inferred keys first; fall back per-message so mixed
+                # ShareGPT/ChatML rows keep valid turns.
+                original_role = message.get(role_key)
+                original_content = message.get(content_key)
+                if original_role is None:
+                    original_role = message.get("role") or message.get("from") or ""
+                if original_content is None:
+                    original_content = message.get("content") or message.get("value") or ""
 
-                # Map to standard role name
                 standard_role = aliases_mapping.get(original_role, original_role)
 
-                # Handle VLM format
                 if is_vlm:
                     original_content = [{"type": "text", "text": original_content}]
 
-                # Create dict with EXPLICIT ORDER
+                # Keep EXPLICIT key order
                 new_message = {"role": standard_role, "content": original_content}
                 new_convo.append(new_message)
 
@@ -144,10 +169,10 @@ def convert_chatml_to_alpaca(
     dataset,
     batch_size = 1000,
     num_proc = None,
+    chat_column: str | None = None,
 ):
     """
-    Converts ChatML format (messages OR conversations) to Alpaca format.
-    Handles both standardized and ShareGPT formats.
+    Convert ChatML (messages OR conversations) to Alpaca format.
 
     Supports:
     - "messages" or "conversations" column
@@ -160,10 +185,11 @@ def convert_chatml_to_alpaca(
         _is_torch_iterable = False
 
     def _convert(examples):
-        # Auto-detect which column name is used
-        chatml_data = (
-            examples.get("messages") or examples.get("conversations") or examples.get("texts")
-        )
+        chatml_data = examples.get(chat_column) if chat_column else None
+        if chatml_data is None:
+            chatml_data = (
+                examples.get("messages") or examples.get("conversations") or examples.get("texts")
+            )
 
         if chatml_data is None:
             raise ValueError("No 'messages' or 'conversations' or 'texts' column found.")
@@ -177,20 +203,20 @@ def convert_chatml_to_alpaca(
             output = ""
 
             for msg in convo:
-                # Handle both standard and ShareGPT formats
+                # Standard and ShareGPT key names
                 role = msg.get("role") or msg.get("from")
                 content = msg.get("content") or msg.get("value")
 
-                # Get first user message as instruction
+                # First user message -> instruction
                 if role in ["user", "human", "input"] and not instruction:
                     instruction = content
-                # Get first assistant message as output
+                # First assistant message -> output
                 elif role in ["assistant", "gpt", "output"] and not output:
                     output = content
                     break  # Stop after first assistant response
 
             instructions.append(instruction)
-            inputs.append("")  # Alpaca typically has empty input
+            inputs.append("")  # Alpaca input usually empty
             outputs.append(output)
 
         return {"instruction": instructions, "input": inputs, "output": outputs}
@@ -220,9 +246,9 @@ def convert_alpaca_to_chatml(
     num_proc = None,
 ):
     """
-    Converts Alpaca format to ChatML format.
+    Convert Alpaca format to ChatML format.
 
-    Output format: Uses 'conversations' column with standard 'role'/'content' structure.
+    Output: 'conversations' column with standard 'role'/'content' dicts.
     """
     try:
         from torch.utils.data import IterableDataset
@@ -238,13 +264,12 @@ def convert_alpaca_to_chatml(
             input_text = examples.get("input", [""] * len(examples["instruction"]))[i]
             output = examples["output"][i]
 
-            # Combine instruction and input (if exists) for user message
+            # User message = instruction + input (if any)
             if input_text and input_text.strip():
                 user_content = f"{instruction}\n\n{input_text}".strip()
             else:
                 user_content = instruction
 
-            # Build conversation in standard ChatML format
             convo = [
                 {"role": "user", "content": user_content},
                 {"role": "assistant", "content": output},
@@ -294,17 +319,14 @@ def convert_to_vlm_format(
     progress_callback = None,
 ):
     """
-    Converts simple {image, text} format to VLM messages format.
+    Convert simple {image, text} format to VLM messages format.
 
     Returns a LIST, not a HuggingFace Dataset (to preserve PIL Images).
-
-    For URL-based image datasets, runs a 200-sample parallel probe first to
-    estimate download speed and failure rate, then reports time estimate or
-    warning through progress_callback before proceeding with the full conversion.
+    For URL-based datasets, runs a 200-sample parallel probe first to
+    estimate speed/failure rate via progress_callback.
 
     Args:
-        progress_callback: Optional callable(status_message=str) to report
-                          progress to the training overlay.
+        progress_callback: Optional callable(status_message=str) for progress.
 
     Returns:
         list: List of dicts with 'messages' field
@@ -313,11 +335,11 @@ def convert_to_vlm_format(
     from .vlm_processing import generate_smart_vlm_instruction
 
     def _notify(msg):
-        """Send status update to the training overlay if callback is available."""
+        """Send a status update to the training overlay if callback set."""
         if progress_callback:
             progress_callback(status_message = msg)
 
-    # Generate smart instruction if not provided
+    # Generate a smart instruction if none provided
     if instruction is None:
         instruction_info = generate_smart_vlm_instruction(
             dataset,
@@ -342,7 +364,7 @@ def convert_to_vlm_format(
 
     def _convert_single_sample(sample):
         """Convert a single sample to VLM format."""
-        # Get image (might be PIL Image, local path, URL, or bare filename)
+        # Image may be a PIL Image, local path, URL, or bare filename
         image_data = sample[image_column]
 
         if isinstance(image_data, str):
@@ -363,19 +385,18 @@ def convert_to_vlm_format(
             else:
                 image_data = Image.open(image_data).convert("RGB")
 
-        # Get text (if list of strings, pick a random one — e.g. multiple captions)
+        # Text: if a list (e.g. multiple captions), pick one at random
         text_data = sample[text_column]
         if isinstance(text_data, list) and len(text_data) > 0:
             import random
             text_data = random.choice(text_data)
 
-        # Get instruction (static or dynamic)
+        # Instruction: static or dynamic
         if uses_dynamic and instruction_column:
             current_instruction = sample[instruction_column]
         else:
             current_instruction = instruction
 
-        # Build VLM messages - simple structure
         messages = [
             {
                 "role": "user",
@@ -387,16 +408,14 @@ def convert_to_vlm_format(
             {"role": "assistant", "content": [{"type": "text", "text": text_data}]},
         ]
 
-        # Return dict with messages
         return {"messages": messages}
 
     total = len(dataset)
     first_image = next(iter(dataset))[image_column]
     has_urls = isinstance(first_image, str) and first_image.startswith(("http://", "https://"))
 
-    # ── Bare-filename detection: images stored as filenames (e.g. "img_001.png")
-    #    that don't exist locally.  Build a basename→repo_path lookup so we can
-    #    resolve them via hf_hub_download during conversion.
+    # ── Bare-filename detection: build a basename→repo_path lookup so
+    #    filename-only images resolve via hf_hub_download during conversion.
     _image_lookup = None
     _IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff")
     if (
@@ -431,7 +450,7 @@ def convert_to_vlm_format(
             logger.info(f"⚠️ Failed to build HF repo image lookup: {e}")
             _image_lookup = None
 
-    # ── URL probe: 200 samples with parallel workers to estimate speed + failure rate ──
+    # ── URL probe: 200 parallel samples to estimate speed + failure rate ──
     PROBE_SIZE = 200
     MAX_FAIL_RATE = 0.3
 
@@ -468,7 +487,7 @@ def convert_to_vlm_format(
                 f"{fail_rate:.0%} of the first {PROBE_SIZE} image URLs failed to download ({probe_fail}/{probe_total})",
                 "Images are external URLs, not embedded in the dataset",
             ]
-            # Try LLM-friendly warning
+            # LLM-friendly warning
             friendly = None
             try:
                 from .llm_assist import llm_generate_dataset_warning
@@ -490,7 +509,7 @@ def convert_to_vlm_format(
             _notify(msg)
             raise ValueError(msg)
 
-        # Estimate total time for remaining samples
+        # Estimate time for remaining samples
         remaining = total - PROBE_SIZE
         estimated_seconds = remaining / throughput if throughput > 0 else 0
         eta_str = _format_eta(estimated_seconds)
@@ -546,7 +565,7 @@ def convert_to_vlm_format(
 
             converted_list.extend(r for r in batch_results if r is not None)
 
-            # Progress update every batch
+            # Per-batch progress update
             elapsed = time.time() - start_time
             done = batch_end
             rate = done / elapsed if elapsed > 0 else 0
@@ -558,7 +577,7 @@ def convert_to_vlm_format(
             )
             _notify(progress_msg)
     else:
-        # Sequential conversion for local/embedded images (fast, no I/O bottleneck)
+        # Sequential conversion for local/embedded images (no I/O bottleneck)
         pbar = tqdm(dataset, total = total, desc = "Converting VLM samples", unit = "sample")
         for sample in pbar:
             try:
@@ -576,7 +595,7 @@ def convert_to_vlm_format(
         logger.info(
             f"⚠️ Skipped {failed_count}/{total} ({fail_rate:.0%}) samples with broken/unreachable images"
         )
-        # For datasets that skipped the probe (small URL datasets), check fail rate now
+        # Small URL datasets skip the probe; check fail rate here
         if has_urls and fail_rate >= MAX_FAIL_RATE:
             issues = [
                 f"{fail_rate:.0%} of images failed to download ({failed_count}/{total})",
@@ -629,7 +648,7 @@ def convert_to_vlm_format(
     logger.info(f"✅ Converted {len(converted_list)}/{total} samples")
     _notify(f"Converted {len(converted_list):,}/{total:,} images successfully")
 
-    # Return list, NOT Dataset
+    # Return list, NOT a Dataset
     return converted_list
 
 
@@ -641,8 +660,8 @@ def convert_sharegpt_with_images_to_vlm_format(
     progress_callback = None,
 ):
     """
-    Converts ShareGPT/ChatML datasets that have a separate image column and
-    ```` placeholders inside the conversation text.
+    Convert ShareGPT/ChatML datasets with a separate image column and
+    ```` placeholders in the conversation text.
 
     Example input::
 
@@ -672,7 +691,7 @@ def convert_sharegpt_with_images_to_vlm_format(
         if progress_callback:
             progress_callback(status_message = msg)
 
-    # ── Resolve image loading strategy (same 3-tier as convert_to_vlm_format) ──
+    # ── Resolve image loading (same 3-tier as convert_to_vlm_format) ──
     total = len(dataset)
     first_image = next(iter(dataset))[image_column]
 
@@ -696,7 +715,7 @@ def convert_sharegpt_with_images_to_vlm_format(
                 for f in repo_files
                 if any(f.lower().endswith(ext) for ext in _IMAGE_EXTS)
             }
-            # Also add the full relative paths as keys (for paths like "sam/images/sa_545504.jpg")
+            # Also key by full relative path (e.g. "sam/images/sa_545504.jpg")
             for f in repo_files:
                 if any(f.lower().endswith(ext) for ext in _IMAGE_EXTS):
                     _image_lookup[f] = f
@@ -714,7 +733,7 @@ def convert_sharegpt_with_images_to_vlm_format(
             _image_lookup = None
 
     def _resolve_image(image_data):
-        """Resolve image data to a PIL Image object."""
+        """Resolve image data to a PIL Image."""
         if hasattr(image_data, "size") and hasattr(image_data, "mode"):
             return image_data  # Already PIL
         if isinstance(image_data, str):
@@ -742,7 +761,7 @@ def convert_sharegpt_with_images_to_vlm_format(
         raise ValueError(f"Cannot resolve image: {type(image_data)}")
 
     def _convert_single_sample(sample):
-        """Convert a single ShareGPT+image sample to standard VLM format."""
+        """Convert one ShareGPT+image sample to standard VLM format."""
         pil_image = _resolve_image(sample[image_column])
         conversation = sample[messages_column]
 
@@ -752,7 +771,7 @@ def convert_sharegpt_with_images_to_vlm_format(
             role = _ROLE_MAP.get(role_raw.lower(), role_raw.lower())
             text = msg.get("value") or msg.get("content") or ""
 
-            # Split on  to interleave text and image content blocks
+            # Interleave text and image blocks around 
             if "" in text:
                 parts = text.split("")
                 content = []
@@ -762,7 +781,7 @@ def convert_sharegpt_with_images_to_vlm_format(
                         content.append({"type": "text", "text": part})
                     if i < len(parts) - 1:
                         content.append({"type": "image", "image": pil_image})
-                # If  was the entire text, content might just be the image
+                # If text was only , content is just the image
                 if not content:
                     content.append({"type": "image", "image": pil_image})
             else:
@@ -804,7 +823,7 @@ def convert_sharegpt_with_images_to_vlm_format(
 
 def convert_llava_to_vlm_format(dataset):
     """
-    Converts Llava format to standard VLM format.
+    Convert Llava format to standard VLM format.
 
     Llava format:
     - messages: [{'content': [{'type': 'image', 'index': 0}, {'type': 'text', 'text': '...'}]}]
@@ -818,23 +837,22 @@ def convert_llava_to_vlm_format(dataset):
     logger.info(f"🔄 Converting {len(dataset)} samples from Llava format to standard VLM format...")
 
     def _convert_single_sample(sample):
-        """Convert a single llava sample to standard VLM format."""
+        """Convert one llava sample to standard VLM format."""
         messages = sample["messages"]
         images = sample.get("images", [])
 
-        # Process each message
         new_messages = []
         for msg in messages:
             new_content = []
 
             for item in msg["content"]:
                 if item["type"] == "image":
-                    # Replace index with actual PIL image
+                    # Replace index with the actual PIL image
                     if "index" in item and item["index"] is not None:
                         img_idx = item["index"]
                         if img_idx < len(images):
                             pil_image = images[img_idx]
-                            # Ensure it's PIL
+                            # Ensure PIL
                             if isinstance(pil_image, str):
                                 pil_image = Image.open(pil_image).convert("RGB")
 
@@ -845,7 +863,7 @@ def convert_llava_to_vlm_format(dataset):
                                 }
                             )
                     else:
-                        # No index, try to use first image
+                        # No index: use the first image
                         if len(images) > 0:
                             pil_image = images[0]
                             if isinstance(pil_image, str):
@@ -854,14 +872,12 @@ def convert_llava_to_vlm_format(dataset):
                             new_content.append({"type": "image", "image": pil_image})
 
                 elif item["type"] == "text":
-                    # Keep text as-is (only type + text)
                     new_content.append({"type": "text", "text": item.get("text", "")})
 
             new_messages.append({"role": msg["role"], "content": new_content})
 
         return {"messages": new_messages}
 
-    # Convert using list comprehension
     converted_list = [_convert_single_sample(sample) for sample in dataset]
 
     logger.info(f"✅ Converted {len(converted_list)} samples")
diff --git a/studio/backend/utils/datasets/format_detection.py b/studio/backend/utils/datasets/format_detection.py
index 829838064a..f5ea5ca138 100644
--- a/studio/backend/utils/datasets/format_detection.py
+++ b/studio/backend/utils/datasets/format_detection.py
@@ -1,12 +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
 
-"""
-Format detection utilities for dataset processing.
-
-This module contains functions for detecting dataset formats (Alpaca, ShareGPT, ChatML),
-detecting multimodal/VLM dataset structures, and heuristic-based column mapping.
-"""
+"""Dataset format detection: Alpaca/ShareGPT/ChatML, multimodal/VLM structures, heuristic column mapping."""
 
 import re
 
@@ -16,23 +11,146 @@ def _keyword_in_column(keyword: str, col_name: str) -> bool:
     return re.search(r"\b" + re.escape(keyword) + r"\b", col_name, re.IGNORECASE) is not None
 
 
+CONVERSATION_COLUMNS = ("messages", "conversations", "texts")
+_CHATML_KEYS = frozenset({"role", "content"})
+_SHAREGPT_KEYS = frozenset({"from", "value"})
+_TRACE_SUFFIXES = ("__trace", "_trace")
+
+
+def _sample_dataset_rows(dataset, limit: int = 100) -> list[dict]:
+    try:
+        total = min(len(dataset), limit)
+        return [dataset[index] for index in range(total)]
+    except Exception:
+        rows = []
+        try:
+            for index, row in enumerate(dataset):
+                if index >= limit:
+                    break
+                rows.append(row)
+        except Exception:
+            return []
+        return rows
+
+
+def _get_dataset_column_names(dataset, sample: dict) -> list[str]:
+    column_names = getattr(dataset, "column_names", None)
+    if isinstance(column_names, list):
+        return [str(column) for column in column_names]
+    return [str(column) for column in sample.keys()]
+
+
+def _is_trace_conversation_name(column_name: str) -> bool:
+    return column_name.lower().endswith(_TRACE_SUFFIXES)
+
+
+def _inspect_conversation_column(rows: list[dict], column_name: str) -> dict | None:
+    turn_keys: set[str] = set()
+    has_chatml = False
+    has_sharegpt = False
+
+    for row in rows:
+        if not isinstance(row, dict) or column_name not in row:
+            continue
+        chat_data = row[column_name]
+        if not isinstance(chat_data, list) or len(chat_data) == 0:
+            continue
+        for turn in chat_data:
+            if not isinstance(turn, dict):
+                continue
+            keys = {str(key) for key in turn.keys()}
+            turn_keys.update(keys)
+            if _SHAREGPT_KEYS.issubset(keys):
+                has_sharegpt = True
+            if _CHATML_KEYS.issubset(keys):
+                has_chatml = True
+
+    if has_sharegpt:
+        return {
+            "format": "sharegpt",
+            "chat_column": column_name,
+            "needs_standardization": True,
+            "sample_keys": sorted(turn_keys),
+        }
+    if has_chatml:
+        return {
+            "format": "chatml",
+            "chat_column": column_name,
+            "needs_standardization": False,
+            "sample_keys": sorted(turn_keys),
+        }
+    if turn_keys:
+        return {
+            "format": "unknown",
+            "chat_column": column_name,
+            "needs_standardization": None,
+            "sample_keys": sorted(turn_keys),
+        }
+    return None
+
+
+def _detect_conversation_column(rows: list[dict], column_names: list[str]) -> dict | None:
+    column_name_set = set(column_names)
+    unknown_exact = None
+    for column_name in CONVERSATION_COLUMNS:
+        if column_name not in column_name_set:
+            continue
+        inspected = _inspect_conversation_column(rows, column_name)
+        if inspected and inspected["format"] in {"sharegpt", "chatml"}:
+            return inspected
+        if inspected and unknown_exact is None:
+            unknown_exact = inspected
+
+    structural_candidates = []
+    for column_name in column_names:
+        if column_name in CONVERSATION_COLUMNS:
+            continue
+        inspected = _inspect_conversation_column(rows, column_name)
+        if inspected and inspected["format"] in {"sharegpt", "chatml"}:
+            structural_candidates.append(inspected)
+
+    trace_candidates = [
+        candidate
+        for candidate in structural_candidates
+        if _is_trace_conversation_name(candidate["chat_column"])
+    ]
+    if len(trace_candidates) == 1:
+        return trace_candidates[0]
+    if len(trace_candidates) > 1:
+        return unknown_exact
+    if len(structural_candidates) == 1:
+        return structural_candidates[0]
+    if unknown_exact is not None:
+        return unknown_exact
+    return None
+
+
 def detect_dataset_format(dataset):
-    """
-    Detects dataset format by inspecting structure.
+    """Detect dataset format by inspecting structure.
 
     Returns:
         dict: {
             "format": "alpaca" | "sharegpt" | "chatml" | "unknown",
-            "chat_column": "messages" | "conversations" | None,
+            "chat_column": str | None,
             "needs_standardization": bool,
             "sample_keys": list of keys found in messages (for debugging)
         }
     """
-    column_names = set(next(iter(dataset)).keys())
+    sample_rows = _sample_dataset_rows(dataset)
+    if not sample_rows:
+        return {
+            "format": "unknown",
+            "chat_column": None,
+            "needs_standardization": None,
+            "sample_keys": [],
+        }
 
-    # Check for Alpaca
+    column_names = _get_dataset_column_names(dataset, sample_rows[0])
+    column_name_set = set(column_names)
+
+    # Alpaca
     alpaca_columns = {"instruction", "output"}
-    if alpaca_columns.issubset(column_names):
+    if alpaca_columns.issubset(column_name_set):
         return {
             "format": "alpaca",
             "chat_column": None,
@@ -40,61 +158,10 @@ def detect_dataset_format(dataset):
             "sample_keys": [],
         }
 
-    # Check for chat-based formats (messages or conversations)
-    chat_column = None
-    if "messages" in column_names:
-        chat_column = "messages"
-    elif "conversations" in column_names:
-        chat_column = "conversations"
-    elif "texts" in column_names:
-        chat_column = "texts"
+    conversation = _detect_conversation_column(sample_rows, column_names)
+    if conversation:
+        return conversation
 
-    if chat_column:
-        # Inspect the structure to determine if ShareGPT or ChatML
-        try:
-            sample = next(iter(dataset))
-            chat_data = sample[chat_column]
-
-            if chat_data and len(chat_data) > 0:
-                first_msg = chat_data[0]
-                msg_keys = set(first_msg.keys())
-
-                # ShareGPT uses "from" and "value"
-                if "from" in msg_keys or "value" in msg_keys:
-                    return {
-                        "format": "sharegpt",
-                        "chat_column": chat_column,
-                        "needs_standardization": True,
-                        "sample_keys": list(msg_keys),
-                    }
-
-                # ChatML uses "role" and "content"
-                elif "role" in msg_keys and "content" in msg_keys:
-                    return {
-                        "format": "chatml",
-                        "chat_column": chat_column,
-                        "needs_standardization": False,
-                        "sample_keys": list(msg_keys),
-                    }
-
-                # Unknown structure but has chat column
-                else:
-                    return {
-                        "format": "unknown",
-                        "chat_column": chat_column,
-                        "needs_standardization": None,
-                        "sample_keys": list(msg_keys),
-                    }
-        except Exception as e:
-            return {
-                "format": "unknown",
-                "chat_column": chat_column,
-                "needs_standardization": None,
-                "sample_keys": [],
-                "error": str(e),
-            }
-
-    # No recognized format
     return {
         "format": "unknown",
         "chat_column": None,
@@ -104,8 +171,7 @@ def detect_dataset_format(dataset):
 
 
 def detect_custom_format_heuristic(dataset):
-    """
-    Smart detection with priority scoring.
+    """Detection with priority scoring.
 
     Strategy for ambiguous keywords like 'task':
     1. Detect assistant first (unambiguous)
@@ -118,7 +184,6 @@ def detect_custom_format_heuristic(dataset):
 
     mapping = {}
 
-    # Keywords
     assistant_words = [
         "output",
         "answer",
@@ -135,7 +200,6 @@ def detect_custom_format_heuristic(dataset):
         "solve",
     ]
 
-    # Split into high/low priority
     user_words_high_priority = [
         "input",
         "question",
@@ -159,10 +223,10 @@ def detect_custom_format_heuristic(dataset):
         "persona",
         "role",
         "template",
-        "task",  # Also in system
+        "task",  # also a system keyword
     ]
 
-    # Metadata columns to ignore
+    # Metadata columns to ignore.
     metadata_exact_match = {
         "id",
         "idx",
@@ -197,7 +261,7 @@ def detect_custom_format_heuristic(dataset):
     }
 
     def has_keyword(col_name, keywords):
-        """Check if any keyword appears in column name."""
+        """True if any keyword appears in the column name."""
         col_lower = col_name.lower()
         col_normalized = col_lower.replace("_", "").replace("-", "").replace(" ", "")
 
@@ -207,7 +271,7 @@ def detect_custom_format_heuristic(dataset):
         return False
 
     def is_metadata(col_name):
-        """Check if column is likely metadata."""
+        """True if the column is likely metadata."""
         col_lower = col_name.lower()
 
         if col_lower in metadata_exact_match:
@@ -229,7 +293,7 @@ def detect_custom_format_heuristic(dataset):
         return False
 
     def get_priority_score(col_name):
-        """Calculate priority score based on column name patterns."""
+        """Priority score from column-name patterns."""
         col_lower = col_name.lower()
         score = 0
 
@@ -240,7 +304,7 @@ def detect_custom_format_heuristic(dataset):
         return score
 
     def get_content_length(col_name):
-        """Get average content length for this column."""
+        """Average content length for this column."""
         try:
             if col_name in sample and sample[col_name]:
                 content = str(sample[col_name])
@@ -250,19 +314,18 @@ def detect_custom_format_heuristic(dataset):
             return 0
 
     def score_column(col_name, keywords, role_type, num_candidates):
-        """Score a column for how likely it is to be a particular role."""
+        """Score how likely a column is to be a given role."""
         if not has_keyword(col_name, keywords):
             return 0
 
         score = 0
         score += 10
 
-        # Penalize ambiguous keywords when scoring for user
+        # Penalize ambiguous "task" so other user columns win.
         if role_type == "user":
             col_lower = col_name.lower()
-            # If column is ONLY "task" (or task_xxx), give it lower priority for user role
             if "task" in col_lower and not any(kw in col_lower for kw in user_words_high_priority):
-                score -= 15  # Significant penalty so other user columns win
+                score -= 15
 
         priority_bonus = get_priority_score(col_name)
         score += priority_bonus
@@ -289,14 +352,12 @@ def detect_custom_format_heuristic(dataset):
 
         return score
 
-    # Filter out metadata columns
     content_columns = [col for col in all_columns if not is_metadata(col)]
 
-    # Count candidates first
     assistant_potential = [col for col in content_columns if has_keyword(col, assistant_words)]
     user_potential = [col for col in content_columns if has_keyword(col, user_words)]
 
-    # STEP 1: Find best ASSISTANT column
+    # STEP 1: best ASSISTANT column
     assistant_candidates = []
     for col in assistant_potential:
         score = score_column(col, assistant_words, "assistant", len(assistant_potential))
@@ -310,7 +371,7 @@ def detect_custom_format_heuristic(dataset):
     else:
         assistant_col = None
 
-    # STEP 2: Find best USER column (with penalty for ambiguous keywords)
+    # STEP 2: best USER column (penalizing ambiguous keywords)
     user_candidates = []
     for col in user_potential:
         if col == assistant_col:
@@ -326,35 +387,32 @@ def detect_custom_format_heuristic(dataset):
     else:
         user_col = None
 
-    # STEP 3: Check ALL remaining columns for SYSTEM matches (priority check)
+    # STEP 3: check remaining columns for SYSTEM matches
     remaining_columns = [col for col in content_columns if col not in mapping]
 
     system_col = None
     for col in remaining_columns:
         if has_keyword(col, system_words):
-            # Found a system match in remaining columns
             mapping[col] = "system"
             system_col = col
             break
 
-    # STEP 4: Handle any additional remaining columns
+    # STEP 4: handle any additional remaining columns
     if system_col:
         remaining_columns = [col for col in remaining_columns if col != system_col]
 
     if len(remaining_columns) >= 1:
         remaining_col = remaining_columns[0]
 
-        # If no strong keyword match, decide based on what's missing
+        # No strong keyword match: decide by what's missing.
         if not has_keyword(remaining_col, user_words + assistant_words):
             mapping[remaining_col] = "system"
         elif user_col is None:
-            # No user column yet, assign this as user
             mapping[remaining_col] = "user"
         else:
-            # Already have user + assistant, treat as system context
             mapping[remaining_col] = "system"
 
-    # VALIDATION: Ensure we have at least user + assistant
+    # Ensure at least user + assistant.
     has_user = any(role == "user" for role in mapping.values())
     has_assistant = any(role == "assistant" for role in mapping.values())
 
@@ -372,28 +430,15 @@ def detect_custom_format_heuristic(dataset):
 
 
 def detect_multimodal_dataset(dataset):
-    """
-    Detects if dataset contains multimodal data (images and/or audio).
+    """Detect multimodal data (images and/or audio) in a dataset.
 
-    Two-pass approach for each modality:
-      1. Column-name heuristic (fast): checks for keywords.
-      2. Value-type inspection (reliable): checks actual sample values.
-
-    Returns:
-        dict: {
-            "is_image": bool,
-            "multimodal_columns": list of column names containing image data,
-            "modality_types": list of detected types (e.g., ["image", "audio"]),
-            "is_audio": bool,
-            "audio_columns": list of column names containing audio data,
-            "detected_audio_column": str or None,
-            "detected_text_column": str or None,
-        }
+    Two passes per modality: column-name keyword heuristic, then value-type
+    inspection. Returns a dict with is_image/is_audio flags, detected columns,
+    modality types, and detected audio/text/speaker columns.
     """
     sample = next(iter(dataset))
     column_names = list(sample.keys())
 
-    # Keywords that indicate image data
     image_keywords = [
         "image",
         "img",
@@ -414,7 +459,6 @@ def detect_multimodal_dataset(dataset):
         "filename",
     ]
 
-    # Keywords that indicate audio data
     audio_keywords = ["audio", "speech", "wav", "waveform", "sound"]
 
     multimodal_columns = []
@@ -422,8 +466,7 @@ def detect_multimodal_dataset(dataset):
     modality_types = set()
 
     # ── Image detection ─────────────────────────────────────
-    # Pass 1: column-name heuristic (word-boundary match to avoid
-    #          false positives like 'pic' in 'topic')
+    # Pass 1: column-name heuristic (word-boundary match)
     for col_name in column_names:
         for keyword in image_keywords:
             if _keyword_in_column(keyword, col_name):
@@ -460,13 +503,13 @@ def detect_multimodal_dataset(dataset):
             audio_columns.append(col_name)
             modality_types.add("audio")
 
-    # Filter out columns that are actually audio from the image list
-    # (e.g. a column named "audio" with {"bytes", "path"} could match _is_image_value)
+    # Drop audio columns from the image list (a {"bytes","path"} audio column
+    # can match _is_image_value).
     if audio_columns:
         audio_set = set(audio_columns)
         multimodal_columns = [c for c in multimodal_columns if c not in audio_set]
 
-    # Detect text column for audio datasets
+    # Text column for audio datasets.
     detected_text_col = None
     if audio_columns:
         text_keywords = ["text", "sentence", "transcript", "transcription", "label"]
@@ -477,7 +520,7 @@ def detect_multimodal_dataset(dataset):
 
     is_audio = len(audio_columns) > 0
 
-    # Detect speaker_id column for TTS datasets (CSM, Orpheus, Spark)
+    # speaker_id column for TTS datasets (CSM, Orpheus, Spark)
     detected_speaker_col = None
     if audio_columns:
         speaker_keywords = ["source", "speaker", "speaker_id"]
@@ -503,7 +546,6 @@ def _is_image_value(value) -> bool:
     if value is None:
         return False
 
-    # PIL Image instance
     try:
         from PIL.Image import Image as PILImage
         if isinstance(value, PILImage):
@@ -511,14 +553,13 @@ def _is_image_value(value) -> bool:
     except ImportError:
         pass
 
-    # HF datasets Image feature stores decoded images as PIL or dicts with
-    # {"bytes": b"...", "path": "..."} when not yet decoded.
+    # HF Image feature: decoded as PIL, or {"bytes", "path"} when undecoded.
     # Exclude audio dicts (decoded audio has "array" + "sampling_rate").
     if isinstance(value, dict):
         if "array" in value and "sampling_rate" in value:
-            return False  # This is audio, not image
+            return False  # audio, not image
         if "bytes" in value and "path" in value:
-            # Check path extension to exclude audio files
+            # Use path extension to exclude audio files.
             path = value.get("path") or ""
             if isinstance(path, str) and any(
                 path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS
@@ -526,20 +567,17 @@ def _is_image_value(value) -> bool:
                 return False
             return True
 
-    # Raw bytes with a known image magic header
     if isinstance(value, (bytes, bytearray)):
         return _has_image_header(value)
 
-    # String that looks like an image file path or URL
+    # String that looks like an image file path or URL.
     _IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".svg")
     if isinstance(value, str) and len(value) < 1000:
         lower = value.strip().lower()
-        # Image URL (http://... ending in image extension)
         if lower.startswith(("http://", "https://")) and any(
             lower.split("?")[0].endswith(ext) for ext in _IMAGE_EXTS
         ):
             return True
-        # Image file path (relative or absolute path ending in image extension)
         if any(lower.endswith(ext) for ext in _IMAGE_EXTS):
             return True
 
@@ -564,11 +602,10 @@ def _is_audio_value(value) -> bool:
     if value is None:
         return False
 
-    # HF datasets Audio feature: decoded → {"array": np.ndarray, "sampling_rate": int}
+    # HF Audio feature: decoded -> {"array", "sampling_rate"}; undecoded -> {"bytes", "path"}.
     if isinstance(value, dict):
         if "array" in value and "sampling_rate" in value:
             return True
-        # Undecoded/streaming → {"bytes": b"...", "path": "some.wav"}
         if "bytes" in value or "path" in value:
             path = value.get("path") or ""
             if isinstance(path, str) and any(
@@ -583,28 +620,22 @@ def _has_image_header(data: bytes) -> bool:
     """Quick magic-byte check for common image formats."""
     if len(data) < 4:
         return False
-    # JPEG
-    if data[:2] == b"\xff\xd8":
+    if data[:2] == b"\xff\xd8":  # JPEG
         return True
-    # PNG
-    if data[:4] == b"\x89PNG":
+    if data[:4] == b"\x89PNG":  # PNG
         return True
-    # GIF
-    if data[:3] == b"GIF":
+    if data[:3] == b"GIF":  # GIF
         return True
-    # WebP
-    if data[:4] == b"RIFF" and len(data) >= 12 and data[8:12] == b"WEBP":
+    if data[:4] == b"RIFF" and len(data) >= 12 and data[8:12] == b"WEBP":  # WebP
         return True
-    # BMP
-    if data[:2] == b"BM":
+    if data[:2] == b"BM":  # BMP
         return True
     return False
 
 
 def detect_vlm_dataset_structure(dataset):
-    """
-    Detects if VLM dataset is:
-    - Standard VLM messages format (image objects in content)
+    """Detect which VLM dataset shape this is:
+    - Standard VLM messages (image objects in content)
     - Llava format (image indices + separate images column)
     - Simple format needing conversion (image + text columns)
     """
@@ -621,7 +652,6 @@ def detect_vlm_dataset_structure(dataset):
 
     column_names = set(sample.keys())
 
-    # Check if has messages column
     if "messages" in column_names:
         messages = sample["messages"]
 
@@ -632,7 +662,7 @@ def detect_vlm_dataset_structure(dataset):
 
                 if isinstance(content, list) and len(content) > 0:
                     if isinstance(content[0], dict) and "type" in content[0]:
-                        # Check for llava format
+                        # Llava format?
                         has_index = any(
                             "index" in item for item in content if isinstance(item, dict)
                         )
@@ -660,8 +690,8 @@ def detect_vlm_dataset_structure(dataset):
                                 "text_column": None,
                             }
 
-    # Check for ShareGPT/ChatML conversations with  placeholder + companion image column
-    # (e.g. Lin-Chen/ShareGPT4V, LLaVA-style datasets)
+    # ShareGPT/ChatML conversations with  placeholder + companion
+    # image column (e.g. Lin-Chen/ShareGPT4V, LLaVA-style datasets)
     for chat_col in ("conversations", "messages"):
         if chat_col not in column_names:
             continue
@@ -671,11 +701,10 @@ def detect_vlm_dataset_structure(dataset):
         first_msg = chat_data[0]
         if not isinstance(first_msg, dict):
             continue
-        # Detect ShareGPT (from/value) or ChatML (role/content) keys
+        # ShareGPT (from/value) or ChatML (role/content).
         msg_text = first_msg.get("value") or first_msg.get("content")
         if not isinstance(msg_text, str):
             continue
-        # Check for  placeholder anywhere in the conversation
         has_image_placeholder = any(
             "" in str(m.get("value", "") or m.get("content", ""))
             for m in chat_data
@@ -683,7 +712,7 @@ def detect_vlm_dataset_structure(dataset):
         )
         if not has_image_placeholder:
             continue
-        # Find companion image column
+        # Find companion image column.
         image_col = None
         for col in column_names:
             if col == chat_col:
@@ -700,9 +729,7 @@ def detect_vlm_dataset_structure(dataset):
                 "messages_column": chat_col,
             }
 
-    # Find image and text columns using metadata filtering
-
-    # Define metadata patterns to EXCLUDE
+    # Find image and text columns, filtering out metadata patterns
     metadata_patterns = {
         "suffixes": [
             "_id",
@@ -726,7 +753,6 @@ def detect_vlm_dataset_structure(dataset):
         ],
     }
 
-    # Image-related keywords
     image_keywords = [
         "image",
         "img",
@@ -739,7 +765,6 @@ def detect_vlm_dataset_structure(dataset):
         "filename",
     ]
 
-    # Text-related keywords
     text_keywords = [
         "text",
         "caption",
@@ -752,14 +777,11 @@ def detect_vlm_dataset_structure(dataset):
     ]
 
     def is_metadata_column(col_name):
-        """Check if column name looks like metadata."""
+        """True if the column name looks like metadata."""
         col_lower = col_name.lower()
 
-        # Check suffixes
         if any(col_lower.endswith(suffix) for suffix in metadata_patterns["suffixes"]):
             return True
-
-        # Check prefixes
         if any(col_lower.startswith(prefix) for prefix in metadata_patterns["prefixes"]):
             return True
 
@@ -767,39 +789,36 @@ def detect_vlm_dataset_structure(dataset):
 
     def _score_image_candidate(col, sample_value):
         """Score a candidate image column by how resolvable its value is."""
-        # PIL Image object (highest priority - already loaded)
+        # PIL Image (already loaded) -> highest.
         if hasattr(sample_value, "size") and hasattr(sample_value, "mode"):
             return 100
 
-        # Dict with image data (bytes/path from HF Image feature)
+        # HF Image feature dict.
         if isinstance(sample_value, dict) and ("bytes" in sample_value or "path" in sample_value):
             return 75
 
         if isinstance(sample_value, str):
-            # URL strings
-            if sample_value.startswith(("http://", "https://")):
+            if sample_value.startswith(("http://", "https://")):  # URL
                 return 70 if not is_metadata_column(col) else 55
-            # Bare file path
-            if is_metadata_column(col):
+            if is_metadata_column(col):  # bare file path
                 return 30
             return 50
 
         return 0
 
     def _probe_image_candidate(col, sample_value):
-        """Quick probe to check if an image candidate is actually reachable.
-        Returns True if likely valid, False if definitely broken."""
+        """Probe whether an image candidate is reachable (True unless definitely broken)."""
         import os
 
-        # PIL / dict — already loaded, always valid
+        # PIL / dict — already loaded.
         if not isinstance(sample_value, str):
             return True
 
-        # Local file — check it exists
+        # Local file — check it exists.
         if not sample_value.startswith(("http://", "https://")):
-            return os.path.exists(sample_value)  # bare filenames return False here, that's OK
+            return os.path.exists(sample_value)  # bare filenames return False, that's OK
 
-        # URL — quick HEAD request with short timeout
+        # URL — quick HEAD with short timeout.
         try:
             import urllib.request
 
@@ -810,11 +829,10 @@ def detect_vlm_dataset_structure(dataset):
             return False
 
     def find_image_column():
-        """Find image column by keyword match + value-based fallback.
-        When multiple candidates exist, probes them to find one that works."""
+        """Find image column by keyword match + value-based fallback, probing for one that works."""
         candidates = []
 
-        # Pass 1: keyword-matched columns
+        # Pass 1: keyword-matched columns.
         for col in column_names:
             if any(_keyword_in_column(keyword, col) for keyword in image_keywords):
                 sample_value = sample[col]
@@ -822,8 +840,8 @@ def detect_vlm_dataset_structure(dataset):
                 if score > 0:
                     candidates.append((col, score))
 
-        # Pass 2: value-based fallback — find columns with image URLs/paths
-        # even if the column name doesn't match image keywords
+        # Pass 2: value-based fallback for image URLs/paths even when the name
+        # doesn't match keywords.
         already = {c[0] for c in candidates}
         for col in column_names:
             if col in already:
@@ -831,7 +849,7 @@ def detect_vlm_dataset_structure(dataset):
             sample_value = sample[col]
             if _is_image_value(sample_value):
                 score = _score_image_candidate(col, sample_value)
-                # Slightly penalise non-keyword columns so keyword matches win on ties
+                # Penalise non-keyword columns so keyword matches win on ties.
                 candidates.append((col, max(score - 5, 1)))
 
         if not candidates:
@@ -839,48 +857,43 @@ def detect_vlm_dataset_structure(dataset):
 
         candidates.sort(key = lambda x: x[1], reverse = True)
 
-        # Single candidate or top candidate is PIL/dict — no probing needed
+        # Single candidate or top is PIL/dict — no probing needed.
         if len(candidates) == 1 or candidates[0][1] >= 75:
             return candidates[0][0]
 
-        # Multiple string-based candidates — probe to find one that actually works
+        # Multiple string candidates — probe for one that works.
         for col, score in candidates:
             sample_value = sample[col]
             if _probe_image_candidate(col, sample_value):
                 return col
 
-        # Nothing probed successfully — return highest-scored anyway and let
-        # conversion handle the error (it may still resolve via hf_hub_download)
+        # None probed OK — return highest-scored; conversion may still resolve it.
         return candidates[0][0]
 
     def find_text_column():
-        """Find text column by filtering out metadata and checking keywords."""
+        """Find text column: skip metadata, match keywords."""
         candidates = []
 
         for col in column_names:
-            # Skip metadata columns
             if is_metadata_column(col):
                 continue
 
-            # Check if contains text keywords (word-boundary match)
             if any(_keyword_in_column(keyword, col) for keyword in text_keywords):
-                # Verify it's actually text
                 sample_value = sample[col]
 
                 if isinstance(sample_value, str) and len(sample_value) > 0:
-                    # Longer text = higher priority (likely content, not just a label)
-                    priority = min(len(sample_value), 1000)  # Cap at 1000
+                    # Longer text = higher priority (content, not a label).
+                    priority = min(len(sample_value), 1000)
                     candidates.append((col, priority))
                 elif (
                     isinstance(sample_value, list)
                     and len(sample_value) > 0
                     and isinstance(sample_value[0], str)
                 ):
-                    # List of strings (e.g. captions list) — lower priority than plain strings
+                    # List of strings (e.g. captions) — lower priority than plain str.
                     priority = min(len(sample_value[0]), 1000) // 2
                     candidates.append((col, priority))
 
-        # Return highest priority candidate
         if candidates:
             candidates.sort(key = lambda x: x[1], reverse = True)
             return candidates[0][0]
diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py
index 10004ce3db..f7b35e2869 100644
--- a/studio/backend/utils/datasets/llm_assist.py
+++ b/studio/backend/utils/datasets/llm_assist.py
@@ -1,16 +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
 
-"""
-LLM-assisted dataset analysis using an ephemeral GGUF helper model.
+"""LLM-assisted dataset analysis using an ephemeral GGUF helper model.
 
-Complements heuristic-based detection in format_detection.py and
-vlm_processing.py.  Only invoked when heuristics are uncertain.
-
-Architecture:
-  - Instantiates LlamaCppBackend, loads model, runs completion(s), unloads.
-  - Not kept warm — VRAM is freed immediately after use.
-  - Gracefully degrades: returns None when unavailable (no binary, OOM, disabled).
+Complements heuristic detection (format_detection.py, vlm_processing.py); only
+invoked when heuristics are uncertain. Loads LlamaCppBackend, runs completion(s),
+unloads (VRAM freed immediately). Degrades gracefully to None when unavailable.
 """
 
 import json
@@ -33,22 +28,15 @@ README_MAX_CHARS = 1500
 
 
 def _strip_think_tags(text: str) -> str:
-    """Strip ... reasoning blocks emitted by some models.
-
-    If the model places its actual answer OUTSIDE the think block, we
-    discard the think block and keep the rest.  If the entire response
-    is INSIDE a think block (nothing useful outside), we extract and
-    return the inner content instead of discarding everything.
-    """
+    """Strip ... blocks, keeping content outside; if all inside, return the inner."""
     if "" not in text:
         return text
 
-    # Try stripping think blocks — keep content outside them
     stripped = re.sub(r".*?\s*", "", text, flags = re.DOTALL).strip()
     if stripped:
         return stripped
 
-    # Everything was inside  tags — extract the inner content of the last block
+    # Everything was inside  tags: return the last block's inner content.
     matches = re.findall(r"(.*?)", text, flags = re.DOTALL)
     if matches:
         return matches[-1].strip()
@@ -57,12 +45,9 @@ def _strip_think_tags(text: str) -> str:
 
 
 def precache_helper_gguf():
-    """
-    Pre-download the helper GGUF to HF cache.
+    """Pre-download the helper GGUF to HF cache (on startup, background thread).
 
-    Called on FastAPI startup in a background thread so subsequent
-    ``_run_with_helper()`` calls skip the download and only pay for
-    llama-server startup.  No-op if already cached or disabled.
+    Lets later ``_run_with_helper()`` calls skip the download. No-op if cached or disabled.
     """
     if os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in ("1", "true"):
         return
@@ -77,12 +62,11 @@ def precache_helper_gguf():
         disable_progress_bars()
         logging.getLogger("huggingface_hub").setLevel(logging.WARNING)
 
-        # Find the GGUF file matching the variant
         api = HfApi()
         files = api.list_repo_files(repo, repo_type = "model")
         gguf_files = [f for f in files if f.endswith(".gguf")]
 
-        # Find all GGUF files matching the variant (may be split into shards)
+        # GGUF files matching the variant (may be split into shards).
         variant_lower = variant.lower().replace("-", "_")
         matching = sorted(f for f in gguf_files if variant_lower in f.lower().replace("-", "_"))
 
@@ -106,11 +90,7 @@ def precache_helper_gguf():
 
 
 def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
-    """
-    Load helper model, run one chat completion, unload.
-
-    Returns the completion text, or None on any failure.
-    """
+    """Load helper model, run one chat completion, unload. Returns text or None on failure."""
     if os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in ("1", "true"):
         return None
 
@@ -150,7 +130,7 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
         ):
             if isinstance(chunk, dict):
                 continue  # skip metadata events
-            cumulative = chunk  # cumulative — last value is full text
+            cumulative = chunk  # last value is full text
 
         result = cumulative.strip()
         result = _strip_think_tags(result)
@@ -178,21 +158,10 @@ def llm_generate_vlm_instruction(
     samples: list[dict],
     dataset_name: Optional[str] = None,
 ) -> Optional[dict]:
+    """Ask a helper LLM for a task-specific VLM instruction (when heuristics are low-confidence).
+
+    Returns {"instruction": str, "confidence": 0.85} or None.
     """
-    Ask a helper LLM to generate a task-specific VLM instruction.
-
-    Called when heuristic instruction generation returns low confidence
-    or falls back to generic.
-
-    Args:
-        column_names: Column names in the dataset.
-        samples: 3-5 sample rows with text values (images replaced by "").
-        dataset_name: Optional HF dataset identifier for context.
-
-    Returns:
-        {"instruction": str, "confidence": 0.85} or None.
-    """
-    # Format samples for the prompt
     formatted = ""
     for i, row in enumerate(samples[:5], 1):
         parts = []
@@ -218,9 +187,8 @@ def llm_generate_vlm_instruction(
     if not result:
         return None
 
-    # Clean up: strip quotes, ensure it's a single sentence
     instruction = result.strip().strip('"').strip("'").strip()
-    # Reject obviously bad outputs (too short, too long, or multi-line)
+    # Reject bad outputs (too short, too long, or multi-line).
     if len(instruction) < 10 or len(instruction) > 200 or "\n" in instruction:
         logger.warning(f"Helper model returned unusable instruction: {instruction!r}")
         return None
@@ -233,18 +201,9 @@ def llm_generate_vlm_instruction(
 
 
 def llm_classify_columns(column_names: list[str], samples: list[dict]) -> Optional[dict[str, str]]:
-    """
-    Ask a helper LLM to classify dataset columns into roles.
+    """Ask a helper LLM to classify columns into roles (when heuristic detection fails).
 
-    Called when heuristic column detection fails (returns None).
-
-    Args:
-        column_names: Column names in the dataset.
-        samples: 3-5 sample rows with values truncated to 200 chars.
-
-    Returns:
-        Dict mapping column_name → role ("user"|"assistant"|"system"|"metadata"),
-        or None on failure.
+    Returns {column_name: role} for roles user|assistant|system|metadata, or None.
     """
     formatted = ""
     for i, row in enumerate(samples[:5], 1):
@@ -270,10 +229,9 @@ def llm_classify_columns(column_names: list[str], samples: list[dict]) -> Option
     if not result:
         return None
 
-    # Parse JSON from response (may have markdown fences)
+    # Parse JSON from response (may have markdown fences).
     text = result.strip()
     if text.startswith("```"):
-        # Strip markdown code fence
         lines = text.split("\n")
         text = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:])
         text = text.strip()
@@ -281,7 +239,6 @@ def llm_classify_columns(column_names: list[str], samples: list[dict]) -> Option
     try:
         mapping = json.loads(text)
     except json.JSONDecodeError:
-        # Try to find JSON object in the response
         import re
         match = re.search(r"\{[^}]+\}", text)
         if match:
@@ -297,7 +254,7 @@ def llm_classify_columns(column_names: list[str], samples: list[dict]) -> Option
     if not isinstance(mapping, dict):
         return None
 
-    # Validate: all values must be valid roles
+    # Keep only valid roles.
     valid_roles = {"user", "assistant", "system", "metadata"}
     cleaned = {}
     for col, role in mapping.items():
@@ -307,7 +264,7 @@ def llm_classify_columns(column_names: list[str], samples: list[dict]) -> Option
     if not cleaned:
         return None
 
-    # Must have at least user + assistant
+    # Must have at least user + assistant.
     roles_present = set(cleaned.values())
     if "user" not in roles_present or "assistant" not in roles_present:
         logger.warning(f"Helper model mapping missing user/assistant: {cleaned}")
@@ -323,19 +280,9 @@ def llm_generate_dataset_warning(
     modality: str = "text",
     column_names: Optional[list[str]] = None,
 ) -> Optional[str]:
-    """
-    Ask the helper LLM to turn technical dataset issues into a user-friendly warning.
+    """Ask the helper LLM to turn technical dataset issues into a friendly warning (any modality).
 
-    Works for all modalities (text, vision, audio).
-
-    Args:
-        issues: List of technical issue descriptions found during analysis.
-        dataset_name: Optional HF dataset name.
-        modality: "text", "vision", or "audio".
-        column_names: Optional list of column names for context.
-
-    Returns:
-        A human-friendly warning string, or None on failure.
+    Returns a human-friendly warning string, or None on failure.
     """
     if not issues:
         return None
@@ -359,7 +306,6 @@ def llm_generate_dataset_warning(
         return None
 
     warning = result.strip()
-    # Reject obviously bad outputs
     if len(warning) < 10 or len(warning) > 500:
         return None
 
@@ -419,7 +365,7 @@ def _generate_with_backend(
         top_k = 20,
         max_tokens = max_tokens,
         repetition_penalty = 1.0,
-        enable_thinking = False,  # Always disable thinking for AI Assist
+        enable_thinking = False,  # disable thinking for AI Assist
     ):
         if isinstance(chunk, dict):
             continue  # skip metadata events
@@ -432,12 +378,7 @@ def _generate_with_backend(
 def fetch_hf_dataset_card(
     dataset_name: str, hf_token: Optional[str] = None
 ) -> tuple[Optional[str], Optional[dict]]:
-    """
-    Fetch HF dataset card (README) and metadata.
-
-    Returns:
-        (readme_text, metadata_dict) or (None, None) on failure.
-    """
+    """Fetch HF dataset card (README) and metadata. Returns (readme, metadata) or (None, None)."""
     try:
         from huggingface_hub import DatasetCard
 
@@ -452,7 +393,7 @@ def fetch_hf_dataset_card(
             else:
                 readme = readme[:README_MAX_CHARS] + "\n[...truncated]"
 
-        # Extract metadata from YAML frontmatter
+        # Extract metadata from YAML frontmatter.
         metadata = {}
         if card.data:
             for key in (
@@ -486,10 +427,9 @@ def _run_multi_pass_advisor(
     model_type: Optional[str] = None,
     hf_token: Optional[str] = None,
 ) -> Optional[dict[str, Any]]:
-    """
-    Multi-pass LLM analysis: classify → convert → validate.
+    """Multi-pass LLM analysis (classify -> convert -> validate), model loaded across passes.
 
-    Keeps model loaded across all passes. Returns combined result dict or None.
+    Returns combined result dict or None.
     """
     if os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in ("1", "true"):
         return None
@@ -619,7 +559,7 @@ def _run_multi_pass_advisor(
             logger.warning(f"Advisor Pass 1 failed to produce JSON: {raw1[:200]}")
             return None
 
-        # If dataset is already conversational, skip passes 2-3
+        # Already conversational: skip passes 2-3.
         if pass1.get("is_conversational") and not pass1.get("needs_conversion"):
             return {
                 "success": True,
@@ -724,11 +664,11 @@ def _run_multi_pass_advisor(
         column_roles = pass2.get("column_roles", {})
         label_map = pass2.get("label_mapping") or {}  # may be null
 
-        # Validate: must have at least one user AND one assistant
+        # Must have at least one user AND one assistant
         roles_present = set(column_roles.values())
         if "user" not in roles_present or "assistant" not in roles_present:
             logger.warning(f"Pass 2 sanity fail: missing user or assistant role: {column_roles}")
-            return None  # triggers fallback to simple classification
+            return None  # falls back to simple classification
 
         # ── Pass 3: System prompt (non-conversational datasets only) ──
         sys_prompt = ""
@@ -739,7 +679,7 @@ def _run_multi_pass_advisor(
             logger.info("Pass 3: Generating system prompt...")
             t3 = time.monotonic()
 
-            # Format label mapping info for the prompt
+            # Format label mapping for the prompt.
             label_info = ""
             if label_map:
                 for col, mapping in label_map.items():
@@ -747,7 +687,6 @@ def _run_multi_pass_advisor(
                         pairs = ", ".join(f"{k} = {v}" for k, v in mapping.items())
                         label_info += f"\nLabel mapping for '{col}': {pairs}"
 
-            # Describe the role assignments for context
             user_cols = [c for c, r in column_roles.items() if r == "user"]
             asst_cols = [c for c, r in column_roles.items() if r == "assistant"]
             task_desc = pass1.get("task_description") or pass1.get("description", "")
@@ -782,18 +721,18 @@ def _run_multi_pass_advisor(
             )
 
             if raw3:
-                # Pass 3 returns raw text, not JSON — clean it up
+                # Pass 3 returns raw text, not JSON.
                 cleaned = raw3.strip().strip('"').strip("'").strip()
                 if len(cleaned) >= 20 and cleaned.lower() not in ("null", "none", ""):
                     sys_prompt = cleaned
 
-        # Build suggested_mapping (column → role, for the frontend dropdowns)
+        # Build suggested_mapping (column -> role) for the frontend dropdowns.
         suggested_mapping = {}
         for col, role in column_roles.items():
             if col in columns and role in ("user", "assistant", "system"):
                 suggested_mapping[col] = role
 
-        # Build user notification from Pass 1 classification
+        # Build user notification from Pass 1 classification.
         desc = pass1.get("task_description") or pass1.get("description", "")
         note_parts = [f"This is a {dtype} dataset (not conversational)."]
         if desc:
@@ -839,23 +778,18 @@ def llm_conversion_advisor(
     model_name: Optional[str] = None,
     model_type: Optional[str] = None,
 ) -> Optional[dict[str, Any]]:
-    """
-    Full conversion advisor: fetch HF card → multi-pass LLM analysis.
+    """Full conversion advisor: fetch HF card -> multi-pass LLM analysis.
 
     Falls back to simple llm_classify_columns() if the multi-pass advisor fails.
-
-    Returns:
-        Dict with keys: success, suggested_mapping, system_prompt, user_template,
-        assistant_template, label_mapping, dataset_type, is_conversational,
-        user_notification. Or None on complete failure.
+    Returns a result dict (success, suggested_mapping, system_prompt, label_mapping,
+    dataset_type, is_conversational, user_notification, ...) or None.
     """
-    # Fetch HF dataset card if this looks like a HF dataset (has a slash)
+    # Fetch HF dataset card if this looks like a HF dataset (has a slash).
     dataset_card = None
     dataset_metadata = None
     if dataset_name and "/" in dataset_name:
         dataset_card, dataset_metadata = fetch_hf_dataset_card(dataset_name, hf_token)
 
-    # Try multi-pass advisor
     result = _run_multi_pass_advisor(
         columns = column_names,
         samples = samples,
@@ -871,7 +805,7 @@ def llm_conversion_advisor(
         logger.info(f"Conversion advisor succeeded: type={result.get('dataset_type')}")
         return result
 
-    # Fallback: simple column classification
+    # Fallback: simple column classification.
     logger.info("Advisor failed, falling back to simple column classification")
     simple_mapping = llm_classify_columns(column_names, samples)
     if simple_mapping:
diff --git a/studio/backend/utils/datasets/model_mappings.py b/studio/backend/utils/datasets/model_mappings.py
index eb2e5482c9..463d26a692 100644
--- a/studio/backend/utils/datasets/model_mappings.py
+++ b/studio/backend/utils/datasets/model_mappings.py
@@ -1,11 +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
 
-"""
-Model and template mappings for dataset processing.
+"""Model and template mappings for dataset processing.
 
-This module contains the mapping dictionaries that associate model names
-with their corresponding chat templates and response markers.
+Maps model names to their chat templates and response markers.
 """
 
 TEMPLATE_TO_MODEL_MAPPER = {
@@ -436,7 +434,7 @@ for key, values in TEMPLATE_TO_MODEL_MAPPER.items():
     for value in values:
         MODEL_TO_TEMPLATE_MAPPER[value] = key
 
-    # Get lowercased
+    # Also map lowercased names.
     lowered_key = key.lower()
     for value in values:
         MODEL_TO_TEMPLATE_MAPPER[value.lower()] = lowered_key
@@ -445,8 +443,8 @@ for key, values in TEMPLATE_TO_MODEL_MAPPER.items():
 def is_gpt_oss_model_name(name: str) -> bool:
     """Name-based check for gpt-oss / harmony models.
 
-    Used by both the in-process backend and the parent-process
-    orchestrator to detect harmony models without an IPC round-trip.
+    Used by the in-process backend and the parent orchestrator to detect
+    harmony models without an IPC round-trip.
     """
     name = (name or "").lower()
     if not name:
diff --git a/studio/backend/utils/datasets/raw_text.py b/studio/backend/utils/datasets/raw_text.py
index 86b1963fc1..03315fb287 100644
--- a/studio/backend/utils/datasets/raw_text.py
+++ b/studio/backend/utils/datasets/raw_text.py
@@ -1,9 +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
 
-"""
-Shared helpers for raw-text dataset preparation.
-"""
+"""Shared helpers for raw-text dataset preparation."""
 
 from dataclasses import dataclass
 from typing import Literal
diff --git a/studio/backend/utils/datasets/vlm_processing.py b/studio/backend/utils/datasets/vlm_processing.py
index a0f1fd9f99..f018913fa8 100644
--- a/studio/backend/utils/datasets/vlm_processing.py
+++ b/studio/backend/utils/datasets/vlm_processing.py
@@ -4,8 +4,8 @@
 """
 VLM (Vision-Language Model) processing utilities.
 
-This module contains functions for generating smart instructions
-for VLM datasets based on content analysis and heuristics.
+Generates smart instructions for VLM datasets via content analysis and
+heuristics.
 """
 
 import re
@@ -19,19 +19,19 @@ def generate_smart_vlm_instruction(
     dataset_name = None,
 ):
     """
-    Generate smart, context-aware instruction for VLM datasets using heuristics.
+    Generate a smart, context-aware instruction for VLM datasets via heuristics.
 
     Strategy:
-    1. Check for explicit question/instruction columns → use that
+    1. Explicit question/instruction column → use that
     2. Infer from text column name + sample content
     3. Analyze dataset name for task hints
-    4. Fall back to generic instruction
+    4. Generic fallback
 
     Returns:
         dict: {
             "instruction": str or None,  # None means use column content
             "instruction_type": "explicit" | "inferred" | "generic",
-            "uses_dynamic_instruction": bool,  # True if instruction varies per sample
+            "uses_dynamic_instruction": bool,  # True if it varies per sample
             "confidence": float,  # 0.0 to 1.0
         }
     """
@@ -39,16 +39,16 @@ def generate_smart_vlm_instruction(
     sample = next(iter(dataset))
 
     # ===== LEVEL 1: Explicit Instruction Columns =====
-    # Check for columns that contain per-sample instructions
+    # Columns that hold per-sample instructions
     question_columns = ["question", "query", "prompt", "instruction", "user_prompt"]
 
     for col in question_columns:
         if col in column_names:
-            # Check if this column has varied content (not just empty/same)
+            # Use it only if it has non-empty content
             sample_content = sample[col]
             if sample_content and str(sample_content).strip():
                 return {
-                    "instruction": None,  # Signal to use column content
+                    "instruction": None,  # use column content
                     "instruction_column": col,
                     "instruction_type": "explicit",
                     "uses_dynamic_instruction": True,
@@ -58,7 +58,6 @@ def generate_smart_vlm_instruction(
     # ===== LEVEL 2: Infer from Column Names + Content =====
     text_col_lower = text_column.lower()
 
-    # Sample the text content to detect patterns
     text_sample = str(sample.get(text_column, ""))[:500]  # First 500 chars
 
     # Task-specific keywords and their instructions
@@ -66,7 +65,7 @@ def generate_smart_vlm_instruction(
         # OCR / Transcription
         "ocr": {
             "keywords": ["ocr", "transcribe", "transcript"],
-            "content_hints": [r"[A-Za-z\u0600-\u06FF]{10,}"],  # Long text passages (Latin/Arabic)
+            "content_hints": [r"[A-Za-z\u0600-\u06FF]{10,}"],  # Long Latin/Arabic passages
             "instruction": "Transcribe all the text shown in this image.",
             "confidence": 0.9,
         },
@@ -122,24 +121,21 @@ def generate_smart_vlm_instruction(
         },
     }
 
-    # Check column name matches
+    # Score each task by column/dataset name and content matches
     best_match = None
     best_score = 0.0
 
     for task_name, task_info in task_patterns.items():
         score = 0.0
 
-        # Check column name
         if any(keyword in text_col_lower for keyword in task_info["keywords"]):
             score += 0.5
 
-        # Check dataset name if provided
         if dataset_name and any(
             keyword in dataset_name.lower() for keyword in task_info["keywords"]
         ):
             score += 0.3
 
-        # Check content patterns
         for pattern in task_info["content_hints"]:
             if re.search(pattern, text_sample, re.IGNORECASE):
                 score += 0.4
@@ -162,7 +158,6 @@ def generate_smart_vlm_instruction(
     if dataset_name:
         name_lower = dataset_name.lower()
 
-        # Common dataset name patterns
         if "vqa" in name_lower or "question" in name_lower:
             return {
                 "instruction": "Answer the question about this image.",
diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py
index 400b5dd066..5f2b2abbcf 100644
--- a/studio/backend/utils/hardware/__init__.py
+++ b/studio/backend/utils/hardware/__init__.py
@@ -1,9 +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
 
-"""
-Hardware detection and GPU utilities
-"""
+"""Hardware detection and GPU utilities."""
 
 from . import hardware as _hardware
 from .hardware import (
@@ -86,8 +84,8 @@ __all__ = [
 
 
 def __getattr__(name: str):
-    """Resolve IS_ROCM at access time so callers always see the live value
-    after detect_hardware() runs (it flips the flag in hardware.py)."""
+    """Resolve IS_ROCM lazily so callers see the live value detect_hardware()
+    sets in hardware.py."""
     if name == "IS_ROCM":
         return getattr(_hardware, "IS_ROCM")
     raise AttributeError(name)
diff --git a/studio/backend/utils/hardware/amd.py b/studio/backend/utils/hardware/amd.py
index 04b9494a29..39aafe0489 100644
--- a/studio/backend/utils/hardware/amd.py
+++ b/studio/backend/utils/hardware/amd.py
@@ -3,9 +3,8 @@
 
 """AMD GPU monitoring via amd-smi.
 
-Mirrors the nvidia.py module structure so hardware.py can swap backends
-based on IS_ROCM. All functions return the same dict shapes as their
-nvidia.py counterparts.
+Mirrors nvidia.py so hardware.py can swap backends based on IS_ROCM.
+All functions return the same dict shapes as their nvidia.py counterparts.
 """
 
 import json
@@ -23,20 +22,19 @@ from utils.subprocess_compat import windows_hidden_subprocess_kwargs
 
 logger = get_logger(__name__)
 
-# amd-smi on Windows must initialise the full ROCm runtime on first call, which
-# can take 15-25 s on cold hardware.  Linux is consistently < 2 s.
+# amd-smi on Windows initialises the full ROCm runtime on first call, which
+# can take 15-25 s on cold hardware. Linux is consistently < 2 s.
 _AMD_SMI_DEFAULT_TIMEOUT = 30 if platform.system() == "Windows" else 10
 
-# Circuit breaker: stop calling amd-smi after this many consecutive failures.
-# On Windows, each failed call spawns a process that may show a UAC/DiskPart
-# elevation prompt.  Once we know amd-smi doesn't work we stop polling it.
+# Circuit breaker: stop polling amd-smi after this many consecutive failures
+# (each Windows failure may pop a UAC/DiskPart elevation prompt).
 _AMD_SMI_FAILURE_LIMIT = 3
 _amd_smi_consecutive_failures = 0
 _amd_smi_disabled = False
 
 
 def _run_amd_smi(*args: str, timeout: int = _AMD_SMI_DEFAULT_TIMEOUT) -> Optional[Any]:
-    """Run amd-smi with the given arguments and return parsed JSON, or None."""
+    """Run amd-smi with the given args and return parsed JSON, or None."""
     global _amd_smi_consecutive_failures, _amd_smi_disabled
     if _amd_smi_disabled:
         return None
@@ -51,8 +49,8 @@ def _run_amd_smi(*args: str, timeout: int = _AMD_SMI_DEFAULT_TIMEOUT) -> Optiona
         )
     except (OSError, subprocess.TimeoutExpired) as e:
         if isinstance(e, FileNotFoundError):
-            # amd-smi ships with Adrenalin, not the HIP SDK -- absence is
-            # expected on HIP SDK-only Windows setups.  Log at debug only.
+            # amd-smi ships with Adrenalin, not the HIP SDK; absence is expected
+            # on HIP SDK-only Windows setups.
             logger.debug("amd-smi not found (not in PATH): %s", e)
         else:
             logger.warning("amd-smi query failed: %s", e)
@@ -75,9 +73,8 @@ def _run_amd_smi(*args: str, timeout: int = _AMD_SMI_DEFAULT_TIMEOUT) -> Optiona
             _amd_smi_disabled = True
         return None
     if not result.stdout.strip():
-        # amd-smi exited successfully but produced no output (e.g. no GPUs
-        # visible on this query, or a version that emits nothing for --json).
-        # This is not a tool failure, so don't count against the circuit breaker.
+        # Exit 0 with no output (no GPUs visible, or a version emitting nothing
+        # for --json). Not a tool failure, so don't trip the circuit breaker.
         logger.debug("amd-smi exited 0 but returned no output")
         return None
     _amd_smi_consecutive_failures = 0  # reset on success
@@ -89,7 +86,7 @@ def _run_amd_smi(*args: str, timeout: int = _AMD_SMI_DEFAULT_TIMEOUT) -> Optiona
 
 
 def _parse_numeric(value: Any) -> Optional[float]:
-    """Extract a numeric value from amd-smi output (may be str, int, float, or dict)."""
+    """Extract a numeric value from amd-smi output (str, int, float, or dict)."""
     if value is None:
         return None
     # Newer amd-smi versions emit {"value": 10, "unit": "W"}
@@ -114,9 +111,8 @@ def _parse_memory_mb(value: Any) -> Optional[float]:
     """Parse a memory value from amd-smi output and return MB.
 
     Handles bare numbers (assumed MB -- the amd-smi convention on every
-    version we have seen), dict-shaped values with explicit units
-    (``{"value": 192, "unit": "GiB"}`` on newer releases), and plain
-    strings like ``"8192 MiB"``.
+    version seen), dict values with explicit units (``{"value": 192,
+    "unit": "GiB"}`` on newer releases), and strings like ``"8192 MiB"``.
     """
     unit = ""
     raw_value = value
@@ -134,8 +130,8 @@ def _parse_memory_mb(value: Any) -> Optional[float]:
     if num is None:
         return None
 
-    # Unit conversion -- GPU tools (including amd-smi) use binary units even
-    # when labeling them "GB" or "MB", so treat GB/GiB and MB/MiB the same.
+    # GPU tools use binary units even when labeled "GB"/"MB", so treat GB/GiB
+    # and MB/MiB the same.
     if "gib" in unit or "gb" in unit:
         return num * 1024
     if "mib" in unit or "mb" in unit:
@@ -146,27 +142,23 @@ def _parse_memory_mb(value: Any) -> Optional[float]:
         # Plain bytes
         return num / (1024 * 1024)
 
-    # No explicit unit -- default to MB, which is the amd-smi convention
-    # for bare numeric values. A previous heuristic assumed values above
-    # ~10M were bytes, but that misclassifies small VRAM allocations
-    # (e.g. 5 MB = 5,242,880 reported without a unit) as ~5 TB. Modern
-    # amd-smi always ships explicit units, so the heuristic branch only
-    # fired for legacy output where MB was already the convention.
+    # No explicit unit: default to MB (the amd-smi convention for bare numbers).
+    # A bytes-above-~10M heuristic was dropped because it misclassified small
+    # VRAM allocations; modern amd-smi always ships explicit units.
     return num
 
 
 def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]:
     """Extract standardized metrics from a single GPU's amd-smi data."""
-    # amd-smi metric output structure varies by version; try common paths
+    # Output structure varies by version; try common paths
     usage = gpu_data.get("usage", gpu_data.get("gpu_activity", {}))
     if isinstance(usage, dict):
         gpu_util = _parse_numeric(usage.get("gfx_activity", usage.get("gpu_use_percent")))
     else:
         gpu_util = _parse_numeric(usage)
 
-    # Temperature -- try multiple keys in priority order.
-    # dict.get() returns "N/A" strings rather than falling through,
-    # so we must try each key and check if it parses to a real number.
+    # Temperature: try keys in priority order, checking each parses to a real
+    # number (dict.get() can return "N/A" strings rather than falling through).
     temp_data = gpu_data.get("temperature", {})
     temp = None
     if isinstance(temp_data, dict):
@@ -191,10 +183,9 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]:
         power_draw = None
         power_limit = None
 
-    # VRAM -- unit-aware parsing to handle varying amd-smi output formats.
-    # Newer amd-smi versions may return {"value": 192, "unit": "GiB"}.
-    # Newer amd-smi uses "mem_usage" with "total_vram" / "used_vram" keys;
-    # older versions use "vram" or "fb_memory_usage" with "used" / "total".
+    # VRAM: unit-aware parsing across amd-smi formats. Newer versions use
+    # "mem_usage" with "total_vram"/"used_vram"; older use "vram" or
+    # "fb_memory_usage" with "used"/"total".
     vram_data = gpu_data.get(
         "mem_usage",
         gpu_data.get("vram", gpu_data.get("fb_memory_usage", {})),
@@ -237,13 +228,11 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]:
 
 
 def _has_real_metrics(metrics: dict[str, Any]) -> bool:
-    """Return True when ``metrics`` contains at least one non-None value.
+    """Return True when ``metrics`` has at least one non-None value.
 
-    ``amd-smi`` can return a zero-exit JSON envelope that is missing every
-    expected field (error response, unsupported card, hipless container).
-    In that case ``_extract_gpu_metrics`` produces a dict where every value
-    is ``None`` -- callers must surface this as ``available: False`` rather
-    than ``available: True`` with empty data.
+    amd-smi can return a zero-exit envelope missing every field (error,
+    unsupported card, hipless container), yielding an all-None dict; callers must
+    surface that as ``available: False``.
     """
     return any(value is not None for value in metrics.values())
 
@@ -255,9 +244,8 @@ def get_physical_gpu_count() -> Optional[int]:
         return None
     if isinstance(data, list):
         return len(data)
-    # Some versions return a dict with a "gpu" / "gpus" key. Guard the
-    # .get() access with an isinstance check so a malformed scalar /
-    # string response from amd-smi cannot raise AttributeError.
+    # Some versions return a dict with a "gpu"/"gpus" key; guard with isinstance
+    # so a malformed scalar/string response can't raise AttributeError.
     if not isinstance(data, dict):
         return None
     gpus = data.get("gpu", data.get("gpus", []))
@@ -267,12 +255,12 @@ def get_physical_gpu_count() -> Optional[int]:
 
 
 def _first_visible_amd_gpu_id() -> Optional[str]:
-    """Return the physical AMD GPU id that should be treated as 'primary'.
+    """Return the physical AMD GPU id treated as 'primary'.
 
     Honours HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES / CUDA_VISIBLE_DEVICES
-    in that order (HIP respects all three). Returns ``"0"`` when none are
-    set, and ``None`` when the env var explicitly narrows to zero GPUs
-    ("" or "-1"), so callers can short-circuit to "available: False".
+    in that order (HIP respects all three). Returns ``"0"`` when none are set,
+    and ``None`` when the env var narrows to zero GPUs ("" or "-1"), so callers
+    can short-circuit to "available: False".
     """
     for env_name in (
         "HIP_VISIBLE_DEVICES",
@@ -285,11 +273,8 @@ def _first_visible_amd_gpu_id() -> Optional[str]:
         raw = raw.strip()
         if raw == "" or raw == "-1":
             return None
-        # Filter out empty tokens after splitting. This tolerates minor
-        # typos like ``HIP_VISIBLE_DEVICES=",1"`` (leading comma, user
-        # clearly meant to narrow to device 1) while still falling
-        # through to the next env var when every token is empty
-        # (e.g. ``,,,``).
+        # Drop empty tokens, tolerating typos like ``",1"`` while still falling
+        # through to the next env var when every token is empty (``,,,``).
         tokens = [t.strip() for t in raw.split(",") if t.strip()]
         if tokens:
             return tokens[0]
@@ -320,10 +305,8 @@ def get_primary_gpu_utilization() -> dict[str, Any]:
 
     metrics = _extract_gpu_metrics(gpu_data)
     if not _has_real_metrics(metrics):
-        # amd-smi returned a JSON envelope with no usable fields (error
-        # response or unsupported card). Surface as unavailable rather
-        # than available-with-empty-data so the UI does not render a
-        # ghost device.
+        # Envelope with no usable fields: surface as unavailable so the UI
+        # doesn't render a ghost device.
         return {"available": False}
     metrics["available"] = True
     return metrics
@@ -352,15 +335,11 @@ def get_visible_gpu_utilization(
             "index_kind": "physical",
         }
 
-    # Extract a device list from amd-smi's envelope. Newer versions return
-    # a JSON array directly, older versions return a dict with a "gpus" /
-    # "gpu" key wrapping the list. Guard non-dict / non-list envelopes
-    # (scalar / string fallbacks from malformed output) so the .get()
-    # access cannot raise AttributeError on an unexpected shape.
+    # Extract a device list across envelope shapes: a JSON array, a dict under
+    # "gpu_data"/"gpus"/"gpu", or a guarded scalar/string fallback.
     if isinstance(data, list):
         gpu_list = data
     elif isinstance(data, dict):
-        # Newer amd-smi wraps output in {"gpu_data": [...]}
         gpu_list = data.get("gpu_data", data.get("gpus", data.get("gpu", [data])))
     else:
         gpu_list = [data]
@@ -369,17 +348,11 @@ def get_visible_gpu_utilization(
 
     devices = []
     for fallback_idx, gpu_data in enumerate(gpu_list):
-        # Skip non-dict entries defensively: if amd-smi ever ships a
-        # scalar inside its "gpus" array (observed on some malformed
-        # output), _extract_gpu_metrics would raise AttributeError on
-        # the first .get() call.
+        # Skip non-dict entries (a scalar in the array would raise AttributeError).
         if not isinstance(gpu_data, dict):
             continue
-        # Use AMD-reported GPU ID when available, fall back to enumeration
-        # index. Newer amd-smi versions wrap scalars as ``{"value": 0,
-        # "unit": "none"}``, so route raw_id through ``_parse_numeric``
-        # which already handles bare ints, floats, strings, and that
-        # dict shape uniformly.
+        # Use the AMD-reported GPU ID, else the enumeration index. _parse_numeric
+        # handles bare ints/floats/strings and the {"value", "unit"} dict shape.
         raw_id = gpu_data.get("gpu", gpu_data.get("gpu_id", gpu_data.get("id", fallback_idx)))
         parsed_id = _parse_numeric(raw_id)
         if parsed_id is None:
@@ -403,10 +376,8 @@ def get_visible_gpu_utilization(
             continue
         metrics = _extract_gpu_metrics(gpu_data)
         if not _has_real_metrics(metrics):
-            # Skip ghost entries: an amd-smi response that decodes to a
-            # dict but contains no usable fields (error envelope, etc.)
-            # would otherwise show up as a device row with all-None
-            # numbers in the UI.
+            # Skip ghost entries (no usable fields) so the UI doesn't show an
+            # all-None device row.
             continue
         metrics["index"] = idx
         metrics["index_kind"] = "physical"
diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py
index 3f2823302a..7292fa185f 100644
--- a/studio/backend/utils/hardware/hardware.py
+++ b/studio/backend/utils/hardware/hardware.py
@@ -39,7 +39,7 @@ logger = get_logger(__name__)
 
 
 class DeviceType(str, Enum):
-    """Supported compute backends. Inherits from str so it serializes cleanly in JSON."""
+    """Supported compute backends. str subclass for clean JSON serialization."""
 
     CUDA = "cuda"
     XPU = "xpu"
@@ -57,14 +57,8 @@ IS_ROCM: bool = False  # True when running on AMD ROCm (HIP) -- routes GPU monit
 def _backend_label(device: DeviceType) -> str:
     """Return the user-facing backend name for API responses.
 
-    Internally we still represent ROCm hosts as ``DeviceType.CUDA`` because
-    ROCm torch sets ``torch.cuda.is_available() = True`` and reuses the whole
-    ``torch.cuda.*`` API surface, so branching on ``DeviceType`` stays
-    consistent with the rest of the codebase. For the JSON responses served
-    to the Studio frontend and other clients, however, "cuda" is misleading
-    on an AMD machine. This helper swaps the label to ``"rocm"`` when the
-    module-level ``IS_ROCM`` flag is set so the UI can render the correct
-    backend name without every caller having to duplicate the check.
+    ROCm hosts stay ``DeviceType.CUDA`` internally (ROCm reuses ``torch.cuda.*``),
+    but "cuda" is misleading in JSON, so swap to ``"rocm"`` when ``IS_ROCM`` is set.
     """
     if IS_ROCM and device == DeviceType.CUDA:
         return "rocm"
@@ -75,12 +69,12 @@ def _backend_label(device: DeviceType) -> str:
 
 
 def is_apple_silicon() -> bool:
-    """Check if running on Apple Silicon hardware (pure platform check, no ML imports)."""
+    """True on Apple Silicon (pure platform check, no ML imports)."""
     return platform.system() == "Darwin" and platform.machine() == "arm64"
 
 
 def _has_torch() -> bool:
-    """Check if PyTorch is importable."""
+    """True if PyTorch is importable."""
     try:
         import torch
         return True
@@ -89,7 +83,7 @@ def _has_torch() -> bool:
 
 
 def _has_mlx() -> bool:
-    """Check if MLX is importable."""
+    """True if MLX is importable."""
     try:
         import mlx.core
         return True
@@ -99,10 +93,9 @@ def _has_mlx() -> bool:
 
 def detect_hardware() -> DeviceType:
     """
-    Detect the best available compute device and set the module-level DEVICE global.
+    Detect the best compute device and set the module-level DEVICE global.
 
-    Should be called exactly once during FastAPI lifespan startup.
-    Safe to call multiple times (idempotent).
+    Call once at FastAPI lifespan startup; idempotent.
 
     Detection order:
       1. CUDA  (NVIDIA GPU, requires torch)
@@ -121,10 +114,8 @@ def detect_hardware() -> DeviceType:
             CHAT_ONLY = False
             device_name = torch.cuda.get_device_properties(0).name
 
-            # Distinguish AMD ROCm (HIP) from NVIDIA CUDA for display purposes.
-            # DeviceType stays CUDA since torch.cuda.* works on ROCm via HIP.
-            # AMD's repo.radeon.com SDK wheels (e.g. 2.9.0+rocmsdk20251116) do
-            # not set torch.version.hip, so fall back to checking __version__.
+            # Distinguish ROCm from CUDA for display only (DeviceType stays CUDA).
+            # AMD SDK wheels don't set torch.version.hip, so fall back to __version__.
             _hip_ver = getattr(torch.version, "hip", None)
             if _hip_ver is not None or "rocm" in torch.__version__.lower():
                 IS_ROCM = True
@@ -148,9 +139,8 @@ def detect_hardware() -> DeviceType:
     if is_apple_silicon() and _has_mlx():
         DEVICE = DeviceType.MLX
         CHAT_ONLY = False
-        # platform.processor() runs `uname -p` which returns "i386" on most
-        # universal2 / Rosetta-shaped Python builds even on native arm64.
-        # platform.machine() is "arm64" once is_apple_silicon() has gated us.
+        # Use platform.machine() ("arm64"); platform.processor() returns "i386"
+        # on universal2 / Rosetta builds even on native arm64.
         chip = platform.machine() or "arm64"
         print(f"Hardware detected: MLX — Apple Silicon ({chip})")
         return DEVICE
@@ -166,8 +156,8 @@ def detect_hardware() -> DeviceType:
 
 def get_device() -> DeviceType:
     """
-    Return the detected device. Auto-detects if detect_hardware() hasn't been called yet.
-    Prefer calling detect_hardware() explicitly at startup instead.
+    Return the detected device, auto-detecting if detect_hardware() hasn't run.
+    Prefer calling detect_hardware() explicitly at startup.
     """
     global DEVICE
     if DEVICE is None:
@@ -178,7 +168,7 @@ def get_device() -> DeviceType:
 def clear_gpu_cache():
     """
     Clear GPU memory cache for the current device.
-    Safe to call on any platform — no-ops gracefully.
+    Safe on any platform — no-ops gracefully.
     """
     gc.collect()
 
@@ -195,15 +185,14 @@ def clear_gpu_cache():
         torch.xpu.synchronize()
         torch.xpu.empty_cache()
     elif device == DeviceType.MLX:
-        # MLX manages memory automatically; no explicit cache clear needed.
-        # mlx.core has no empty_cache equivalent — gc.collect() above is enough.
+        # MLX manages memory automatically; gc.collect() above is enough.
         pass
 
 
 def get_gpu_memory_info() -> Dict[str, Any]:
     """
-    Get GPU memory information.
-    Supports CUDA (NVIDIA), MLX (Apple Silicon), and CPU-only environments.
+    Get GPU memory info.
+    Supports CUDA (NVIDIA), MLX (Apple Silicon), and CPU-only.
     """
     device = get_device()
 
@@ -275,16 +264,14 @@ def get_gpu_memory_info() -> Dict[str, Any]:
             import mlx.core as mx
             import psutil
 
-            # MLX uses unified memory. Total = system RAM. GPU memory used
-            # comes from IORegistry's AGXAccelerator (system-wide, no sudo).
+            # Unified memory: total = system RAM, GPU used from IORegistry AGX.
             total = psutil.virtual_memory().total
             agx = _read_apple_gpu_stats()
             allocated = agx.get("vram_used_bytes", 0) if agx else 0
 
             try:
                 info = mx.device_info()
-                # See detect_hardware(): platform.processor() can return "i386"
-                # on native arm64 Python builds, so prefer machine() as fallback.
+                # prefer machine(); processor() can return "i386" on native arm64.
                 gpu_name = info.get("device_name") or platform.machine() or "arm64"
             except Exception:
                 gpu_name = platform.machine() or "arm64"
@@ -352,13 +339,11 @@ def get_gpu_summary() -> Dict[str, Any]:
 
 def get_package_versions() -> Dict[str, Optional[str]]:
     """
-    Return the installed versions of key ML packages.
+    Return installed versions of key ML packages.
 
-    Uses importlib.metadata (stdlib) so no subprocess is needed.
-    CUDA version comes from torch.version.cuda.
-
-    Returns dict with keys: unsloth, torch, transformers, cuda.
-    Missing packages yield None.
+    Uses importlib.metadata (stdlib), no subprocess. CUDA version from
+    torch.version.cuda. Returns dict keyed unsloth/torch/transformers/cuda;
+    missing packages yield None.
     """
     packages = ("unsloth", "torch", "transformers")
     versions: Dict[str, Optional[str]] = {}
@@ -415,11 +400,10 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any]
     devices = []
     for ordinal, phys_idx in enumerate(device_indices):
         try:
-            # torch uses 0-based ordinals relative to CUDA_VISIBLE_DEVICES
+            # torch ordinals are 0-based relative to CUDA_VISIBLE_DEVICES.
             props = mod.get_device_properties(ordinal)
             total_bytes = props.total_memory
-            # Prefer mem_get_info (reports system-wide usage, not just this
-            # process) so auto-selection accounts for other GPU consumers.
+            # Prefer mem_get_info (system-wide) so auto-select sees other consumers.
             if hasattr(mod, "mem_get_info"):
                 free_bytes, total_bytes = mod.mem_get_info(ordinal)
                 used_bytes = total_bytes - free_bytes
@@ -443,9 +427,9 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any]
 
 
 def _smi_query(func_name: str, *args, **kwargs) -> Optional[Dict[str, Any]]:
-    """Run a query against the appropriate SMI backend (amd-smi or nvidia-smi).
+    """Query the appropriate SMI backend (amd-smi or nvidia-smi).
 
-    Returns the result dict if available, or None on failure/unavailability.
+    Returns the result dict if available, else None.
     """
     if IS_ROCM:
         backend_name = "amd-smi"
@@ -474,8 +458,8 @@ def _smi_query(func_name: str, *args, **kwargs) -> Optional[Dict[str, Any]]:
 def _read_apple_gpu_stats() -> Dict[str, Any]:
     """Query macOS IORegistry for AGX (Apple GPU) live stats. No sudo needed.
 
-    Returns dict with utilization_pct, vram_used_bytes (system-wide GPU memory).
-    Returns empty dict on failure.
+    Returns dict with utilization_pct, vram_used_bytes (system-wide GPU
+    memory), or empty dict on failure.
     """
     try:
         result = subprocess.run(
@@ -574,7 +558,7 @@ def _rocm_windows_perf_counter_gpu_util_pct() -> Optional[float]:
 def _rocm_linux_sysfs_vram_gb() -> tuple[Optional[float], Optional[float]]:
     """Query system-wide AMD GPU VRAM via Linux DRM sysfs.
 
-    Reads /sys/class/drm/card*/device/mem_info_vram_* which the kernel
+    Reads /sys/class/drm/card*/device/mem_info_vram_*, which the kernel
     updates in real-time across all processes. No tools required.
     Returns (used_gb, total_gb) or (None, None) on failure.
     """
@@ -597,8 +581,8 @@ def _rocm_linux_sysfs_vram_gb() -> tuple[Optional[float], Optional[float]]:
 def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[float]]:
     """Query system-wide dedicated GPU VRAM via Windows Performance Counters.
 
-    Uses the same data source as Task Manager so it reflects cross-process
-    usage accurately. Works for any GPU vendor without amd-smi or nvidia-smi.
+    Same data source as Task Manager, so cross-process usage is accurate.
+    Works for any GPU vendor without amd-smi or nvidia-smi.
     Returns (used_gb, total_gb) or (None, None) on failure.
     """
     if platform.system() != "Windows":
@@ -637,13 +621,11 @@ def get_gpu_utilization() -> Dict[str, Any]:
         if result is not None:
             result["backend"] = _backend_label(device)
             if IS_ROCM:
-                # Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.)
+                # Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.).
                 _reconcile_primary_rocm_unified_memory(result, _get_parent_visible_gpu_spec())
             return result
-        # SMI tool unavailable or returned no usable data. On Windows, query
-        # the Performance Counter API (same source as Task Manager) for
-        # system-wide dedicated VRAM — covers cross-process usage that
-        # torch.cuda.mem_get_info cannot see from the Studio server process.
+        # SMI unavailable. On Windows, use Performance Counters (Task Manager
+        # source) for system-wide VRAM, covering cross-process usage torch can't see.
         if IS_ROCM and platform.system() == "Windows":
             _win_used, _win_total = _rocm_windows_perf_counter_vram_gb()
             if _win_used is not None and _win_total is not None:
@@ -705,8 +687,7 @@ def get_gpu_utilization() -> Dict[str, Any]:
                 "power_utilization_pct": None,
             }
 
-    # MLX path: single _read_apple_gpu_stats() call carries both VRAM-used
-    # bytes and GPU utilization %. psutil for unified-memory total is cheap.
+    # MLX: _read_apple_gpu_stats() carries both VRAM-used and GPU util%.
     if device == DeviceType.MLX:
         try:
             import psutil
@@ -772,8 +753,8 @@ def _apply_unified_memory_correction(
     """Per-device reconciliation: when torch reports a larger memory total
     than amd-smi, overwrite the smi VRAM fields in place.
 
-    Used by both the multi-device and primary-device reconciliation helpers
-    so the two endpoints stay in sync on AMD iGPUs with unified memory.
+    Used by both the multi-device and primary-device reconcilers so the two
+    endpoints stay in sync on AMD iGPUs with unified memory.
     """
     torch_total_gb = torch_info["total_gb"]
     smi_total_gb = device_metrics.get("vram_total_gb") or 0.0
@@ -796,9 +777,8 @@ def _apply_unified_memory_correction(
 def _reconcile_rocm_unified_memory(utilization: Dict[str, Any], device_indices: list[int]) -> None:
     """Fix amd-smi VRAM for ROCm unified-memory GPUs (e.g. Strix Halo).
 
-    amd-smi reports only the dedicated slice (~512 MB); torch sees the full
-    GTT pool (~128 GB). When torch total > smi total, overwrite per-device
-    VRAM fields so GPU selection uses the real available memory.
+    amd-smi reports only the dedicated slice; torch sees the full GTT pool. When
+    torch total > smi total, overwrite per-device VRAM fields with the real value.
     """
     torch_devices = _torch_get_per_device_info(device_indices)
     if not torch_devices:
@@ -820,10 +800,8 @@ def _reconcile_primary_rocm_unified_memory(
         # No visibility env var set: torch ordinal 0 is the primary device.
         primary_idx = [0]
     elif len(numeric_ids) == 0:
-        # Empty mask (HIP_VISIBLE_DEVICES="" or "-1"): no GPU is visible to
-        # this process. Querying torch device 0 would raise a RuntimeError or
-        # return stale/wrong data, so bail out rather than writing bad values
-        # into the utilization dict.
+        # Empty mask: no GPU visible. Querying torch device 0 would raise or
+        # return stale data, so bail rather than write bad values.
         return
     else:
         primary_idx = [int(numeric_ids[0])]
@@ -847,15 +825,14 @@ def get_visible_gpu_utilization() -> Dict[str, Any]:
             result["backend"] = _backend_label(device)
             numeric_ids = parent_visible_spec.get("numeric_ids")
             if IS_ROCM and numeric_ids is not None:
-                # Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.)
+                # Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.).
                 _reconcile_rocm_unified_memory(result, numeric_ids)
             return result
 
     # Torch-based fallback for CUDA (nvidia-smi unavailable, AMD ROCm) and XPU (Intel)
     if device in (DeviceType.CUDA, DeviceType.XPU):
         parent_ids = get_parent_visible_gpu_ids()
-        # When parent_visible_ids is empty (UUID/MIG mask or no CVD set),
-        # enumerate torch-visible ordinals so the UI still shows devices.
+        # Empty parent_ids (UUID/MIG mask or no CVD): enumerate torch ordinals.
         if parent_ids:
             torch_indices = parent_ids
             index_kind = "physical"
@@ -942,14 +919,11 @@ _visible_gpu_count: Optional[int] = None
 
 
 def _get_parent_visible_gpu_spec() -> Dict[str, Any]:
-    # ROCm uses HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES in addition to
-    # CUDA_VISIBLE_DEVICES (which HIP also respects).  Check ROCm-specific
-    # env vars first so multi-GPU AMD setups are handled correctly.
-    # Use explicit None checks (not `or`) so empty string "" is honoured
-    # as "no visible GPUs" rather than falling through to CUDA_VISIBLE_DEVICES.
+    # ROCm uses HIP/ROCR_VISIBLE_DEVICES on top of CUDA_VISIBLE_DEVICES; check
+    # them first. Explicit None checks (not `or`) so "" reads as "no visible GPUs".
     cuda_visible = None
-    # Prefer ROCm masks only on a ROCm host, or when no CUDA mask is set, so a
-    # stale HIP_VISIBLE_DEVICES on an NVIDIA host can't override CUDA_VISIBLE_DEVICES.
+    # Prefer ROCm masks only on a ROCm host or when no CUDA mask is set, so a
+    # stale HIP_VISIBLE_DEVICES on NVIDIA can't override CUDA_VISIBLE_DEVICES.
     _is_rocm_spec = IS_ROCM or (
         "CUDA_VISIBLE_DEVICES" not in os.environ
         and ("HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ)
@@ -1027,7 +1001,7 @@ def resolve_requested_gpu_ids(gpu_ids: Optional[list[int]]) -> list[int]:
             f"Parent-visible GPUs: {parent_visible_ids}"
         )
 
-    # Reject negative IDs unconditionally.
+    # Reject negative IDs.
     negative_ids = [gpu_id for gpu_id in requested_ids if gpu_id < 0]
     if negative_ids:
         raise ValueError(
@@ -1035,16 +1009,13 @@ def resolve_requested_gpu_ids(gpu_ids: Optional[list[int]]) -> list[int]:
             f"Rejected IDs: {negative_ids}. Parent-visible GPUs: {parent_visible_ids}"
         )
 
-    # Only enforce the physical upper bound when we have a reliable count
-    # from nvidia-smi. When the count comes from torch, it reflects visible
-    # devices (filtered by CUDA_VISIBLE_DEVICES), not the physical total,
-    # so high physical indices like 3 would be falsely rejected on a
-    # CUDA_VISIBLE_DEVICES="2,3" machine that reports device_count()=2.
-    # The parent-visible check below is authoritative in all cases.
+    # Only enforce the physical upper bound when the count is reliable (nvidia-smi).
+    # A torch count reflects only visible devices, so it could falsely reject valid
+    # physical indices. The parent-visible check below is always authoritative.
     if physical_gpu_count > 0 and parent_visible_ids:
         max_parent_id = max(parent_visible_ids)
         if physical_gpu_count > max_parent_id:
-            # Count is plausibly physical (not just visible), so enforce it
+            # Count is plausibly physical, so enforce it.
             out_of_range = [gpu_id for gpu_id in requested_ids if gpu_id >= physical_gpu_count]
             if out_of_range:
                 raise ValueError(
@@ -1123,13 +1094,11 @@ def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = Non
 
 
 def _determine_attention_impl_for_gpu_estimate(config) -> str:
-    # torch.distributed is incomplete on Windows ROCm — torch._C is a C
-    # extension (not a package), so Python cannot import the submodule
-    # torch._C._distributed_c10d that torch.distributed depends on.
-    # Inject an empty stub into sys.modules BEFORE importing torch.distributed
-    # so the import succeeds, then patch the missing process-group helpers.
+    # torch.distributed is incomplete on Windows ROCm (torch._C._distributed_c10d
+    # can't be imported). Inject stubs into sys.modules before importing
+    # torch.distributed, then patch the missing process-group helpers.
     if sys.platform == "win32" and IS_ROCM:
-        # Dummy class for any name torch.distributed tries to import from these stubs
+        # Dummy for any name torch.distributed imports from these stubs.
         class _Dummy:
             pass
 
@@ -1140,8 +1109,7 @@ def _determine_attention_impl_for_gpu_estimate(config) -> str:
         ):
             if _c10d_name not in sys.modules:
                 _stub = types.ModuleType(_c10d_name)
-                # torch.distributed imports these names from _distributed_c10d;
-                # provide no-op dummies so the import doesn't raise AttributeError.
+                # No-op dummies for names torch.distributed imports from _distributed_c10d.
                 for _sym in (
                     "FakeProcessGroup",
                     "ProcessGroup",
@@ -1177,11 +1145,9 @@ def _determine_attention_impl_for_gpu_estimate(config) -> str:
     from unsloth.models._utils import resolve_attention_implementation
     from transformers import AutoModel, AutoModelForCausalLM
 
-    # why: resolve_attention_implementation calls _set_attn_impl which writes
-    # _attn_implementation onto the config; PreTrainedConfig's setter walks
-    # `sub_configs` and propagates to nested text_config / sub-configs, so a
-    # shallow copy still mutates those shared inner objects on the cached
-    # config returned by _load_config_for_gpu_estimate. Deepcopy isolates them.
+    # why: resolve_attention_implementation writes _attn_implementation onto the
+    # config and propagates to nested sub-configs; a shallow copy would still
+    # mutate the cached config's shared inner objects. Deepcopy isolates them.
     config_copy = copy.deepcopy(config)
 
     model_class = None
@@ -1357,31 +1323,29 @@ def estimate_required_model_memory_gb(
                 config
             )
         except Exception as e:
-            # Log at debug: on Windows ROCm the torch.distributed stub does
-            # not implement Store, so this fires on every estimate call.
-            # It is expected and non-actionable -- eager is the safe fallback.
+            # Debug-level: fires every estimate on Windows ROCm (stub lacks Store);
+            # expected and non-actionable -- eager is the safe fallback.
             logger.debug(
                 "Could not resolve attention implementation for '%s': %s",
                 estimate_model,
                 e,
             )
-            # why: if we cannot prove flash attention is usable, charge the
-            # quadratic non-flash activation path so GPU selection stays
-            # conservative.
+            # why: charge the quadratic non-flash activation path so GPU
+            # selection stays conservative when flash attn isn't proven usable.
             vram_config.attention_implementation = "eager"
     arch = extract_arch_config(config) if config is not None else None
 
     if arch is not None:
         breakdown = estimate_training_vram(arch, vram_config)
-        # why: extract_arch_config only sees text_config; safetensors include
-        # vision/audio tower bytes that the text-arch fp16 total misses.
+        # why: extract_arch_config only sees text_config; add the vision/audio
+        # tower bytes that the text-arch fp16 total misses.
         arch_fp16_bytes = compute_total_params(arch) * 2
         extra_bytes = max(0, int(model_size_bytes) - arch_fp16_bytes)
         if extra_bytes > 0:
             breakdown.model_weights += extra_bytes
             if training_method == "full":
-                # why: full fine-tuning makes the extra (vision/audio) params
-                # trainable; optimizer + gradient bytes scale with them too.
+                # why: full fine-tuning makes extra params trainable; optimizer +
+                # gradient bytes scale with them.
                 extra_params = extra_bytes // 2
                 breakdown.optimizer_states += compute_optimizer_bytes(
                     extra_params,
@@ -1400,7 +1364,7 @@ def estimate_required_model_memory_gb(
             )
         return required_gb, metadata
 
-    # Fallback when model config is unavailable
+    # Fallback when model config is unavailable.
     overhead_gb = CUDA_OVERHEAD_BYTES / (1024**3)
     if training_method == "full":
         required_gb = model_size_gb * 3.5 + overhead_gb
@@ -1460,9 +1424,7 @@ def auto_select_gpu_ids(
         return None, metadata
 
     if required_gb is None:
-        # Cannot estimate model size -- fall back to all visible GPUs
-        # rather than risk loading on a single GPU that may not have
-        # enough memory.
+        # Can't estimate size -- use all visible GPUs rather than risk one too small.
         parent_ids = get_parent_visible_gpu_ids()
         metadata["selection_mode"] = "fallback_all"
         metadata["selected_gpu_ids"] = parent_ids
@@ -1500,17 +1462,13 @@ def auto_select_gpu_ids(
     free_by_index = {item["index"]: item["free_gb"] for item in ranked}
     selected: list[int] = []
     usable_gb = 0.0
-    # Multi-GPU sharding has overhead from inter-GPU communication (NCCL
-    # all-reduce, PCIe/NVLink transfers, synchronization barriers), so each
-    # additional GPU contributes less than its raw free memory. The first GPU
-    # keeps its full capacity (no cross-device overhead). 0.85 was calibrated
-    # empirically on 2-8 GPU setups with NVLink and PCIe topologies -- the
-    # 15% discount accounts for NCCL buffers (~2-5% of VRAM), pipeline bubble
-    # overhead, and memory fragmentation from non-uniform shard sizes.
+    # Sharding has inter-GPU overhead, so each extra GPU contributes less than
+    # its raw free memory (first GPU keeps full capacity). 0.85 is empirical on
+    # 2-8 GPU setups: covers NCCL buffers, pipeline bubbles, fragmentation.
     multi_gpu_overhead = 0.85
 
-    # Per-GPU check: activations don't shard, so each GPU needs its weight
-    # shard + full activation cost. Use precomputed min_per_gpu_N values.
+    # Per-GPU check: activations don't shard, so each GPU needs its weight shard
+    # + full activation cost. Uses precomputed min_per_gpu_N values.
     vram_breakdown = estimate_metadata.get("vram_breakdown", {})
 
     for candidate in ranked:
@@ -1547,7 +1505,7 @@ def auto_select_gpu_ids(
             )
             return selected, metadata
 
-    # Use only GPUs with verified VRAM data (from gpu_candidates, not raw devices)
+    # Use only GPUs with verified VRAM data.
     fallback_all = [c["index"] for c in gpu_candidates] if gpu_candidates else parent_ids
     metadata["selection_mode"] = "fallback_all"
     if ranked:
@@ -1586,18 +1544,17 @@ def prepare_gpu_selection(
     """Resolve which physical GPUs to use for a model load.
 
     GPU selection modes:
-      - **Explicit** (``gpu_ids=[5, 6, 7]``): the caller chooses exact GPUs.
-        All listed GPUs are used and the model is sharded across them via
-        ``device_map="balanced"``, regardless of whether the model would fit
-        on fewer GPUs.  IDs are validated against the parent-visible set.
-      - **Auto** (``gpu_ids=None`` or ``[]``): ``auto_select_gpu_ids`` estimates
-        VRAM requirements and picks the *minimum* number of GPUs needed,
-        preferring GPUs with the most free memory.
+      - **Explicit** (``gpu_ids=[5, 6, 7]``): caller chooses exact GPUs.
+        All listed GPUs are used and the model is sharded via
+        ``device_map="balanced"``, even if it would fit on fewer. IDs are
+        validated against the parent-visible set.
+      - **Auto** (``gpu_ids=None`` or ``[]``): ``auto_select_gpu_ids``
+        estimates VRAM needs and picks the *minimum* GPUs needed,
+        preferring those with the most free memory.
 
-    The returned ``gpu_ids`` list is later passed to ``get_device_map()`` which
-    maps it to a Hugging Face ``device_map`` string, and to ``apply_gpu_ids()``
-    in the worker subprocess which narrows ``CUDA_VISIBLE_DEVICES`` before any
-    torch/CUDA initialisation.
+    The returned ``gpu_ids`` is later passed to ``get_device_map()`` (maps it
+    to a Hugging Face ``device_map`` string) and to ``apply_gpu_ids()`` in the
+    worker subprocess (narrows ``CUDA_VISIBLE_DEVICES`` before torch/CUDA init).
     """
     if gpu_ids and get_device() != DeviceType.CUDA:
         raise ValueError(
@@ -1633,8 +1590,7 @@ def get_physical_gpu_count() -> int:
     Return the number of physical GPUs on the machine.
 
     Uses ``nvidia-smi -L`` on NVIDIA (unaffected by CUDA_VISIBLE_DEVICES),
-    with a torch-based fallback for AMD ROCm and Intel XPU.
-    Result is cached after the first call.
+    with a torch fallback for AMD ROCm and Intel XPU. Cached after first call.
     """
     global _physical_gpu_count
     if _physical_gpu_count is not None:
@@ -1654,7 +1610,7 @@ def get_physical_gpu_count() -> int:
                 return _physical_gpu_count
         except Exception:
             pass
-        # SMI tool unavailable or failed -- fall back to torch
+        # SMI unavailable -- fall back to torch.
         count = _torch_get_physical_gpu_count()
         _physical_gpu_count = count if count is not None else 1
         return _physical_gpu_count
@@ -1676,10 +1632,10 @@ def get_physical_gpu_count() -> int:
 def _backend_visible_devices_env() -> Optional[str]:
     """Return the raw visibility env string that applies to this backend.
 
-    On ROCm, HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES take precedence
-    over CUDA_VISIBLE_DEVICES; the helper mirrors the resolution logic in
-    ``_get_parent_visible_gpu_spec`` so ``backend_cuda_visible_devices``
-    reports the value that is actually narrowing the visible device set.
+    On ROCm, HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES take precedence over
+    CUDA_VISIBLE_DEVICES; this mirrors ``_get_parent_visible_gpu_spec`` so
+    ``backend_cuda_visible_devices`` reports the value actually narrowing the
+    visible device set.
     """
     if IS_ROCM:
         return _get_parent_visible_gpu_spec().get("raw")
@@ -1690,7 +1646,7 @@ def get_backend_visible_gpu_info() -> Dict[str, Any]:
     device = get_device()
     if device in (DeviceType.CUDA, DeviceType.XPU):
         parent_visible_ids = get_parent_visible_gpu_ids()
-        # Try native SMI tool first (nvidia-smi for NVIDIA, skipped for ROCm)
+        # Try native SMI first (nvidia-smi; skipped for ROCm).
         if device == DeviceType.CUDA and not IS_ROCM:
             try:
                 from . import nvidia
@@ -1706,9 +1662,8 @@ def get_backend_visible_gpu_info() -> Dict[str, Any]:
             except Exception as e:
                 logger.warning("Backend GPU visibility query failed: %s", e)
 
-        # Torch fallback (AMD ROCm, Intel XPU, nvidia-smi missing/failed)
-        # When parent_visible_ids is empty (UUID/MIG mask), enumerate by
-        # torch ordinal so the UI still shows devices.
+        # Torch fallback (ROCm, XPU, nvidia-smi missing). Empty parent_visible_ids
+        # (UUID/MIG mask) -> enumerate by torch ordinal so the UI shows devices.
         if parent_visible_ids:
             torch_indices = parent_visible_ids
             index_kind = "physical"
@@ -1789,15 +1744,15 @@ def get_visible_gpu_count() -> int:
     Return the number of GPUs visible to this process.
 
     Respects ``CUDA_VISIBLE_DEVICES`` -- if set, only those GPUs count.
-    Falls back to physical count if the env var is unset or torch is
-    unavailable.  Result is cached after the first call.
+    Falls back to physical count if unset or torch is unavailable.
+    Cached after the first call.
     """
     global _visible_gpu_count
     if _visible_gpu_count is not None:
         return _visible_gpu_count
 
-    # Use _get_parent_visible_gpu_spec() which already handles
-    # HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES on ROCm.
+    # _get_parent_visible_gpu_spec() already handles HIP_VISIBLE_DEVICES /
+    # ROCR_VISIBLE_DEVICES on ROCm.
     visible_spec = _get_parent_visible_gpu_spec()
     if visible_spec["raw"] is not None:
         raw = visible_spec["raw"].strip()
@@ -1809,7 +1764,7 @@ def get_visible_gpu_count() -> int:
             _visible_gpu_count = len([x for x in raw.split(",") if x.strip()])
         return _visible_gpu_count
 
-    # No visibility env var set -- try torch, fall back to physical count
+    # No visibility env var set -- try torch, else physical count
     try:
         import torch
         if get_device() == DeviceType.XPU and hasattr(torch, "xpu"):
@@ -1826,8 +1781,7 @@ def apply_gpu_ids(gpu_ids) -> None:
     if gpu_ids is None:
         return
 
-    # Empty list means "no GPUs visible" -- treat the same as None
-    # (inherit parent) to avoid setting CUDA_VISIBLE_DEVICES="" which
+    # Empty list -> treat like None (inherit parent); setting CUDA_VISIBLE_DEVICES=""
     # disables CUDA entirely and crashes downstream torch calls.
     if isinstance(gpu_ids, (list, tuple)) and len(gpu_ids) == 0:
         return
@@ -1840,24 +1794,16 @@ def apply_gpu_ids(gpu_ids) -> None:
         value = str(gpu_ids)
 
     os.environ["CUDA_VISIBLE_DEVICES"] = value
-    # Keep ROCm visibility env vars in sync so _get_parent_visible_gpu_spec()
-    # picks up the narrowed set on AMD systems. Workers can call
-    # apply_gpu_ids() before detect_hardware() runs (so IS_ROCM is still
-    # its default False), so also mirror the selection whenever the
-    # parent process already set a ROCm visibility variable -- that
-    # way a downstream ROCm process inherits the narrowed mask even
-    # before Studio's hardware detection has classified the host.
-    # Final fallback: probe torch.version.hip so AMD workers without
-    # HIP_VISIBLE_DEVICES still get the correct ROCm visibility mask.
+    # Keep ROCm visibility env vars in sync. Workers may call apply_gpu_ids()
+    # before detect_hardware() (IS_ROCM still False), so also mirror when the
+    # parent set a ROCm visibility var, with a torch.version.hip probe fallback.
     _inherits_rocm_visibility = (
         "HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ
     )
     _is_rocm = IS_ROCM or _inherits_rocm_visibility
     if not _is_rocm:
-        # torch.version.hip is a non-empty string on ROCm, None on CUDA.
-        # AMD SDK / Radeon ROCm wheels can leave torch.version.hip unset but
-        # still encode "rocm" in torch.__version__, matching detect_hardware().
-        # Broad except: a probe failure must never crash a training worker.
+        # torch.version.hip is set on ROCm, None on CUDA; AMD SDK wheels may leave
+        # it unset but encode "rocm" in __version__. Broad except: never crash a worker.
         try:
             import torch as _torch
             _is_rocm = (
@@ -1886,23 +1832,22 @@ def get_device_map(gpu_ids: Optional[list[int]] = None) -> str:
     Returns ``"balanced"`` (shard evenly across GPUs) when:
       - ``gpu_ids`` explicitly lists >1 GPU, **or**
       - ``CUDA_VISIBLE_DEVICES`` uses UUID/MIG identifiers (non-numeric) and
-        more than one GPU is visible (fallback: we cannot resolve numeric IDs,
-        so we assume the caller intends multi-GPU).
+        >1 GPU is visible (fallback: numeric IDs unresolvable, so assume
+        multi-GPU is intended).
 
-    Returns ``"sequential"`` (single device) in all other cases, including
-    non-CUDA backends (CPU, MLX).
+    Returns ``"sequential"`` (single device) otherwise, including non-CUDA
+    backends (CPU, MLX).
 
-    Callers should use ``prepare_gpu_selection()`` upstream to determine the
-    ``gpu_ids`` list -- that function handles the smart auto-selection of the
-    minimum number of GPUs needed for a given model.
+    Use ``prepare_gpu_selection()`` upstream to determine ``gpu_ids`` -- it
+    handles auto-selecting the minimum GPUs needed for a model.
     """
     device = get_device()
     if device == DeviceType.CUDA:
         multi_gpu = gpu_ids is not None and len(gpu_ids) > 1
 
         if not multi_gpu:
-            # UUID/MIG masks cannot be split into numeric IDs, so if multiple
-            # GPUs are visible we assume multi-GPU sharding is intended.
+            # UUID/MIG masks can't be split into numeric IDs; >1 visible GPU
+            # means multi-GPU sharding is intended.
             parent_visible_spec = _get_parent_visible_gpu_spec()
             if parent_visible_spec["numeric_ids"] is None and get_visible_gpu_count() > 1:
                 multi_gpu = True
@@ -1944,18 +1889,14 @@ def safe_num_proc(desired: Optional[int] = None) -> int:
     """
     Return a safe ``num_proc`` for ``dataset.map()`` calls.
 
-    On Windows, always returns 1 because Python uses ``spawn`` instead of
-    ``fork`` for multiprocessing -- the overhead of re-importing torch,
-    transformers, unsloth etc. per worker is typically slower than
-    single-process for normal dataset sizes.
+    On Windows always returns 1: Python uses ``spawn`` not ``fork``, so
+    re-importing torch/transformers/unsloth per worker is typically slower
+    than single-process for normal dataset sizes.
 
-    On multi-GPU machines (where multiple GPUs are *visible* to this
-    process) the NVIDIA driver spawns extra background threads, making
-    ``os.fork()`` prone to deadlocks when many workers are created.
-    This helper caps ``num_proc`` to 4 on such machines.
-
-    When ``CUDA_VISIBLE_DEVICES`` restricts to a single GPU, the cap
-    does not apply.
+    On multi-GPU machines (multiple GPUs *visible* to this process) the
+    NVIDIA driver spawns extra background threads, making ``os.fork()``
+    deadlock-prone with many workers, so this caps ``num_proc`` to 4.
+    The cap does not apply when ``CUDA_VISIBLE_DEVICES`` restricts to one GPU.
 
     Args:
         desired: The num_proc you *want*. If None, auto-computes from
@@ -1964,9 +1905,8 @@ def safe_num_proc(desired: Optional[int] = None) -> int:
     Returns:
         A safe integer ≥ 1.
     """
-    # Windows and macOS use 'spawn' for multiprocessing -- the overhead of
-    # re-importing torch/transformers/unsloth per worker is typically slower
-    # than single-process.
+    # Windows/macOS use 'spawn'; re-importing torch/transformers/unsloth per
+    # worker is typically slower than single-process.
     if sys.platform in ("win32", "darwin"):
         return 1
 
@@ -1989,9 +1929,8 @@ def safe_thread_num_proc(desired: Optional[int] = None) -> int:
     """
     Return a safe worker count for ``ThreadPoolExecutor`` calls.
 
-    Unlike ``safe_num_proc()``, this does NOT cap to 1 on macOS/Windows.
-    Threads share the parent process address space and are unaffected by
-    the ``spawn`` vs ``fork`` distinction.
+    Unlike ``safe_num_proc()``, does NOT cap to 1 on macOS/Windows: threads
+    share the parent address space, unaffected by ``spawn`` vs ``fork``.
 
     Args:
         desired: The thread count you *want*. If None, auto-computes
@@ -2010,9 +1949,9 @@ def dataset_map_num_proc(desired: Optional[int] = None) -> Optional[int]:
     """
     Return a safe ``num_proc`` for ``Dataset.map()`` and ``Dataset.filter()``.
 
-    Returns ``None`` on spawn-based platforms (Windows, macOS) because
-    ``datasets`` treats ``num_proc=1`` as multiprocessing (creates ``Pool(1)``).
-    Only ``num_proc=None`` guarantees in-process execution.
+    Returns ``None`` on spawn platforms (Windows, macOS) because ``datasets``
+    treats ``num_proc=1`` as multiprocessing (creates ``Pool(1)``); only
+    ``num_proc=None`` guarantees in-process execution.
     """
     if sys.platform in ("win32", "darwin"):
         return None
diff --git a/studio/backend/utils/hardware/nvidia.py b/studio/backend/utils/hardware/nvidia.py
index 6cead61f08..f98ca4343e 100644
--- a/studio/backend/utils/hardware/nvidia.py
+++ b/studio/backend/utils/hardware/nvidia.py
@@ -110,9 +110,8 @@ def get_primary_gpu_utilization() -> dict[str, Any]:
 def get_visible_gpu_utilization(
     parent_visible_ids: Optional[list[int]], parent_cuda_visible_devices: Optional[str] = None
 ) -> dict[str, Any]:
-    # When parent_visible_ids is None (UUID/MIG mask), we cannot safely
-    # map nvidia-smi rows to the process's visible devices. Return empty
-    # instead of exposing all physical GPUs.
+    # parent_visible_ids None (UUID/MIG mask): can't map nvidia-smi rows to
+    # visible devices, so return empty rather than exposing all physical GPUs.
     if parent_visible_ids is None:
         return {
             "available": False,
@@ -196,8 +195,8 @@ def get_visible_gpu_utilization(
 def get_backend_visible_gpu_info(
     parent_visible_ids: Optional[list[int]], backend_cuda_visible_devices: Optional[str]
 ) -> dict[str, Any]:
-    # When parent_visible_ids is None (UUID/MIG mask), we cannot safely
-    # map nvidia-smi rows to the process's visible devices.
+    # parent_visible_ids None (UUID/MIG mask): can't map nvidia-smi rows to
+    # visible devices.
     if parent_visible_ids is None:
         return {
             "available": False,
@@ -249,7 +248,7 @@ def get_backend_visible_gpu_info(
             continue
         if visible_ordinals is not None and idx not in visible_ordinals:
             continue
-        # Use split with limit to handle GPU names containing commas
+        # Rejoin in case the GPU name contains commas
         name = parts[1] if len(parts) == 3 else ", ".join(parts[1:-1])
         try:
             mem_total_mb = int(parts[-1])
diff --git a/studio/backend/utils/hardware/vram_estimation.py b/studio/backend/utils/hardware/vram_estimation.py
index ddc39733e3..86069ead3d 100644
--- a/studio/backend/utils/hardware/vram_estimation.py
+++ b/studio/backend/utils/hardware/vram_estimation.py
@@ -59,8 +59,7 @@ OPTIMIZER_BYTES_PER_PARAM: Dict[str, int] = {
 }
 
 # (full_ft_multiplier, lora_multiplier) — fraction of num_layers.
-# LoRA: frozen base layers skip activation storage, but you always need
-# at least ~1 layer in flight during backprop recomputation.
+# LoRA: frozen layers skip activation storage, but ~1 is in flight during backprop.
 GC_LAYER_MULTIPLIERS = {
     "none": (None, None),
     "true": (2.0, 1.0),
@@ -125,8 +124,7 @@ class VramBreakdown:
     gradients: int
     activations: int
     cuda_overhead: int
-    # Equals `activations`; retained for backward compatibility with
-    # consumers that read this field.
+    # Equals `activations`; kept for backward compat with field consumers.
     activations_computed: int = 0
 
     @property
@@ -141,10 +139,10 @@ class VramBreakdown:
         )
 
     def min_gpu_vram(self, n_gpus: int) -> int:
-        """Minimum VRAM a single GPU needs: its shard + non-shardable costs.
+        """Min VRAM one GPU needs: its shard + non-shardable costs.
 
-        Weights/LoRA/optimizer/gradients shard across GPUs.
-        Activations do NOT shard (the GPU running a layer holds them).
+        Weights/LoRA/optimizer/gradients shard across GPUs; activations do
+        NOT (the GPU running a layer holds them).
         """
         shardable = self.model_weights + self.lora_adapters + self.optimizer_states + self.gradients
         per_gpu_fixed = self.activations + self.cuda_overhead
@@ -163,7 +161,7 @@ class VramBreakdown:
 
 
 def _first_scalar(value):
-    # why: ERNIE MoE configs ship moe_intermediate_size / moe_num_experts as
+    # ERNIE MoE ships moe_intermediate_size / moe_num_experts as
     # [routed, shared] lists; downstream arithmetic needs the routed scalar.
     if isinstance(value, (list, tuple)):
         return value[0] if value else None
@@ -171,8 +169,8 @@ def _first_scalar(value):
 
 
 def _max_scalar(value):
-    # why: Hunyuan-V1-MoE moe_topk can be a per-layer list; activation
-    # accounting uses the max top-k as a conservative upper bound.
+    # Hunyuan-V1-MoE moe_topk can be a per-layer list; activation accounting
+    # uses max top-k as a conservative upper bound.
     if isinstance(value, (list, tuple)):
         items = [v for v in value if v is not None]
         return max(items) if items else None
@@ -181,16 +179,15 @@ def _max_scalar(value):
 
 def _compute_dense_layer_indices(text_config, total_layers: int) -> tuple:
     """Layer indices that use dense MLP instead of MoE. Position matters."""
-    # why: transformers Exaone-MoE / Laguna / Hy_v3 / GLM-MoE-DSA / GLM4-MoE-Lite /
-    # Ernie4_5_VL_MoE prefer per-position `mlp_layer_types` over the prefix-style
-    # `first_k_dense_replace` and may omit `decoder_sparse_step` entirely.
+    # Exaone-MoE / Laguna / Hy_v3 / GLM-MoE-DSA / GLM4-MoE-Lite / Ernie4_5_VL_MoE
+    # prefer per-position `mlp_layer_types` over prefix `first_k_dense_replace`.
     layer_types = getattr(text_config, "mlp_layer_types", None)
     if layer_types:
         return tuple(
             i for i, t in enumerate(layer_types[:total_layers]) if str(t).lower() == "dense"
         )
 
-    # why: Llama4TextConfig.__init__ auto-populates self.moe_layers from
+    # Llama4TextConfig.__init__ auto-populates self.moe_layers from
     # interleave_moe_layer_step; Llama4TextDecoderLayer dispatches via
     # `layer_idx in config.moe_layers` (modeling_llama4.py).
     llama4_moe_layers = getattr(text_config, "moe_layers", None)
@@ -198,10 +195,9 @@ def _compute_dense_layer_indices(text_config, total_layers: int) -> tuple:
         moe_indices = {int(i) for i in llama4_moe_layers}
         return tuple(i for i in range(total_layers) if i not in moe_indices)
 
-    # why: transformers ERNIE 4.5 MoE / ERNIE 4.5 VL MoE declare MoE layers
-    # via moe_layer_start_index / moe_layer_end_index / moe_layer_interval;
-    # the model's per-layer guard is `(layer_idx + 1) % interval == 0` with
-    # start <= layer_idx <= end (modeling_ernie4_5_moe.py).
+    # ERNIE 4.5 (VL) MoE: layers via moe_layer_start/end_index + interval;
+    # per-layer guard `(layer_idx+1) % interval == 0` within [start, end]
+    # (modeling_ernie4_5_moe.py).
     moe_start = getattr(text_config, "moe_layer_start_index", None)
     moe_interval = getattr(text_config, "moe_layer_interval", None)
     if moe_start is not None and moe_interval is not None and int(moe_interval) > 0:
@@ -261,8 +257,8 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]:
 
     num_kv_heads = getattr(text_config, "num_key_value_heads", num_heads)
 
-    # why: DBRX places its MoE attrs on the DbrxFFNConfig sub-config; probe
-    # ffn_config as a secondary source so DBRX is not misclassified as dense.
+    # DBRX places its MoE attrs on the DbrxFFNConfig sub-config; probe
+    # ffn_config as a secondary source so DBRX isn't misclassified as dense.
     ffn_config = getattr(text_config, "ffn_config", None)
 
     def _moe_attr(name):
@@ -286,8 +282,8 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]:
     if moe_intermediate_raw is None:
         moe_intermediate_raw = _moe_attr("ffn_hidden_size")
     moe_intermediate = _first_scalar(moe_intermediate_raw)
-    # why: Exaone-MoE / ERNIE families alias num_shared_experts /
-    # moe_num_shared_experts to the canonical n_shared_experts.
+    # Exaone-MoE / ERNIE alias num_shared_experts / moe_num_shared_experts
+    # to the canonical n_shared_experts.
     n_shared_experts = (
         _first_scalar(_moe_attr("n_shared_experts"))
         or _first_scalar(_moe_attr("num_shared_experts"))
@@ -297,9 +293,8 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]:
     shared_expert_intermediate_size = _moe_attr("shared_expert_intermediate_size")
     if shared_expert_intermediate_size and n_shared_experts == 0:
         n_shared_experts = 1
-    # why: DBRX exposes moe_top_k, Hunyuan-V1-MoE exposes moe_topk (which can
-    # be a per-layer list); _max_scalar normalizes list values to the worst
-    # case so int(...) below cannot crash on the canonical attribute_map path.
+    # DBRX moe_top_k; Hunyuan-V1-MoE moe_topk (may be a per-layer list).
+    # _max_scalar normalizes lists to the worst case so int(...) can't crash.
     num_experts_per_tok = (
         _max_scalar(_moe_attr("num_experts_per_tok"))
         or _max_scalar(_moe_attr("top_k_experts"))
@@ -313,9 +308,8 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]:
         dense_layer_indices = _compute_dense_layer_indices(text_config, num_layers)
     num_dense_layers = len(dense_layer_indices)
 
-    # why: Llama4 dense layers use intermediate_size_mlp; routed and shared
-    # experts use intermediate_size. Llama4TextMoe builds one shared_expert
-    # per MoE layer (modeling_llama4.py).
+    # Llama4 dense layers use intermediate_size_mlp; experts use
+    # intermediate_size. One shared_expert per MoE layer (modeling_llama4.py).
     intermediate_size_mlp_raw = _first_scalar(_moe_attr("intermediate_size_mlp"))
     dense_intermediate_size = (
         int(intermediate_size_mlp_raw) if intermediate_size_mlp_raw is not None else None
@@ -386,8 +380,8 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]:
 
 
 def _targets_all_linear(target_modules) -> bool:
-    # why: peft LoraConfig accepts target_modules="all-linear" as a bare
-    # string; iterating a string yields chars and never matches the set.
+    # peft LoraConfig accepts target_modules="all-linear" as a bare string;
+    # iterating a string yields chars and never matches the set.
     if isinstance(target_modules, str):
         target_modules = [target_modules]
     normalized = {str(module).lower().replace("_", "-") for module in target_modules}
@@ -425,10 +419,9 @@ def _is_kv_shared_layer(arch: ModelArchConfig, layer_idx: int) -> bool:
     if arch.num_kv_shared_layers <= 0:
         return False
     first_shared = arch.num_hidden_layers - arch.num_kv_shared_layers
-    # why: transformers Gemma4 (modeling_gemma4.py:1031, modular_gemma4.py:863)
-    # uses the same `> 0` guard so a fully-shared config raises during model
-    # construction; matching upstream avoids producing a detailed estimate
-    # for a shape the actual model code rejects.
+    # Gemma4 (modeling_gemma4.py:1031, modular_gemma4.py:863) uses the same
+    # `> 0` guard so a fully-shared config raises at model construction;
+    # matching upstream avoids estimating a shape the model code rejects.
     return layer_idx >= first_shared > 0
 
 
@@ -439,9 +432,9 @@ def _is_dense_mlp_layer(arch: ModelArchConfig, layer_idx: int) -> bool:
 
 
 def _per_layer_input_quantizable(arch: ModelArchConfig) -> int:
-    # why: Gemma4 PLE block adds per_layer_model_projection (single Linear),
+    # Gemma4 PLE block adds per_layer_model_projection (single Linear),
     # per_layer_input_gate (per layer), and per_layer_projection (per layer);
-    # see transformers gemma4/modular_gemma4.py:1077-1083 and :1247-1253.
+    # see gemma4/modular_gemma4.py:1077-1083 and :1247-1253.
     pli = arch.hidden_size_per_layer_input
     if pli <= 0:
         return 0
@@ -460,10 +453,8 @@ def _per_layer_input_norm_elements(arch: ModelArchConfig) -> int:
 
 
 def _per_layer_input_lora_params(arch: ModelArchConfig, r: int, target_modules) -> int:
-    # why: Unsloth's get_peft_regex (unsloth_zoo/peft_utils.py) requires module
-    # names to contain a component tag (mlp/attn/...); PLE module names lack
-    # any tag, so all-linear training does NOT attach LoRA to them. Only count
-    # PLE LoRA when the user explicitly names PLE modules.
+    # get_peft_regex requires a component tag (mlp/attn/...); PLE names lack
+    # one, so all-linear skips them. Count PLE LoRA only when named explicitly.
     pli = arch.hidden_size_per_layer_input
     if pli <= 0:
         return 0
@@ -543,18 +534,16 @@ def _module_path_matches(skip_module: str, alias: str) -> bool:
     if alias_parts[0] == "layers":
         return skip_parts == alias_parts
     if len(skip_parts) <= len(alias_parts):
-        # why: transformers BNB quantizer suffix-matches short skip entries
-        # like ["q_proj"] / ["lm_head"] against full module paths, so a skip
-        # shorter than the alias is a tail match.
+        # BNB suffix-matches short skip entries (["q_proj"], ["lm_head"]) so a
+        # skip shorter than the alias is a tail match.
         return alias_parts[-len(skip_parts) :] == skip_parts
     if skip_parts[-len(alias_parts) :] != alias_parts:
         return False
     prefix_parts = skip_parts[: len(skip_parts) - len(alias_parts)]
     if not prefix_parts:
         return True
-    # why: bound the prefix to known text-tower roots so VLM skip names like
-    # vision_tower.model.layers..self_attn.q_proj do not shadow the text
-    # alias model.layers..self_attn.q_proj.
+    # Bound the prefix to text-tower roots so VLM skips like
+    # vision_tower.model.layers... don't shadow the text alias.
     return ".".join(prefix_parts) in _SKIP_MODULE_TEXT_PREFIXES
 
 
@@ -587,9 +576,8 @@ def _build_text_module_elements(arch: ModelArchConfig) -> tuple[Dict[str, int],
         mlp_dims = {name: dim for name, dim in dims.items() if name in MLP_TARGET_MODULES}
 
         if is_mla:
-            # why: _text_linear_dims uses (hd, hd) for q/o; MLA actually splits
-            # into q_a/q_b/kv_a/kv_b, so emit a single self_attn aggregate at
-            # the authoritative MLA per-layer total.
+            # MLA splits q/o into q_a/q_b/kv_a/kv_b; emit a single self_attn
+            # aggregate at the authoritative MLA per-layer total.
             layer_modules["self_attn"] = _compute_attn_elements(arch)
         else:
             for name, (in_dim, out_dim) in attn_dims.items():
@@ -607,17 +595,13 @@ def _build_text_module_elements(arch: ModelArchConfig) -> tuple[Dict[str, int],
                 layer_modules["mlp.experts"] = _compute_routed_moe_elements(arch)
                 shared_moe = _compute_shared_moe_elements(arch)
                 if shared_moe:
-                    # why: Qwen3.5-MoE exposes shared expert as
-                    # mlp.shared_expert; Exaone-MoE/Laguna/GLM-style configs use
-                    # mlp.shared_experts. Register both names so child-path
-                    # llm_int8_skip_modules entries match the right shared block.
+                    # Qwen3.5-MoE: mlp.shared_expert; Exaone-MoE/Laguna/GLM:
+                    # mlp.shared_experts. Register both so skip_modules match.
                     layer_modules["mlp.shared_expert"] = shared_moe
                 if arch.moe_has_dense_mlp:
-                    # why: enable_moe_block runs the dense MLP and the MoE
-                    # experts in parallel; register both for skip matching.
-                    # Non-structured _text_linear_dims returns mlp_size from
-                    # _get_mlp_size which prefers moe_intermediate_size, so
-                    # rebuild dense dims from arch.intermediate_size directly.
+                    # enable_moe_block runs dense MLP and experts in parallel;
+                    # register both. Non-structured _get_mlp_size prefers
+                    # moe_intermediate_size, so rebuild dense dims directly.
                     if _uses_structured_layer_shapes(arch):
                         dense_dims = mlp_dims
                     else:
@@ -640,8 +624,8 @@ def _build_text_module_elements(arch: ModelArchConfig) -> tuple[Dict[str, int],
             )
 
         if pli > 0:
-            # why: register PLE per-layer linears so llm_int8_skip_modules
-            # entries like model.layers.0.per_layer_input_gate match.
+            # Register PLE per-layer linears so llm_int8_skip_modules entries
+            # like model.layers.0.per_layer_input_gate match.
             layer_modules["per_layer_input_gate"] = hd_global * pli
             layer_modules["per_layer_projection"] = pli * hd_global
 
@@ -650,10 +634,9 @@ def _build_text_module_elements(arch: ModelArchConfig) -> tuple[Dict[str, int],
             for name, value in layer_modules.items()
             if name == "self_attn" or name.startswith("self_attn.")
         )
-        # why: gemma4 enable_moe_block puts routed experts at the sibling
-        # layers..experts attribute, not under self.mlp; the layer's "mlp"
-        # aggregate must reflect only the dense MLP path so a skip module
-        # `model.layers.0.mlp` does not over-skip into the experts block.
+        # gemma4 enable_moe_block puts routed experts at sibling
+        # layers..experts, not under self.mlp; keep the "mlp" aggregate to
+        # the dense path so a `model.layers.0.mlp` skip doesn't over-skip.
         is_sibling_experts = bool(arch.moe_has_dense_mlp)
         mlp_total = sum(
             value
@@ -683,12 +666,10 @@ def _build_text_module_elements(arch: ModelArchConfig) -> tuple[Dict[str, int],
             elements[canonical] = value
             _add_module_aliases(aliases, canonical, canonical.removeprefix("text."))
             if name == "mlp.experts" and arch.moe_has_dense_mlp:
-                # why: gemma4 enable_moe_block exposes routed experts at
-                # layers..experts (sibling of self.mlp), not under mlp.
+                # gemma4: routed experts at sibling layers..experts, not mlp.
                 _add_module_aliases(aliases, canonical, f"layers.{layer_idx}.experts")
             elif name == "mlp.shared_expert":
-                # why: Exaone-MoE / Laguna / GLM-style configs use the plural
-                # `shared_experts` attribute name; register both spellings.
+                # Exaone-MoE/Laguna/GLM use plural `shared_experts`; add both.
                 _add_module_aliases(
                     aliases,
                     canonical,
@@ -733,8 +714,8 @@ def _get_mlp_size(arch: ModelArchConfig) -> int:
 
 
 def _dense_mlp_size(arch: ModelArchConfig) -> int:
-    # why: Llama4 dense layers use intermediate_size_mlp; routed/shared
-    # experts use intermediate_size. Other configs leave the field None.
+    # Llama4 dense layers use intermediate_size_mlp; routed/shared experts use
+    # intermediate_size. Other configs leave the field None.
     return arch.dense_intermediate_size or arch.intermediate_size
 
 
@@ -764,7 +745,7 @@ def _compute_dense_mlp_elements(arch: ModelArchConfig) -> int:
 
 
 def _shared_expert_size(arch: ModelArchConfig) -> int:
-    # why: Qwen3.5-MoE shared expert has its own intermediate_size (default 512)
+    # Qwen3.5-MoE shared expert has its own intermediate_size (default 512)
     # distinct from moe_intermediate_size; fall back to routed mlp_size for
     # families that share it (deepseek-style configs).
     return arch.shared_expert_intermediate_size or _get_mlp_size(arch)
@@ -782,10 +763,8 @@ def _compute_shared_moe_elements(arch: ModelArchConfig) -> int:
     hd = arch.hidden_size
     shared_size = _shared_expert_size(arch)
     total = hd * shared_size * 3 * arch.n_shared_experts
-    # why: only Qwen2-MoE / Qwen3.5-MoE define a shared_expert_gate Linear
-    # (hidden_size→1); other families (Exaone-MoE, HY-V3, GLM4-MoE-Lite, Laguna)
-    # have shared_experts without a gate. shared_expert_intermediate_size is the
-    # Qwen-style discriminator.
+    # Only Qwen2/Qwen3.5-MoE add a shared_expert_gate Linear (hidden_size->1);
+    # shared_expert_intermediate_size is the Qwen-style discriminator.
     if arch.shared_expert_intermediate_size:
         total += arch.n_shared_experts * hd
     return total
@@ -824,8 +803,8 @@ def _compute_layer_elements(arch: ModelArchConfig):
             n_moe = n_layers - n_dense
             moe_mlp_total = _compute_moe_mlp_elements(arch) * n_moe
             if arch.moe_has_dense_mlp:
-                # why: enable_moe_block runs dense MLP and MoE experts in
-                # parallel; count dense for every layer alongside MoE.
+                # enable_moe_block runs dense MLP and MoE experts in parallel;
+                # count dense for every layer alongside MoE.
                 mlp_total = sum(per_layer_dense_mlp) + moe_mlp_total
             else:
                 dense_only_total = sum(
@@ -958,11 +937,10 @@ def compute_lora_params(arch: ModelArchConfig, lora_rank: int, target_modules: l
         if n_experts > 1:
             n_dense = arch.num_dense_layers
             n_moe = n_layers - n_dense
-            # why: peft "all-linear" attaches LoRA to nn.Linear only;
-            # routed experts are nn.Parameter and need explicit
-            # gate_proj/up_proj/down_proj naming via Unsloth's
-            # get_moe_target_parameters. Shared experts are nn.Linear and
-            # are picked up by get_peft_regex.
+            # peft "all-linear" attaches LoRA to nn.Linear only; routed experts
+            # are nn.Parameter and need explicit gate_proj/up_proj/down_proj
+            # naming via Unsloth's get_moe_target_parameters. Shared experts are
+            # nn.Linear, picked up by get_peft_regex.
             routed_moe = (
                 0
                 if all_linear
@@ -983,7 +961,7 @@ def compute_lora_params(arch: ModelArchConfig, lora_rank: int, target_modules: l
             )
             moe_mlp = routed_moe + shared_moe
             if arch.moe_has_dense_mlp:
-                # why: parallel dense MLP coexists with MoE on every layer.
+                # Parallel dense MLP coexists with MoE on every layer.
                 mlp_total = structured_dense_mlp + moe_mlp * n_moe
             else:
                 dense_only = sum(
@@ -999,7 +977,7 @@ def compute_lora_params(arch: ModelArchConfig, lora_rank: int, target_modules: l
         attn_total = _lora_attn_elements(arch, r, selected_modules) * n_layers
         n_dense = arch.num_dense_layers
         n_moe = n_layers - n_dense
-        # why: routed and shared experts may use different intermediate sizes
+        # Routed and shared experts may use different intermediate sizes
         # (Qwen3.5-MoE: routed mlp_size != shared_expert_intermediate_size).
         # See structured branch for the all-linear exclusion rationale; only
         # routed (nn.Parameter) experts are excluded under all-linear.
@@ -1064,7 +1042,7 @@ def compute_gradient_bytes(trainable_params: int) -> int:
 
 
 def _is_linear_attention(attention_implementation: Optional[str]) -> bool:
-    # why: PyTorch SDPA dispatches to flash/memory-efficient O(n) backends; only
+    # PyTorch SDPA dispatches to flash/memory-efficient O(n) backends; only
     # eager (and other non-flash impls) need the quadratic correction.
     return attention_implementation in LINEAR_ATTENTION_IMPLS
 
@@ -1081,17 +1059,16 @@ def _layer_qkv_mlp_sizes(arch: ModelArchConfig, layer_idx: int) -> tuple:
     is_moe_layer = n_experts > 1 and not _is_dense_mlp_layer(arch, layer_idx)
     if _uses_structured_layer_shapes(arch):
         q_size, kv_size, _has_k, _has_v = _layer_attention_dims(arch, layer_idx)
-        # why: KV-shared layers (Gemma4/Gemma3n) drop k_proj/v_proj WEIGHTS but
-        # the donor layer's K/V tensors stay alive across the shared range, so
-        # activation memory still pays for kv_size; only the weight path uses
-        # has_k/has_v.
+        # KV-shared layers (Gemma4/Gemma3n) drop k/v WEIGHTS but the donor's
+        # K/V tensors stay alive, so activations still pay kv_size; only the
+        # weight path uses has_k/has_v.
         layer_type = _layer_types(arch)[layer_idx]
         use_alt_attention = arch.attention_k_eq_v and layer_type != "sliding_attention"
         kv_count = 1 if use_alt_attention else 2
         qkv_size = q_size + kv_size * kv_count
         if is_moe_layer:
-            # why: each token routes through `num_experts_per_tok` experts; their
-            # gate/up/down intermediates are all live during MLP forward.
+            # Each token routes through num_experts_per_tok experts; all their
+            # gate/up/down intermediates are live during MLP forward.
             mlp_size = _get_mlp_size(arch) * arch.num_experts_per_tok
             if arch.n_shared_experts:
                 mlp_size += _shared_expert_size(arch) * arch.n_shared_experts
@@ -1119,9 +1096,8 @@ def _per_layer_activation_bytes(
     activation_qkv = seq_len * batch_size * qkv_size
     residual_memory = (seq_len * batch_size) * 2
     activation_mlp = seq_len * batch_size * (mlp_size + mlp_size)
-    # why: per_layer_input_gate (hd-sized) and per_layer_projection (pli-sized)
-    # outputs materialize once per decoder layer when hidden_size_per_layer_input
-    # is set; see gemma4/modular_gemma4.py:1141-1145.
+    # PLE gate (hd) + projection (pli) outputs materialize once per decoder
+    # layer when hidden_size_per_layer_input is set (gemma4 modular:1141-1145).
     pli = arch.hidden_size_per_layer_input
     activation_ple = seq_len * batch_size * (arch.hidden_size + pli) if pli > 0 else 0
     return int((activation_qkv + residual_memory + activation_mlp + activation_ple) * 2 * 1.25)
@@ -1154,8 +1130,8 @@ def compute_activation_bytes(
         )
         linear_bytes = int(max_layer_bytes * effective_layers)
 
-    # why: gemma4 per_layer_model_projection runs once outside the per-decoder
-    # loop and materializes a [B, S, L, PLI] tensor; see modular_gemma4.py:1247.
+    # gemma4 per_layer_model_projection runs once outside the per-decoder loop
+    # and materializes a [B, S, L, PLI] tensor; see modular_gemma4.py:1247.
     pli = arch.hidden_size_per_layer_input
     if pli > 0:
         linear_bytes += int(seq_len * batch_size * n_layers * pli * 2 * 1.25)
diff --git a/studio/backend/utils/helper_precache_settings.py b/studio/backend/utils/helper_precache_settings.py
new file mode 100644
index 0000000000..db19a2d028
--- /dev/null
+++ b/studio/backend/utils/helper_precache_settings.py
@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Persisted opt-in controls for Helper LLM startup pre-cache."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+HELPER_PRECACHE_SETTING_KEY = "helper_model_preload_on_startup"
+DEFAULT_HELPER_PRECACHE_ENABLED = False
+
+
+def _coerce_bool(value: Any) -> bool | None:
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        normalized = value.strip().lower()
+        if normalized in {"1", "true", "yes", "on"}:
+            return True
+        if normalized in {"0", "false", "no", "off", ""}:
+            return False
+    return None
+
+
+def helper_model_disabled_by_env() -> bool:
+    """Return True when existing broad helper-disable env var is active."""
+    return os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in {"1", "true"}
+
+
+def get_helper_precache_enabled() -> bool:
+    """Read the persisted startup pre-cache preference.
+
+    Missing or unreadable settings default to False so Studio startup never
+    performs optional network work unless the user explicitly opted in.
+    """
+    try:
+        from storage.studio_db import get_app_setting
+        stored = get_app_setting(HELPER_PRECACHE_SETTING_KEY, None)
+    except Exception:
+        stored = None
+    parsed = _coerce_bool(stored)
+    return parsed if parsed is not None else DEFAULT_HELPER_PRECACHE_ENABLED
+
+
+def set_helper_precache_enabled(value: Any) -> bool:
+    """Persist whether Studio should pre-cache the Helper LLM at startup."""
+    parsed = _coerce_bool(value)
+    if parsed is None:
+        raise ValueError("Helper LLM startup pre-cache must be true or false.")
+
+    from storage.studio_db import upsert_app_settings
+
+    upsert_app_settings({HELPER_PRECACHE_SETTING_KEY: parsed})
+    return parsed
+
+
+def should_preload_helper_on_startup() -> bool:
+    """Gate the startup pre-cache thread.
+
+    The persisted setting is opt-in and the existing broad disable env var wins.
+    Explicit AI Assist calls do not use this gate; they remain user-triggered.
+    """
+    return get_helper_precache_enabled() and not helper_model_disabled_by_env()
diff --git a/studio/backend/utils/inference/inference_config.py b/studio/backend/utils/inference/inference_config.py
index e07bc62cfd..05eb08067c 100644
--- a/studio/backend/utils/inference/inference_config.py
+++ b/studio/backend/utils/inference/inference_config.py
@@ -1,13 +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
 
-"""
-Inference configuration loading utilities.
-
-This module provides functions to load inference parameters (temperature, top_p, top_k, min_p)
-from model YAML configuration files, with fallback to default.yaml.
-Includes family-based lookup from inference_defaults.json for GGUF models.
-"""
+"""Load inference params (temperature, top_p, top_k, min_p) from model YAML, family defaults, or default.yaml."""
 
 from pathlib import Path
 from typing import Dict, Any, Optional
@@ -48,29 +42,22 @@ def _load_family_defaults():
 
 
 def get_family_inference_params(model_id: str) -> Dict[str, Any]:
-    """
-    Look up recommended inference parameters by model family.
+    """Look up recommended inference params by model family.
 
-    Extracts the model family from the identifier (e.g. "unsloth/Qwen3.5-9B-GGUF" -> "qwen3.5")
-    and returns the matching parameters from inference_defaults.json.
-
-    Args:
-        model_id: Model identifier (e.g. "unsloth/Qwen3.5-9B-GGUF")
-
-    Returns:
-        Dict with inference params, or empty dict if no family match.
+    Extracts the family from the identifier (e.g. "unsloth/Qwen3.5-9B-GGUF" ->
+    "qwen3.5") and returns matching params from inference_defaults.json, or {}.
     """
     _load_family_defaults()
 
     if not _FAMILY_PATTERNS or not _FAMILY_DEFAULTS:
         return {}
 
-    # Normalize: lowercase, strip org prefix
+    # Normalize: lowercase, strip org prefix.
     normalized = model_id.lower()
     if "/" in normalized:
         normalized = normalized.split("/", 1)[1]
 
-    # Match against patterns (ordered longest-match-first in the JSON)
+    # Match patterns (ordered longest-match-first in the JSON).
     for pattern in _FAMILY_PATTERNS:
         if pattern in normalized:
             params = _FAMILY_DEFAULTS.get(pattern, {})
@@ -87,14 +74,11 @@ def _has_specific_yaml(model_identifier: str) -> bool:
     script_dir = Path(__file__).parent.parent.parent
     defaults_dir = script_dir / "assets" / "configs" / "model_defaults"
 
-    # Check the mapping
     if model_identifier.lower() in _REVERSE_MODEL_MAPPING:
         return True
 
-    # For local filesystem paths (e.g. C:\Users\...\model on Windows),
-    # normalize backslashes so Path().parts splits correctly on POSIX/WSL,
-    # then try matching the last 1-2 path components against the registry
-    # (mirrors the logic in load_model_defaults).
+    # For local paths, normalize backslashes so Path().parts splits correctly,
+    # then match the last 1-2 components against the registry (mirrors load_model_defaults).
     _is_local = is_local_path(model_identifier)
     _normalized = normalize_path(model_identifier) if _is_local else model_identifier
 
@@ -109,9 +93,7 @@ def _has_specific_yaml(model_identifier: str) -> bool:
     else:
         _lookup = model_identifier
 
-    # Check for exact filename match (basename for local paths to avoid
-    # passing absolute paths into rglob which raises
-    # "Non-relative patterns are unsupported" on Windows).
+    # Exact filename match (basename for local paths; absolute paths break rglob on Windows).
     model_filename = _lookup.replace("/", "_") + ".yaml"
     for config_path in defaults_dir.rglob(model_filename):
         if config_path.is_file():
@@ -121,30 +103,14 @@ def _has_specific_yaml(model_identifier: str) -> bool:
 
 
 def load_inference_config(model_identifier: str) -> Dict[str, Any]:
+    """Load inference params for a model.
+
+    Priority: model-specific YAML, then family defaults (inference_defaults.json),
+    then default.yaml. Returns a dict of temperature/top_p/top_k/min_p/etc.
     """
-    Load inference configuration parameters for a model.
-
-    Priority chain:
-    1. Model-specific YAML (if it exists and has inference params)
-    2. Family-based defaults from inference_defaults.json
-    3. default.yaml fallback
-
-    Args:
-        model_identifier: Model identifier (e.g., "unsloth/llama-3-8b-bnb-4bit")
-
-    Returns:
-        Dictionary containing inference parameters:
-        {
-            "temperature": float,
-            "top_p": float,
-            "top_k": int,
-            "min_p": float
-        }
-    """
-    # Load model defaults to get inference parameters
     model_defaults = load_model_defaults(model_identifier)
 
-    # Load default.yaml for fallback values
+    # default.yaml for fallback values.
     script_dir = Path(__file__).parent.parent.parent
     defaults_dir = script_dir / "assets" / "configs" / "model_defaults"
     default_config_path = defaults_dir / "default.yaml"
@@ -158,18 +124,18 @@ def load_inference_config(model_identifier: str) -> Dict[str, Any]:
         except Exception as e:
             logger.warning(f"Failed to load default.yaml: {e}")
 
-    # Family-based defaults from inference_defaults.json
+    # Family-based defaults from inference_defaults.json.
     family_params = get_family_inference_params(model_identifier)
 
     model_inference = model_defaults.get("inference", {})
 
-    # If the model has its own YAML config, those values take priority over family defaults.
-    # If it only fell back to default.yaml, family defaults take priority.
+    # Model's own YAML beats family defaults; if it only fell back to
+    # default.yaml, family defaults win.
     has_own_yaml = _has_specific_yaml(model_identifier)
 
     def _get_param(key, hardcoded_default):
         if has_own_yaml:
-            # Model-specific YAML wins, then family fills gaps, then default.yaml
+            # Model-specific YAML wins, then family fills gaps, then default.yaml.
             val = model_inference.get(key)
             if val is not None and isinstance(val, (int, float)):
                 return val
@@ -177,7 +143,7 @@ def load_inference_config(model_identifier: str) -> Dict[str, Any]:
                 return family_params[key]
             return default_inference.get(key, hardcoded_default)
         else:
-            # No model-specific YAML: family wins, then default.yaml
+            # No model-specific YAML: family wins, then default.yaml.
             if key in family_params:
                 return family_params[key]
             return default_inference.get(key, hardcoded_default)
diff --git a/studio/backend/utils/llama_cpp_freshness.py b/studio/backend/utils/llama_cpp_freshness.py
index f0e642a38e..3e5066ca2d 100644
--- a/studio/backend/utils/llama_cpp_freshness.py
+++ b/studio/backend/utils/llama_cpp_freshness.py
@@ -45,7 +45,7 @@ def _cache_dir() -> Path:
 
 def read_install_marker(binary_path: Optional[str]) -> Optional[dict]:
     """Walk up from binary_path to find UNSLOTH_PREBUILT_INFO.json.
-    None means no marker (source build / custom path) or invalid JSON."""
+    None = no marker (source build / custom path) or invalid JSON."""
     if not binary_path:
         return None
     cached = _marker_cache.get(binary_path)
@@ -53,10 +53,7 @@ def read_install_marker(binary_path: Optional[str]) -> Optional[dict]:
         return cached
     p = Path(binary_path)
     marker: Optional[dict] = None
-    # Cover all _find_llama_server_binary layouts:
-    #   /llama-server                          (1 up)
-    #   /build/bin/llama-server                (3 up, Linux/macOS cmake)
-    #   /build/bin/Release/llama-server.exe   (4 up, Windows cmake)
+    # Cover all _find_llama_server_binary layouts (binary is 1-4 dirs deep):
     for parent in p.parents[:5]:
         candidate = parent / _INSTALL_MARKER_NAME
         if candidate.is_file():
diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py
index 808e2b012e..74d08ac116 100644
--- a/studio/backend/utils/models/__init__.py
+++ b/studio/backend/utils/models/__init__.py
@@ -1,9 +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
 
-"""
-Model and LoRA configuration handling
-"""
+"""Model and LoRA configuration handling."""
 
 from .model_config import (
     ModelConfig,
diff --git a/studio/backend/utils/models/checkpoints.py b/studio/backend/utils/models/checkpoints.py
index 99770ce7a6..5a992926ec 100644
--- a/studio/backend/utils/models/checkpoints.py
+++ b/studio/backend/utils/models/checkpoints.py
@@ -1,9 +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
 
-"""
-Checkpoint scanning utilities for discovering training runs and their checkpoints.
-"""
+"""Checkpoint scanning utilities for discovering training runs and checkpoints."""
 
 import json
 import structlog
@@ -16,11 +14,7 @@ logger = get_logger(__name__)
 
 
 def _read_checkpoint_loss(checkpoint_path: Path) -> Optional[float]:
-    """
-    Read the training loss from a checkpoint's trainer_state.json.
-
-    Returns the loss from the last log_history entry, or None if unavailable.
-    """
+    """Read loss from the last log_history entry of trainer_state.json, or None."""
     trainer_state = checkpoint_path / "trainer_state.json"
     if not trainer_state.exists():
         return None
@@ -38,14 +32,13 @@ def _read_checkpoint_loss(checkpoint_path: Path) -> Optional[float]:
 def scan_checkpoints(
     outputs_dir: str = str(outputs_root()),
 ) -> List[Tuple[str, List[Tuple[str, str, Optional[float]]], dict]]:
-    """
-    Scan outputs folder for training runs and their checkpoints.
+    """Scan outputs folder for training runs and their checkpoints.
 
     Returns:
-        List of tuples: [(model_name, [(display_name, checkpoint_path, loss), ...], metadata), ...]
-        metadata keys: base_model, peft_type, lora_rank (all optional)
-        The first entry in each checkpoint list is the main adapter; its loss is
-        set to the loss of the last (highest-step) intermediate checkpoint.
+        [(model_name, [(display_name, checkpoint_path, loss), ...], metadata), ...]
+        metadata keys (optional): base_model, peft_type, lora_rank.
+        First checkpoint entry is the main adapter; its loss mirrors the last
+        (highest-step) intermediate checkpoint.
     """
     models = []
     outputs_path = resolve_output_dir(outputs_dir)
@@ -65,7 +58,7 @@ def scan_checkpoints(
             if not (config_file.exists() or adapter_config.exists()):
                 continue
 
-            # Extract training metadata from adapter_config.json / config.json
+            # Training metadata from adapter_config.json / config.json
             metadata: dict = {}
             try:
                 if adapter_config.exists():
@@ -77,7 +70,7 @@ def scan_checkpoints(
                     cfg = json.loads(config_file.read_text())
                     metadata["base_model"] = cfg.get("_name_or_path")
 
-                # Detect BNB quantization from config.json (present in both cases)
+                # Detect BNB quantization from config.json
                 if config_file.exists():
                     if "cfg" not in dir():
                         cfg = json.loads(config_file.read_text())
@@ -91,8 +84,8 @@ def scan_checkpoints(
             except Exception:
                 pass
 
-            # Fallback: extract base model name from folder name
-            # e.g. "unsloth_Llama-3.2-3B-Instruct_1771227800" → "unsloth/Llama-3.2-3B-Instruct"
+            # Fallback: extract base model name from the folder name, e.g.
+            # "unsloth_Llama-3.2-3B-Instruct_1771227800" → "unsloth/Llama-3.2-3B-Instruct"
             if not metadata.get("base_model"):
                 parts = item.name.rsplit("_", 1)
                 if len(parts) == 2 and parts[1].isdigit():
@@ -103,13 +96,13 @@ def scan_checkpoints(
                     else:
                         metadata["base_model"] = name_part
 
-            # This is a valid training run
+            # Valid training run.
             checkpoints = []
 
-            # Placeholder for the main adapter — loss filled from last checkpoint below
+            # Main adapter placeholder — loss filled from the last checkpoint below.
             checkpoints.append((item.name, str(item), None))
 
-            # Scan for intermediate checkpoints (checkpoint-N subdirs)
+            # Scan for intermediate checkpoints (checkpoint-N subdirs).
             for sub in sorted(item.iterdir()):
                 if not sub.is_dir() or not sub.name.startswith("checkpoint-"):
                     continue
@@ -119,7 +112,7 @@ def scan_checkpoints(
                     loss = _read_checkpoint_loss(sub)
                     checkpoints.append((sub.name, str(sub), loss))
 
-            # Assign the last checkpoint's loss to the main adapter entry
+            # Assign the last checkpoint's loss to the main adapter entry.
             if len(checkpoints) > 1:
                 last_checkpoint_loss = checkpoints[-1][2]
                 checkpoints[0] = (
diff --git a/studio/backend/utils/models/gguf_metadata.py b/studio/backend/utils/models/gguf_metadata.py
index 32618773b7..a2912cc843 100644
--- a/studio/backend/utils/models/gguf_metadata.py
+++ b/studio/backend/utils/models/gguf_metadata.py
@@ -1,10 +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
 
-"""Free-function ``general.*`` reader for GGUF headers, used by
-``detect_mmproj_file`` to pair weights and projectors via
-``general.base_model.0.repo_url``. ~30 ms per file, cached by
-(path, mtime, size)."""
+"""``general.*`` reader for GGUF headers, used by ``detect_mmproj_file`` to
+pair weights and projectors via ``general.base_model.0.repo_url``. ~30 ms
+per file, cached by (path, mtime, size)."""
 
 from __future__ import annotations
 
@@ -65,9 +64,9 @@ def _cache_key(path: str) -> Optional[_CacheKey]:
 
 
 def read_gguf_general_metadata(path: str) -> Optional[Dict[str, str]]:
-    """Return ``general.*`` strings from a GGUF header, or ``None`` if
-    the file is missing, unreadable, or not a GGUF. ``{}`` means the
-    file is valid but carries none of the wanted keys."""
+    """Return ``general.*`` strings from a GGUF header, or ``None`` if the
+    file is missing, unreadable, or not a GGUF. ``{}`` means valid but
+    carrying none of the wanted keys."""
     key = _cache_key(path)
     if key is None:
         return None
@@ -156,9 +155,9 @@ _FIXED_VTYPE_SIZES: Dict[int, int] = {
 
 
 def _skip_gguf_value(f, vtype: int) -> bool:
-    """Advance past one GGUF value. ``f.seek(.., 1)`` past EOF is legal
-    on a regular file so truncation is detected on the next read; we
-    only return False for unknown types or sanity-bound overflow."""
+    """Advance past one GGUF value. ``f.seek(.., 1)`` past EOF is legal on a
+    regular file, so truncation is caught on the next read; return False only
+    for unknown types or sanity-bound overflow."""
     if vtype == 8:  # STRING
         slen_bytes = f.read(8)
         if len(slen_bytes) < 8:
@@ -265,9 +264,9 @@ def _read_gguf_bool(path: str, wanted_key: str) -> Optional[bool]:
 
 
 def read_mmproj_audio_capability(path: str) -> Optional[bool]:
-    """``clip.has_audio_encoder`` from an mmproj GGUF (e.g. Gemma 4's gemma4ua):
-    ``True``/``False`` if present, ``None`` if absent / unreadable. Flags
-    audio-input models independently of tokenizer token names."""
+    """``clip.has_audio_encoder`` from an mmproj GGUF (e.g. Gemma 4's
+    gemma4ua): ``True``/``False`` if present, ``None`` if absent/unreadable.
+    Flags audio-input models independently of tokenizer token names."""
     return _read_gguf_bool(path, "clip.has_audio_encoder")
 
 
diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py
index 92960485a4..0df7a1a477 100644
--- a/studio/backend/utils/models/model_config.py
+++ b/studio/backend/utils/models/model_config.py
@@ -1,9 +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
 
-"""
-Model and LoRA configuration handling
-"""
+"""Model and LoRA configuration handling."""
 
 from dataclasses import dataclass
 from typing import Optional, Dict, Any
@@ -58,7 +56,7 @@ def _env_offline() -> bool:
 import re as _re
 
 _MODEL_SIZE_RE = _re.compile(r"(?:^|[-_/])(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE)
-# MoE active-parameter pattern: matches "A3B", "A3.5B", etc.
+# MoE active-parameter pattern: "A3B", "A3.5B", etc.
 _ACTIVE_SIZE_RE = _re.compile(r"(?:^|[-_/])a(\d+\.?\d*)\s*([bm])(?:$|[-_/])", _re.IGNORECASE)
 
 
@@ -66,8 +64,8 @@ def extract_model_size_b(model_id: str) -> float | None:
     """Extract model size in billions from a model identifier.
 
     Prefers MoE active-parameter notation (e.g. ``A3B`` in
-    ``Qwen3.5-35B-A3B``) over the total parameter count.
-    Handles both ``B`` (billions) and ``M`` (millions) suffixes.
+    ``Qwen3.5-35B-A3B``) over total params. Handles ``B`` (billions)
+    and ``M`` (millions) suffixes.
     """
     mid = (model_id or "").lower()
     active = _ACTIVE_SIZE_RE.search(mid)
@@ -81,9 +79,9 @@ def extract_model_size_b(model_id: str) -> float | None:
     return val / 1000.0 if size.group(2).lower() == "m" else val
 
 
-# Model name mapping: maps all equivalent model names to their canonical YAML config file
-# Format: "canonical_model_name.yaml": [list of all equivalent model names]
-# Based on the model mapper provided - canonical filename is based on the first model name in the mapper
+# Maps equivalent model names to their canonical YAML config file.
+# Format: "canonical_model_name.yaml": [equivalent model names].
+# Canonical filename derives from the first model name in each list.
 MODEL_NAME_MAPPING = {
     # ── Embedding models ──
     "unsloth_all-MiniLM-L6-v2.yaml": [
@@ -457,7 +455,7 @@ MODEL_NAME_MAPPING = {
     ],
 }
 
-# Reverse mapping for quick lookup: model_name -> canonical_filename
+# Reverse lookup: model_name -> canonical_filename
 _REVERSE_MODEL_MAPPING = {}
 for canonical_file, model_names in MODEL_NAME_MAPPING.items():
     for model_name in model_names:
@@ -470,19 +468,16 @@ def load_model_config(
     token: Optional[str] = None,
     trust_remote_code: bool = True,
 ):
-    """
-    Load model config with optional authentication control.
-    """
+    """Load model config with optional authentication control."""
     from transformers import AutoConfig
 
     if token:
-        # Explicit token provided - use it
         return AutoConfig.from_pretrained(
             model_name, trust_remote_code = trust_remote_code, token = token
         )
 
     if not use_auth:
-        # Load without any authentication (for public model checks)
+        # No auth, for public model checks
         with without_hf_auth():
             return AutoConfig.from_pretrained(
                 model_name,
@@ -490,7 +485,7 @@ def load_model_config(
                 token = None,
             )
 
-    # Use default authentication (cached tokens)
+    # Default auth (cached tokens)
     return AutoConfig.from_pretrained(
         model_name,
         trust_remote_code = trust_remote_code,
@@ -507,8 +502,13 @@ _VLM_MODEL_TYPES = {
     "internvl_chat",
     "cogvlm2",
     "minicpmv",
+    "gemma4",
 }
 
+# Audio-only models that share the ForConditionalGeneration suffix
+# (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration).
+_AUDIO_ONLY_MODEL_TYPES = {"csm", "whisper"}
+
 # Pre-computed .venv_t5 paths and backend dir for subprocess version switching.
 # Vision check uses 5.5.0 (newest, recognizes all architectures).
 from utils.paths.storage_roots import studio_root as _studio_root  # noqa: E402
@@ -516,9 +516,77 @@ from utils.paths.storage_roots import studio_root as _studio_root  # noqa: E402
 _VENV_T5_DIR = str(_studio_root() / ".venv_t5_550")
 _BACKEND_DIR = str(Path(__file__).resolve().parent.parent.parent)
 
-# Inline script executed in a subprocess with transformers 5.x activated.
-# Receives model_name and token via argv, prints JSON result to stdout.
-_VISION_CHECK_SCRIPT = r"""
+
+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")
+        or hasattr(config, "img_processor")
+        or hasattr(config, "image_token_index")
+        or model_type in _VLM_MODEL_TYPES
+    )
+
+
+def _raw_config_has_vision_config(
+    model_name: str, hf_token: Optional[str] = None
+) -> Optional[bool]:
+    try:
+        if is_local_path(model_name):
+            config_path = Path(normalize_path(model_name)).expanduser() / "config.json"
+        else:
+            from huggingface_hub import hf_hub_download
+            config_path = Path(
+                hf_hub_download(
+                    repo_id = model_name,
+                    filename = "config.json",
+                    token = hf_token,
+                )
+            )
+        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
+            or "img_processor" in config
+            or "image_token_index" in config
+            or model_type in _VLM_MODEL_TYPES
+        )
+    except Exception as exc:
+        logger.warning("Could not read config.json for '%s': %s", model_name, exc)
+        return None
+
+
+# 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"
+    "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"
+    "        or hasattr(config, 'img_processor')\n"
+    "        or hasattr(config, 'image_token_index')\n"
+    "        or model_type in _VLM_MODEL_TYPES\n"
+    "    )\n"
+)
+
+# Subprocess script run with transformers 5.x active. Takes model_name and
+# token via argv, prints JSON result to stdout.
+_VISION_CHECK_SCRIPT = (
+    r"""
 import sys, os, json
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
@@ -532,32 +600,20 @@ sys.path.insert(0, venv_t5)
 if backend_dir not in sys.path:
     sys.path.insert(0, backend_dir)
 
+"""
+    + _VISION_CHECK_INLINE_HELPERS
+    + r"""
 try:
     from transformers import AutoConfig
+
     kwargs = {"trust_remote_code": True}
     if token:
         kwargs["token"] = token
     config = AutoConfig.from_pretrained(model_name, **kwargs)
 
-    is_vlm = False
-    if hasattr(config, "architectures"):
-        is_vlm = any(
-            x.endswith(("ForConditionalGeneration", "ForVisionText2Text"))
-            for x in config.architectures
-        )
-    if not is_vlm and hasattr(config, "vision_config"):
-        is_vlm = True
-    if not is_vlm and hasattr(config, "img_processor"):
-        is_vlm = True
-    if not is_vlm and hasattr(config, "image_token_index"):
-        is_vlm = True
-    if not is_vlm and hasattr(config, "model_type"):
-        vlm_types = {"phi3_v","llava","llava_next","llava_onevision",
-                      "internvl_chat","cogvlm2","minicpmv"}
-        if config.model_type in vlm_types:
-            is_vlm = True
+    is_vlm = _is_vlm(config)
 
-    model_type = getattr(config, "model_type", "unknown")
+    model_type = getattr(config, "model_type", None)
     archs = getattr(config, "architectures", [])
     print(json.dumps({"is_vision": is_vlm, "model_type": model_type,
                        "architectures": archs}))
@@ -565,19 +621,16 @@ except Exception as exc:
     print(json.dumps({"error": str(exc)}))
     sys.exit(1)
 """
+)
 
 
 def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None) -> Optional[bool]:
-    """Run is_vision_model check in a subprocess with transformers 5.x.
+    """Run is_vision_model in a subprocess with transformers 5.x.
 
-    Same pattern as training/inference workers: spawn a clean subprocess
-    with .venv_t5/ prepended to sys.path so AutoConfig recognizes newer
-    architectures (glm4_moe_lite, etc.).
-
-    Returns True/False for definitive results, or None for transient failures
-    (timeouts, subprocess errors) so callers can decide whether to cache
-    the result. Subprocess failures are treated as transient because they
-    can be caused by temporary HF/auth/network issues.
+    Spawns a clean subprocess with .venv_t5/ on sys.path so AutoConfig
+    recognizes newer architectures. Returns True/False for definitive results,
+    or None for transient failures (timeouts, subprocess errors), which are not
+    cached so they can be retried.
     """
     token_arg = hf_token or ""
 
@@ -637,42 +690,35 @@ def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None)
 
 
 def _token_fingerprint(token: Optional[str]) -> Optional[str]:
-    """Return a SHA256 digest of the token for use as a cache key.
-
-    Avoids storing the raw bearer token in process memory as a dict key.
-    """
+    """SHA256 digest of the token for use as a cache key (avoids storing the
+    raw bearer token in process memory)."""
     if token is None:
         return None
     return hashlib.sha256(token.encode("utf-8")).hexdigest()
 
 
-# Cache vision detection results per session to avoid repeated subprocess spawns.
-# Keyed by (normalized_model_name, token_fingerprint) to handle gated models correctly.
-# Only definitive results (True/False from successful detection) are cached;
-# transient failures (network errors, timeouts) are NOT cached so they can be retried.
+# Cache vision detection per session to avoid repeated subprocess spawns.
+# Keyed by (normalized_model_name, token_fingerprint) to handle gated models.
+# Only definitive results are cached; transient failures (network, timeouts)
+# are NOT cached so they can be retried.
 _vision_detection_cache: Dict[Tuple[str, Optional[str]], bool] = {}
 _vision_cache_lock = threading.Lock()
 
 
 def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
     """
-    Detect vision-language models (VLMs) by checking architecture in config.
-    Works for fine-tuned models since they inherit the base architecture.
+    Detect vision-language models (VLMs) via architecture in config. Works for
+    fine-tuned models since they inherit the base architecture.
 
-    For models that require transformers 5.x (e.g. GLM-4.7-Flash), the check
-    runs in a subprocess with .venv_t5/ activated -- same pattern as the
-    training and inference workers.
-
-    Results are cached per (model_name, token_fingerprint) for the lifetime of
-    the process to avoid repeated subprocess spawns and HuggingFace API calls.
-    Transient failures are not cached so they can be retried on the next call.
+    Models needing transformers 5.x are checked in a .venv_t5/ subprocess.
+    Results are cached per (model_name, token_fingerprint) for the process
+    lifetime; transient failures are not cached so they can be retried.
 
     Args:
         model_name: Model identifier (HF repo or local path)
-        hf_token: Optional HF token for accessing gated/private models
+        hf_token: Optional HF token for gated/private models
     """
-    # Normalize model name for cache key to avoid duplicate entries for
-    # different casings of the same HF repo (e.g. "Org/Model" vs "org/model").
+    # Normalize model name so different casings of the same repo share a key
     try:
         if is_local_path(model_name):
             resolved_name = normalize_path(model_name)
@@ -687,21 +733,17 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
         resolved_name = model_name
     cache_key = (resolved_name, _token_fingerprint(hf_token))
 
-    # Lock-free fast path for cache hits. Uses a sentinel to distinguish
-    # "key not found" from "value is False" in a single atomic dict.get() call.
+    # Lock-free fast path for cache hits. Sentinel distinguishes "key not found"
+    # from "value is False" in a single atomic dict.get() call.
     _MISS = object()
     cached = _vision_detection_cache.get(cache_key, _MISS)
     if cached is not _MISS:
         return cached
 
-    # Compute outside the lock to avoid serializing long-running detection
-    # (subprocess spawns with 60s timeout, HF API calls) across all models.
-    # The tradeoff: two concurrent calls for the same uncached model may
-    # both run detection, but they produce the same result and the second
-    # write is a benign no-op.
+    # Compute outside the lock so long-running detection isn't serialized across
+    # models. Two concurrent calls may both run, but produce the same result.
     result = _is_vision_model_uncached(resolved_name, hf_token)
-    # Only cache definitive results; None means a transient failure occurred
-    # and we should retry on the next call instead of locking in a wrong answer.
+    # Only cache definitive results; None is a transient failure, retry later.
     if result is not None:
         with _vision_cache_lock:
             _vision_detection_cache[cache_key] = result
@@ -710,17 +752,13 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
 
 
 def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) -> Optional[bool]:
-    """Uncached vision model detection -- called by is_vision_model().
+    """Uncached vision detection; use is_vision_model() instead.
 
-    Returns True/False for definitive results, or None when detection failed
-    due to a transient error (network, timeout, subprocess failure) so the
-    caller knows not to cache the result.
-
-    Do not call directly; use is_vision_model() instead.
+    Returns True/False for definitive results, or None on transient errors
+    (network, timeout, subprocess failure) so the caller knows not to cache.
     """
-    # Models that need transformers 5.x must be checked in a subprocess
-    # because AutoConfig in the main process (transformers 4.57.x) doesn't
-    # recognize their architectures.
+    # Models needing transformers 5.x must be checked in a subprocess: the main
+    # process (transformers 4.57.x) doesn't recognize their architectures.
     from utils.transformers_version import needs_transformers_5
 
     if needs_transformers_5(model_name):
@@ -728,54 +766,36 @@ 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,
         )
-        return _is_vision_model_subprocess(model_name, hf_token = hf_token)
+        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)
 
     try:
         config = load_model_config(model_name, use_auth = True, token = hf_token)
 
-        # Exclude audio-only models that share ForConditionalGeneration suffix
+        # Exclude audio-only models sharing the ForConditionalGeneration suffix
         # (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration)
-        _audio_only_model_types = {"csm", "whisper"}
         model_type = getattr(config, "model_type", None)
-        if model_type in _audio_only_model_types:
+        if model_type in _AUDIO_ONLY_MODEL_TYPES:
             return False
 
-        # Check 1: Architecture class name patterns
-        if hasattr(config, "architectures"):
-            is_vlm = any(x.endswith(_VLM_ARCH_SUFFIXES) for x in config.architectures)
-            if is_vlm:
-                logger.info(
-                    f"Model {model_name} detected as VLM: architecture {config.architectures}"
-                )
-                return True
-
-        # Check 2: Has vision_config (most VLMs: LLaVA, Gemma-3, Qwen2-VL, etc.)
-        if hasattr(config, "vision_config"):
-            logger.info(f"Model {model_name} detected as VLM: has vision_config")
+        if _is_vlm(config):
+            archs = getattr(config, "architectures", None) or []
+            logger.info(
+                "Model %s detected as VLM (model_type=%s, architectures=%s)",
+                model_name,
+                model_type,
+                archs,
+            )
             return True
 
-        # Check 3: Has img_processor (Phi-3.5 Vision uses this instead of vision_config)
-        if hasattr(config, "img_processor"):
-            logger.info(f"Model {model_name} detected as VLM: has img_processor")
-            return True
-
-        # Check 4: Has image_token_index (common in VLMs for image placeholder tokens)
-        if hasattr(config, "image_token_index"):
-            logger.info(f"Model {model_name} detected as VLM: has image_token_index")
-            return True
-
-        # Check 5: Known VLM model_type values that may not match above checks
-        if hasattr(config, "model_type"):
-            if config.model_type in _VLM_MODEL_TYPES:
-                logger.info(f"Model {model_name} detected as VLM: model_type={config.model_type}")
-                return True
-
         return False
 
     except Exception as e:
         logger.warning(f"Could not determine if {model_name} is vision model: {e}")
-        # Permanent failures (model not found, gated, bad config) should be
-        # cached as False. Transient failures (network, timeout) should not.
+        # Permanent failures (not found, gated, bad config) cache as False;
+        # transient ones (network, timeout) should not.
         try:
             from huggingface_hub.errors import RepositoryNotFoundError, GatedRepoError
         except ImportError:
@@ -797,10 +817,10 @@ 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 results per session to avoid repeated API calls
+# Cache detection per session to avoid repeated API calls
 _audio_detection_cache: Dict[str, Optional[str]] = {}
 
-# Tokenizer token patterns → audio_type (all 6 types detected from tokenizer_config.json)
+# Tokenizer token patterns → audio_type (all 6 types from tokenizer_config.json)
 _AUDIO_TOKEN_PATTERNS = {
     "csm": lambda tokens: "<|AUDIO|>" in tokens and "<|audio_eos|>" in tokens,
     "whisper": lambda tokens: "<|startoftranscript|>" in tokens,
@@ -818,13 +838,11 @@ _AUDIO_TOKEN_PATTERNS = {
 
 
 def detect_audio_type(model_name: str, hf_token: Optional[str] = None) -> Optional[str]:
-    """
-    Dynamically detect if a model is an audio model and return its type.
+    """Detect if a model is an audio model and return its type.
 
-    Fully dynamic — works for any model, not just known ones.
-    Uses tokenizer_config.json special tokens to detect all 6 audio types.
-
-    Returns: audio_type string ('snac', 'csm', 'bicodec', 'dac', 'whisper', 'audio_vlm') or None.
+    Works for any model via tokenizer_config.json special tokens.
+    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]
@@ -838,10 +856,10 @@ def detect_audio_type(model_name: str, hf_token: Optional[str] = None) -> Option
 
 
 def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None) -> Optional[str]:
-    """Detect audio type from tokenizer special tokens (for LLM-based audio models).
+    """Detect audio type from tokenizer special tokens.
 
-    First checks local HF cache, then fetches tokenizer_config.json from HuggingFace.
-    Checks added_tokens_decoder for distinctive patterns.
+    Checks local HF cache first, then fetches tokenizer_config.json from HF;
+    examines added_tokens_decoder for distinctive patterns.
     """
 
     def _check_token_patterns(tok_config: dict) -> Optional[str]:
@@ -854,7 +872,7 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
                 return audio_type
         return None
 
-    # 1) Check local HF cache first (works for gated/offline models)
+    # 1) Local HF cache first (works for gated/offline models)
     try:
         repo_dir = get_cache_path(model_name)
         if repo_dir is not None and repo_dir.exists():
@@ -880,7 +898,6 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
         import os
 
         paths_to_try = ["tokenizer_config.json", "LLM/tokenizer_config.json"]
-        # Use provided token, or fall back to env
         token = hf_token or os.environ.get("HF_TOKEN")
         headers = {}
         if token:
@@ -904,10 +921,7 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
 
 
 def is_audio_input_type(audio_type: Optional[str]) -> bool:
-    """Check if an audio_type accepts audio input (ASR/speech understanding).
-
-    Whisper (ASR) and audio_vlm (Gemma3n) accept audio input.
-    """
+    """True if an audio_type accepts audio input: whisper (ASR), audio_vlm (Gemma3n)."""
     return audio_type in ("whisper", "audio_vlm")
 
 
@@ -916,8 +930,7 @@ def _is_mmproj(filename: str) -> bool:
     return "mmproj" in filename.lower()
 
 
-# Family tokens for #5347's filename fallback. Lowercase. Order does not
-# matter (see ``_detect_family_token``).
+# Family tokens for #5347's filename fallback. Lowercase; order irrelevant.
 _MODEL_FAMILY_TOKENS: tuple[str, ...] = (
     "qwen",
     "gemma",
@@ -950,8 +963,8 @@ _MODEL_FAMILY_TOKENS: tuple[str, ...] = (
 )
 
 
-# Word-bounded match: any letter on either side disqualifies. Stops
-# ``phi`` matching ``sapphire``, ``yi`` matching ``tiny``, etc.
+# Word-bounded match: a letter on either side disqualifies (stops ``phi``
+# matching ``sapphire``, ``yi`` matching ``tiny``).
 _FAMILY_TOKEN_RE_CACHE: Dict[str, "_re.Pattern[str]"] = {}
 
 
@@ -978,8 +991,8 @@ def _detect_family_token(filename: str) -> Optional[str]:
 
 
 def mmproj_matches_model_family(model_path: str, mmproj_path: str) -> bool:
-    """Defense-in-depth guard for the launcher: True unless both filenames
-    carry recognised family tokens that disagree."""
+    """Launcher guard: True unless both filenames carry recognised family
+    tokens that disagree."""
     model_fam = _detect_family_token(Path(model_path).name)
     mmproj_fam = _detect_family_token(Path(mmproj_path).name)
     if model_fam is None or mmproj_fam is None:
@@ -1072,7 +1085,7 @@ def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional
                 continue
             if resolved in seen_resolved:
                 continue
-            # Prefer ``general.type=='mmproj'``; fall back to filename.
+            # Prefer ``general.type=='mmproj'``, else filename.
             meta = read_gguf_general_metadata(str(resolved))
             by_meta = is_mmproj_by_metadata(meta)
             if by_meta is True or (by_meta is None and _is_mmproj(f.name)):
@@ -1125,18 +1138,11 @@ def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional
 
 
 def detect_gguf_model(path: str) -> Optional[str]:
-    """
-    Check if the given local path is or contains a GGUF model file.
+    """Check if a local path is or contains a GGUF model file.
 
-    Handles two cases:
-    1. path is a direct .gguf file path
-    2. path is a directory containing .gguf files
-
-    Skips mmproj (vision projection) files — those must be passed via
-    ``--mmproj``, not ``-m``.  Use :func:`detect_mmproj_file` instead.
-
-    Returns the full path to the .gguf file if found, None otherwise.
-    For HuggingFace repo detection, use detect_gguf_model_remote() instead.
+    Handles a direct .gguf path or a directory of .gguf files. Skips mmproj
+    files (pass those via ``--mmproj``; see :func:`detect_mmproj_file`). Returns
+    the .gguf path or None. For HF repos, use detect_gguf_model_remote().
     """
     p = Path(path)
 
@@ -1167,12 +1173,9 @@ def detect_gguf_model(path: str) -> Optional[str]:
     return None
 
 
-# Preferred GGUF quantization levels, in descending priority.
-# Q4_K_M is a good default: small, fast, acceptable quality.
-# UD (Unsloth Dynamic) variants are always preferred over standard quants
-# because they provide better quality per bit. If the repo has no UD variants
-# (e.g., bartowski repos), the standard quants are used as fallback.
-# Ordered by best size/quality tradeoff, not raw quality.
+# Preferred GGUF quant levels, descending priority. UD (Unsloth Dynamic)
+# variants beat standard quants on quality per bit; repos without UD fall back
+# to standard quants. Ordered by size/quality tradeoff, not raw quality.
 _GGUF_QUANT_PREFERENCE = [
     # UD variants (best quality per bit) -- Q4 is the sweet spot
     "UD-Q4_K_XL",
@@ -1216,23 +1219,16 @@ _GGUF_QUANT_PREFERENCE = [
 
 
 def _pick_best_gguf(filenames: list[str]) -> Optional[str]:
-    """
-    Pick the best GGUF file from a list of filenames.
-
-    Prefers quantization levels in _GGUF_QUANT_PREFERENCE order.
-    Falls back to the first .gguf file found.
-    """
+    """Pick the best GGUF file: quant levels in _GGUF_QUANT_PREFERENCE order, else first .gguf."""
     gguf_files = [f for f in filenames if f.lower().endswith(".gguf")]
     if not gguf_files:
         return None
 
-    # Try preferred quantization levels
     for quant in _GGUF_QUANT_PREFERENCE:
         for f in gguf_files:
             if quant in f:
                 return f
 
-    # Fallback: first GGUF file
     return gguf_files[0]
 
 
@@ -1247,7 +1243,7 @@ class GgufVariantInfo:
 
 def _extract_quant_label(filename: str) -> str:
     """
-    Extract quantization label like Q4_K_M, IQ4_XS, BF16 from a GGUF filename.
+    Extract quant label like Q4_K_M, IQ4_XS, BF16 from a GGUF filename.
 
     Examples:
         "gemma-3-4b-it-Q4_K_M.gguf"          → "Q4_K_M"
@@ -1274,8 +1270,8 @@ def _extract_quant_label(filename: str) -> str:
     )
     match = re.search(quant_re, stem, re.IGNORECASE)
     # Subdir layouts like ``BF16/foo.gguf`` keep the quant in the directory,
-    # not the basename. Look at the parent dirs too so the variant label
-    # matches the snapshot-relative path produced elsewhere.
+    # not the basename. Check parent dirs too so the label matches the
+    # snapshot-relative path produced elsewhere.
     if not match and "/" in filename:
         parents = filename.rsplit("/", 1)[0]
         for segment in reversed(parents.split("/")):
@@ -1286,16 +1282,16 @@ def _extract_quant_label(filename: str) -> str:
     if match:
         prefix = match.group(1) or ""
         return f"{prefix}{match.group(2)}"
-    # Fallback: last segment after hyphen
+    # Fallback: last hyphen-separated segment
     return stem.split("-")[-1]
 
 
 def _iter_hf_cache_snapshots(repo_id: str):
     """Yield HF cache snapshot dirs for *repo_id*, newest first.
 
-    Empty generator if HF_HUB_CACHE is missing, the repo isn't cached,
-    or has no snapshots. Repo name match is case-insensitive to handle
-    casing drift between download time and lookup.
+    Empty if HF_HUB_CACHE is missing, the repo isn't cached, or has no
+    snapshots. Repo name match is case-insensitive to handle casing drift
+    between download time and lookup.
     """
     try:
         from huggingface_hub import constants as hf_constants
@@ -1342,18 +1338,17 @@ def _list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufV
 def list_gguf_variants(
     repo_id: str, hf_token: Optional[str] = None
 ) -> tuple[list[GgufVariantInfo], bool]:
-    """
-    List all GGUF quantization variants in a HuggingFace repo.
+    """List all GGUF quant variants in a HF repo.
 
-    Separates main model files from mmproj (vision projection) files.
-    The presence of mmproj files indicates a vision-capable model.
+    Separates main model files from mmproj (vision projection) files; mmproj
+    presence flags a vision-capable model.
 
     Returns:
-        (variants, has_vision): list of non-mmproj GGUF variants + vision flag.
+        (variants, has_vision): non-mmproj GGUF variants + vision flag.
     """
     from huggingface_hub import model_info as hf_model_info
 
-    # Offline: skip the API and serve from cache.
+    # Offline: skip the API and serve from cache
     if _env_offline():
         cached = _list_gguf_variants_from_hf_cache(repo_id)
         if cached is not None:
@@ -1362,9 +1357,9 @@ def list_gguf_variants(
     try:
         info = hf_model_info(repo_id, token = hf_token, files_metadata = True)
     except Exception as e:
-        # Permanent errors (deleted/gated/bad revision) must surface to
-        # the caller; serving stale cache here would mask the real cause.
-        # Matches the early-return in ``detect_gguf_model_remote``.
+        # Permanent errors (deleted/gated/bad revision) must surface to the
+        # caller; serving stale cache would mask the real cause. Matches the
+        # early-return in ``detect_gguf_model_remote``.
         if type(e).__name__ in (
             "RepositoryNotFoundError",
             "GatedRepoError",
@@ -1386,7 +1381,7 @@ def list_gguf_variants(
     has_vision = False
 
     quant_totals: dict[str, int] = {}  # quant -> total bytes
-    quant_first_file: dict[str, str] = {}  # quant -> first filename (for display)
+    quant_first_file: dict[str, str] = {}  # quant -> first filename (display)
 
     for sibling in info.siblings:
         fname = sibling.rfilename
@@ -1394,7 +1389,7 @@ def list_gguf_variants(
             continue
         size = sibling.size or 0
 
-        # mmproj files are vision projection models, not main model files
+        # mmproj files are vision projections, not main model files
         if "mmproj" in fname.lower():
             has_vision = True
             continue
@@ -1413,9 +1408,8 @@ def list_gguf_variants(
             )
         )
 
-    # Sort by size descending (largest = best quality first).
-    # Recommended pinning and OOM demotion are handled client-side
-    # where GPU VRAM info is available.
+    # Sort by size descending (largest = best quality first); pinning and OOM
+    # demotion happen client-side where GPU VRAM info exists.
     variants.sort(key = lambda v: -v.size_bytes)
 
     return variants, has_vision
@@ -1424,11 +1418,10 @@ def list_gguf_variants(
 def _resolve_gguf_dir(p: Path) -> Optional[Path]:
     """Resolve a path to the directory containing GGUF variants.
 
-    If *p* is already a directory, returns it directly.  If *p* is a ``.gguf``
-    file whose parent directory has model metadata (``config.json`` or
-    ``adapter_config.json``), returns the parent -- all GGUFs in that
-    directory belong to the same model.  Returns ``None`` for loose standalone
-    GGUFs (no config) to avoid cross-wiring unrelated models.
+    Directory *p* returns directly. A ``.gguf`` file whose parent dir has
+    model metadata (``config.json`` or ``adapter_config.json``) returns the
+    parent -- all GGUFs there belong to the same model. Returns ``None`` for
+    loose standalone GGUFs (no config) to avoid cross-wiring unrelated models.
     """
     if p.is_dir():
         return p
@@ -1444,14 +1437,13 @@ def _resolve_gguf_dir(p: Path) -> Optional[Path]:
 
 
 def list_local_gguf_variants(directory: str) -> tuple[list[GgufVariantInfo], bool]:
-    """List GGUF quantization variants in a local directory.
+    """List GGUF quant variants in a local directory.
 
-    Mirrors :func:`list_gguf_variants` but reads from the filesystem
-    instead of the HuggingFace API.  Aggregates shard sizes by quant
-    label so that split GGUFs appear as a single variant.
+    Like :func:`list_gguf_variants` but reads the filesystem. Aggregates shard
+    sizes by quant label so split GGUFs appear as one variant.
 
     Returns:
-        (variants, has_vision): list of non-mmproj GGUF variants + vision flag.
+        (variants, has_vision): non-mmproj GGUF variants + vision flag.
     """
     p = _resolve_gguf_dir(Path(directory))
     if p is None:
@@ -1461,10 +1453,10 @@ def list_local_gguf_variants(directory: str) -> tuple[list[GgufVariantInfo], boo
     quant_first_file: dict[str, str] = {}
     has_vision = False
 
-    # Recurse so variant-specific subdirectories (e.g. ``BF16/...gguf``
-    # used by some HF GGUF repos for the largest quants) are picked up.
-    # Filenames in the result preserve the relative subpath so that
-    # ``_find_local_gguf_by_variant`` can locate the file again.
+    # Recurse so variant-specific subdirs (e.g. ``BF16/...gguf`` used by
+    # some HF GGUF repos for the largest quants) are picked up. Result
+    # filenames keep the relative subpath so ``_find_local_gguf_by_variant``
+    # can locate the file again.
     for f in sorted(_iter_gguf_files(p, recursive = True)):
         if _is_mmproj(f.name):
             has_vision = True
@@ -1473,8 +1465,8 @@ def list_local_gguf_variants(directory: str) -> tuple[list[GgufVariantInfo], boo
             size = f.stat().st_size
         except OSError:
             size = 0
-        # Pass the relative path so ``BF16/foo.gguf`` and ``Q4_K_M/foo.gguf``
-        # produce distinct quant labels instead of collapsing on basename.
+        # Use the relative path so ``BF16/foo.gguf`` and ``Q4_K_M/foo.gguf``
+        # get distinct quant labels instead of collapsing on basename.
         rel = f.relative_to(p).as_posix()
         quant = _extract_quant_label(rel)
         quant_totals[quant] = quant_totals.get(quant, 0) + size
@@ -1496,8 +1488,8 @@ def list_local_gguf_variants(directory: str) -> tuple[list[GgufVariantInfo], boo
 def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]:
     """Find the GGUF file in *directory* matching a quantization *variant*.
 
-    For sharded GGUFs (multiple files with the same quant label), returns
-    the first shard (sorted by name) which is what ``llama-server -m`` expects.
+    For sharded GGUFs (multiple files sharing a quant label), returns the
+    first shard (sorted by name), which is what ``llama-server -m`` expects.
 
     Returns the resolved absolute path, or ``None`` if no match.
     """
@@ -1505,10 +1497,10 @@ def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]:
     if p is None:
         return None
 
-    # Recurse into subdirectories so variants stored under a quant-named
-    # subdir (e.g. ``BF16/foo-BF16-00001-of-00002.gguf``) are found.
-    # Match against the relative path so the quant label can come from
-    # the directory name when the basename omits it.
+    # Recurse so variants under a quant-named subdir (e.g.
+    # ``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)
@@ -1522,8 +1514,8 @@ def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]:
 def _detect_gguf_from_hf_cache(repo_id: str) -> Optional[str]:
     """Best GGUF filename for *repo_id* from the local HF cache, or None.
 
-    Excludes mmproj (vision projector) files so a partial cache that
-    only has the projector cannot route the projector as the main model.
+    Excludes mmproj (vision projector) files so a partial cache holding only
+    the projector cannot route it as the main model.
     """
     for snap in _iter_hf_cache_snapshots(repo_id):
         rel_files = [
@@ -1537,21 +1529,11 @@ def _detect_gguf_from_hf_cache(repo_id: str) -> Optional[str]:
 
 
 def detect_gguf_model_remote(repo_id: str, hf_token: Optional[str] = None) -> Optional[str]:
-    """
-    Check if a HuggingFace repo contains GGUF files.
+    """Return the best GGUF filename in a HF repo, or None.
 
-    Returns the filename of the best GGUF file in the repo, or None.
-
-    Retries on transient HF Hub failures (network hiccups, 5xx, slow
-    cold-start of the API). Without retry, a single transient failure
-    here returns None silently and the caller treats the repo as
-    non-GGUF -- which on Apple Silicon (Mac UI route) means falling
-    through to the MLX backend, which then fails opening a non-existent
-    config.json on the GGUF-only repo. Three attempts with 1s/2s/4s
-    backoff covers the typical free-runner HF Hub flakiness.
-
-    When offline, falls back to the local HF cache so a downloaded
-    repo is still routed to llama-server (not MLX/Unsloth).
+    Retries (3 attempts, 1s/2s/4s backoff) on transient HF Hub failures: a
+    silent None would make the caller treat a GGUF-only repo as non-GGUF and
+    fall through to MLX on Apple Silicon. Offline falls back to the local cache.
     """
     import time
     from huggingface_hub import model_info as hf_model_info
@@ -1569,7 +1551,7 @@ def detect_gguf_model_remote(repo_id: str, hf_token: Optional[str] = None) -> Op
             return _pick_best_gguf(repo_files)
         except Exception as e:
             last_err = e
-            # 404 / RepoNotFound is permanent -- don't waste attempts.
+            # 404 / RepoNotFound is permanent -- don't retry
             err_name = type(e).__name__
             if err_name in (
                 "RepositoryNotFoundError",
@@ -1601,11 +1583,7 @@ def download_gguf_file(
     filename: str,
     hf_token: Optional[str] = None,
 ) -> str:
-    """
-    Download a specific GGUF file from a HuggingFace repo.
-
-    Returns the local path to the downloaded file.
-    """
+    """Download a specific GGUF file from a HF repo; returns the local path."""
     from huggingface_hub import hf_hub_download
 
     local_path = hf_hub_download(
@@ -1616,35 +1594,29 @@ def download_gguf_file(
     return local_path
 
 
-# Cache embedding detection results per session to avoid repeated HF API calls
+# Cache embedding detection per session to avoid repeated HF API calls
 _embedding_detection_cache: Dict[tuple, bool] = {}
 
 
 def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool:
-    """
-    Detect embedding/sentence-transformer models using HuggingFace model metadata.
+    """Detect embedding/sentence-transformer models via HF metadata.
 
-    Uses a belt-and-suspenders approach combining three signals:
-      1. "sentence-transformers" in model tags
-      2. "feature-extraction" in model tags
-      3. pipeline_tag is "sentence-similarity" or "feature-extraction"
-
-    This catches all known embedding models including those like gte-modernbert
-    whose library_name is "transformers" rather than "sentence-transformers".
+    Combines three signals: "sentence-transformers" or "feature-extraction" in
+    tags, or pipeline_tag in {"sentence-similarity", "feature-extraction"}.
+    Catches models like gte-modernbert whose library_name is "transformers".
 
     Args:
         model_name: Model identifier (HF repo or local path)
-        hf_token: Optional HF token for accessing gated/private models
+        hf_token: Optional HF token for gated/private models
 
     Returns:
-        True if the model is an embedding model, False otherwise.
-        Defaults to False for local paths or on errors.
+        True if embedding model, else False (default for local paths or errors).
     """
     cache_key = (model_name, hf_token)
     if cache_key in _embedding_detection_cache:
         return _embedding_detection_cache[cache_key]
 
-    # Local paths: check for sentence-transformer marker file (modules.json)
+    # Local paths: check for sentence-transformer marker (modules.json)
     if is_local_path(model_name):
         local_dir = normalize_path(model_name)
         is_emb = os.path.isfile(os.path.join(local_dir, "modules.json"))
@@ -1726,12 +1698,11 @@ def _looks_like_lora_adapter(model_dir: Path) -> bool:
 
 
 def scan_trained_models(outputs_dir: str = str(outputs_root())) -> List[Tuple[str, str, str]]:
-    """
-    Scan outputs folder for trained Studio models.
+    """Scan outputs folder for trained Studio models.
 
     Returns:
-        List of tuples: [(display_name, model_path, model_type), ...]
-        model_type is "lora" for adapter runs and "merged" for full finetunes.
+        List of (display_name, model_path, model_type), where model_type is
+        "lora" for adapter runs or "merged" for full finetunes.
     """
     trained_models = []
     outputs_path = resolve_output_dir(outputs_dir)
@@ -1752,7 +1723,7 @@ def scan_trained_models(outputs_dir: str = str(outputs_root())) -> List[Tuple[st
                 trained_models.append((display_name, model_path, model_type))
                 logger.debug("Found trained model: %s (%s)", display_name, model_type)
 
-        # Sort by modification time (newest first)
+        # Sort by mtime, newest first
         trained_models.sort(key = lambda x: Path(x[1]).stat().st_mtime, reverse = True)
 
         logger.info(
@@ -1770,16 +1741,14 @@ def scan_trained_models(outputs_dir: str = str(outputs_root())) -> List[Tuple[st
 def scan_exported_models(
     exports_dir: str = str(exports_root()),
 ) -> List[Tuple[str, str, str, Optional[str]]]:
-    """
-    Scan exports folder for exported models (merged, LoRA, GGUF).
+    """Scan exports folder for exported models (merged, LoRA, GGUF).
 
-    Supports two directory layouts:
-      - Two-level: {run}/{checkpoint}/  (merged & LoRA exports)
-      - Flat:      {name}-finetune-gguf/  (GGUF exports)
+    Supports two layouts: two-level {run}/{checkpoint}/ (merged & LoRA) and
+    flat {name}-finetune-gguf/ (GGUF).
 
     Returns:
-        List of tuples: [(display_name, model_path, export_type, base_model), ...]
-        export_type: "lora" | "merged" | "gguf"
+        List of (display_name, model_path, export_type, base_model), where
+        export_type is "lora" | "merged" | "gguf".
     """
     results = []
     exports_path = resolve_export_dir(exports_dir)
@@ -1792,8 +1761,8 @@ def scan_exported_models(
             if not run_dir.is_dir():
                 continue
 
-            # Check for flat GGUF export (e.g. exports/gemma-3-4b-it-finetune-gguf/)
-            # Filter out mmproj (vision projection) files — they aren't loadable as main models
+            # Flat GGUF export (e.g. exports/gemma-3-4b-it-finetune-gguf/).
+            # Skip mmproj (vision projection) files — not loadable as main models.
             gguf_files = [f for f in _iter_gguf_files(run_dir) if not _is_mmproj(f.name)]
             if gguf_files:
                 base_model = None
@@ -1806,7 +1775,7 @@ def scan_exported_models(
                     pass
 
                 display_name = run_dir.name
-                model_path = str(gguf_files[0])  # path to the .gguf file
+                model_path = str(gguf_files[0])
                 results.append((display_name, model_path, "gguf", base_model))
                 logger.debug(f"Found GGUF export: {display_name}")
                 continue
@@ -1845,8 +1814,8 @@ def scan_exported_models(
                 elif has_gguf:
                     export_type = "gguf"
                     gguf_list = list(_iter_gguf_files(checkpoint_dir))
-                    # Check checkpoint_dir first, then fall back to parent run_dir
-                    # (export.py writes metadata to the top-level export directory)
+                    # checkpoint_dir first, then run_dir (export.py writes
+                    # metadata to the top-level export dir)
                     for meta_dir in (checkpoint_dir, run_dir):
                         export_meta = meta_dir / "export_metadata.json"
                         try:
@@ -1866,8 +1835,7 @@ def scan_exported_models(
                 else:
                     continue
 
-                # Fallback: read base model from the original training run's
-                # adapter_config.json in ./outputs/{run_name}/
+                # Fallback: base model from ./outputs/{run_name}/adapter_config.json
                 if not base_model:
                     outputs_adapter_cfg = resolve_output_dir(run_dir.name) / "adapter_config.json"
                     try:
@@ -1953,22 +1921,14 @@ def get_base_model_from_checkpoint(checkpoint_path: str) -> Optional[str]:
 
 
 def get_base_model_from_lora(lora_path: str) -> Optional[str]:
-    """
-    Read the base model name from a LoRA adapter's config.
-
-    Args:
-        lora_path: Path to the LoRA adapter directory
-
-    Returns:
-        Base model identifier or None if not found
-    """
+    """Read the base model name from a LoRA adapter's config, or None."""
     try:
         lora_path_obj = Path(lora_path)
 
         if not _looks_like_lora_adapter(lora_path_obj):
             return None
 
-        # Try adapter_config.json first
+        # adapter_config.json first
         adapter_config_path = lora_path_obj / "adapter_config.json"
         if adapter_config_path.exists():
             with open(adapter_config_path, "r") as f:
@@ -1995,13 +1955,10 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]:
         #     except Exception as e:
         #         logger.warning(f"Could not load training_args.bin: {e}")
 
-        # Last resort: parse from directory name
-        # Format: unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit_timestamp
+        # Last resort: parse from dir name (unsloth__)
         dir_name = lora_path_obj.name
         if dir_name.startswith("unsloth_"):
-            # Remove timestamp suffix (usually _1234567890)
             parts = dir_name.split("_")
-            # Reconstruct model name
             if len(parts) >= 2:
                 model_parts = parts[1:-1]  # Skip "unsloth" and timestamp
                 base_model = "unsloth/" + "_".join(model_parts)
@@ -2021,28 +1978,19 @@ UI_STATUS_INDICATORS = [" (Ready)", " (Loading...)", " (Active)", "↓ "]
 
 
 def load_model_defaults(model_name: str) -> Dict[str, Any]:
-    """
-    Load default training parameters for a model from YAML file.
+    """Load default training parameters for a model from a YAML file.
 
-    Args:
-        model_name: Model identifier (e.g., "unsloth/Meta-Llama-3.1-8B-bnb-4bit")
-
-    Returns:
-        Dictionary with default parameters from YAML file, or empty dict if not found
-
-    The function looks for a YAML file in configs/model_defaults/ (including subfolders)
-    based on the model name or its aliases from MODEL_NAME_MAPPING.
-    If no specific file exists, it falls back to default.yaml.
+    Looks in configs/model_defaults/ (incl. subfolders) by model name or its
+    MODEL_NAME_MAPPING aliases, else falls back to default.yaml. Returns the
+    parameter dict, or {} if none found.
     """
     try:
-        # Get the script directory to locate configs
         script_dir = Path(__file__).parent.parent.parent
         defaults_dir = script_dir / "assets" / "configs" / "model_defaults"
 
-        # First, check if model is in the mapping
+        # Check the mapping first
         if model_name.lower() in _REVERSE_MODEL_MAPPING:
             canonical_file = _REVERSE_MODEL_MAPPING[model_name.lower()]
-            # Search in subfolders and root
             for config_path in defaults_dir.rglob(canonical_file):
                 if config_path.is_file():
                     with open(config_path, "r", encoding = "utf-8") as f:
@@ -2050,10 +1998,9 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]:
                         logger.info(f"Loaded model defaults from {config_path} (via mapping)")
                         return config
 
-        # If model_name is a local path (e.g. /home/.../Spark-TTS-0.5B/LLM from
-        # adapter_config.json, or C:\Users\...\model on Windows), try matching
-        # the last 1-2 path components against the registry
-        # (e.g. "Spark-TTS-0.5B/LLM").
+        # For local paths (e.g. /home/.../Spark-TTS-0.5B/LLM from
+        # adapter_config.json, or C:\Users\...\model on Windows), match the
+        # last 1-2 path components against the registry (e.g. "Spark-TTS-0.5B/LLM").
         _is_local_path = is_local_path(model_name)
         # Normalize Windows backslash paths so Path().parts splits correctly
         # on POSIX/WSL hosts (pathlib treats backslashes as literals on Linux).
@@ -2074,13 +2021,13 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]:
                                     )
                                     return config
 
-        # Try exact model name match (for backward compatibility).
-        # For local filesystem paths, use only the directory basename to
-        # avoid passing absolute paths (e.g. C:\...) into rglob which
-        # raises "Non-relative patterns are unsupported" on Windows.
+        # Exact model name match (backward compatibility). For local paths,
+        # use only the dir basename to avoid passing absolute paths (e.g.
+        # C:\...) into rglob, which raises "Non-relative patterns are
+        # unsupported" on Windows.
         _lookup_name = Path(_normalized).name if _is_local_path else model_name
         model_filename = _lookup_name.replace("/", "_") + ".yaml"
-        # Search in subfolders and root
+        # Search subfolders and root
         for config_path in defaults_dir.rglob(model_filename):
             if config_path.is_file():
                 with open(config_path, "r", encoding = "utf-8") as f:
@@ -2106,17 +2053,17 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]:
 
 @dataclass
 class ModelConfig:
-    """Configuration for a model to load"""
+    """Configuration for a model to load."""
 
     identifier: str  # Clean model identifier (org/name or path)
     display_name: str  # Original UI display name
     path: str  # Normalized filesystem path
-    is_local: bool  # Is this a local file vs HF model?
-    is_cached: bool  # Is this already in HF cache?
-    is_vision: bool  # Is this a vision model?
-    is_lora: bool  # Is this a lora adapter?
-    is_gguf: bool = False  # Is this a GGUF model?
-    is_audio: bool = False  # Is this a TTS audio model?
+    is_local: bool  # Local file vs HF model?
+    is_cached: bool  # Already in HF cache?
+    is_vision: bool  # Vision model?
+    is_lora: bool  # LoRA adapter?
+    is_gguf: bool = False  # GGUF model?
+    is_audio: bool = False  # TTS audio model?
     audio_type: Optional[str] = None  # Audio codec type: 'snac', 'csm', 'bicodec', 'dac'
     has_audio_input: bool = False  # Accepts audio input (ASR/speech understanding)
     gguf_file: Optional[str] = None  # Full path to the .gguf file (local mode)
@@ -2133,17 +2080,12 @@ class ModelConfig:
         lora_path: str,
         hf_token: Optional[str] = None,
     ) -> Optional["ModelConfig"]:
-        """
-        Create ModelConfig from a local LoRA adapter path.
-
-        Automatically detects the base model from adapter config.
+        """Create ModelConfig from a local LoRA adapter path, auto-detecting the
+        base model from adapter config.
 
         Args:
-            lora_path: Path to LoRA adapter (e.g., "./outputs/unsloth_Meta-Llama-3.1_.../")
+            lora_path: Path to the LoRA adapter directory
             hf_token: HF token for vision detection
-
-        Returns:
-            ModelConfig for the LoRA adapter
         """
         try:
             lora_path_obj = Path(lora_path)
@@ -2152,27 +2094,23 @@ class ModelConfig:
                 logger.error(f"LoRA path does not exist: {lora_path}")
                 return None
 
-            # Get base model
             base_model = get_base_model_from_lora(lora_path)
             if not base_model:
                 logger.error(f"Could not determine base model for LoRA: {lora_path}")
                 return None
 
-            # Check if base model is vision
             is_vision = is_vision_model(base_model, hf_token = hf_token)
-
-            # Check if base model is audio
             audio_type = detect_audio_type(base_model, hf_token = hf_token)
 
             display_name = lora_path_obj.name
-            identifier = lora_path  # Use path as identifier for local LoRAs
+            identifier = lora_path  # path is the identifier for local LoRAs
 
             return cls(
                 identifier = identifier,
                 display_name = display_name,
                 path = lora_path,
                 is_local = True,
-                is_cached = True,  # Local LoRAs are always "cached"
+                is_cached = True,  # local LoRAs are always cached
                 is_vision = is_vision,
                 is_lora = True,
                 is_audio = audio_type is not None and audio_type != "audio_vlm",
@@ -2193,25 +2131,18 @@ class ModelConfig:
         is_lora: bool = False,
         gguf_variant: Optional[str] = None,
     ) -> Optional["ModelConfig"]:
-        """
-        Create ModelConfig from a clean model identifier.
-
-        For FastAPI routes where the frontend sends sanitized model paths.
-        No Gradio dropdown parsing - expects clean identifiers like:
-        - "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit"
-        - "./outputs/my_lora_adapter"
-        - "/absolute/path/to/model"
+        """Create ModelConfig from a clean model identifier (HF repo or local
+        path), for FastAPI routes that send sanitized paths.
 
         Args:
             model_id: Clean model identifier (HF repo name or local path)
             hf_token: Optional HF token for vision detection on gated models
             is_lora: Whether this is a LoRA adapter
-            gguf_variant: Optional GGUF quantization variant (e.g. "Q4_K_M").
-                For remote GGUF repos, specifies which quant to load via -hf.
-                If None, auto-selects using _pick_best_gguf().
+            gguf_variant: Optional GGUF quant variant (e.g. "Q4_K_M") to load
+                via -hf for remote repos; None auto-selects via _pick_best_gguf().
 
         Returns:
-            ModelConfig or None if configuration cannot be created
+            ModelConfig or None if it cannot be created.
         """
         if not model_id or not model_id.strip():
             return None
@@ -2225,8 +2156,8 @@ class ModelConfig:
             identifier = f"unsloth/{identifier}"
             path = identifier
 
-        # Preserve requested casing, but if a case-variant already exists in local HF cache,
-        # reuse that exact repo_id spelling to avoid one-time re-downloads after #2592.
+        # Reuse a cached case-variant's exact repo_id spelling to avoid
+        # one-time re-downloads after #2592.
         if not is_local:
             resolved_identifier = resolve_cached_repo_id_case(identifier)
             if resolved_identifier != identifier:
@@ -2248,12 +2179,12 @@ class ModelConfig:
                 display_name = Path(gguf_file).stem
                 logger.info(f"Detected local GGUF model: {gguf_file}")
 
-                # Detect vision: check if base model is vision, then look for mmproj
+                # Vision: check base model, then look for mmproj
                 mmproj_file = None
                 gguf_is_vision = False
                 gguf_dir = Path(gguf_file).parent
 
-                # Determine if this is a vision model from export metadata
+                # Is this a vision model, per export metadata?
                 base_is_vision = False
                 meta_path = gguf_dir / "export_metadata.json"
                 if meta_path.exists():
@@ -2266,15 +2197,9 @@ class ModelConfig:
                     except Exception as e:
                         logger.debug(f"Could not read export metadata: {e}")
 
-                # If vision (or mmproj happens to exist), find the mmproj
-                # file. The recursive variant scan in
-                # ``_find_local_gguf_by_variant`` may have returned a
-                # weight file inside a quant-named subdir (e.g.
-                # ``.../BF16/foo.gguf``) while ``mmproj-*.gguf`` lives
-                # at the snapshot root. Pass ``search_root=path`` so
-                # ``detect_mmproj_file`` walks up to the snapshot root
-                # instead of seeing only the weight file's immediate
-                # parent.
+                # Pass search_root=path so detect_mmproj_file walks up to the
+                # snapshot root: the weight may sit in a quant subdir while
+                # mmproj-*.gguf lives at the root.
                 mmproj_file = detect_mmproj_file(gguf_file, search_root = path)
                 if mmproj_file:
                     gguf_is_vision = True
@@ -2295,11 +2220,10 @@ class ModelConfig:
                     gguf_mmproj_file = mmproj_file,
                 )
         else:
-            # Check if the HF repo contains GGUF files
+            # Does the HF repo contain GGUF files?
             gguf_filename = detect_gguf_model_remote(identifier, hf_token = hf_token)
             if gguf_filename:
-                # Preflight: verify llama-server binary exists BEFORE user waits
-                # for a multi-GB download that llama-server handles natively
+                # Preflight: verify llama-server binary exists before a multi-GB download
                 from core.inference.llama_cpp import LlamaCppBackend
 
                 if not LlamaCppBackend._find_llama_server_binary():
@@ -2308,11 +2232,10 @@ class ModelConfig:
                         "Run setup.sh to build it, or set LLAMA_SERVER_PATH."
                     )
 
-                # Use list_gguf_variants() to detect vision & resolve variant
+                # list_gguf_variants() detects vision & resolves the variant
                 variants, has_vision = list_gguf_variants(identifier, hf_token = hf_token)
                 variant = gguf_variant
-                if not variant:
-                    # Auto-select best quantization
+                if not variant:  # auto-select best quant
                     variant_filenames = [v.filename for v in variants]
                     best = _pick_best_gguf(variant_filenames)
                     if best:
@@ -2339,7 +2262,7 @@ class ModelConfig:
                     gguf_variant = variant,
                 )
 
-        # Auto-detect LoRA for local paths (check adapter_config.json on disk)
+        # Auto-detect LoRA for local paths (adapter_config.json on disk)
         if not is_lora and is_local:
             detected_base = (
                 get_base_model_from_lora(path) if _looks_like_lora_adapter(Path(path)) else None
@@ -2362,7 +2285,7 @@ class ModelConfig:
             except Exception as e:
                 logger.debug(f"Could not check remote LoRA status for '{identifier}': {e}")
 
-            # API may have failed; adapter_config.json may still be cached.
+            # API may have failed; adapter_config.json could still be cached.
             if not is_lora:
                 for snap in _iter_hf_cache_snapshots(identifier):
                     if (snap / "adapter_config.json").is_file():
@@ -2377,7 +2300,7 @@ class ModelConfig:
                 # Local LoRA: read adapter_config.json from disk
                 base_model = get_base_model_from_lora(path)
             else:
-                # Remote LoRA: download adapter_config.json from HF
+                # Remote LoRA: fetch adapter_config.json from HF
                 try:
                     from huggingface_hub import hf_hub_download
 
@@ -2428,10 +2351,7 @@ class ModelConfig:
         hf_token: Optional[str] = None,
         is_lora: bool = False,
     ) -> Optional["ModelConfig"]:
-        """
-        Create a universal ModelConfig from UI dropdown/search selections.
-        Handles base models and LoRA adapters.
-        """
+        """Create a ModelConfig from UI dropdown/search selections (base models and LoRAs)."""
         selected = None
         if search_value and search_value.strip():
             selected = search_value.strip()
@@ -2443,7 +2363,7 @@ class ModelConfig:
 
         display_name = selected
 
-        #  Use the correct 'local_models' parameter to resolve display names
+        # Resolve display names via the 'local_models' parameter
         if " (Active)" in selected or " (Ready)" in selected:
             clean_display_name = selected.replace(" (Active)", "").replace(" (Ready)", "")
             if local_models:
@@ -2452,7 +2372,7 @@ class ModelConfig:
                         selected = local_path
                         break
 
-        # Clean all UI status indicators to get the final identifier
+        # Strip all UI status indicators to get the final identifier
         identifier = selected
         for status in UI_STATUS_INDICATORS:
             identifier = identifier.replace(status, "")
@@ -2472,23 +2392,23 @@ class ModelConfig:
                 identifier = resolved_identifier
                 path = resolved_identifier
 
-        # --- Logic for Base Model and Vision Detection ---
+        # --- Base Model and Vision Detection ---
         base_model = None
         is_vision = False
 
         if is_lora:
-            # For a LoRA, we MUST find its base model.
+            # A LoRA MUST have a base model.
             base_model = get_base_model_from_lora(path)
             if not base_model:
                 logger.warning(
                     f"Could not determine base model for LoRA '{path}'. Cannot create config."
                 )
-                return None  # Cannot proceed without a base model
+                return None  # cannot proceed without a base model
 
-            # A LoRA's vision capability is determined by its base model.
+            # A LoRA's vision capability comes from its base model.
             is_vision = is_vision_model(base_model, hf_token = hf_token)
         else:
-            # For a base model, just check its own vision status.
+            # Base model: check its own vision status.
             is_vision = is_vision_model(identifier, hf_token = hf_token)
 
         from utils.paths import is_model_cached
@@ -2503,5 +2423,5 @@ class ModelConfig:
             is_cached = is_cached,
             is_vision = is_vision,
             is_lora = is_lora,
-            base_model = base_model,  # This will be None for base models, and populated for LoRAs
+            base_model = base_model,  # None for base models, set for LoRAs
         )
diff --git a/studio/backend/utils/paths/__init__.py b/studio/backend/utils/paths/__init__.py
index eba7a9de6c..1dcc9191fd 100644
--- a/studio/backend/utils/paths/__init__.py
+++ b/studio/backend/utils/paths/__init__.py
@@ -1,9 +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
 
-"""
-Path utilities for model and dataset handling
-"""
+"""Path utilities for model and dataset handling."""
 
 from .path_utils import (
     normalize_path,
@@ -46,8 +44,8 @@ from .storage_roots import (
     resolve_dataset_path,
 )
 
-# Re-export shim: name-load the project-path helpers so the import-hoist
-# safety net sees them used here, not just listed in __all__ as strings.
+# Re-export shim: mark project-path helpers as used so the import-hoist
+# safety net does not flag them as unused.
 _REEXPORTED = (documents_root, project_workspaces_root)
 
 __all__ = [
diff --git a/studio/backend/utils/paths/path_utils.py b/studio/backend/utils/paths/path_utils.py
index 22c6c46ee1..e8dabc8954 100644
--- a/studio/backend/utils/paths/path_utils.py
+++ b/studio/backend/utils/paths/path_utils.py
@@ -17,7 +17,7 @@ logger = get_logger(__name__)
 # Per-process cache to avoid repeated cache-dir scans for the same identifier.
 _CACHE_CASE_RESOLUTION_MEMO: dict[str, str] = {}
 
-# Lightweight instrumentation counters for operational visibility.
+# Instrumentation counters for operational visibility.
 _CACHE_CASE_RESOLUTION_STATS: dict[str, int] = {
     "calls": 0,
     "memo_hits": 0,
@@ -44,27 +44,17 @@ _IS_WSL: bool = _is_wsl()
 
 
 def normalize_path(path: str) -> str:
-    """
-    Normalize filesystem paths for cross-platform use.
+    """Normalize filesystem paths for cross-platform use.
 
-    On WSL, converts Windows drive-letter paths to ``/mnt//...``.
-    On native Windows, keeps the drive letter and normalizes separators.
-    On Linux/macOS (non-WSL), paths are returned with forward slashes.
-
-    Examples (WSL):
-        C:\\Users\\... -> /mnt/c/Users/...
-    Examples (native Windows):
-        C:\\Users\\... -> C:/Users/...
-    Examples (Linux/macOS):
-        /home/user/... -> /home/user/... (unchanged)
+    WSL maps drive-letter paths to ``/mnt//...``; native Windows keeps
+    the drive and normalizes separators; elsewhere slashes are forward-only.
     """
     if not path:
         return path
 
     # Handle Windows drive letters (C:\\ or c:\\)
     if len(path) >= 3 and path[1] == ":" and path[2] in ("\\", "/"):
-        # Only map to /mnt// when running under WSL;
-        # on native Windows the drive letter must be preserved.
+        # Map to /mnt// only under WSL; native Windows keeps the drive letter.
         if _IS_WSL:
             drive = path[0].lower()
             rest = path[3:].replace("\\", "/")
@@ -86,7 +76,7 @@ def is_local_path(path: str) -> bool:
     if not path:
         return False
 
-    # If it exists on disk, treat as local (covers relative paths like "outputs/foo").
+    # Exists on disk → local (covers relative paths like "outputs/foo").
     try:
         if Path(normalize_path(path)).expanduser().exists():
             return True
@@ -122,7 +112,7 @@ def is_model_cached(model_name: str) -> bool:
     if not cache_path:
         return False
 
-    # Check for actual model files
+    # Check for model files
     for suffix in [".safetensors", ".bin", ".json"]:
         if list(cache_path.rglob(f"*{suffix}")):
             return True
@@ -146,9 +136,9 @@ def _hf_hub_cache_dir() -> Path:
 def resolve_cached_repo_id_case(model_name: str, use_memo: bool = True) -> str:
     """Resolve repo_id to the exact casing already present in local HF cache.
 
-    Policy: prefer the requested/canonical repo_id, but if a case-variant already
-    exists in local HF cache, reuse that exact cached spelling. This avoids
-    duplicate downloads while preserving user intent whenever possible.
+    Policy: prefer the requested/canonical repo_id, but reuse a case-variant's
+    exact cached spelling if one already exists in local HF cache. Avoids
+    duplicate downloads while preserving user intent where possible.
     """
     _CACHE_CASE_RESOLUTION_STATS["calls"] += 1
 
@@ -163,8 +153,7 @@ def resolve_cached_repo_id_case(model_name: str, use_memo: bool = True) -> str:
 
     expected_dir = f"models--{model_name.replace('/', '--')}"
 
-    # Always check the exact-case path first so a newly-appeared exact match
-    # wins over any previously memoized variant.
+    # Exact-case path first so a new exact match beats a memoized variant.
     exact_path = cache_dir / expected_dir
     if exact_path.is_dir():
         if use_memo:
@@ -172,8 +161,7 @@ def resolve_cached_repo_id_case(model_name: str, use_memo: bool = True) -> str:
         _CACHE_CASE_RESOLUTION_STATS["exact_hits"] += 1
         return model_name
 
-    # Validate memoized entries still exist on disk before returning them.
-    # This prevents stale results when cache dirs are deleted/recreated.
+    # Revalidate memoized entries on disk to avoid stale results.
     if use_memo:
         cached = _CACHE_CASE_RESOLUTION_MEMO.get(model_name)
         if cached is not None:
@@ -181,7 +169,7 @@ def resolve_cached_repo_id_case(model_name: str, use_memo: bool = True) -> str:
             if cached_path.is_dir():
                 _CACHE_CASE_RESOLUTION_STATS["memo_hits"] += 1
                 return cached
-            # Stale entry -- drop it and re-scan below.
+            # Stale entry -- drop it and re-scan below
             _CACHE_CASE_RESOLUTION_MEMO.pop(model_name, None)
 
     expected_lower = expected_dir.lower()
@@ -200,7 +188,7 @@ def resolve_cached_repo_id_case(model_name: str, use_memo: bool = True) -> str:
             candidates.append(repo_part.replace("--", "/"))
 
         if candidates:
-            # Deterministic tie-break if multiple case variants coexist.
+            # Deterministic tie-break if multiple case variants coexist
             resolved = sorted(candidates)[0]
             if len(candidates) > 1:
                 _CACHE_CASE_RESOLUTION_STATS["tie_breaks"] += 1
diff --git a/studio/backend/utils/paths/storage_roots.py b/studio/backend/utils/paths/storage_roots.py
index b254c20f97..63d24afd0c 100644
--- a/studio/backend/utils/paths/storage_roots.py
+++ b/studio/backend/utils/paths/storage_roots.py
@@ -11,9 +11,9 @@ import tempfile
 
 
 def _infer_studio_home_from_venv() -> Path | None:
-    """Return parent dir of sys.prefix as STUDIO_HOME if running from an
+    """Return parent of sys.prefix as STUDIO_HOME when running from an
     installer-managed unsloth_studio venv. Sentinel-gated (share/studio.conf
-    or bin shim) so a developer venv named unsloth_studio is not misidentified.
+    or bin shim) so a dev venv named unsloth_studio isn't misidentified.
     """
     try:
         prefix = Path(sys.prefix).resolve()
@@ -38,8 +38,8 @@ def studio_root() -> Path:
     """Studio install root.
 
     Priority: UNSLOTH_STUDIO_HOME, then STUDIO_HOME alias, then sys.prefix
-    inference, then legacy ~/.unsloth/studio. UNSLOTH_STUDIO_HOME wins when
-    both are set (the more specific signal beats the generic alias).
+    inference, then legacy ~/.unsloth/studio. UNSLOTH_STUDIO_HOME wins if
+    both are set (specific signal beats generic alias).
     """
     override = (os.environ.get("UNSLOTH_STUDIO_HOME") or "").strip()
     if not override:
@@ -56,7 +56,7 @@ def studio_root() -> Path:
 
 
 def cache_root() -> Path:
-    """Central cache directory for all studio downloads (models, datasets, etc.)."""
+    """Central cache dir for all studio downloads (models, datasets, etc.)."""
     return studio_root() / "cache"
 
 
@@ -158,16 +158,15 @@ def ensure_dir(path: Path) -> Path:
 
 
 def legacy_hf_cache_dir() -> Path:
-    """Old Unsloth-specific HF hub cache, kept for backward-compat scanning."""
+    """Old Unsloth-specific HF hub cache, kept for backward-compat scans."""
     return cache_root() / "huggingface" / "hub"
 
 
 def hf_default_cache_dir() -> Path:
-    """Return the platform default HuggingFace hub cache (ignoring env overrides).
+    """Platform default HuggingFace hub cache (ignoring env overrides).
 
-    This is the location HF uses when no ``HF_HUB_CACHE`` / ``HF_HOME``
-    env var is set.  We scan it so that models a user downloaded *before*
-    installing Unsloth Studio are still discovered.
+    Where HF caches when no ``HF_HUB_CACHE`` / ``HF_HOME`` is set. Scanned
+    so models downloaded *before* installing Unsloth Studio are discovered.
     """
     return Path.home() / ".cache" / "huggingface" / "hub"
 
@@ -183,7 +182,7 @@ def lmstudio_model_dirs() -> list[Path]:
             seen.add(resolved)
             dirs.append(p)
 
-    # 1. Check LM Studio settings.json for custom downloads folder
+    # LM Studio settings.json custom downloads folder
     settings_path = Path.home() / ".lmstudio" / "settings.json"
     if settings_path.is_file():
         try:
@@ -195,10 +194,10 @@ def lmstudio_model_dirs() -> list[Path]:
         except Exception:
             pass
 
-    # 2. LM Studio current default models directory (all platforms)
+    # LM Studio default models directory (all platforms)
     _add(Path.home() / ".lmstudio" / "models")
 
-    # 3. Legacy LM Studio cache location
+    # Legacy LM Studio cache location
     _add(Path.home() / ".cache" / "lm-studio" / "models")
 
     return dirs
@@ -207,17 +206,17 @@ def lmstudio_model_dirs() -> list[Path]:
 def well_known_model_dirs() -> list[Path]:
     """Return directories commonly used by other local LLM tools.
 
-    Used by the folder browser to offer quick-pick chips. Returns only
-    paths that exist on disk, so the UI never shows dead chips. Order
-    reflects a rough "likelihood the user has models here" -- LM Studio
-    and Ollama first, then the generic fallbacks.
+    Backs the folder browser's quick-pick chips. Returns only paths that
+    exist on disk, so the UI never shows dead chips. Order reflects rough
+    likelihood of models being there -- LM Studio and Ollama first, then
+    generic fallbacks.
     """
     candidates: list[Path] = []
 
     # LM Studio (reuses the logic above, including settings.json override)
     candidates.extend(lmstudio_model_dirs())
 
-    # Ollama -- both the user-level and common system-wide install paths
+    # Ollama -- user-level and common system-wide install paths
     # (https://github.com/ollama/ollama/issues/733).
     ollama_env = os.environ.get("OLLAMA_MODELS")
     if ollama_env:
@@ -229,11 +228,11 @@ def well_known_model_dirs() -> list[Path]:
     # HF hub cache root (separate from the explicit HF cache chip)
     candidates.append(Path.home() / ".cache" / "huggingface" / "hub")
 
-    # Generic "my models" spots users tend to drop things into
+    # Generic "my models" spots users drop things into
     for name in ("models", "Models"):
         candidates.append(Path.home() / name)
 
-    # Deduplicate while preserving order; keep only extant dirs
+    # Dedupe preserving order; keep only extant dirs
     out: list[Path] = []
     seen: set[str] = set()
     for p in candidates:
@@ -250,17 +249,11 @@ def well_known_model_dirs() -> list[Path]:
 
 
 def _setup_cache_env() -> None:
-    """Set cache environment variables for HuggingFace, uv, and vLLM.
+    """Set cache env vars for HuggingFace, uv, and vLLM.
 
-    Respects the standard HF cache resolution chain: explicit ``HF_HOME``
-    / ``HF_HUB_CACHE`` env vars take priority, then ``XDG_CACHE_HOME``,
-    then the platform default (``~/.cache/huggingface``).  The legacy
-    Unsloth cache is still *scanned* for models but is never set as the
-    active download target.
-
-    Only sets variables that are not already set by the user, so
-    explicit overrides (e.g. HF_HOME=/data/hf) are respected.
-    Works on Linux, macOS, and Windows.
+    Respects the standard HF cache chain (explicit HF_HOME / HF_HUB_CACHE,
+    then XDG_CACHE_HOME, then ~/.cache/huggingface) and only sets vars the
+    user hasn't, so explicit overrides are honored.
     """
     root = cache_root()
     xdg_cache = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")).expanduser()
@@ -327,8 +320,8 @@ def resolve_under_root(
 ) -> Path:
     """Resolve ``path_value`` and assert the result is under ``root``.
 
-    Absolutes are accepted only if already contained (so internal pre-resolved
-    paths re-enter idempotently); user-facing schemas reject absolutes upstream.
+    Absolutes are accepted only if already contained (so pre-resolved
+    internal paths re-enter idempotently); schemas reject absolutes upstream.
     """
     if not path_value or not str(path_value).strip():
         return root
diff --git a/studio/backend/utils/studio_version.py b/studio/backend/utils/studio_version.py
index e59439a2cc..98c48fe45c 100644
--- a/studio/backend/utils/studio_version.py
+++ b/studio/backend/utils/studio_version.py
@@ -74,8 +74,8 @@ def _exact_git_studio_tag(repo_root: Path) -> str | None:
 def get_studio_version(repo_root: Path | None = None) -> str:
     """Return the installed Studio release tag for display, or ``dev``.
 
-    This value is intentionally separate from the PyPI ``unsloth`` package
-    version used by update checks. It never performs network requests.
+    Intentionally separate from the PyPI ``unsloth`` package version used by
+    update checks. Never performs network requests.
     """
     resolved_repo_root = repo_root or _repo_root()
 
diff --git a/studio/backend/utils/subprocess_compat.py b/studio/backend/utils/subprocess_compat.py
index bedf8cf2e6..f2fa6eadc7 100644
--- a/studio/backend/utils/subprocess_compat.py
+++ b/studio/backend/utils/subprocess_compat.py
@@ -8,10 +8,9 @@ import sys
 
 
 def windows_hidden_subprocess_kwargs() -> dict[str, object]:
-    """Return Windows-only subprocess kwargs that suppress console windows.
+    """Windows-only subprocess kwargs that suppress console windows.
 
-    On non-Windows platforms returns an empty dict so callers can always
-    unpack the result into ``subprocess.run`` / ``subprocess.Popen`` via
+    Empty dict off Windows, so callers can always unpack via
     ``**windows_hidden_subprocess_kwargs()``.
     """
     if sys.platform != "win32":
diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py
index 16f964d628..bf8d14b7cb 100644
--- a/studio/backend/utils/transformers_version.py
+++ b/studio/backend/utils/transformers_version.py
@@ -1,29 +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
 
-"""
-Automatic transformers version switching.
+"""Automatic transformers version switching.
 
-Some newer model architectures (Ministral-3, GLM-4.7-Flash, Qwen3-30B-A3B MoE,
-tiny_qwen3_moe) require transformers>=5.3.0, while Gemma 4 models require
-transformers>=5.5.0.  Everything else needs the default 4.57.x that ships
-with Unsloth.
+Some newer architectures need transformers>=5.3.0 (.venv_t5_530/); Gemma 4
+needs >=5.5.0 (.venv_t5_550/). Everything else uses the default 4.57.x. A
+custom-named LoRA adapter's base model is resolved from adapter_config.json.
 
-Two separate target directories are maintained:
-  - .venv_t5_530/  — transformers 5.3.0 (Ministral-3, GLM, Qwen3 MoE, etc.)
-  - .venv_t5_550/  — transformers 5.5.0 (Gemma 4)
-
-When loading a LoRA adapter with a custom name, we resolve the base model from
-``adapter_config.json`` and check *that* against the model list.
-
-Strategy:
-  Training and inference run in subprocesses that activate the correct version
-  via sys.path (prepending the appropriate .venv_t5_*/ directory). See:
-    - core/training/worker.py
-    - core/inference/worker.py
-
-  For export (still in-process), ensure_transformers_version() does a lightweight
-  sys.path swap using the same directories pre-installed by setup.sh.
+Training/inference run in subprocesses that activate the right version via
+sys.path; export (in-process) uses ensure_transformers_version() for the swap.
 """
 
 import importlib
@@ -57,8 +42,7 @@ def _env_offline() -> bool:
 # Detection
 # ---------------------------------------------------------------------------
 
-# Lowercase substrings — if ANY appears anywhere in the lowered model name,
-# we need transformers 5.3.0.
+# Lowercase substrings — any match in the lowered model name needs transformers 5.3.0.
 TRANSFORMERS_5_MODEL_SUBSTRINGS: tuple[str, ...] = (
     "ministral-3-",  # Ministral-3-{3,8,14}B-{Instruct,Reasoning,Base}-2512
     "glm-4.7-flash",  # GLM-4.7-Flash
@@ -85,15 +69,15 @@ _TRANSFORMERS_550_MODEL_TYPES: set[str] = {
     "gemma4",
 }
 
-# Tokenizer classes that only exist in transformers>=5.x
+# Tokenizer classes that only exist in transformers>=5.x.
 _TRANSFORMERS_5_TOKENIZER_CLASSES: set[str] = {
     "TokenizersBackend",
 }
 
-# Cache for dynamic tokenizer_config.json lookups to avoid repeated fetches
+# 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)
+# Cache for dynamic config.json lookups (architecture/model_type checks).
 _config_needs_550_cache: dict[str, bool] = {}
 
 # Versions
@@ -116,12 +100,10 @@ _VENV_T5_DIR = _VENV_T5_550_DIR
 def activate_transformers_for_subprocess(model_name: str) -> None:
     """Activate the correct transformers version in a subprocess worker.
 
-    Call this BEFORE any ML imports. Resolves LoRA adapters to their base
-    model, determines the required tier, and prepends the appropriate
-    ``.venv_t5_*`` directory to ``sys.path``.  Also propagates the path
-    via ``PYTHONPATH`` for child processes (e.g. GGUF converter).
-
-    Used by training, inference, and export workers.
+    Call BEFORE any ML imports. Resolves LoRA adapters to their base model,
+    determines the required tier, prepends the appropriate ``.venv_t5_*`` dir to
+    ``sys.path``, and propagates it via ``PYTHONPATH`` for child processes
+    (e.g. GGUF converter). Used by training, inference, and export workers.
     """
     resolved = _resolve_base_model(model_name)
     tier = get_transformers_tier(resolved)
@@ -155,11 +137,10 @@ def activate_transformers_for_subprocess(model_name: str) -> None:
 def _resolve_base_model(model_name: str) -> str:
     """If *model_name* points to a LoRA adapter, return its base model.
 
-    Checks for ``adapter_config.json`` locally first.  Only calls the heavier
-    ``get_base_model_from_lora`` for paths that are actual local directories
-    (avoids noisy warnings for plain HF model IDs).
-
-    Returns the original *model_name* unchanged if it is not a LoRA adapter.
+    Checks ``adapter_config.json`` locally first. Only calls the heavier
+    ``get_base_model_from_lora`` for real local directories (avoids noisy
+    warnings for plain HF model IDs). Returns *model_name* unchanged if not a
+    LoRA adapter.
     """
     # --- Fast local check ---------------------------------------------------
     local_path = Path(model_name)
@@ -221,11 +202,11 @@ def _resolve_base_model(model_name: str) -> str:
 
 
 def _check_tokenizer_config_needs_v5(model_name: str) -> bool:
-    """Fetch tokenizer_config.json from HuggingFace and check if the
-    tokenizer_class requires transformers 5.x.
+    """True if the model's tokenizer_class requires transformers 5.x.
 
-    Results are cached in ``_tokenizer_class_cache`` to avoid repeated fetches.
-    Returns False on any network/parse error (fail-open to default version).
+    Checks local tokenizer_config.json, else fetches from HuggingFace. Cached in
+    ``_tokenizer_class_cache``. Returns False on any network/parse error
+    (fail-open to default version).
     """
     if model_name in _tokenizer_class_cache:
         return _tokenizer_class_cache[model_name]
@@ -280,12 +261,11 @@ def _check_tokenizer_config_needs_v5(model_name: str) -> bool:
 
 
 def _check_config_needs_550(model_name: str) -> bool:
-    """Check ``config.json`` for architectures or model_type that require
-    transformers 5.5.0 (e.g. Gemma 4).
+    """True if ``config.json`` has architectures/model_type needing transformers
+    5.5.0 (e.g. Gemma 4).
 
-    Checks locally first, then falls back to fetching from HuggingFace.
-    Results are cached in ``_config_needs_550_cache``.
-    Returns False on any error (fail-open to lower tier).
+    Checks locally first, else fetches from HuggingFace. Cached in
+    ``_config_needs_550_cache``. Returns False on any error (fail-open to lower tier).
     """
     if model_name in _config_needs_550_cache:
         return _config_needs_550_cache[model_name]
@@ -352,10 +332,8 @@ def _check_config_needs_550(model_name: str) -> bool:
 def get_transformers_tier(model_name: str) -> str:
     """Return the transformers tier required for *model_name*.
 
-    Returns ``"550"`` for models needing transformers 5.5.0 (e.g. Gemma 4),
-    ``"530"`` for models needing transformers 5.3.0 (e.g. Ministral-3, Qwen3 MoE),
-    or ``"default"`` for everything else (4.57.x).
-
+    ``"550"`` for transformers 5.5.0 (e.g. Gemma 4), ``"530"`` for 5.3.0
+    (e.g. Ministral-3, Qwen3 MoE), or ``"default"`` for everything else (4.57.x).
     The 5.5.0 check runs first, then 5.3.0.
     """
     lowered = model_name.lower()
@@ -406,12 +384,11 @@ _PURGE_PREFIXES = (
     "trl",
     "accelerate",
     "auto_gptq",
-    # NOTE: bitsandbytes is intentionally EXCLUDED — it registers torch custom
-    # operators at import time via torch.library.define(). Those registrations
-    # live in torch's global operator registry which survives module purge.
-    # Re-importing bitsandbytes after purge → duplicate registration → crash.
-    # Our own modules that import from transformers at module level
-    # (e.g. model_config.py: `from transformers import AutoConfig`)
+    # NOTE: bitsandbytes is intentionally EXCLUDED -- it registers torch custom
+    # operators via torch.library.define() into torch's global registry, which
+    # survives module purge; re-importing after purge -> duplicate registration
+    # -> crash.
+    # Our own modules that import from transformers at module level.
     "utils.models",
     "core.training",
     "core.inference",
@@ -462,15 +439,15 @@ def _venv_dir_is_valid(venv_dir: str, packages: tuple[str, ...]) -> bool:
         pkg_name = parts[0]
         pkg_version = parts[1] if len(parts) > 1 else None
         pkg_name_norm = pkg_name.replace("-", "_")
-        # Check directory exists
+        # Directory must exist.
         if not any(
             (Path(venv_dir) / d).is_dir() for d in (pkg_name_norm, pkg_name_norm.replace("_", "-"))
         ):
             return False
-        # For unpinned packages, existence is enough
+        # Unpinned packages: existence is enough.
         if pkg_version is None:
             continue
-        # Check version via .dist-info metadata
+        # Check version via .dist-info metadata.
         dist_info_found = False
         for di in Path(venv_dir).glob(f"{pkg_name_norm}-*.dist-info"):
             metadata = di / "METADATA"
@@ -504,7 +481,7 @@ def _venv_t5_is_valid() -> bool:
 
 def _install_to_dir(pkg: str, target_dir: str) -> bool:
     """Install a single package into *target_dir*, preferring uv then pip."""
-    # Try uv first (faster) if already on PATH -- do NOT install uv at runtime
+    # Try uv first (faster) if on PATH -- do NOT install uv at runtime.
     if shutil.which("uv"):
         result = subprocess.run(
             [
@@ -529,7 +506,7 @@ def _install_to_dir(pkg: str, target_dir: str) -> bool:
             return True
         logger.warning("uv install of %s failed, falling back to pip", pkg)
 
-    # Fallback to pip
+    # Fallback to pip.
     result = subprocess.run(
         [
             sys.executable,
@@ -621,13 +598,13 @@ def ensure_transformers_version(model_name: str) -> None:
       • Need 5.3.0 → prepend .venv_t5_530/ to sys.path, purge modules.
       • Need 4.x  → remove all .venv_t5_*/ from sys.path, purge modules.
 
-    For LoRA adapters with custom names, the base model is resolved from
+    For custom-named LoRA adapters, the base model is resolved from
     ``adapter_config.json`` before checking.
 
-    NOTE: Training and inference use subprocess isolation instead of this
-    function. This is only used by the export path (routes/export.py).
+    NOTE: Training and inference use subprocess isolation instead. Used only by
+    the export path (routes/export.py).
     """
-    # Resolve LoRA adapters to their base model for accurate detection
+    # Resolve LoRA adapters to their base model for accurate detection.
     resolved = _resolve_base_model(model_name)
     tier = get_transformers_tier(resolved)
 
@@ -666,10 +643,10 @@ def ensure_transformers_version(model_name: str) -> None:
                 model_name,
             )
             return
-        # Different 5.x → need to switch (e.g. 5.3.0 loaded but need 5.5.0)
+        # Different 5.x → must switch (e.g. 5.3.0 loaded but need 5.5.0).
         in_memory_major = int(in_memory.split(".")[0])
         if in_memory_major == target_major and venv_dir is None:
-            # Both are default (4.x) — close enough
+            # Both are default (4.x) — close enough.
             logger.info(
                 "transformers %s already loaded — correct for '%s'",
                 in_memory,
@@ -679,7 +656,7 @@ def ensure_transformers_version(model_name: str) -> None:
 
     # --- Switch version -----------------------------------------------------
     if venv_dir is not None:
-        # First remove any other 5.x venv from sys.path
+        # First remove any other 5.x venv from sys.path.
         _deactivate_5x()
         if not ensure_fn():
             raise RuntimeError(
diff --git a/studio/backend/utils/update_status.py b/studio/backend/utils/update_status.py
index 9b71ff31a0..ad9dabcf36 100644
--- a/studio/backend/utils/update_status.py
+++ b/studio/backend/utils/update_status.py
@@ -3,9 +3,8 @@
 
 """Web update status helpers for browser-served Unsloth Studio.
 
-This module is intentionally side-effect light: no network work happens at
-import time or from /api/health. The PyPI check is lazy, cached, and only used
-for normal PyPI-managed installs.
+Side-effect light: no network work at import time or from /api/health.
+The PyPI check is lazy, cached, and only for PyPI-managed installs.
 """
 
 from __future__ import annotations
@@ -66,9 +65,9 @@ def reset_update_status_cache() -> None:
 def detect_install_source() -> str:
     """Return a coarse install source without exposing local paths.
 
-    Sources are intentionally conservative. PEP 610 local/vcs metadata wins.
-    Legacy source installs are treated as local only when package files resolve
-    outside site-packages/dist-packages and under a Git checkout.
+    Conservative: PEP 610 local/vcs metadata wins. Legacy source
+    installs count as local only when package files resolve outside
+    site-packages/dist-packages and under a Git checkout.
     """
     try:
         dist = distribution(PACKAGE_NAME)
diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py
index 7a7774be40..3818253ac9 100644
--- a/studio/backend/utils/utils.py
+++ b/studio/backend/utils/utils.py
@@ -1,9 +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
 
-"""
-Shared backend utilities
-"""
+"""Shared backend utilities."""
 
 import os
 import structlog
@@ -18,16 +16,12 @@ logger = get_logger(__name__)
 
 
 # ── Client-safe error helpers ───────────────────────────────────
-# Never return raw exception text to clients (it can leak paths/internals);
-# log the full exception server-side and return a generic message.
+# Never return raw exception text to clients; log server-side, return generic.
 
 
 def safe_error_detail(error: Exception, fallback: str = "An internal error occurred") -> str:
-    """Map a caught exception to a generic, client-safe message.
-
-    Never includes raw ``str(error)`` (which can leak internal paths or stack
-    detail); known transient conditions get a friendlier hint. Always log the
-    real exception server-side (e.g. via ``log_and_http_error``) for diagnosis.
+    """Map an exception to a generic, client-safe message (never raw
+    ``str(error)``, which can leak paths). Log the real exception server-side.
     """
     text = str(error).lower()
     if (
@@ -43,10 +37,10 @@ def safe_error_detail(error: Exception, fallback: str = "An internal error occur
 
 
 def safe_curated_detail(error: Exception, fallback: str = "An internal error occurred") -> str:
-    """Client-safe text for curated domain/validation exceptions meant for the user.
+    """Client-safe text for curated domain/validation exceptions.
 
-    Keeps the message (paths stripped) instead of a generic fallback; use for known
-    exception types, keep ``safe_error_detail`` for generic ``Exception``.
+    Keeps the message (paths stripped) instead of a generic fallback; for known
+    exception types only (use ``safe_error_detail`` for generic ``Exception``).
     """
     from utils.native_path_leases import redact_native_paths
 
@@ -69,7 +63,7 @@ def log_and_http_error(
     """
     from fastapi import HTTPException
 
-    # Works for both structlog and stdlib loggers; exc_info=error logs its traceback.
+    # exc_info=error works for both structlog and stdlib loggers.
     (log or logger).error(f"{event}: {error}", exc_info = error)
     return HTTPException(status_code = status_code, detail = public_message)
 
@@ -77,14 +71,13 @@ def log_and_http_error(
 @contextmanager
 def without_hf_auth():
     """
-    Context manager to temporarily disable HuggingFace authentication.
+    Temporarily disable HuggingFace authentication.
 
     Usage:
         with without_hf_auth():
             # Code that should run without cached tokens
             model_info(model_name, token=None)
     """
-    # Save environment variables
     saved_env = {}
     env_vars = ["HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_HOME"]
     for var in env_vars:
@@ -92,11 +85,10 @@ def without_hf_auth():
             saved_env[var] = os.environ[var]
             del os.environ[var]
 
-    # Save disable flag
     saved_disable = os.environ.get("HF_HUB_DISABLE_IMPLICIT_TOKEN")
     os.environ["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = "1"
 
-    # Move token files temporarily
+    # Move token files aside temporarily
     token_files = []
     token_locations = [
         Path.home() / ".cache" / "huggingface" / "token",
@@ -121,7 +113,7 @@ def without_hf_auth():
             except Exception as e:
                 logger.error(f"Failed to restore token {original}: {e}")
 
-        # Restore environment
+        # Restore env
         for var, value in saved_env.items():
             os.environ[var] = value
 
@@ -133,14 +125,11 @@ def without_hf_auth():
 
 def format_error_message(error: Exception, model_name: str) -> str:
     """
-    Format user-friendly error messages for common issues.
+    Format a user-friendly error message for common load issues.
 
     Args:
         error: The exception that occurred
         model_name: Name of the model being loaded
-
-    Returns:
-        User-friendly error string
     """
     error_str = str(error).lower()
     model_short = model_name.split("/")[-1] if "/" in model_name else model_name
@@ -171,5 +160,4 @@ def format_error_message(error: Exception, model_name: str) -> str:
         )
         return f"Not enough {device_label} memory to load '{model_short}'. Try a smaller model or free memory."
 
-    # Generic fallback
     return str(error)
diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py
index cca30bfd44..98697df83c 100644
--- a/studio/backend/utils/wheel_utils.py
+++ b/studio/backend/utils/wheel_utils.py
@@ -24,16 +24,12 @@ FLASH_ATTN_RELEASE_BASE_URL = "https://github.com/Dao-AILab/flash-attention/rele
 
 @functools.lru_cache(maxsize = 1)
 def has_blackwell_gpu() -> bool:
-    """Return True if any visible NVIDIA GPU has compute capability >= 10.0
-    (Blackwell: sm_100, sm_120, sm_121, ...).
+    """Return True if any visible NVIDIA GPU has compute capability >= 10.0 (Blackwell).
 
-    Dao-AILab does not publish prebuilt flash-attention wheels for these
-    architectures, and the older-arch wheels fail to load on Blackwell, so
-    callers use this gate to skip the flash-attn install/upgrade path.
-
-    Result is cached for the process lifetime since GPU hardware does not
-    change. Tests that mock subprocess/nvidia-smi must call
-    ``has_blackwell_gpu.cache_clear()`` before each invocation.
+    Dao-AILab ships no flash-attention wheels for these archs and older-arch wheels
+    fail to load, so callers use this to skip the flash-attn install path. Cached
+    for the process lifetime; tests mocking nvidia-smi must call
+    ``has_blackwell_gpu.cache_clear()`` first.
     """
     exe = shutil.which("nvidia-smi")
     if not exe:
diff --git a/studio/frontend/package-lock.json b/studio/frontend/package-lock.json
index f36f2d8e79..7d1283e94d 100644
--- a/studio/frontend/package-lock.json
+++ b/studio/frontend/package-lock.json
@@ -31,6 +31,7 @@
         "@tailwindcss/vite": "^4.2.2",
         "@tanstack/react-router": "1.169.2",
         "@tanstack/react-table": "^8.21.3",
+        "@tanstack/react-virtual": "3.13.25",
         "@tauri-apps/api": "^2.10.1",
         "@tauri-apps/plugin-clipboard-manager": "^2.3.2",
         "@tauri-apps/plugin-notification": "^2.3.3",
@@ -6112,6 +6113,23 @@
         "react-dom": ">=16.8"
       }
     },
+    "node_modules/@tanstack/react-virtual": {
+      "version": "3.13.25",
+      "resolved": "https://registry.npmjs.org/@tanstack/react-virtual/-/react-virtual-3.13.25.tgz",
+      "integrity": "sha512-bmNoqMu6gcAW9JGrKVB0Q1tN1i5RONZF8r1fW0bbE4Oyf3DwEGnzzQJ2OW+Ozg1P4s8PyugkHg2ULZoFQN+cqw==",
+      "license": "MIT",
+      "dependencies": {
+        "@tanstack/virtual-core": "3.15.0"
+      },
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/tannerlinsley"
+      },
+      "peerDependencies": {
+        "react": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0",
+        "react-dom": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0"
+      }
+    },
     "node_modules/@tanstack/router-core": {
       "version": "1.169.2",
       "resolved": "https://registry.npmjs.org/@tanstack/router-core/-/router-core-1.169.2.tgz",
@@ -6154,6 +6172,16 @@
         "url": "https://github.com/sponsors/tannerlinsley"
       }
     },
+    "node_modules/@tanstack/virtual-core": {
+      "version": "3.15.0",
+      "resolved": "https://registry.npmjs.org/@tanstack/virtual-core/-/virtual-core-3.15.0.tgz",
+      "integrity": "sha512-0AwPGx0I8QxPYjAxShT/+z+ZOe9u8mW5rsXvivCTjRfRmz9a43+3mRyi4wwlyoUqOC56q/jatKa0Bh9M99BEHQ==",
+      "license": "MIT",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/tannerlinsley"
+      }
+    },
     "node_modules/@tauri-apps/api": {
       "version": "2.11.0",
       "resolved": "https://registry.npmjs.org/@tauri-apps/api/-/api-2.11.0.tgz",
diff --git a/studio/frontend/package.json b/studio/frontend/package.json
index 5ba0db143f..8537b0c076 100644
--- a/studio/frontend/package.json
+++ b/studio/frontend/package.json
@@ -40,6 +40,7 @@
     "@tailwindcss/vite": "^4.2.2",
     "@tanstack/react-router": "1.169.2",
     "@tanstack/react-table": "^8.21.3",
+    "@tanstack/react-virtual": "3.13.25",
     "@tauri-apps/api": "^2.10.1",
     "@tauri-apps/plugin-clipboard-manager": "^2.3.2",
     "@tauri-apps/plugin-notification": "^2.3.3",
diff --git a/studio/frontend/public/hub/profile/logo/anthropic.svg b/studio/frontend/public/hub/profile/logo/anthropic.svg
new file mode 100644
index 0000000000..7545cc8f3e
--- /dev/null
+++ b/studio/frontend/public/hub/profile/logo/anthropic.svg
@@ -0,0 +1,6 @@
+
+
+
+
+
+
\ No newline at end of file
diff --git a/studio/frontend/public/hub/profile/logo/cohere.png b/studio/frontend/public/hub/profile/logo/cohere.png
new file mode 100644
index 0000000000..99eabbb54f
Binary files /dev/null and b/studio/frontend/public/hub/profile/logo/cohere.png differ
diff --git a/studio/frontend/public/hub/profile/logo/deepseek.svg b/studio/frontend/public/hub/profile/logo/deepseek.svg
new file mode 100644
index 0000000000..d1ba06b942
--- /dev/null
+++ b/studio/frontend/public/hub/profile/logo/deepseek.svg
@@ -0,0 +1,14 @@
+
+
+  
+    
+  
+  
+    
+  
+
\ No newline at end of file
diff --git a/studio/frontend/public/hub/profile/logo/google.png b/studio/frontend/public/hub/profile/logo/google.png
new file mode 100644
index 0000000000..01bb81206e
Binary files /dev/null and b/studio/frontend/public/hub/profile/logo/google.png differ
diff --git a/studio/frontend/public/hub/profile/logo/hf.svg b/studio/frontend/public/hub/profile/logo/hf.svg
new file mode 100644
index 0000000000..ab959d165f
--- /dev/null
+++ b/studio/frontend/public/hub/profile/logo/hf.svg
@@ -0,0 +1,8 @@
+
+
+
+
+
+
+
+
diff --git a/studio/frontend/public/hub/profile/logo/ibm.png b/studio/frontend/public/hub/profile/logo/ibm.png
new file mode 100644
index 0000000000..31c965f0b3
Binary files /dev/null and b/studio/frontend/public/hub/profile/logo/ibm.png differ
diff --git a/studio/frontend/public/hub/profile/logo/meta.svg b/studio/frontend/public/hub/profile/logo/meta.svg
new file mode 100644
index 0000000000..9fa656bd6b
--- /dev/null
+++ b/studio/frontend/public/hub/profile/logo/meta.svg
@@ -0,0 +1,19 @@
+
+
+Logo of Meta Platforms -- Graphic created by Detmar Owen
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/studio/frontend/public/hub/profile/logo/microsoft.svg b/studio/frontend/public/hub/profile/logo/microsoft.svg
new file mode 100644
index 0000000000..5334aa7ca6
--- /dev/null
+++ b/studio/frontend/public/hub/profile/logo/microsoft.svg
@@ -0,0 +1 @@
+
\ No newline at end of file
diff --git a/studio/frontend/public/hub/profile/logo/minimax-color.png b/studio/frontend/public/hub/profile/logo/minimax-color.png
new file mode 100644
index 0000000000..e9472c676d
Binary files /dev/null and b/studio/frontend/public/hub/profile/logo/minimax-color.png differ
diff --git a/studio/frontend/public/hub/profile/logo/mistral.svg b/studio/frontend/public/hub/profile/logo/mistral.svg
new file mode 100644
index 0000000000..40c2591b31
--- /dev/null
+++ b/studio/frontend/public/hub/profile/logo/mistral.svg
@@ -0,0 +1,19 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
diff --git a/studio/frontend/public/hub/profile/logo/moonshot.jpg b/studio/frontend/public/hub/profile/logo/moonshot.jpg
new file mode 100644
index 0000000000..956a5b58b1
Binary files /dev/null and b/studio/frontend/public/hub/profile/logo/moonshot.jpg differ
diff --git a/studio/frontend/public/hub/profile/logo/nvidia.svg b/studio/frontend/public/hub/profile/logo/nvidia.svg
new file mode 100644
index 0000000000..ae65b09a2b
--- /dev/null
+++ b/studio/frontend/public/hub/profile/logo/nvidia.svg
@@ -0,0 +1 @@
+
\ No newline at end of file
diff --git a/studio/frontend/public/hub/profile/logo/openai.svg b/studio/frontend/public/hub/profile/logo/openai.svg
new file mode 100644
index 0000000000..74d9b1b44b
--- /dev/null
+++ b/studio/frontend/public/hub/profile/logo/openai.svg
@@ -0,0 +1,5 @@
+
+
+  
+  
+
\ No newline at end of file
diff --git a/studio/frontend/public/hub/profile/logo/qwen.png b/studio/frontend/public/hub/profile/logo/qwen.png
new file mode 100644
index 0000000000..67d2258f40
Binary files /dev/null and b/studio/frontend/public/hub/profile/logo/qwen.png differ
diff --git a/studio/frontend/public/hub/profile/logo/xai.svg b/studio/frontend/public/hub/profile/logo/xai.svg
new file mode 100644
index 0000000000..0c83eb3d9b
--- /dev/null
+++ b/studio/frontend/public/hub/profile/logo/xai.svg
@@ -0,0 +1 @@
+
\ No newline at end of file
diff --git a/studio/frontend/public/hub/profile/logo/zai.svg b/studio/frontend/public/hub/profile/logo/zai.svg
new file mode 100644
index 0000000000..28ca7280a1
--- /dev/null
+++ b/studio/frontend/public/hub/profile/logo/zai.svg
@@ -0,0 +1,215 @@
+
+
+
+
+
+	
+	
+		
+			
+				
+					
+					
+					
+				
+			
+		
+	
+
+
\ No newline at end of file
diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx
index 6f3c7618be..25dbfdc780 100644
--- a/studio/frontend/src/app/provider.tsx
+++ b/studio/frontend/src/app/provider.tsx
@@ -9,7 +9,9 @@ import {
   shouldUseCustomWindowTitlebar,
 } from "@/components/tauri/window-titlebar";
 import { Toaster } from "@/components/ui/sonner";
+import { TooltipProvider } from "@/components/ui/tooltip";
 import { WebUpdateBanner } from "@/components/web/update-banner";
+import { DownloadManagerPanel } from "@/features/hub/download-manager";
 import { getTauriAuthFailure, tauriAutoAuth } from "@/features/auth";
 import { NativeIntentDrain } from "@/features/native-intents/native-intent-drain";
 import { useTauriBackend, type BackendStatus } from "@/hooks/use-tauri-backend";
@@ -48,11 +50,10 @@ async function applyAppWindowLayout(isCurrent: WindowLayoutGuard): Promise
 
   const win = getCurrentWindow();
   // Decide first-launch vs restore from the on-disk state file BEFORE touching the
-  // window. Probing the window itself after restoreStateCurrent is unreliable:
-  // on GTK, set_size against a hidden window is deferred until show(), so
-  // innerSize() reads a stale value and any baseline fallback would overwrite the
-  // queued restore. On macOS the same probe works, hence the inconsistency
-  // between previous iterations of this code.
+  // window. Probing the window after restoreStateCurrent is unreliable: on GTK,
+  // set_size on a hidden window is deferred until show(), so innerSize() reads a
+  // stale value and a baseline fallback would overwrite the queued restore. On
+  // macOS the same probe works, hence the inconsistency between prior iterations.
   const hasSavedState = await invoke("has_saved_window_state");
   if (!isCurrent()) return;
 
@@ -60,9 +61,8 @@ async function applyAppWindowLayout(isCurrent: WindowLayoutGuard): Promise
   if (!isCurrent()) return;
 
   if (hasSavedState) {
-    // Subsequent launch: the plugin handles size, position, and maximized,
-    // with built-in off-screen protection (monitor-intersection check) for
-    // positions saved on a now-disconnected display.
+    // Subsequent launch: plugin restores size/position/maximized, with built-in
+    // off-screen protection for positions saved on a now-disconnected display.
     await restoreStateCurrent(
       StateFlags.SIZE | StateFlags.POSITION | StateFlags.MAXIMIZED,
     );
@@ -87,8 +87,8 @@ async function applyAppWindowLayout(isCurrent: WindowLayoutGuard): Promise
   if (!isCurrent()) return;
   await win.show();
   if (!isCurrent()) return;
-  // Apply constraints after restore/show. Setting constraints before plugin restore
-  // can emit a Resized event and overwrite the plugin's cached saved size.
+  // Apply constraints after restore/show: doing so before plugin restore can emit
+  // a Resized event and overwrite the plugin's cached saved size.
   await win.setSizeConstraints({ minWidth: MIN_WINDOW_WIDTH, minHeight: MIN_WINDOW_HEIGHT });
 }
 
@@ -257,6 +257,7 @@ function TauriWrapper({ children }: { children: ReactNode }) {
     return (
       <>
         {children}
+        
         
       
     );
@@ -274,6 +275,7 @@ function TauriWrapper({ children }: { children: ReactNode }) {
       
       
       {children}
+      
     
   ) : (
     
-      
-        {children}
-      
-      
+      
+        
+          {children}
+        
+        
+      
     
   );
 }
diff --git a/studio/frontend/src/app/router.tsx b/studio/frontend/src/app/router.tsx
index a0ca1e8cdb..dbb74ee1a1 100644
--- a/studio/frontend/src/app/router.tsx
+++ b/studio/frontend/src/app/router.tsx
@@ -12,6 +12,7 @@ import { Route as exportRoute } from "./routes/export";
 import { Route as gridTestRoute } from "./routes/grid-test";
 import { Route as indexRoute } from "./routes/index";
 import { Route as loginRoute } from "./routes/login";
+import { Route as hubRoute } from "./routes/hub";
 import { Route as onboardingRoute } from "./routes/onboarding";
 import { Route as projectsRoute } from "./routes/projects";
 import { Route as changePasswordRoute } from "./routes/change-password";
@@ -24,6 +25,7 @@ const routeTree = rootRoute.addChildren([
   loginRoute,
   changePasswordRoute,
   gridTestRoute,
+  hubRoute,
   settingsRoute,
   studioRoute,
   chatRoute,
diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx
index da74a9e8a7..c8e47902b9 100644
--- a/studio/frontend/src/app/routes/__root.tsx
+++ b/studio/frontend/src/app/routes/__root.tsx
@@ -42,6 +42,7 @@ const CHAT_ONLY_ALLOWED = new Set([
   "/",
   "/chat",
   "/projects",
+  "/hub",
   "/login",
   "/signup",
   "/change-password",
@@ -55,8 +56,8 @@ function isChatOnlyAllowed(pathname: string): boolean {
 
 export const Route = createRootRoute({
   beforeLoad: async ({ location }) => {
-    // Ensure platform info is fetched before checking chat-only guard.
-    // fetchDeviceType caches after first call, so subsequent navigations are instant.
+    // Fetch platform info before the chat-only guard. fetchDeviceType caches,
+    // so later navigations are instant.
     await fetchDeviceType();
     const chatOnly = usePlatformStore.getState().isChatOnly();
     if (chatOnly && !isChatOnlyAllowed(location.pathname)) {
diff --git a/studio/frontend/src/app/routes/hub.tsx b/studio/frontend/src/app/routes/hub.tsx
new file mode 100644
index 0000000000..dcd6617ec8
--- /dev/null
+++ b/studio/frontend/src/app/routes/hub.tsx
@@ -0,0 +1,29 @@
+// 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 { createRoute } from "@tanstack/react-router";
+import { lazy } from "react";
+import { requireAuth } from "../auth-guards";
+import { Route as rootRoute } from "./__root";
+
+const ModelsPage = lazy(() =>
+  import("@/features/hub/hub-page").then((m) => ({
+    default: m.ModelsPage,
+  })),
+);
+
+export interface ModelsSearch {
+  tab?: "discover" | "downloaded";
+}
+
+export const Route = createRoute({
+  getParentRoute: () => rootRoute,
+  path: "/hub",
+  beforeLoad: () => requireAuth(),
+  component: ModelsPage,
+  validateSearch: (search: Record): ModelsSearch => {
+    const raw = search.tab;
+    if (raw === "discover" || raw === "downloaded") return { tab: raw };
+    return {};
+  },
+});
diff --git a/studio/frontend/src/app/routes/settings.tsx b/studio/frontend/src/app/routes/settings.tsx
index fa97a450f7..84650bdb8b 100644
--- a/studio/frontend/src/app/routes/settings.tsx
+++ b/studio/frontend/src/app/routes/settings.tsx
@@ -7,10 +7,10 @@ import { useSettingsDialogStore } from "@/features/settings";
 import { requireAuth } from "../auth-guards";
 import { Route as rootRoute } from "./__root";
 
-// /settings is a deep link to the modal. Open it, then redirect home.
-// Tab title is driven by useSettingsDialogStore in __root.tsx since the
-// redirect means /settings never stays matched; staticData is just a
-// safety net if beforeLoad ever stops throwing.
+// /settings deep-links the modal: open it, then redirect home. Tab title is
+// driven by useSettingsDialogStore in __root.tsx since the redirect means
+// /settings never stays matched; staticData is a safety net if beforeLoad
+// ever stops throwing.
 export const Route = createRoute({
   getParentRoute: () => rootRoute,
   path: "/settings",
diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx
index 941c2e8a74..2abc4d190b 100644
--- a/studio/frontend/src/components/app-sidebar.tsx
+++ b/studio/frontend/src/components/app-sidebar.tsx
@@ -47,6 +47,7 @@ import { cn } from "@/lib/utils";
 import {
   ChefHatIcon,
   CursorInfo02Icon,
+  DashboardCircleIcon,
   Delete02Icon,
   DownloadSquare01Icon,
   Edit03Icon,
@@ -138,11 +139,8 @@ function getTourId(pathname: string): string | null {
   return null;
 }
 
-// Hugeicons' TestTube01Icon ships with two interior bubbles (paths #4
-// and #5 of the 5-path definition). Slicing to the first three paths
-// keeps the test-tube outline + horizontal cap + liquid line, dropping
-// the bubbles. The original export stays untouched, and HugeiconsIcon
-// renders this trimmed array exactly the same way.
+// TestTube01Icon's last 2 paths are interior bubbles; slice to the first
+// 3 (outline + cap + liquid line) to drop them. Original export untouched.
 const TestTubeOutlineIcon = TestTube01Icon.slice(
   0,
   3,
@@ -261,11 +259,11 @@ export function AppSidebar() {
 
   const scrollRef = useRef(null);
   const [scrolled, setScrolled] = useState(false);
-  // Bottom fade hides at the very bottom (and for short, non-scrolling lists)
-  // so the last row isn't washed out - Gemini-style.
+  // Bottom fade hides at the very bottom / for short lists so the last row
+  // isn't washed out (Gemini-style).
   const [canScrollDown, setCanScrollDown] = useState(false);
-  // Driven only from onScroll + a content-change effect below. Deliberately NO
-  // ResizeObserver: its callback-driven setState created a render loop (React
+  // Driven only from onScroll + a content-change effect below. No
+  // ResizeObserver: its callback-driven setState caused a render loop (React
   // #185). Both setters bail out when unchanged, so neither path can loop.
   const syncScrollState = (el: HTMLDivElement) => {
     const nextScrolled = el.scrollTop > 0;
@@ -311,10 +309,9 @@ export function AppSidebar() {
   const selectedHistoryRunId = useTrainingRuntimeStore((s) => s.selectedHistoryRunId);
   const setSelectedHistoryRunId = useTrainingRuntimeStore((s) => s.setSelectedHistoryRunId);
 
-  // Recompute the bottom-fade state on mount and whenever the list height can
-  // change (items load, sections collapse/expand, route switches the visible
-  // list) - onScroll never fires for short, non-scrolling lists. Guarded
-  // setState below means this can't loop even if a dep is a fresh reference.
+  // Recompute bottom-fade on mount and whenever list height can change
+  // (items load, sections toggle, route switch) - onScroll never fires for
+  // short, non-scrolling lists. Guarded setState below can't loop.
   useEffect(() => {
     const el = scrollRef.current;
     if (!el) return;
@@ -761,9 +758,8 @@ export function AppSidebar() {
               label={t("shell.navigation.search")}
               active={false}
               onClick={() => {
-                // Search is read-only over chat history and never runs
-                // inference, so it stays available while training (unlike
-                // New chat, which is gated on `chatDisabled`).
+                // Search is read-only and never runs inference, so it stays
+                // available while training (unlike New chat, gated on chatDisabled).
                 useChatSearchStore.getState().open();
                 closeMobileIfOpen();
               }}
@@ -796,8 +792,16 @@ export function AppSidebar() {
                   closeMobileIfOpen();
                 }}
               />
-              {/* Train has its own labelled section when expanded; surface it as
-                  a plain icon here only while the sidebar is collapsed. */}
+               {
+                  navigate({ to: "/hub" });
+                  closeMobileIfOpen();
+                }}
+              />
+              {/* Train has a labelled section when expanded; plain icon here only when collapsed. */}
               
                 
                   {runItems.map((run) => {
-                    // An explicit sidebar selection wins. Otherwise highlight
-                    // the active job only while the "Current Run" tab is the
-                    // view - that covers a live run (it auto-switches there) and
-                    // a just-finished/errored run you're still viewing, while
-                    // keeping the Configure tab unhighlighted even though
-                    // `activeJobId` stays pinned to the last job.
+                    // Explicit selection wins. Otherwise highlight the active
+                    // job only while the "Current Run" tab is the view, keeping
+                    // the Configure tab unhighlighted even though activeJobId
+                    // stays pinned to the last job.
                     const isActiveRun =
                       selectedHistoryRunId != null
                         ? run.id === selectedHistoryRunId
@@ -988,10 +990,9 @@ export function AppSidebar() {
       
 
       
-        {/* Fade above the profile box, shown only while there's more list below
-            the fold; at the very bottom (or for short lists) it fades out so the
-            last row shows fully (Gemini-style). `right-2` keeps it clear of the
-            8px scrollbar gutter so the scrollbar isn't faded out. */}
+        {/* Fade above the profile box, shown only when there's more list below
+            the fold; at the bottom (or short lists) it fades so the last row
+            shows fully (Gemini-style). right-2 keeps it clear of the 8px scrollbar gutter. */}