From 3b235895bdf08410e0a9032e663e82c0de60a6a4 Mon Sep 17 00:00:00 2001 From: Anas Khan Date: Wed, 15 Jul 2026 02:26:12 +0530 Subject: [PATCH 01/17] Fix DeepScaleR-1.5B mapper entry pointing its 16bit repo at DeepHermes-3-8B (#7088) --- tests/test_get_model_name.py | 17 +++++++++++++++++ unsloth/models/mapper.py | 2 +- 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/tests/test_get_model_name.py b/tests/test_get_model_name.py index 993931c0ef..33ad316d88 100644 --- a/tests/test_get_model_name.py +++ b/tests/test_get_model_name.py @@ -102,6 +102,19 @@ class TestGetModelName(unittest.TestCase): ), ("unsloth/Kimi-K2-Instruct", True, "unsloth/Kimi-K2-Instruct-BF16", True), ("unsloth/Kimi-K2-Instruct", False, "unsloth/Kimi-K2-Instruct", False), + # DeepScaleR-1.5B must resolve to its own 16bit repo, not another model + ( + "agentica-org/DeepScaleR-1.5B-Preview", + False, + "unsloth/DeepScaleR-1.5B-Preview", + True, + ), + ( + "agentica-org/DeepScaleR-1.5B-Preview", + True, + "unsloth/DeepScaleR-1.5B-Preview-unsloth-bnb-4bit", + True, + ), # Fallback-to-original behavior "nonexistent-user/nonexistent-model-123", "google/gemma-3-random-prototype-123", @@ -157,6 +170,10 @@ class TestGetModelName(unittest.TestCase): with self.subTest(src = src): self.assertEqual(FLOAT_TO_INT_MAPPER[src], expected) self.assertEqual(MAP_TO_UNSLOTH_16bit["qwen/qwen3-8b-fp8"], "unsloth/Qwen3-8B-FP8") + self.assertEqual( + MAP_TO_UNSLOTH_16bit["agentica-org/deepscaler-1.5b-preview"], + "unsloth/DeepScaleR-1.5B-Preview", + ) if __name__ == "__main__": diff --git a/unsloth/models/mapper.py b/unsloth/models/mapper.py index f3a0e1f9bb..4558bb0f28 100644 --- a/unsloth/models/mapper.py +++ b/unsloth/models/mapper.py @@ -701,7 +701,7 @@ __INT_TO_FLOAT_MAPPER = \ "unsloth/Qwen2.5-VL-72B-Instruct-bnb-4bit", ), "unsloth/DeepScaleR-1.5B-Preview-unsloth-bnb-4bit" : ( - "unsloth/DeepHermes-3-Llama-3-8B-Preview", + "unsloth/DeepScaleR-1.5B-Preview", "agentica-org/DeepScaleR-1.5B-Preview", "unsloth/DeepScaleR-1.5B-Preview-bnb-4bit", ), From 387b2f28e38f44e0586dff31161f10c96a70fdec Mon Sep 17 00:00:00 2001 From: Anas Khan Date: Wed, 15 Jul 2026 07:03:54 +0530 Subject: [PATCH 02/17] fix(dataprep): guard smart_chunk_text against stride >= chunk_size (#7126) RawTextDataLoader.smart_chunk_text takes chunk_size and stride as its own arguments, so a direct call with stride >= chunk_size bypasses the constructor validation. In that case `start_idx += chunk_size - stride` is non-positive, so start_idx never advances past the first window and the chunking loop never terminates (hangs). Re-add the chunk_size/stride guard at the top of smart_chunk_text so direct callers fail fast with a clear ValueError. The constructor keeps its own guard for the internal callers (defense in depth). Add a regression test that calls smart_chunk_text directly with stride == chunk_size and stride > chunk_size and asserts it raises instead of hanging. Signed-off-by: Anas Khan <83116240+anxkhn@users.noreply.github.com> --- tests/test_raw_text.py | 20 ++++++++++++++++++++ unsloth/dataprep/raw_text.py | 7 +++++++ 2 files changed, 27 insertions(+) diff --git a/tests/test_raw_text.py b/tests/test_raw_text.py index d7f6c317fe..c0f5d5f398 100644 --- a/tests/test_raw_text.py +++ b/tests/test_raw_text.py @@ -141,6 +141,26 @@ def test_raw_text_loader(): except ValueError as e: assert "stride" in str(e) and "chunk_size" in str(e) + # smart_chunk_text validation: called directly, chunk_size/stride are its own + # arguments and bypass the constructor guard, so it must guard itself or an + # invalid stride makes `start_idx += chunk_size - stride` non-positive and the + # chunking loop never terminates (hangs). + long_text = "This is a test file for raw text training. " * 10 + valid_chunks = loader.smart_chunk_text(long_text, chunk_size = 5, stride = 2) + assert len(valid_chunks) > 0, "Valid stride should produce chunks" + + try: + loader.smart_chunk_text(long_text, chunk_size = 5, stride = 5) + assert False, "Should raise ValueError for stride == chunk_size" + except ValueError as e: + assert "stride" in str(e) and "chunk_size" in str(e) + + try: + loader.smart_chunk_text(long_text, chunk_size = 5, stride = 10) + assert False, "Should raise ValueError for stride > chunk_size" + except ValueError as e: + assert "stride" in str(e) and "chunk_size" in str(e) + # Preprocessor. preprocessor = TextPreprocessor() clean_text = preprocessor.clean_text(" messy text \n\n\n ") diff --git a/unsloth/dataprep/raw_text.py b/unsloth/dataprep/raw_text.py index 7d18b1ff29..0993f0b4e9 100644 --- a/unsloth/dataprep/raw_text.py +++ b/unsloth/dataprep/raw_text.py @@ -132,6 +132,13 @@ class RawTextDataLoader: 3. Maintains context with stride overlap 4. Returns tokenized chunks directly (more efficient) or text chunks """ + if chunk_size <= 0: + raise ValueError(f"chunk_size must be positive, got {chunk_size}") + if stride >= chunk_size: + raise ValueError( + f"stride ({stride}) must be smaller than chunk_size ({chunk_size}) to progress the chunking loop" + ) + # Tokenize the whole text once for accurate token counts tokenized = self.tokenizer(text, return_tensors = "pt", add_special_tokens = False) tokens = tokenized["input_ids"] From 2b52da98cca76c6b8d849c20e7a7d624902b36e6 Mon Sep 17 00:00:00 2001 From: Nilay <118994073+NilayYadav@users.noreply.github.com> Date: Wed, 15 Jul 2026 07:05:38 +0530 Subject: [PATCH 03/17] Fix Hub offline status (#7129) Co-authored-by: Etherl <61019402+Etherll@users.noreply.github.com> --- .../frontend/src/features/hub/lib/network.ts | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/studio/frontend/src/features/hub/lib/network.ts b/studio/frontend/src/features/hub/lib/network.ts index 48d165c92a..a0bdf6598d 100644 --- a/studio/frontend/src/features/hub/lib/network.ts +++ b/studio/frontend/src/features/hub/lib/network.ts @@ -7,10 +7,10 @@ export function isBrowserOffline(): boolean { const NETWORK_STATUS_EVENT = "unsloth-network-status"; const REMOTE_OFFLINE_TTL_MS = 30_000; -const HUGGING_FACE_REMOTE_ORIGINS = [ - "https://huggingface.co", - "https://datasets-server.huggingface.co", -] as const; +// Discovery and repository pages are served by the main Hugging Face origin. +// Keep optional services such as datasets-server separate so an outage there +// cannot make the whole Hub appear offline. +const HUGGING_FACE_ORIGIN = "https://huggingface.co"; const noopUnsubscribe = () => undefined; type RemoteNetworkScope = string | readonly string[]; @@ -33,7 +33,7 @@ export function getBrowserOfflineRetryDelayMs(): number { // recovery doesn't stall on platforms where navigator.onLine is stuck false. return Math.max( 0, - getRemoteOfflineUntil(HUGGING_FACE_REMOTE_ORIGINS) - Date.now(), + getRemoteOfflineUntil(HUGGING_FACE_ORIGIN) - Date.now(), ); } @@ -56,7 +56,7 @@ function getRemoteOfflineUntil(scope: RemoteNetworkScope): number { } export function isRemoteNetworkOffline( - scope: RemoteNetworkScope = HUGGING_FACE_REMOTE_ORIGINS, + scope: RemoteNetworkScope = HUGGING_FACE_ORIGIN, ): boolean { return getRemoteOfflineUntil(scope) > Date.now(); } @@ -66,7 +66,7 @@ export function isHuggingFaceOffline(): boolean { // WebKitGTK/Tauri webviews). The authoritative signal is the empirical // remote-offline TTL, set when a real fetch fails and cleared on next success; // navigator's online/offline events still drive re-evaluation. - return isRemoteNetworkOffline(HUGGING_FACE_REMOTE_ORIGINS); + return isRemoteNetworkOffline(HUGGING_FACE_ORIGIN); } export function markRemoteNetworkOnline(origin?: string): void { @@ -85,13 +85,13 @@ export function markRemoteNetworkOnline(origin?: string): void { } export function markRemoteNetworkOffline( - originOrTtl: string | number = HUGGING_FACE_REMOTE_ORIGINS[0], + originOrTtl: string | number = HUGGING_FACE_ORIGIN, ttlMs = REMOTE_OFFLINE_TTL_MS, ): void { const origin = typeof originOrTtl === "string" ? originOrTtl - : HUGGING_FACE_REMOTE_ORIGINS[0]; + : HUGGING_FACE_ORIGIN; const ttl = typeof originOrTtl === "number" ? originOrTtl : ttlMs; const nextUntil = Date.now() + ttl; if (nextUntil <= (remoteOfflineUntilByOrigin.get(origin) ?? 0)) { From f9aa818ca8b87fe3bc8f8e786faa2171ff6f67a4 Mon Sep 17 00:00:00 2001 From: Anas Khan Date: Wed, 15 Jul 2026 12:53:06 +0530 Subject: [PATCH 04/17] fix: name unsloth_vllm_standby parameter in vLLM standby error (#7089) * fix: name unsloth_vllm_standby parameter in vLLM standby error FastBaseModel.from_pretrained's vLLM-standby guard raised "UNSLOTH_VLLM_STANDBY is True, but UNSLOTH_VLLM_STANDBY is not set to 1", naming the environment variable in both clauses. The value that is True is the unsloth_vllm_standby parameter, not the env var, so the message was self-contradictory. Name the parameter in the first clause, matching the sibling guard in FastLlamaModel.from_pretrained. Signed-off-by: Anas Khan <83116240+anxkhn@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove vLLM standby error message test --------- Signed-off-by: Anas Khan <83116240+anxkhn@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen --- unsloth/models/vision.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 3b5af39e69..729c191e83 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -739,7 +739,7 @@ class FastBaseModel: if unsloth_vllm_standby and os.environ.get("UNSLOTH_VLLM_STANDBY", "0") != "1": raise RuntimeError( - "Unsloth: UNSLOTH_VLLM_STANDBY is True, but UNSLOTH_VLLM_STANDBY is not set to 1!" + "Unsloth: `unsloth_vllm_standby` is True, but environment variable `UNSLOTH_VLLM_STANDBY` is not set to 1!" ) if model_types is None: From dc65638b7d1e2686cb2c16d9fed13c7ea5df6fb2 Mon Sep 17 00:00:00 2001 From: Gaurav Dubey Date: Wed, 15 Jul 2026 12:54:11 +0530 Subject: [PATCH 05/17] Studio: expose Windows drive roots in the folder browser (#7082) * Studio: expose Windows drive roots in the folder browser The model-selection folder browser bounds navigation to the roots returned by _build_browse_allowlist(), which exposed Linux removable-media mounts via linux_run_media_mount_roots() but had no Windows analog. As a result a user on C: could not browse to D:/E: to pick a model directory. Add windows_drive_roots(), a Windows-only companion to linux_run_media_mount_roots() that lists readable logical drive roots, and wire it into both browse-allowlist builders and their suggestion chips so other drives are both navigable and offered as quick-picks. The helper is a no-op on Linux/macOS, so existing platforms are unaffected. Closes #6368 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: cover the Windows drive-root browse wiring with an integration test Add an allowlist integration test mirroring the Linux side's test_legacy_browse_allowlist_includes_linux_run_media_mounts: it extracts _build_browse_allowlist from routes/models.py, stubs external_media so windows_drive_roots() yields a fake drive root, and asserts that root becomes browsable through the built allowlist. Proves the wiring, not just the helper. * Studio: skip inactive drives via GetLogicalDrives before probing Resolve active logical drives from GetLogicalDrives() before probing each letter with os.path.isdir. Probing a drive letter mapped to a disconnected network share can otherwise block the async backend for tens of seconds per letter. The call degrades gracefully (falls back to probing all letters) when ctypes/windll is unavailable, so behavior is unchanged on Linux/macOS. Tests override the bitmask source to stay deterministic on real Windows hosts. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: allow browsing descendants of a drive-root allowlist entry routes/models.py _is_path_inside_allowlist() checked descendants with startswith(root_real + os.sep). A drive root ("D:\") already ends in a separator, so the prefix became "D:\\" and a child like "D:\models" was rejected with 403 after the browser opened the drive root. Only append a separator when the root does not already end in one. folder_browser.py already uses commonpath and was unaffected. Adds a regression test covering the separator-terminated-root descendant case. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: enforce the system-directory denylist during folder browsing Exposing whole Windows drive roots (and any legacy-registered filesystem root) widened the browse allowlist above system directories, but the browse resolvers only re-applied the credential/config denylist, not the _denied_path_prefixes() system-dir denylist that scan-folder registration enforces. That let browse-folders enumerate C:\Windows, C:\Program Files, /etc and /proc. - Add is_denied_system_path() to both storage modules and enforce it in both browse resolvers (legacy routes/models.py and hub folder_browser.py), on each resolved child and on the final target, keeping the /run/media carve-out. - Rework the legacy _is_path_inside_allowlist to use splitdrive + commonpath so a Windows drive root authorizes its descendants while a bare POSIX / does not, and to compare case-insensitively like the hub browser. - Reject the filesystem root in the legacy add_scan_folder, matching the hub. - Hide denied system dirs from browse listings and suggestion chips. - Add tests/test_browse_denylist.py and update the external-media path tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: make browse-denylist tests OS-portable The browse-time denylist tests used real /etc and tmp_path locations; on macOS tmp lives under the (legitimately denied) /private/var and /etc resolves to /private/etc, so three tests failed there. Pin the platform / use a tmp-based denied prefix so they assert the same behavior on Linux, macOS and Windows. * Studio: apply the bare POSIX-root guard to the hub folder browser too The _is_path_inside_allowlist guard that stops a legacy-registered '/' scan folder from authorizing every absolute path lived only in the legacy browser. The hub browser used commonpath without it, so a stale '/' row let it descend into /var, /root, /home -- which the system-directory denylist (/proc /sys /dev /etc /boot /run) does not cover, while the legacy browser blocked them. Mirror the legacy guard so both browsers treat '/' identically. Also resolve each directory entry before the denylist check in both listing loops, so a symlink or junction pointing into a denied dir is hidden instead of rendered as a row that 403s on descent. Adds legacy-vs-hub parity tests. * Studio: bound Windows drive probing so a disconnected mapping can't stall the browser GetLogicalDrives includes mapped network drives, so a disconnected but still mapped drive (e.g. Z: -> \\nas\share) stays set in the bitmask and reaches os.path.isdir, which can block for tens of seconds while Windows tries to reconnect. Because windows_drive_roots() runs synchronously while building both folder-browser responses, one stale mapping stalled every browse request. Probe each surviving drive in a daemon thread bounded by a short timeout and skip it if it does not answer in time, so a hung mapping is dropped instead of blocking the caller. Connected drives (local or network) still respond well within the timeout, so drive discovery is unchanged. Corrects the GetLogicalDrives docstring, which claimed the bitmask alone prevented the stall. * Studio: probe drive/media roots once per browse request, not twice Both folder browsers called windows_drive_roots() (and linux_run_media_mount_roots()) twice per browse request: once to seed the allowlist in _build_browse_allowlist() and again to build the suggestion chips. With the bounded drive probe, a disconnected mapped network drive then paid the timeout twice per folder click. Probe both once in the request handler and pass the results into _build_browse_allowlist(), reusing them for the chips, in both the legacy and hub browsers. Adds a test asserting the roots are reused, not re-probed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: run the legacy browse endpoint in the threadpool, fix its stale test Two follow-ups from review of the drive-probe changes: - browse_folders was 'async def' but does only blocking filesystem I/O (the timeout-bounded drive probe, iterdir, realpath). On the event loop a disconnected mapped drive waiting out its probe timeout stalled every other request. Declare it sync 'def' so FastAPI runs it in the threadpool, matching the hub browse endpoint. No await was used in the body. - test_browse_folders_hides_sensitive_dirs monkeypatched _build_browse_allowlist with a zero-arg lambda; the once-per-request refactor now calls it with (media_roots, drive_roots), so the lambda raised TypeError. Accept and ignore the args. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: probe Windows drive roots concurrently so multiple dead mappings don't stack timeouts windows_drive_roots() probed each candidate serially, so N disconnected-but-mapped network drives each paid the full per-drive timeout in turn (e.g. four stale mappings added ~8s to every folder-browser request). Collect the candidate roots first, then probe them all at once under a single overall deadline, so the added delay stays at ~one timeout regardless of how many drives are disconnected. _readable_dir_within stays as a thin single-path wrapper for its existing callers/tests. * Studio: tighten comments in the folder-browser drive-root changes Condense the comments and docstrings added by the Windows drive-root and system-directory denylist work to be shorter and clearer while keeping the security and correctness rationale intact. Comment and docstring text only; no code changes. * Studio: iterate the input, not the results dict, when collecting readable drive probes _readable_dirs_within returned {path for path, ok in results.items()...}, but a probe thread that exceeded the join deadline is still alive and can insert its key into results during that iteration, raising 'dictionary changed size during iteration' -- reachable exactly in the disconnected-mapped-drive case the probe exists for. Iterate the fixed input list and read results.get(path) (an atomic read) instead. * Studio: keep the browse-route containment tests denylist-inert so they pass on macOS test_browse_folders_route.py exercises allowlist containment and the file-vs-directory guard, not the system-directory denylist. On macOS pytest tmp_path resolves under /private/var, a denied prefix, so _resolve_browse_target 403s the fixture dirs before the containment logic runs (4 failures). Add an autouse fixture that makes is_denied_system_path inert in this file; the denylist keeps its own coverage in test_browse_denylist.py. * Studio: keep the hub browse tests denylist-inert so they pass on macOS * Studio: register a UNC share root; only reject local filesystem roots * Studio: reject device drive roots and browse a registered UNC share root * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: treat device-namespace volume GUID roots as local filesystem roots --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen --- .../hub/services/models/folder_browser.py | 87 +++- studio/backend/hub/storage/scan_folders.py | 26 +- .../backend/hub/tests/test_model_services.py | 17 +- studio/backend/routes/models.py | 113 +++++- studio/backend/storage/studio_db.py | 27 +- studio/backend/tests/test_browse_denylist.py | 371 ++++++++++++++++++ .../tests/test_browse_folders_route.py | 12 + .../tests/test_linux_external_media_paths.py | 9 +- .../test_windows_external_drive_paths.py | 354 +++++++++++++++++ studio/backend/utils/paths/external_media.py | 129 ++++++ 10 files changed, 1117 insertions(+), 28 deletions(-) create mode 100644 studio/backend/tests/test_browse_denylist.py create mode 100644 studio/backend/tests/test_windows_external_drive_paths.py diff --git a/studio/backend/hub/services/models/folder_browser.py b/studio/backend/hub/services/models/folder_browser.py index eb137127fb..d56b62c318 100644 --- a/studio/backend/hub/services/models/folder_browser.py +++ b/studio/backend/hub/services/models/folder_browser.py @@ -15,6 +15,7 @@ from loggers import get_logger from hub.schemas.inventory import BrowseEntry, BrowseFoldersResponse from hub.storage.scan_folders import ( contains_sensitive_path_component, + is_denied_system_path, list_scan_folders, ) from hub.utils.paths import ( @@ -27,7 +28,10 @@ from hub.utils.paths import ( studio_root, well_known_model_dirs, ) -from utils.paths.external_media import linux_run_media_mount_roots +from utils.paths.external_media import ( + linux_run_media_mount_roots, + windows_drive_roots, +) from hub.services.models.common import _safe_is_dir from hub.services.models.local_inventory import _resolve_hf_cache_dir @@ -158,8 +162,14 @@ def _looks_like_model_dir(directory: Path) -> bool: return False -def _build_browse_allowlist() -> list[Path]: - """Root directories the browser may walk (also seeds the suggestion chips): HOME, resolved HF cache dirs, Studio outputs/exports/root, registered scan folders, and well-known local-LLM dirs. Each is added only if it resolves to a real directory so the sandbox has no dead boundary.""" +def _build_browse_allowlist( + media_roots: Optional[list[Path]] = None, drive_roots: Optional[list[Path]] = None +) -> list[Path]: + """Root directories the browser may walk (also seeds the suggestion chips): HOME, resolved HF cache dirs, Studio outputs/exports/root, registered scan folders, and well-known local-LLM dirs. Each is added only if it resolves to a real directory so the sandbox has no dead boundary. + + *media_roots* / *drive_roots* let the caller pass already-probed + removable-media and Windows drive roots so they aren't scanned again (a + disconnected mapped drive can make each probe slow); probed here when ``None``.""" from hub.storage.scan_folders import list_scan_folders candidates: list[Path] = [] @@ -176,7 +186,13 @@ def _build_browse_allowlist() -> list[Path]: candidates.append(resolved) _add(Path.home()) - for p in linux_run_media_mount_roots(): + if media_roots is None: + media_roots = linux_run_media_mount_roots() + if drive_roots is None: + drive_roots = windows_drive_roots() + for p in media_roots: + _add(p) + for p in drive_roots: _add(p) _add(_resolve_hf_cache_dir()) try: @@ -218,7 +234,14 @@ def _build_browse_allowlist() -> list[Path]: def _is_path_inside_allowlist(target: Path, allowed_roots: list[Path]) -> bool: - """True if *target* equals or descends from any allowed root; uses ``os.path.realpath`` so symlinks cannot escape the sandbox.""" + """True if *target* equals or descends from any allowed root; uses ``os.path.realpath`` so symlinks cannot escape the sandbox. + + A Windows drive root (``D:\\``) authorizes its descendants, but a bare POSIX + root (``/``) must NOT: a single ``/`` allowlist entry (e.g. a legacy scan + folder) would otherwise authorize every absolute path, reaching ``/var``, + ``/root``, etc. the denylist does not cover. Mirrors the legacy browser so + both treat ``/`` identically. + """ try: target_real = os.path.normcase(os.path.realpath(str(target))) except OSError: @@ -228,13 +251,25 @@ def _is_path_inside_allowlist(target: Path, allowed_roots: list[Path]) -> bool: root_real = os.path.normcase(os.path.realpath(str(root))) except OSError: continue + if target_real == root_real: + return True + drive, tail = os.path.splitdrive(root_real) + if os.path.dirname(root_real) == root_real and not drive: + # Bare POSIX filesystem root ("/"): equality above is the only + # match; do not let it authorize arbitrary descendants. + continue + if drive.startswith(("\\\\", "//")) and not tail: + # Bare UNC share root (\\server\share): os.path.commonpath raises + # "can't mix absolute and relative" on it, so authorize its + # descendants with a boundary-safe prefix test (normcase applied). + if target_real.startswith(root_real.rstrip("\\/") + os.sep): + return True + continue try: if os.path.commonpath([target_real, root_real]) == root_real: return True except ValueError: continue - if target_real == root_real: - return True return False @@ -347,6 +382,11 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa status_code = 403, detail = "Credential or configuration directories are not browseable.", ) + if is_denied_system_path(str(resolved_child)): + raise HTTPException( + status_code = 403, + detail = "System directories are not browseable.", + ) current = resolved_child if contains_sensitive_path_component(str(current)): @@ -354,6 +394,13 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa status_code = 403, detail = "Credential or configuration directories are not browseable.", ) + # Zero-component case: the requested path IS an allowlist root + # (e.g. a legacy-registered "/" or a Windows drive root). + if is_denied_system_path(str(current)): + raise HTTPException( + status_code = 403, + detail = "System directories are not browseable.", + ) if not current.is_dir(): raise HTTPException( status_code = 400, @@ -382,9 +429,13 @@ def browse_folders_response( """ from hub.storage.scan_folders import list_scan_folders + # Probe removable-media and Windows drive roots once; the allowlist and + # chips reuse the result so a disconnected mapped drive isn't scanned twice. + media_roots = linux_run_media_mount_roots() + drive_roots = windows_drive_roots() # Build the allowlist once -- the sandbox check and suggestion chips share # it so chips are always navigable. - allowed_roots = _build_browse_allowlist() + allowed_roots = _build_browse_allowlist(media_roots, drive_roots) try: target = _resolve_browse_target(path, allowed_roots) @@ -440,6 +491,15 @@ def browse_folders_response( # descending into them is refused and registration rejects them. if contains_sensitive_path_component(name): continue + # Same for denied system dirs (C:\Windows, /etc, ...): descent 403s, + # so don't render them as clickable rows. Resolve first so a + # symlink/junction into a denied dir is hidden too, not just a literal name. + try: + resolved_child = os.path.realpath(str(child)) + except (OSError, ValueError): + resolved_child = str(child) + if is_denied_system_path(resolved_child): + continue entries.append( BrowseEntry( name = name, @@ -487,13 +547,22 @@ def browse_folders_response( return if resolved in seen_sug: return + # Drop a denied system dir (e.g. a stale scan-folder row) so it never + # becomes a chip that 403s on click. Drive roots stay: only their + # system subdirectories are denied, not the root itself. + if is_denied_system_path(resolved): + return if _safe_is_dir(resolved): seen_sug.add(resolved) suggestions.append(resolved) # Home first as the safe fallback. _add_sug(Path.home()) - for p in linux_run_media_mount_roots(): + # Reuse the roots probed for the allowlist above (no second drive scan). + for p in media_roots: + _add_sug(p) + # Windows drive roots so the user can hop between C:, D:, E: ... + for p in drive_roots: _add_sug(p) # The HF cache root in use (honors HF_HOME / HF_HUB_CACHE), then the default. try: diff --git a/studio/backend/hub/storage/scan_folders.py b/studio/backend/hub/storage/scan_folders.py index fdb15c7c3c..81623d2759 100644 --- a/studio/backend/hub/storage/scan_folders.py +++ b/studio/backend/hub/storage/scan_folders.py @@ -16,7 +16,7 @@ from datetime import datetime, timezone from storage.studio_db import get_connection from hub.utils.paths import normalize_path -from utils.paths.external_media import is_linux_run_media_path +from utils.paths.external_media import is_linux_run_media_path, is_local_filesystem_root from utils.paths.sensitive import ( contains_sensitive_path_component as _shared_contains_sensitive_path_component, ) @@ -52,6 +52,25 @@ def _denied_path_prefixes() -> list[str]: return [] +def is_denied_system_path(path: str) -> bool: + """True if *path* is, or descends from, a denied system directory. + + Mirrors the denylist add_scan_folder() enforces at registration so the + browser refuses /etc, /proc, C:\\Windows, etc. even when the allowlist holds + a broad root (a Windows drive root C:\\ or a legacy-registered / root). The + /run carve-out keeps Linux removable-media mounts browseable. Expects an + already-resolved (realpath) path so symlinks cannot escape into a denied subtree. + """ + is_win = platform.system() == "Windows" + check = os.path.normcase(path) if is_win else path + for prefix in _denied_path_prefixes(): + if check == prefix or check.startswith(prefix + os.sep): + if prefix == "/run" and is_linux_run_media_path(check): + continue + return True + return False + + def _contains_sensitive_path_component(path: str) -> bool: return _shared_contains_sensitive_path_component(path) @@ -108,8 +127,9 @@ def add_scan_folder(path: str) -> dict: raise ValueError("Path must be a directory, not a file") if not os.access(normalized, os.R_OK | os.X_OK): raise ValueError("Path is not readable") - if os.path.dirname(normalized) == normalized: - # Registering a filesystem root would expose denied system dirs via browse. + if is_local_filesystem_root(normalized): + # A local fs root ("/", "C:\\") would expose denied system dirs via browse; + # a UNC share root (\\server\share) has none under it and stays registerable. raise ValueError("The filesystem root cannot be registered") if _contains_sensitive_path_component(normalized): raise ValueError("Credential or configuration directories are not allowed") diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py index 44701c0b64..2c33e09b2b 100644 --- a/studio/backend/hub/tests/test_model_services.py +++ b/studio/backend/hub/tests/test_model_services.py @@ -36,6 +36,20 @@ from hub.utils import ( from hub.workers import hf_download +@pytest.fixture(autouse = True) +def _denylist_inert(monkeypatch): + # The browse tests here exercise allowlist containment, symlink safety and + # the sensitive-name filter, not the system-directory denylist (which has + # its own suite in tests/test_browse_denylist.py). On macOS tmp_path + # resolves under /private/var, a denied prefix, so _resolve_browse_target + # would 403 the fixture dirs before that logic runs. Keep the denylist inert + # so these assertions hold on every platform. folder_browser binds + # is_denied_system_path at import, so patch it on that module, not on + # scan_folders. The "rejects" cases still 403 via the allowlist/sensitive + # checks, and the non-browse tests never call it. + monkeypatch.setattr(folder_browser, "is_denied_system_path", lambda _p: False) + + def _repo(repo_id: str, files: list[SimpleNamespace], repo_path: Path): return SimpleNamespace( repo_id = repo_id, @@ -228,7 +242,8 @@ def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path): home = tmp_path / "home" (home / ".ssh").mkdir(parents = True) (home / "models").mkdir() - monkeypatch.setattr(folder_browser, "_build_browse_allowlist", lambda: [home]) + # Accept and ignore the optional (media_roots, drive_roots) args the caller now passes. + monkeypatch.setattr(folder_browser, "_build_browse_allowlist", lambda *_a, **_k: [home]) response = folder_browser.browse_folders_response(str(home), show_hidden = True) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index c23ab1d428..b8526c75e7 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -1188,7 +1188,9 @@ def _looks_like_model_dir(directory: Path) -> bool: return False -def _build_browse_allowlist() -> list[Path]: +def _build_browse_allowlist( + media_roots: Optional[list[Path]] = None, drive_roots: Optional[list[Path]] = None +) -> list[Path]: """Return the root directories the folder browser may walk. The same list seeds the sidebar suggestion chips, so chip targets are @@ -1196,13 +1198,20 @@ def _build_browse_allowlist() -> list[Path]: outputs/exports/studio root, registered scan folders, and well-known local-LLM dirs (LM Studio, Ollama, ``~/models``); each added only if it resolves to a real directory. + + *media_roots* / *drive_roots* let the caller pass already-probed + removable-media and Windows drive roots so they aren't scanned again (a + disconnected mapped drive can make each probe slow); probed here when ``None``. """ from utils.paths import ( hf_default_cache_dir, legacy_hf_cache_dir, well_known_model_dirs, ) - from utils.paths.external_media import linux_run_media_mount_roots + from utils.paths.external_media import ( + linux_run_media_mount_roots, + windows_drive_roots, + ) from storage.studio_db import list_scan_folders candidates: list[Path] = [] @@ -1218,7 +1227,13 @@ def _build_browse_allowlist() -> list[Path]: candidates.append(resolved) _add(Path.home()) - for p in linux_run_media_mount_roots(): + if media_roots is None: + media_roots = linux_run_media_mount_roots() + if drive_roots is None: + drive_roots = windows_drive_roots() + for p in media_roots: + _add(p) + for p in drive_roots: _add(p) _add(_resolve_hf_cache_dir()) try: @@ -1269,19 +1284,43 @@ def _build_browse_allowlist() -> list[Path]: def _is_path_inside_allowlist(target: Path, allowed_roots: list[Path]) -> bool: """True if *target* equals or descends from any allowed root. - Uses ``os.path.realpath`` so symlinks can't escape the sandbox. + Uses ``os.path.realpath`` (symlinks can't escape the sandbox) and + ``os.path.commonpath`` for a component-wise containment test, so a string + prefix like ``/home/u`` never matches a sibling ``/home/user2`` while a + drive root ``D:\\`` still contains ``D:\\models``. A Windows drive root + authorizes its descendants, but a bare POSIX root ``/`` must NOT, else one + ``/`` allowlist entry would authorize every absolute path. ``normcase`` keeps + the drive-letter comparison case-insensitive, matching the hub browser. """ try: - target_real = os.path.realpath(str(target)) + target_real = os.path.normcase(os.path.realpath(str(target))) except OSError: return False for root in allowed_roots: try: - root_real = os.path.realpath(str(root)) + root_real = os.path.normcase(os.path.realpath(str(root))) except OSError: continue - if target_real == root_real or target_real.startswith(root_real + os.sep): + if target_real == root_real: return True + drive, tail = os.path.splitdrive(root_real) + if os.path.dirname(root_real) == root_real and not drive: + # Bare POSIX filesystem root ("/"): equality above is the only + # match; do not let it authorize arbitrary descendants. + continue + if drive.startswith(("\\\\", "//")) and not tail: + # Bare UNC share root (\\server\share): os.path.commonpath raises + # "can't mix absolute and relative" on it, so authorize its + # descendants with a boundary-safe prefix test (normcase applied). + if target_real.startswith(root_real.rstrip("\\/") + os.sep): + return True + continue + try: + if os.path.commonpath([target_real, root_real]) == root_real: + return True + except ValueError: + # Different drives / mixed absolute-relative: not contained. + continue return False @@ -1339,7 +1378,10 @@ def _match_browse_child(current: Path, name: str) -> Optional[Path]: def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Path: """Resolve a requested browse path by walking from trusted allowlist roots.""" - from storage.studio_db import contains_sensitive_path_component + from storage.studio_db import ( + contains_sensitive_path_component, + is_denied_system_path, + ) requested_path = _normalize_browse_request_path(path) resolved_roots: list[Path] = [] @@ -1396,6 +1438,11 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa status_code = 403, detail = "Credential or configuration directories are not browseable.", ) + if is_denied_system_path(str(resolved_child)): + raise HTTPException( + status_code = 403, + detail = "System directories are not browseable.", + ) current = resolved_child if contains_sensitive_path_component(str(current)): @@ -1403,6 +1450,13 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa status_code = 403, detail = "Credential or configuration directories are not browseable.", ) + # Zero-component case: the requested path IS an allowlist root + # (e.g. a legacy-registered "/" or a Windows drive root). + if is_denied_system_path(str(current)): + raise HTTPException( + status_code = 403, + detail = "System directories are not browseable.", + ) if not current.is_dir(): raise HTTPException( status_code = 400, @@ -1420,8 +1474,12 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa ) +# Sync (def, not async) so FastAPI runs the blocking filesystem I/O (drive +# probes, iterdir, realpath) in the threadpool: a disconnected mapped drive can +# make the probe wait out its timeout, which on the event loop would stall every +# other request. Matches the hub browse endpoint. @router.get("/browse-folders", response_model = BrowseFoldersResponse) -async def browse_folders( +def browse_folders( path: Optional[str] = Query( None, description = ( @@ -1450,11 +1508,22 @@ async def browse_folders( then hidden (if ``show_hidden=true``). """ from utils.paths import hf_default_cache_dir, well_known_model_dirs - from utils.paths.external_media import linux_run_media_mount_roots - from storage.studio_db import contains_sensitive_path_component, list_scan_folders + from utils.paths.external_media import ( + linux_run_media_mount_roots, + windows_drive_roots, + ) + from storage.studio_db import ( + contains_sensitive_path_component, + is_denied_system_path, + list_scan_folders, + ) + # Probe removable-media and Windows drive roots once; the allowlist and + # chips reuse the result so a disconnected mapped drive isn't scanned twice. + media_roots = linux_run_media_mount_roots() + drive_roots = windows_drive_roots() # Build once; the sandbox check and suggestion chips share it. - allowed_roots = _build_browse_allowlist() + allowed_roots = _build_browse_allowlist(media_roots, drive_roots) try: target = _resolve_browse_target(path, allowed_roots) @@ -1506,6 +1575,15 @@ async def browse_folders( continue if contains_sensitive_path_component(name): continue + # Hide denied system dirs (C:\Windows, /etc, ...) so they don't + # render as clickable rows that then 403 on descent. Resolve first + # so a symlink/junction into a denied dir is hidden too, not just a literal name. + try: + resolved_child = os.path.realpath(str(child)) + except (OSError, ValueError): + resolved_child = str(child) + if is_denied_system_path(resolved_child): + continue entries.append( BrowseEntry( name = name, @@ -1553,13 +1631,22 @@ async def browse_folders( return if resolved in seen_sug: return + # Drop a denied system dir (e.g. a stale scan-folder row) so it never + # becomes a chip that 403s on click. Drive roots stay: only their + # system subdirectories are denied, not the root itself. + if is_denied_system_path(resolved): + return if _safe_is_dir(resolved): seen_sug.add(resolved) suggestions.append(resolved) # Home first -- the safe fallback when everything else is cold. _add_sug(Path.home()) - for p in linux_run_media_mount_roots(): + # Reuse the roots probed for the allowlist above (no second drive scan). + for p in media_roots: + _add_sug(p) + # Windows drive roots so the user can hop between C:, D:, E: ... + for p in drive_roots: _add_sug(p) # The HF cache root the process is actually using. try: diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 87aa50ee26..4e0c711b69 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -27,7 +27,7 @@ from utils.paths import ( project_workspaces_root, studio_db_path, ) -from utils.paths.external_media import is_linux_run_media_path +from utils.paths.external_media import is_linux_run_media_path, is_local_filesystem_root from utils.paths.sensitive import ( contains_sensitive_path_component as _shared_contains_sensitive_path_component, ) @@ -69,6 +69,25 @@ def _denied_path_prefixes() -> list[str]: return [] +def is_denied_system_path(path: str) -> bool: + """True if *path* is, or descends from, a denied system directory. + + Mirrors the denylist add_scan_folder() enforces at registration so the + browser refuses /etc, /proc, C:\\Windows, etc. even when the allowlist holds + a broad root (a Windows drive root C:\\ or a legacy-registered / root). The + /run carve-out keeps Linux removable-media mounts browseable. Expects an + already-resolved (realpath) path so symlinks cannot escape into a denied subtree. + """ + is_win = platform.system() == "Windows" + check = os.path.normcase(path) if is_win else path + for prefix in _denied_path_prefixes(): + if check == prefix or check.startswith(prefix + os.sep): + if prefix == "/run" and is_linux_run_media_path(check): + continue + return True + return False + + def _contains_sensitive_path_component(path: str) -> bool: return _shared_contains_sensitive_path_component(path) @@ -931,6 +950,12 @@ def add_scan_folder(path: str) -> dict: raise ValueError("Path must be a directory, not a file") if not os.access(normalized, os.R_OK | os.X_OK): raise ValueError("Path is not readable") + # Reject a local filesystem root ("/", or a bare Windows drive root "C:\\"): + # registering one seeds the browse allowlist with a root above denied system + # dirs. A UNC share root (\\server\share) has none under it and was + # registerable before this guard, so it stays allowed. Mirrors scan_folders.py. + if is_local_filesystem_root(normalized): + raise ValueError("The filesystem root cannot be registered") if _contains_sensitive_path_component(normalized): raise ValueError("Credential or configuration directories are not allowed") diff --git a/studio/backend/tests/test_browse_denylist.py b/studio/backend/tests/test_browse_denylist.py new file mode 100644 index 0000000000..e21dc5c5a8 --- /dev/null +++ b/studio/backend/tests/test_browse_denylist.py @@ -0,0 +1,371 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""System-directory denylist enforcement for the folder browser. + +Once the allowlist can hold a whole Windows drive root (C:\\) or a legacy / +root, the browse endpoints must re-apply the ``_denied_path_prefixes()`` policy +``add_scan_folder`` enforces, so /etc, /proc, C:\\Windows, C:\\Program Files stay +unbrowseable even under an allowlisted root. Windows/macOS branches run on this +POSIX host by AST-extracting the pure helper with ``ntpath`` / a mocked ``platform``. +""" + +from __future__ import annotations + +import ast +import ntpath +import os +import posixpath +from pathlib import Path +from types import SimpleNamespace +from typing import Optional + +import pytest + +from hub.storage import scan_folders +from storage import studio_db +from utils.paths.external_media import is_local_filesystem_root + + +_BACKEND_ROOT = Path(__file__).resolve().parent.parent + + +class _HTTPException(Exception): + def __init__(self, status_code: int, detail: str): + super().__init__(detail) + self.status_code = status_code + self.detail = detail + + +def _extract_is_denied_windows(): + """is_denied_system_path (+ _denied_path_prefixes) from studio_db.py under Windows semantics (ntpath) on a POSIX host.""" + src = (_BACKEND_ROOT / "storage" / "studio_db.py").read_text(encoding = "utf-8") + tree = ast.parse(src) + funcs = [ + n + for n in tree.body + if isinstance(n, ast.FunctionDef) + and n.name in {"_denied_path_prefixes", "is_denied_system_path"} + ] + module = ast.Module(body = funcs, type_ignores = []) + ast.fix_missing_locations(module) + + win_os = SimpleNamespace( + sep = "\\", + environ = { + "SystemRoot": r"C:\Windows", + "ProgramFiles": r"C:\Program Files", + "ProgramFiles(x86)": r"C:\Program Files (x86)", + }, + path = SimpleNamespace(normcase = ntpath.normcase), + ) + ns = { + "os": win_os, + "platform": SimpleNamespace(system = lambda: "Windows"), + # /run has no Windows analog, so the carve-out is never reached. + "is_linux_run_media_path": lambda _p: False, + } + exec(compile(module, "", "exec"), ns) + return ns["is_denied_system_path"] + + +# is_denied_system_path -- Linux (real helper, this host) +@pytest.mark.parametrize( + "path", + [ + "/etc", + "/etc/ssl/private", + "/proc", + "/proc/1", + "/sys", + "/dev", + "/boot", + "/run", + "/run/systemd/private", + "/run/media", + "/run/media/dspofu", + ], +) +def test_is_denied_system_path_linux_denies_system_dirs(monkeypatch, path): + monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux") + assert studio_db.is_denied_system_path(path) is True + + +@pytest.mark.parametrize( + "path", + ["/run/media/dspofu/nvmeB", "/run/media/dspofu/nvmeB/models"], +) +def test_is_denied_system_path_linux_allows_run_media_mounts(monkeypatch, path): + # The /run/media// carve-out keeps removable media browseable. + monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux") + assert studio_db.is_denied_system_path(path) is False + + +@pytest.mark.parametrize( + "path", + ["/etc-backup", "/etcetera", "/home/u/models", "/mnt/data", "/devices", "/", "/opt/models"], +) +def test_is_denied_system_path_linux_allows_non_system(monkeypatch, path): + monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux") + assert studio_db.is_denied_system_path(path) is False + + +def test_legacy_and_hub_denylist_agree(monkeypatch): + monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux") + monkeypatch.setattr(scan_folders.platform, "system", lambda: "Linux") + for p in ["/etc", "/proc/1", "/home/u", "/boot", "/opt/x"]: + assert studio_db.is_denied_system_path(p) == scan_folders.is_denied_system_path(p) + + +# is_denied_system_path -- Windows (ntpath-backed), case-insensitive + collisions +@pytest.mark.parametrize( + "path", + [ + r"C:\Windows", + r"C:\Windows\System32", + r"c:\windows", + r"C:\WINDOWS\Temp", + r"C:\Program Files", + r"C:\Program Files\x", + r"C:\Program Files (x86)\y", + r"c:\program files", + ], +) +def test_is_denied_system_path_windows_denies_system_dirs(path): + is_denied = _extract_is_denied_windows() + assert is_denied(path) is True + + +@pytest.mark.parametrize( + "path", + [ + r"C:\Models", + r"D:\models", + r"C:\WindowsApps", + r"C:\ProgramData", + r"C:\Program Files Extra", + r"E:\gguf", + r"C:\Users\me\models", + ], +) +def test_is_denied_system_path_windows_allows_non_system(path): + is_denied = _extract_is_denied_windows() + assert is_denied(path) is False + + +# _resolve_browse_target -- real-FS integration (legacy browser) +def _extract_resolver(): + """Extract the legacy browse resolver; its inline imports use the real storage.studio_db policy.""" + src = (_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8") + tree = ast.parse(src) + names = { + "_is_path_inside_allowlist", + "_normalize_browse_request_path", + "_browse_relative_parts", + "_match_browse_child", + "_resolve_browse_target", + } + funcs = [n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name in names] + module = ast.Module(body = funcs, type_ignores = []) + ast.fix_missing_locations(module) + ns = { + "os": os, + "Path": Path, + "Optional": Optional, + "HTTPException": _HTTPException, + "logger": SimpleNamespace(warning = lambda *a, **k: None, debug = lambda *a, **k: None), + } + exec(compile(module, "", "exec"), ns) + return ns["_resolve_browse_target"] + + +def test_resolve_browse_target_blocks_etc_via_root(): + # Registering "/" must not make /etc browsable (Codex #3 regression guard). + resolve = _extract_resolver() + with pytest.raises(_HTTPException) as exc: + resolve("/etc", [Path("/")]) + assert exc.value.status_code == 403 + + +def test_resolve_browse_target_blocks_stale_denied_root(tmp_path, monkeypatch): + # A stale scan-folder row pointing at a denied dir is refused by the + # browse-time denylist even though it is its own allowlist root. A tmp-based + # denied prefix (+ Linux compare) keeps the assertion OS-agnostic: on macOS + # tmp lives under the already-denied /private/var, masking the message. + denied = (tmp_path / "sysfake").resolve() + denied.mkdir() + monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux") + monkeypatch.setattr(studio_db, "_denied_path_prefixes", lambda: [str(denied)]) + resolve = _extract_resolver() + with pytest.raises(_HTTPException) as exc: + resolve(str(denied), [denied]) + assert exc.value.status_code == 403 + assert "System directories" in exc.value.detail + + +def test_resolve_browse_target_allows_root_itself(): + resolve = _extract_resolver() + assert resolve("/", [Path("/")]) == Path("/") + + +def test_resolve_browse_target_allows_legit_nested_dir(tmp_path, monkeypatch): + # Force the Linux denylist so the macOS temp location (under the denied + # /private/var) doesn't reject the tmp fixture; a normal nested dir must not be over-blocked. + monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux") + resolve = _extract_resolver() + base = tmp_path / "allowed" + sub = base / "models" / "gguf" + sub.mkdir(parents = True) + assert resolve(str(sub), [base]) == sub.resolve() + + +def test_resolve_browse_target_symlink_escape_blocked(tmp_path): + resolve = _extract_resolver() + base = tmp_path / "allowed" + base.mkdir() + link = base / "escape" + try: + link.symlink_to("/etc", target_is_directory = True) + except OSError: + pytest.skip("symlinks unsupported on this host") + with pytest.raises(_HTTPException) as exc: + resolve(str(link), [base]) + assert exc.value.status_code == 403 + + +# _is_path_inside_allowlist -- bare POSIX root parity (legacy == hub) +def _extract_is_inside(rel_parts, *, os_module = os): + """Extract a standalone _is_path_inside_allowlist (os/Path only) so both browsers' copies compare without importing their heavy modules.""" + src = _BACKEND_ROOT.joinpath(*rel_parts).read_text(encoding = "utf-8") + tree = ast.parse(src) + funcs = [ + n + for n in tree.body + if isinstance(n, ast.FunctionDef) and n.name == "_is_path_inside_allowlist" + ] + module = ast.Module(body = funcs, type_ignores = []) + ast.fix_missing_locations(module) + ns = {"os": os_module, "Path": Path} + exec(compile(module, f"", "exec"), ns) + return ns["_is_path_inside_allowlist"] + + +# ntpath semantics with a no-FS realpath, so UNC containment can be driven on a +# POSIX CI (the real realpath cannot resolve \\server\share off Windows). +_WIN_OS = SimpleNamespace( + sep = ntpath.sep, + path = SimpleNamespace( + realpath = lambda p: ntpath.normpath(str(p)), + normcase = ntpath.normcase, + splitdrive = ntpath.splitdrive, + dirname = ntpath.dirname, + commonpath = ntpath.commonpath, + ), +) + + +def test_legacy_and_hub_allowlist_agree_on_posix_root(): + # A bare "/" allowlist entry must authorize only "/" itself in BOTH + # browsers, never descend into /var, /root, /home (which the denylist does + # not cover). Guards the hub browser against authorizing every absolute path. + legacy = _extract_is_inside(["routes", "models.py"]) + hub = _extract_is_inside(["hub", "services", "models", "folder_browser.py"]) + roots = [Path("/")] + for tgt in ["/var", "/root", "/home", "/usr", "/opt", "/etc"]: + assert legacy(Path(tgt), roots) is False + assert hub(Path(tgt), roots) is False + # "/" itself stays browseable; only its descendants are withheld. + assert legacy(Path("/"), roots) is True + assert hub(Path("/"), roots) is True + + +def test_hub_allowlist_authorizes_normal_nested_dir(tmp_path): + # The bare-root special case must not over-block a normal allowlist root's descendants. + hub = _extract_is_inside(["hub", "services", "models", "folder_browser.py"]) + base = tmp_path / "allowed" + sub = base / "models" / "gguf" + sub.mkdir(parents = True) + assert hub(sub, [base]) is True + assert hub(base, [base]) is True + + +# add_scan_folder -- filesystem-root rejection parity (legacy == hub) +def test_legacy_add_scan_folder_rejects_filesystem_root(monkeypatch): + monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux") + with pytest.raises(ValueError, match = "filesystem root"): + studio_db.add_scan_folder("/") + + +def test_hub_add_scan_folder_rejects_filesystem_root(monkeypatch): + monkeypatch.setattr(scan_folders.platform, "system", lambda: "Linux") + with pytest.raises(ValueError, match = "filesystem root"): + scan_folders.add_scan_folder("/") + + +# is_local_filesystem_root: reject "/" and "C:\\" (roots above denied system dirs), +# but NOT a UNC share root -- registering \\server\share was allowed before this +# guard and has no system dirs under it. _pathmod drives Windows semantics on POSIX CI. +@pytest.mark.parametrize( + "path, pathmod, expected", + [ + # Local filesystem roots -> rejected (True). + ("/", posixpath, True), + ("C:\\", ntpath, True), + ("c:\\", ntpath, True), + ("D:\\", ntpath, True), + # UNC share roots -> NOT a local root, stay registerable (False). + (r"\\server\share", ntpath, False), + (r"\\nas\models", ntpath, False), + ("//server/share", ntpath, False), + # Device / extended-length volume roots -> still local roots (rejected), + # so neither \\?\C:\ nor a drive-letter-less \\?\Volume{GUID}\ can slip + # past the guard as if it were a share root. + (r"\\?\C:" + "\\", ntpath, True), + (r"\\.\C:" + "\\", ntpath, True), + (r"\\?\C:", ntpath, True), + (r"\\.\C:", ntpath, True), + (r"\\?\Volume{2f8e6d31-0000-0000-0000-100000000000}" + "\\", ntpath, True), + (r"\\.\Volume{2f8e6d31-0000-0000-0000-100000000000}", ntpath, True), + # Device-namespace UNC share root -> stays registerable (False). + (r"\\?\UNC\server\share", ntpath, False), + # Non-root paths (incl. deep device / extended-length) -> not a root (False). + ("C:\\Models", ntpath, False), + (r"\\server\share\models", ntpath, False), + (r"\\?\C:\Users\me\models", ntpath, False), + (r"\\?\Volume{2f8e6d31-0000-0000-0000-100000000000}\models", ntpath, False), + ("/home/user", posixpath, False), + ], +) +def test_is_local_filesystem_root(path, pathmod, expected): + assert is_local_filesystem_root(path, _pathmod = pathmod) is expected + + +def test_both_guards_use_the_shared_local_root_helper(): + # Register-root parity: both browsers reject the same roots via one helper, so a + # UNC-share exemption can never drift between the legacy and hub code paths. + legacy_src = (_BACKEND_ROOT / "storage" / "studio_db.py").read_text(encoding = "utf-8") + hub_src = (_BACKEND_ROOT / "hub" / "storage" / "scan_folders.py").read_text(encoding = "utf-8") + assert "is_local_filesystem_root(normalized)" in legacy_src + assert "is_local_filesystem_root(normalized)" in hub_src + + +# A registered UNC share root must authorize its own descendants in both browsers. +# os.path.commonpath raises "can't mix absolute and relative" on a bare +# \\server\share, so containment falls back to a boundary-safe prefix test; without +# it, registering a UNC share (now allowed) would 403 every folder under it. +@pytest.mark.parametrize( + "rel_parts", + [ + ["routes", "models.py"], + ["hub", "services", "models", "folder_browser.py"], + ], +) +def test_unc_share_root_authorizes_its_descendants(rel_parts): + is_inside = _extract_is_inside(rel_parts, os_module = _WIN_OS) + root = [Path(r"\\server\share")] + assert is_inside(Path(r"\\server\share"), root) is True # the root itself + assert is_inside(Path(r"\\server\share\models"), root) is True # direct child + assert is_inside(Path(r"\\server\share\a\b\c"), root) is True # deep descendant + assert is_inside(Path(r"\\SERVER\SHARE\Models"), root) is True # case-insensitive + assert is_inside(Path(r"\\server\share2\models"), root) is False # sibling share + assert is_inside(Path(r"C:\models"), root) is False # different volume diff --git a/studio/backend/tests/test_browse_folders_route.py b/studio/backend/tests/test_browse_folders_route.py index 3a607e6b10..970057f1e2 100644 --- a/studio/backend/tests/test_browse_folders_route.py +++ b/studio/backend/tests/test_browse_folders_route.py @@ -22,6 +22,18 @@ if "structlog" not in sys.modules: ) import routes.models as models_route +import storage.studio_db as studio_db + + +@pytest.fixture(autouse = True) +def _denylist_inert(monkeypatch): + # These tests exercise allowlist containment and the file-vs-directory guard, + # not the system-directory denylist (which has its own suite in + # test_browse_denylist.py). On macOS tmp_path resolves under /private/var, a + # denied prefix, so _resolve_browse_target would 403 the fixture dirs before + # the containment logic runs. Keep the denylist inert here so these + # assertions hold on every platform. + monkeypatch.setattr(studio_db, "is_denied_system_path", lambda _p: False) def test_resolve_browse_target_returns_allowed_directory(tmp_path): diff --git a/studio/backend/tests/test_linux_external_media_paths.py b/studio/backend/tests/test_linux_external_media_paths.py index c763248f6a..b735bd1132 100644 --- a/studio/backend/tests/test_linux_external_media_paths.py +++ b/studio/backend/tests/test_linux_external_media_paths.py @@ -252,10 +252,17 @@ def test_legacy_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tm outputs_root = lambda: tmp_path / "missing-outputs", exports_root = lambda: tmp_path / "missing-exports", ) - fake_external_media = SimpleNamespace(linux_run_media_mount_roots = lambda: [media_root]) + fake_external_media = SimpleNamespace( + linux_run_media_mount_roots = lambda: [media_root], + windows_drive_roots = lambda: [], + ) fake_studio_db = SimpleNamespace( list_scan_folders = lambda: [], contains_sensitive_path_component = studio_db.contains_sensitive_path_component, + # The media root is a legitimate mount, not denied; the .ssh 403 below + # comes from the credential check. A False stub keeps this OS-independent + # (on macOS tmp_path lives under the denied /private/var). + is_denied_system_path = lambda _p: False, ) monkeypatch.setitem(sys.modules, "utils.paths", fake_paths) monkeypatch.setitem(sys.modules, "utils.paths.external_media", fake_external_media) diff --git a/studio/backend/tests/test_windows_external_drive_paths.py b/studio/backend/tests/test_windows_external_drive_paths.py new file mode 100644 index 0000000000..9686d45c9f --- /dev/null +++ b/studio/backend/tests/test_windows_external_drive_paths.py @@ -0,0 +1,354 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import ast +import os +import sys +from pathlib import Path +from types import SimpleNamespace +from typing import Optional + +from utils.paths import external_media + + +_BACKEND_ROOT = Path(__file__).resolve().parent.parent + + +class _HTTPException(Exception): + def __init__(self, status_code: int, detail: str): + super().__init__(detail) + self.status_code = status_code + self.detail = detail + + +def _extract_routes_function(name: str, ns_extra: Optional[dict] = None) -> dict: + """Exec one top-level function from routes/models.py without importing the module (which pulls in FastAPI).""" + tree = ast.parse((_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8")) + fn = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == name) + module = ast.Module(body = [fn], type_ignores = []) + ast.fix_missing_locations(module) + ns = {"os": os, "Path": Path, "Optional": Optional} + if ns_extra: + ns.update(ns_extra) + exec(compile(module, "", "exec"), ns) + return ns + + +def _stub_windows(monkeypatch, existing_drives): + """Simulate Windows exposing only *existing_drives* (e.g. {"C", "D"}) as readable roots, independent of the host FS. + + Overriding _active_windows_drive_bitmask keeps it deterministic even on a + real Windows host, where live GetLogicalDrives would return the actual layout.""" + monkeypatch.setattr(external_media.platform, "system", lambda: "Windows") + mask = sum(1 << (ord(d.upper()) - ord("A")) for d in existing_drives) + monkeypatch.setattr(external_media, "_active_windows_drive_bitmask", lambda: mask) + present = {f"{d.upper()}:\\" for d in existing_drives} + monkeypatch.setattr(external_media.os.path, "isdir", lambda p: str(p) in present) + monkeypatch.setattr(external_media.os, "access", lambda p, _mode: str(p) in present) + + +def test_windows_drive_roots_empty_off_windows(monkeypatch): + # Regression guard: the helper is a no-op on Linux/macOS so it can't change the allowlist on the platforms CI runs on. + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + assert external_media.windows_drive_roots() == [] + monkeypatch.setattr(external_media.platform, "system", lambda: "Darwin") + assert external_media.windows_drive_roots() == [] + + +def test_windows_drive_roots_lists_readable_drives(monkeypatch): + _stub_windows(monkeypatch, {"C", "D", "E"}) + + roots = external_media.windows_drive_roots(drive_letters = "CDEF") + + # F is absent, so it is skipped; the rest are exposed in order. + assert roots == [Path("C:\\"), Path("D:\\"), Path("E:\\")] + + +def test_windows_drive_roots_skips_absent_and_unreadable(monkeypatch): + _stub_windows(monkeypatch, {"C"}) + + roots = external_media.windows_drive_roots(drive_letters = "CDE") + + assert roots == [Path("C:\\")] + + +def test_windows_drive_roots_ignores_bad_letters_and_dedupes(monkeypatch): + _stub_windows(monkeypatch, {"C", "D"}) + + roots = external_media.windows_drive_roots( + drive_letters = ["c:", "C", "D", "1", "AB", "", " d "], + ) + + assert roots == [Path("C:\\"), Path("D:\\")] + + +def test_readable_dir_within_times_out(monkeypatch): + # A probe that outlives the timeout is reported not-readable, so a hung + # (disconnected mapped network) drive is skipped instead of blocking. + import time + + monkeypatch.setattr(external_media.os.path, "isdir", lambda p: time.sleep(5) or True) + monkeypatch.setattr(external_media.os, "access", lambda p, _mode: True) + start = time.monotonic() + ok = external_media._readable_dir_within("Z:\\", timeout = 0.2) + elapsed = time.monotonic() - start + assert ok is False + assert elapsed < 3.0 # returned on the timeout, did not wait out the 5s stall + + +def test_readable_dir_within_reports_fast_probe(monkeypatch): + monkeypatch.setattr(external_media.os.path, "isdir", lambda p: True) + monkeypatch.setattr(external_media.os, "access", lambda p, _mode: True) + assert external_media._readable_dir_within("C:\\", timeout = 2.0) is True + + +def test_windows_drive_roots_skips_hung_drive(monkeypatch): + # A disconnected mapped drive stays set in the bitmask and its os.path.isdir + # stalls; it must be skipped without stalling enumeration. C answers, D hangs, + # so only C is listed, bounded by the per-drive timeout, not the stall. + import time + + monkeypatch.setattr(external_media.platform, "system", lambda: "Windows") + monkeypatch.setattr( + external_media, + "_active_windows_drive_bitmask", + lambda: sum(1 << (ord(d) - ord("A")) for d in "CD"), + ) + monkeypatch.setattr(external_media, "_DRIVE_PROBE_TIMEOUT_S", 0.2) + + def _isdir(p): + if str(p) == "D:\\": + time.sleep(5) # simulate the reconnect stall + return True + return str(p) == "C:\\" + + monkeypatch.setattr(external_media.os.path, "isdir", _isdir) + monkeypatch.setattr(external_media.os, "access", lambda p, _mode: True) + + start = time.monotonic() + roots = external_media.windows_drive_roots(drive_letters = "CD") + elapsed = time.monotonic() - start + + assert roots == [Path("C:\\")] + assert elapsed < 3.0 # bounded by the per-drive timeout, not the 5s stall + + +def test_windows_drive_roots_probes_hung_drives_in_parallel(monkeypatch): + # Several disconnected mapped drives must add ~one timeout total, not one + # per drive: C answers fast, D/E/F stall. The concurrent probe stays bounded + # by a single deadline where serial probing would cost ~4x the timeout. + import time + + monkeypatch.setattr(external_media.platform, "system", lambda: "Windows") + monkeypatch.setattr( + external_media, + "_active_windows_drive_bitmask", + lambda: sum(1 << (ord(d) - ord("A")) for d in "CDEF"), + ) + timeout = 0.2 + monkeypatch.setattr(external_media, "_DRIVE_PROBE_TIMEOUT_S", timeout) + + def _isdir(p): + if str(p) == "C:\\": + return True + time.sleep(5) # every other drive simulates a reconnect stall + return True + + monkeypatch.setattr(external_media.os.path, "isdir", _isdir) + monkeypatch.setattr(external_media.os, "access", lambda p, _mode: True) + + start = time.monotonic() + roots = external_media.windows_drive_roots(drive_letters = "CDEF") + elapsed = time.monotonic() - start + + assert roots == [Path("C:\\")] + # 3 stalled drives probed in parallel finish within ~1 timeout, well under the ~3*timeout a serial probe would take. + assert elapsed < 3 * timeout + + +def test_browse_allowlist_includes_windows_drive_roots(monkeypatch, tmp_path): + # End-to-end wiring: windows_drive_roots() output flows into the browse + # allowlist built by routes/models.py, mirroring the Linux media-mounts test. + tree = ast.parse((_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8")) + function_names = { + "_build_browse_allowlist", + "_browse_relative_parts", + "_is_path_inside_allowlist", + "_match_browse_child", + "_normalize_browse_request_path", + "_resolve_browse_target", + } + functions = [ + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name in function_names + ] + module = ast.Module(body = functions, type_ignores = []) + ast.fix_missing_locations(module) + + home = tmp_path / "home" + drive_root = tmp_path / "D_drive" + model_dir = drive_root / "modelsAI" / "gguf" + home.mkdir() + model_dir.mkdir(parents = True) + + fake_paths = SimpleNamespace( + hf_default_cache_dir = lambda: tmp_path / "missing-default-hf", + legacy_hf_cache_dir = lambda: tmp_path / "missing-legacy-hf", + well_known_model_dirs = lambda: [], + studio_root = lambda: tmp_path / "missing-studio", + outputs_root = lambda: tmp_path / "missing-outputs", + exports_root = lambda: tmp_path / "missing-exports", + ) + fake_external_media = SimpleNamespace( + linux_run_media_mount_roots = lambda: [], + windows_drive_roots = lambda: [drive_root], + ) + fake_studio_db = SimpleNamespace( + list_scan_folders = lambda: [], + contains_sensitive_path_component = lambda _p: False, + # The simulated D:\ root maps to a tmp_path dir, not a denied system path. + is_denied_system_path = lambda _p: False, + ) + monkeypatch.setitem(sys.modules, "utils.paths", fake_paths) + monkeypatch.setitem(sys.modules, "utils.paths.external_media", fake_external_media) + monkeypatch.setitem(sys.modules, "storage.studio_db", fake_studio_db) + + ns = { + "HTTPException": _HTTPException, + "os": os, + "Path": Path, + "Optional": Optional, + "_safe_is_dir": lambda p: Path(p).is_dir(), + "_resolve_hf_cache_dir": lambda: tmp_path / "missing-hf", + "logger": SimpleNamespace(debug = lambda *_args, **_kwargs: None), + } + exec(compile(module, "", "exec"), ns) + + allowlist = ns["_build_browse_allowlist"]() + + # The simulated Windows drive root is now browsable, and a model dir on it resolves. + assert drive_root.resolve() in allowlist + assert ns["_resolve_browse_target"](str(model_dir), allowlist) == model_dir.resolve() + + +def test_build_browse_allowlist_reuses_passed_roots(monkeypatch, tmp_path): + # Double-probe fix: a browse request probes the drive/media roots once and + # passes them in, so _build_browse_allowlist must NOT scan + # windows_drive_roots() again (a disconnected drive would double the stall). + tree = ast.parse((_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8")) + functions = [ + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "_build_browse_allowlist" + ] + module = ast.Module(body = functions, type_ignores = []) + ast.fix_missing_locations(module) + + drive_root = tmp_path / "D_drive" + drive_root.mkdir() + + calls = {"drive": 0, "media": 0} + + def _drive_roots(): + calls["drive"] += 1 + return [drive_root] + + def _media_roots(): + calls["media"] += 1 + return [] + + fake_paths = SimpleNamespace( + hf_default_cache_dir = lambda: tmp_path / "missing-default-hf", + legacy_hf_cache_dir = lambda: tmp_path / "missing-legacy-hf", + well_known_model_dirs = lambda: [], + studio_root = lambda: tmp_path / "missing-studio", + outputs_root = lambda: tmp_path / "missing-outputs", + exports_root = lambda: tmp_path / "missing-exports", + ) + fake_external_media = SimpleNamespace( + linux_run_media_mount_roots = _media_roots, + windows_drive_roots = _drive_roots, + ) + fake_studio_db = SimpleNamespace(list_scan_folders = lambda: []) + monkeypatch.setitem(sys.modules, "utils.paths", fake_paths) + monkeypatch.setitem(sys.modules, "utils.paths.external_media", fake_external_media) + monkeypatch.setitem(sys.modules, "storage.studio_db", fake_studio_db) + + ns = { + "os": os, + "Path": Path, + "Optional": Optional, + "_safe_is_dir": lambda p: Path(p).is_dir(), + "_resolve_hf_cache_dir": lambda: tmp_path / "missing-hf", + "logger": SimpleNamespace(debug = lambda *_args, **_kwargs: None), + } + exec(compile(module, "", "exec"), ns) + build = ns["_build_browse_allowlist"] + + # Roots passed in -> neither helper is probed, but the roots still flow in. + allowlist = build([], [drive_root]) + assert calls == {"drive": 0, "media": 0} + assert drive_root.resolve() in allowlist + + # No args -> each helper is probed exactly once. + build() + assert calls == {"drive": 1, "media": 1} + + +def test_is_path_inside_allowlist_real_descendants_and_siblings(tmp_path): + # Component-wise containment (commonpath): a genuine descendant is allowed, + # but a sibling sharing only a string prefix ("models_root_evil" vs + # "models_root") is not, which the old startswith check could miss. + ns = _extract_routes_function("_is_path_inside_allowlist") + root = tmp_path / "models_root" + child = root / "gguf" / "qwen" + sibling = tmp_path / "models_root_evil" + child.mkdir(parents = True) + sibling.mkdir() + + is_inside = ns["_is_path_inside_allowlist"] + assert is_inside(root, [root]) is True # the root itself + assert is_inside(child, [root]) is True # a genuine descendant + assert is_inside(sibling, [root]) is False # prefix-collision sibling + + +def test_is_path_inside_allowlist_posix_root_does_not_authorize_descendants(monkeypatch): + # Regression for the reported POSIX "/" unlock: a bare filesystem root may + # match itself but must NOT authorize arbitrary descendants such as /etc. + ns = _extract_routes_function("_is_path_inside_allowlist") + monkeypatch.setattr(os.path, "realpath", lambda p: str(p)) # keep "/" intact + + is_inside = ns["_is_path_inside_allowlist"] + assert is_inside("/", ["/"]) is True # the root itself + assert is_inside("/etc", ["/"]) is False # not a licensed descendant + assert is_inside("/root/models", ["/"]) is False + + +def test_is_path_inside_allowlist_windows_drive_root_descendants(): + # Exercise the Windows drive-root branch on a POSIX host by backing os.path + # with ntpath and an identity realpath (the simulated drives don't exist + # here). A drive root authorizes its descendants; a different drive does not. + import ntpath + + win_os = SimpleNamespace( + sep = "\\", + path = SimpleNamespace( + normcase = ntpath.normcase, + realpath = lambda p: str(p), + splitdrive = ntpath.splitdrive, + dirname = ntpath.dirname, + commonpath = ntpath.commonpath, + ), + ) + ns = _extract_routes_function("_is_path_inside_allowlist", {"os": win_os}) + is_inside = ns["_is_path_inside_allowlist"] + + assert is_inside("D:\\", ["D:\\"]) is True # drive root itself + assert is_inside("D:\\models", ["D:\\"]) is True # descendant on the drive + assert is_inside("D:\\models\\gguf", ["D:\\"]) is True # deeper descendant + assert is_inside("d:\\models", ["D:\\"]) is True # case-insensitive drive letter + assert is_inside("C:\\Users", ["D:\\"]) is False # different drive + assert is_inside("D:\\models", ["E:\\"]) is False diff --git a/studio/backend/utils/paths/external_media.py b/studio/backend/utils/paths/external_media.py index 1f1754664f..0ea0477cc7 100644 --- a/studio/backend/utils/paths/external_media.py +++ b/studio/backend/utils/paths/external_media.py @@ -8,6 +8,10 @@ from __future__ import annotations import getpass import os import platform +import string +import threading +import time +from collections.abc import Iterable from pathlib import Path from utils.paths.sensitive import ( @@ -16,6 +20,33 @@ from utils.paths.sensitive import ( ) +def is_local_filesystem_root(path: str, *, _pathmod = os.path) -> bool: + """True for a bare local filesystem root -- POSIX ``/``, a drive root ``C:\\``, + or a device-namespace volume root like ``\\\\?\\C:\\`` or + ``\\\\?\\Volume{GUID}\\`` -- which sit above denied system dirs, but NOT a UNC + share root (``\\\\server\\share`` or its ``\\\\?\\UNC\\...`` form), which has + none under it and was registerable before this guard. ``splitdrive`` is empty + on POSIX servers, so this reduces to the plain ``dirname == self`` test there. + ``_pathmod`` lets tests drive ``ntpath`` semantics on a POSIX CI. + """ + # Resolve the Windows device / extended-length namespace, where \\?\C:\, + # \\.\C:\ and \\?\Volume{GUID}\ are all bare LOCAL volume roots (rejected) + # while only \\?\UNC\server\share is a UNC share (handled like \\server\share). + if path[:4].lower() in ("\\\\?\\", "\\\\.\\"): + rest = path[4:] + if rest[:4].lower() == "unc\\": + path = "\\\\" + rest[4:] + else: + # A device volume root is just the volume specifier (C:, Volume{GUID}) + # with no further component; a deeper path is an ordinary folder. + core = rest.rstrip("\\/") + return "\\" not in core and "/" not in core + if _pathmod.dirname(path) != path: + return False + drive, _ = _pathmod.splitdrive(path) + return drive[:2] not in ("\\\\", "//") + + def _is_linux_media_mount_path(path: str, media_root: Path | str) -> bool: normalized = os.path.normpath(os.path.realpath(os.path.expanduser(path))) root = os.path.normpath(os.path.realpath(os.path.expanduser(str(media_root)))) @@ -98,3 +129,101 @@ def linux_run_media_mount_roots( seen.add(key) roots.append(resolved) return roots + + +def _active_windows_drive_bitmask() -> int: + """Active-logical-drive bitmask from ``GetLogicalDrives`` (bit 0 = ``A:``), or ``0`` when unavailable. + + A fast non-blocking call that lets :func:`windows_drive_roots` skip the + ``os.path.isdir`` probe on unmapped letters. A disconnected network mapping + stays set here, so it does not guard the reconnect stall on its own; + :func:`windows_drive_roots` bounds each surviving probe too. Returns ``0`` + (probe every letter) when ctypes/``windll`` is missing. + """ + try: + import ctypes + return int(ctypes.windll.kernel32.GetLogicalDrives()) + except Exception: # noqa: BLE001 -- best-effort; fall back to probing all letters + return 0 + + +# A disconnected mapped drive stays set in the GetLogicalDrives bitmask, so +# ``os.path.isdir`` on it can block for tens of seconds. Bound each drive probe +# so one stale mapping cannot stall a whole folder-browser request. +_DRIVE_PROBE_TIMEOUT_S = 2.0 + + +def _readable_dirs_within(paths: Iterable[str], timeout: float) -> set[str]: + """Which of *paths* are readable directories, probed concurrently under one overall *timeout* (seconds). + + Each path is checked (``os.path.isdir`` + ``os.access(R_OK)``) in its own + daemon thread and the call waits at most *timeout* total, not per path, so N + stalled network drives add ~timeout instead of N*timeout. A path not + answering ``True`` by the deadline is treated as unreadable. The daemon + threads are never joined past the deadline, so a stuck OS call cannot delay + interpreter exit or block the caller (``os.path.isdir`` releases the GIL). + """ + paths = list(paths) # fixed input we can iterate twice; one probe per path + results: dict[str, bool] = {} + + def _probe(path: str) -> None: + try: + results[path] = os.path.isdir(path) and os.access(path, os.R_OK) + except OSError: + results[path] = False + + threads: list[threading.Thread] = [] + for path in paths: + thread = threading.Thread(target = _probe, args = (path,), daemon = True) + thread.start() + threads.append(thread) + + deadline = time.monotonic() + timeout + for thread in threads: + thread.join(max(0.0, deadline - time.monotonic())) + + # Iterate the fixed input, not results.items(): a probe that timed out is + # still alive and may insert its key here, which would raise "dictionary + # changed size during iteration". results.get() is an atomic read. + return {path for path in paths if results.get(path)} + + +def _readable_dir_within(path: str, timeout: float) -> bool: + """``os.path.isdir(path) and os.access(path, R_OK)``, bounded by *timeout* seconds; single-path wrapper over :func:`_readable_dirs_within`.""" + return path in _readable_dirs_within((path,), timeout) + + +def windows_drive_roots(drive_letters: Iterable[str] = string.ascii_uppercase) -> list[Path]: + """Readable logical drive roots (``C:\\``, ``D:\\`` ...) for the folder browser; the Windows analog of :func:`linux_run_media_mount_roots`. + + Without it the allowlist and chips only reach the home drive, so a user + cannot navigate from ``C:`` to ``D:``/``E:``. ``GetLogicalDrives`` drops + unmapped letters; the rest are probed concurrently under a single timeout + and kept only if readable in time. A disconnected mapped drive stays active + in the bitmask and its ``os.path.isdir`` can hang for tens of seconds, so + parallel probing bounds the added delay at ~one timeout rather than one per + drive. Returns ``[]`` off Windows. + """ + if platform.system() != "Windows": + return [] + + active_mask = _active_windows_drive_bitmask() + candidates: list[str] = [] + seen: set[str] = set() + for letter in drive_letters: + letter = letter.strip().rstrip(":").upper() + if len(letter) != 1 or letter not in string.ascii_uppercase: + continue + if active_mask and not active_mask & (1 << (ord(letter) - ord("A"))): + continue + root_text = f"{letter}:\\" + key = os.path.normcase(root_text) + if key in seen: + continue + seen.add(key) + candidates.append(root_text) + + # Bounded concurrent probe: an active bitmask bit can still be a + # disconnected mapping whose os.path.isdir blocks, so probe all at once. + readable = _readable_dirs_within(candidates, _DRIVE_PROBE_TIMEOUT_S) + return [Path(root_text) for root_text in candidates if root_text in readable] From 67339b15fd893cc5fe35d57ace6a950ccbd4ea55 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 15 Jul 2026 02:45:05 -0700 Subject: [PATCH 06/17] Studio CI: make tool-calling SSE probes resilient to transport stalls (#7137) * Studio CI: make tool-calling SSE probes resilient to transport stalls * Studio CI: bound tool-probe SSE stalls to the job timeout and stop accepting partial tool events * Studio CI: surface HTTP errors from the seed loop and bound the Mac best-effort probes * Studio CI: keep completed tool results on a stall and cap seed reads by the remaining budget --- .github/workflows/studio-inference-smoke.yml | 137 ++++++++++++++---- .../workflows/studio-mac-inference-smoke.yml | 82 ++++++++--- .../studio-windows-inference-smoke.yml | 91 +++++++++--- 3 files changed, 232 insertions(+), 78 deletions(-) diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml index f540c11da4..cf8c021e38 100644 --- a/.github/workflows/studio-inference-smoke.yml +++ b/.github/workflows/studio-inference-smoke.yml @@ -485,7 +485,7 @@ jobs: print(f"[retry] {path}: {exc!r}", flush = True) time.sleep(15) - def post_sse(path, body, *, timeout = 600): + def post_sse(path, body, *, timeout = 600, retries = 1, complete_on = None): """POST a streaming request and accumulate the assistant text deltas. The server-side agentic loop ALWAYS returns SSE regardless of the request's `stream` field, so any @@ -501,6 +501,22 @@ jobs: invocation markers / tool output, since `delta.content` alone is not evidence that the tool path executed. + + A shared CI runner can stall the stream transport (the + connection opening, or a mid-stream read) even when Studio + is healthy, so retry a stall once with a fresh request + capped at 300s. A stall means the stream did NOT complete, + so partial events are normally NOT returned (an early + tool_start with no tool_end is not proof the tool loop + finished). The one exception is `complete_on`: an optional + predicate over the events collected so far -- when a stall + happens after it is already satisfied (the tool ran and + produced its result before the trailing read timed out), + those events are returned rather than discarded, so the + stall-after-answer case still counts. HTTP status errors + surface immediately; a stall that yields no completed result + across all attempts re-raises so the caller can rotate to + the next seed. """ body = {**body, "stream": True} data = json.dumps(body).encode() @@ -513,26 +529,45 @@ jobs: "Content-Type": "application/json", }, ) - parts = [] - events = [] - with urllib.request.urlopen(req, timeout = timeout) as resp: - for raw in resp: - line = raw.decode().strip() - if not line.startswith("data: "): - continue - payload = line[6:] - if payload == "[DONE]": - break - events.append(payload) - try: - chunk = json.loads(payload) - except json.JSONDecodeError: - continue - for choice in chunk.get("choices", []): - delta = choice.get("delta", {}) or {} - if delta.get("content"): - parts.append(delta["content"]) - return "".join(parts), events + for attempt in range(retries + 1): + parts = [] + events = [] + t = timeout if attempt == 0 else min(timeout, 300) + try: + with urllib.request.urlopen(req, timeout = t) as resp: + for raw in resp: + line = raw.decode().strip() + if not line.startswith("data: "): + continue + payload = line[6:] + if payload == "[DONE]": + break + events.append(payload) + try: + chunk = json.loads(payload) + except json.JSONDecodeError: + continue + for choice in chunk.get("choices", []): + delta = choice.get("delta", {}) or {} + if delta.get("content"): + parts.append(delta["content"]) + return "".join(parts), events + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + # A stall after the tool already produced its result is + # the case this probe exists to tolerate: keep those + # events. But a stall with only an early tool_start (no + # completed output) is not proof the tool loop finished, + # so it must not pass -- retry once, then raise so + # _run_tool_probe rotates to the next seed. + if complete_on is not None and complete_on(events): + print(f"[retry-sse] {path}: {exc!r}; keeping {len(events)} completed events", flush = True) + return "".join(parts), events + if attempt == retries: + raise + print(f"[retry-sse] {path}: {exc!r}", flush = True) + time.sleep(15) _STUDIO_TOOL_TYPES = { "tool_start", "tool_end", "tool_use", "tool_result", @@ -669,17 +704,54 @@ jobs: """ attempts_log = [] best = None + # Cap the wall-clock spent rotating through stalled seeds so a + # persistent no-data wedge fails fast (clean assertion) instead + # of being killed by the job's timeout-minutes. A healthy or + # merely degenerate round answers in seconds, so all seeds still + # run in the normal case; only stalls consume the budget. + probe_deadline = time.monotonic() + 300 for attempt_i in range(max_attempts): + # Cap each read by the budget still remaining (not just a flat + # 180s) and skip an attempt too small to finish, so the whole + # rotation stays within ~300s -- two probes then fit the job's + # timeout-minutes even if every seed stalls. + remaining = int(probe_deadline - time.monotonic()) + if attempt_i and remaining < 30: + print(f"[tools] {label}: seed-rotation budget spent after {attempt_i} attempts", flush = True) + break attempt_seed = SEED + attempt_i - content, events = post_sse("/v1/chat/completions", { - "messages": [{"role": "user", "content": prompt}], - "enable_tools": True, - "enabled_tools": enabled, - "session_id": f"{session}-att{attempt_i}", - "temperature": TOOL_PROBE_TEMP, - "seed": attempt_seed, - "max_tokens": 600, - }) + try: + # Bounded per-attempt timeout, no inner retry -- the seed + # loop IS the retry, so a stall raises quickly and rotates + # rather than spending post_sse's full 600+300s. complete_on + # keeps a stall that already produced the tool result (only + # the trailing read timed out) instead of discarding it. + content, events = post_sse("/v1/chat/completions", { + "messages": [{"role": "user", "content": prompt}], + "enable_tools": True, + "enabled_tools": enabled, + "session_id": f"{session}-att{attempt_i}", + "temperature": TOOL_PROBE_TEMP, + "seed": attempt_seed, + "max_tokens": 600, + }, timeout = min(180, remaining), retries = 0, + complete_on = lambda ev: _tool_invoked(ev) and _tool_output_contains(ev, *needles)) + except urllib.error.HTTPError: + # HTTPError subclasses URLError, so re-raise a real 4xx/5xx + # here instead of letting the transport-stall handler below + # swallow it and rotate seeds -- an endpoint status failure + # must surface, not be masked as missing tool evidence. + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + # A transport stall that outlived post_sse's own retry: + # log it as a failed attempt and rotate to the next seed + # rather than sinking the whole probe on one bad stream. + attempts_log.append({ + "attempt": attempt_i, "seed": attempt_seed, + "transport_error": repr(exc), + }) + print(f"[tools] retry {label} attempt {attempt_i}: transport {exc!r}", flush = True) + continue invoked = _tool_invoked(events) produced = _tool_output_contains(events, *needles) attempts_log.append({ @@ -740,6 +812,9 @@ jobs: # enough that requiring a tool_call marker would create # red-herring failures from infra rather than from Studio. try: + # Best-effort and bounded: a single 180s attempt keeps a stall + # from eating the job's timeout-minutes (it already WARNs, so a + # retry buys nothing). content, events = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}], "enable_tools": True, @@ -748,7 +823,7 @@ jobs: "temperature": 0.0, "seed": SEED, "max_tokens": 400, - }) + }, timeout = 180, retries = 0) print( f"[tools] PASS web_search stream ({len(content)} chars in content, " f"{len(events)} raw events)" diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml index 03c0a8580d..d3d765aa84 100644 --- a/.github/workflows/studio-mac-inference-smoke.yml +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -471,11 +471,22 @@ jobs: print(f"[retry] {path}: {exc!r}", flush = True) time.sleep(15) - def post_sse(path, body, *, timeout = 600): + def post_sse(path, body, *, timeout = 600, retries = 1, soft = False): """POST a streaming request and accumulate the assistant text deltas. The server-side agentic loop ALWAYS returns SSE regardless of the request's `stream` field, so any - call with enable_tools=true must use this helper.""" + call with enable_tools=true must use this helper. + + A shared CI runner can stall the stream transport (the + connection opening, or a mid-stream read) even when Studio + is healthy, so harden the read three ways: retry a stall + once with a fresh request capped at 300s; return any text + already streamed before a stall (a stall on the trailing + tokens, after the answer arrived, still counts); and when + every attempt yields nothing, a hard call re-raises while a + soft call (the best-effort server-side tool probes) returns + None so the caller can WARN instead of sinking the whole + job. HTTP status errors always surface immediately.""" body = {**body, "stream": True} data = json.dumps(body).encode() req = urllib.request.Request( @@ -487,24 +498,43 @@ jobs: "Content-Type": "application/json", }, ) - parts = [] - with urllib.request.urlopen(req, timeout = timeout) as resp: - for raw in resp: - line = raw.decode().strip() - if not line.startswith("data: "): - continue - payload = line[6:] - if payload == "[DONE]": - break - try: - chunk = json.loads(payload) - except json.JSONDecodeError: - continue - for choice in chunk.get("choices", []): - delta = choice.get("delta", {}) or {} - if delta.get("content"): - parts.append(delta["content"]) - return "".join(parts) + for attempt in range(retries + 1): + parts = [] + t = timeout if attempt == 0 else min(timeout, 300) + try: + with urllib.request.urlopen(req, timeout = t) as resp: + for raw in resp: + line = raw.decode().strip() + if not line.startswith("data: "): + continue + payload = line[6:] + if payload == "[DONE]": + break + try: + chunk = json.loads(payload) + except json.JSONDecodeError: + continue + for choice in chunk.get("choices", []): + delta = choice.get("delta", {}) or {} + if delta.get("content"): + parts.append(delta["content"]) + return "".join(parts) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + # Text already streamed is a valid signal -- keep it + # rather than re-running a heavy generation. + if parts: + joined = "".join(parts) + print(f"[retry-sse] {path}: {exc!r}; keeping {len(joined)} partial chars", flush = True) + return joined + if attempt == retries: + if soft: + print(f"[tools] WARN {path}: SSE transport stalled with no data ({exc!r}) -- non-blocking", flush = True) + return None + raise + print(f"[retry-sse] {path}: {exc!r}", flush = True) + time.sleep(15) # ── 1. Standard OpenAI function calling ────────────────────── weather_tool = { @@ -575,6 +605,10 @@ jobs: # macos-14 free runner is ~10 tok/s on Qwen3.5-2B Q4_K_XL; # cap max_tokens tightly so each SSE round stays under ~30s # even when the model stalls in a degenerate output state. + # retries=0 on the best-effort probes: this job's 25-minute cap + # allows a 10-minute model load, so a no-data stall must be a + # single 180s attempt (not 180+15+180s) to leave room for the + # thinking checks. A soft/best-effort probe only WARNs anyway. content = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "What is 123 * 456? Use the python tool to compute it and tell me the number."}], "enable_tools": True, @@ -583,8 +617,10 @@ jobs: "temperature": TEMP, "seed": SEED, "max_tokens": 128, - }, timeout = 180) - if "56088" in content or "56,088" in content: + }, timeout = 180, retries = 0, soft = True) + if content is None: + print("[tools] WARN python tool: SSE transport stalled after retries -- non-blocking") + elif "56088" in content or "56,088" in content: print(f"[tools] PASS python tool ({len(content)} chars, found 56088)") else: # Empty stream is a known Mac-quant degeneracy too; log @@ -616,7 +652,7 @@ jobs: "temperature": TEMP, "seed": SEED, "max_tokens": 96, - }, timeout = 180) + }, timeout = 180, retries = 0) print(f"[tools] PASS web_search stream ({len(content)} chars)") except Exception as exc: print(f"[tools] WARN web_search probe failed (non-blocking): {exc}") diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 0453c9212a..233292f7a3 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -677,7 +677,22 @@ jobs: print(f"[retry] {path}: {exc!r}", flush = True) time.sleep(15) - def post_sse(path, body, *, timeout = 600): + def post_sse(path, body, *, timeout = 600, retries = 1, soft = False): + # The server-side agentic loop always answers over SSE. A + # shared CI runner can stall the stream transport (the + # connection opening, or a mid-stream read) even when Studio + # is healthy, so harden the read three ways: + # * retry a transport stall once with a fresh request, + # capped at 300s (a healthy server answers a retry + # quickly, a wedged one never does); + # * return any text already streamed before a stall, so a + # stall on the trailing tokens -- after the answer + # arrived -- still counts; + # * when every attempt yields nothing, a hard call + # re-raises while a soft call (the best-effort + # server-side tool probes) returns None so the caller + # can WARN instead of sinking the whole job. + # HTTP status errors always surface immediately. body = {**body, "stream": True} data = json.dumps(body).encode() req = urllib.request.Request( @@ -689,24 +704,43 @@ jobs: "Content-Type": "application/json", }, ) - parts = [] - with urllib.request.urlopen(req, timeout = timeout) as resp: - for raw in resp: - line = raw.decode().strip() - if not line.startswith("data: "): - continue - payload = line[6:] - if payload == "[DONE]": - break - try: - chunk = json.loads(payload) - except json.JSONDecodeError: - continue - for choice in chunk.get("choices", []): - delta = choice.get("delta", {}) or {} - if delta.get("content"): - parts.append(delta["content"]) - return "".join(parts) + for attempt in range(retries + 1): + parts = [] + t = timeout if attempt == 0 else min(timeout, 300) + try: + with urllib.request.urlopen(req, timeout = t) as resp: + for raw in resp: + line = raw.decode().strip() + if not line.startswith("data: "): + continue + payload = line[6:] + if payload == "[DONE]": + break + try: + chunk = json.loads(payload) + except json.JSONDecodeError: + continue + for choice in chunk.get("choices", []): + delta = choice.get("delta", {}) or {} + if delta.get("content"): + parts.append(delta["content"]) + return "".join(parts) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + # Text already streamed is a valid signal -- keep it + # rather than re-running a heavy generation. + if parts: + joined = "".join(parts) + print(f"[retry-sse] {path}: {exc!r}; keeping {len(joined)} partial chars", flush = True) + return joined + if attempt == retries: + if soft: + print(f"[tools] WARN {path}: SSE transport stalled with no data ({exc!r}) -- non-blocking", flush = True) + return None + raise + print(f"[retry-sse] {path}: {exc!r}", flush = True) + time.sleep(15) # ── 1. Standard OpenAI function calling ────────────────────── weather_tool = { @@ -749,6 +783,11 @@ jobs: ) # ── 2. Server-side python tool ─────────────────────────────── + # Bound each soft probe to a single 180s attempt (timeout=180, + # retries=0): this job runs two of them back-to-back under a + # 30-minute cap, so the default 600+15+300s per stall could hit + # the workflow timeout before the thinking checks run. A soft + # probe only WARNs anyway, so a retry buys nothing. content = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "What is 123 * 456? Use the python tool to compute it and tell me the number."}], "enable_tools": True, @@ -757,8 +796,10 @@ jobs: "temperature": TEMP, "seed": SEED, "max_tokens": 600, - }) - if "56088" in content or "56,088" in content: + }, timeout = 180, retries = 0, soft = True) + if content is None: + print("[tools] WARN python tool: SSE transport stalled after retries -- non-blocking") + elif "56088" in content or "56,088" in content: print(f"[tools] PASS python tool ({len(content)} chars, found 56088)") else: assert content, "python tool: SSE stream empty" @@ -780,8 +821,10 @@ jobs: "temperature": TEMP, "seed": SEED, "max_tokens": 600, - }) - if "hello-bash-tool" in content: + }, timeout = 180, retries = 0, soft = True) + if content is None: + print("[tools] WARN terminal tool: SSE transport stalled after retries -- non-blocking") + elif "hello-bash-tool" in content: print(f"[tools] PASS terminal tool ({len(content)} chars)") else: assert content, "terminal tool: SSE stream empty" @@ -802,7 +845,7 @@ jobs: "temperature": TEMP, "seed": SEED, "max_tokens": 400, - }) + }, timeout = 180, retries = 0) print(f"[tools] PASS web_search stream ({len(content)} chars)") except Exception as exc: print(f"[tools] WARN web_search probe failed (non-blocking): {exc}") From 4beb0a3a5f30d6a98d26c71d93c22067a2c418b5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 15 Jul 2026 02:46:55 -0700 Subject: [PATCH 07/17] Studio: force-terminate a stuck training stop after a grace period (#7099) * Studio: force-terminate a stuck training stop after a grace period A Stop-with-save only signals the worker and waits for it to save and exit; force_terminate() was reachable only from the /reset cancel path. On Windows + ROCm the worker saves the adapter fine but then wedges in post-save GPU/HIP teardown and never exits, so the run stays in "Stopping..." forever, is_training stays true, and /reset returns 409. Add a stop watchdog: when a stop is requested, a daemon escalates to force_terminate() a short grace after the worker's "complete" (save done), or after an absolute cap covering a hang during save. After escalation the parent state is finalized (is_training=False, "Training stopped.") even if the OS never reaps the wedged worker, so the UI leaves "Stopping..." and a new run can start. No behavior change on a clean quick exit. Grace and timeout are configurable via UNSLOTH_STUDIO_TRAINING_STOP_GRACE_S (15) and UNSLOTH_STUDIO_TRAINING_STOP_TIMEOUT_S (120). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: harden the training stop watchdog per review Address review feedback so a stop can never corrupt a checkpoint or leave the run stuck: - Never force-kill an in-progress save. The absolute cap is now a last-resort backstop: raise the save default to 600s and only kill past that long window; a not-yet-complete save is not treated as a hang. Cancels have nothing to save, so they keep a shorter 120s cap via UNSLOTH_STUDIO_TRAINING_CANCEL_TIMEOUT_S. The save vs cancel path is now explicit and the backstop logs a clear warning. - Always finalize even if force_terminate raises on a wedged child (try/finally), so the watchdog never dies leaving the run in "Stopping...". - Preserve output_dir when the watchdog finalizes so a saved checkpoint is still recorded in run history. - Track the watched process per watchdog: a new run always gets its own watcher, and a stale watchdog on an old proc no longer suppresses it. - Terminate only the captured proc; force_terminate revalidates under the lock that it is still the current worker, so it can never kill a fresh run. - Name the watchdog thread for debuggability. * studio: tighten training-stop watchdog comments Comment-only pass: collapse the watchdog docstrings and inline notes to fewer lines while keeping the rationale. No behavior change. * Studio: make the stop watchdog safe against concurrent runs and the pump Target-scope the escalation finalize so a stale watchdog can never clobber a run that replaced its worker: capture the watched proc and job id, and no-op the finalize (handle, progress, and DB) when a new run has already taken over. Honor a later cancel by tightening an in-flight save watchdog to the shorter cancel cap. Serialize the DB helpers on the lock so the watchdog and pump can no longer double-create, double-finalize, or corrupt the metric buffer when a force-terminate hands off to a still-finalizing pump. Add regression tests: finalize no-ops when superseded, finalize runs for its own worker, a later cancel tightens the cap, finalize is single-winner under concurrency, finalize honors expected_job_id, and concurrent flushes claim each metric exactly once. * Studio: close the remaining stop-watchdog vs start/pump races Guard the escalation finalize by the watched job id in addition to the proc: start_training sets current_job_id before it installs the new _proc, so a stale watchdog entering during that startup window still sees the old dead handle and was not caught by the proc-only guard. Capture the job id when the watchdog starts and require it to still match before touching state. Snapshot the run id and final progress under the finalize lock and thread them through the flush and finish_run calls, so a new run that starts between the finalize claim and the DB writes cannot be flushed or marked stopped under the old run's finalizer. Publish _db_run_created only after create_run commits, gated by a dedicated in-progress flag, so a concurrent finalize can no longer run finish_run against a not-yet-inserted row and leave the run stuck as running. Add regression tests for the startup-window job-id guard, run-id pinned flush, snapshot-based finalize across a new run, and create-not-published-before-insert. * Studio: finalize a force-stopped run by its captured id If a new run starts in the gap after the watchdog clears _proc and marks the backend idle, current_job_id changes, so the previous expected_job_id guard made the finalize skip and left the stopped run recorded as running. Capture the run id, metrics, and final progress under the lock (where current_job_id is still the watched run) and finalize by that captured id via _finish_stopped_run: finish_run is an idempotent UPDATE and insert_metrics_batch upserts, so a concurrent pump finalize of the same run is harmless and a newly started run is never touched. Add a test that the watched run is finalized by id with its buffered metrics, and update the escalation tests to assert finalize goes through _finish_stopped_run. * Studio: keep force-stop finalization retryable and unclaimed until the row exists Only claim _run_finalized in the escalation when the DB row already exists; if an early create failed and the pump is retrying it, claiming would make the pump's later finalize no-op and strand the row as running, so leave the finalize to that create-then-finalize path. On a DB error in _finish_stopped_run (e.g. a transient SQLite lock), unclaim the finalize and requeue the drained metrics when the run is still current, so the pump or a later retry can still record the run stopped instead of leaving history with an active run and lost metrics. A superseded run's state is never touched. Add tests: no claim before the row exists, requeue+unclaim on a DB error, and a superseded run left untouched on error. * Studio: tighten stop-watchdog comments Reduce the wording of the docstrings and inline comments added by this PR without dropping any of the concurrency invariants (dual proc/job-id supersession guard, finalize-by-captured-id, publish-after-commit, snapshot-under-lock, unclaim and requeue on error). Comments and docstrings only; no code change. * Studio: record the stopped run's DB state before dropping _proc A wedged worker still reports alive, so the pump never reaches its own finalize and bails on its _proc-is-None guard once the escalation drops the handle. So the watchdog is the sole finalizer: record the terminal DB state (create the row if a start-time create failed, then finish by captured id) BEFORE dropping _proc. While the handle is held is_training_active() stays true, so no new run can start and current_job_id stays the watched run for the write; _proc is dropped last, guarded on target_proc so a run that did replace the worker keeps its handle. _finish_stopped_run retries a transient DB error a few times (the pump can no longer retry once _proc is gone) and unclaims on final failure only when the run is still current. Add tests for create-then-finalize, retry-then-unclaim, and not dropping a new run's handle. * Studio: job-guard the DB create flags against a racing new run _ensure_db_run_created publishes backend-wide _db_run_created and _db_create_in_progress flags. When the watchdog creates a missing row for an escalated stop, the killed worker lets a new /start proceed mid-create, so the stale create could publish those flags against the new current_job_id, making the new run skip inserting its own row (metric/finalize then target a missing run). Publish the flags only when the captured job id is still current; the row is still created by id, and the new run owns/creates its own. Also reset _db_create_in_progress in start_training so a stale claim can't block a new run. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/training/training.py | 383 ++++++-- .../tests/test_training_stop_watchdog.py | 824 ++++++++++++++++++ 2 files changed, 1153 insertions(+), 54 deletions(-) create mode 100644 studio/backend/tests/test_training_stop_watchdog.py diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 2ddda19951..6b32ec873d 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -39,6 +39,27 @@ from utils.paths import outputs_root logger = get_logger(__name__) + +def _env_int(name: str, default: int) -> int: + try: + raw = (os.environ.get(name) or "").strip() + return int(raw) if raw else default + except ValueError: + return default + + +# Stop-watchdog escalation timeouts. Primary trigger: a short grace once "complete" +# (save done). Absolute cap is a backstop: long for save=True so a slow save is never +# killed mid-write, shorter for a cancel that has nothing to save. +_STOP_GRACE_S = _env_int("UNSLOTH_STUDIO_TRAINING_STOP_GRACE_S", 15) +_STOP_TIMEOUT_S = _env_int("UNSLOTH_STUDIO_TRAINING_STOP_TIMEOUT_S", 600) +_CANCEL_TIMEOUT_S = _env_int("UNSLOTH_STUDIO_TRAINING_CANCEL_TIMEOUT_S", 120) + +# Watchdog DB finalize: a few short retries so a transient SQLite lock doesn't lose the +# terminal state, since the watchdog is the sole finalizer once _proc is dropped. +_DB_FINALIZE_RETRIES = 3 +_DB_FINALIZE_RETRY_S = 0.5 + _pyplot = None _pyplot_failed = False @@ -741,6 +762,13 @@ class TrainingBackend: self._pump_running: bool = False self._lock = threading.Lock() + # Stop watchdog: after a stop is requested, escalates to force_terminate() + # if the worker does not exit on its own within a bounded time. The watched + # proc is tracked so a new run always gets its own watcher. + self._stop_watchdog: Optional[threading.Thread] = None + self._stop_watchdog_proc: Optional[mp.Process] = None + self._complete_seen = threading.Event() + # Progress state (updated by pump thread from subprocess events) self._progress = TrainingProgress() self._should_stop = False @@ -765,6 +793,7 @@ class TrainingBackend: self._metric_buffer: list[dict] = [] self._run_finalized: bool = False self._db_run_created: bool = False + self._db_create_in_progress: bool = False self._db_total_steps_set: bool = False self._db_config: Optional[dict] = None self._db_started_at: Optional[str] = None @@ -896,6 +925,7 @@ class TrainingBackend: self.current_job_id = job_id self._should_stop = False self._cancel_requested = False + self._complete_seen.clear() self._progress = TrainingProgress( is_training = True, status_message = "Initializing training..." ) @@ -911,6 +941,7 @@ class TrainingBackend: self._metric_buffer.clear() self._run_finalized = False self._db_run_created = False + self._db_create_in_progress = False # a stale watchdog create can't block this run self._db_total_steps_set = False self._db_config = _sanitize_db_config(config) self._db_started_at = datetime.now(timezone.utc).isoformat() @@ -953,15 +984,212 @@ class TrainingBackend: self._progress.status_message = ( "Stopping training and saving checkpoint..." if save else "Cancelling training..." ) + # Guarantee the run finalizes even if the worker wedges after saving. + self._start_stop_watchdog(cancel = not save) return True - def force_terminate(self) -> None: - """Force-kill the training subprocess so state can be reset immediately.""" + def _start_stop_watchdog(self, cancel: bool) -> None: + """Start a daemon that force-terminates the worker if a requested stop does not + exit on its own. No-op if no worker is alive or a live watchdog already watches + this proc (a stale watchdog on an old proc never blocks a new run's watcher).""" with self._lock: - if self._proc is not None and self._proc.is_alive(): - logger.info("Force-terminating training subprocess (pid=%s)", self._proc.pid) - self._proc.terminate() proc = self._proc + if proc is None or not proc.is_alive(): + return + if ( + self._stop_watchdog is not None + and self._stop_watchdog.is_alive() + and self._stop_watchdog_proc is proc + ): + return + watchdog = threading.Thread( + target = self._stop_watchdog_loop, + args = (proc, cancel, self.current_job_id), + name = f"stop-watchdog-{self.current_job_id or 'unknown'}", + daemon = True, + ) + self._stop_watchdog = watchdog + self._stop_watchdog_proc = proc + watchdog.start() + + def _stop_watchdog_loop( + self, + target_proc: "mp.Process", + cancel: bool, + watched_job_id: Optional[str] = None, + ) -> None: + """Escalate a stuck stop to force_terminate(): grace after "complete", else the + absolute backstop (see the module timeouts). No-ops on a clean exit; exits + silently if a new run replaces the worker.""" + started = time.monotonic() + complete_at: Optional[float] = None + reason = "" + while True: + with self._lock: + superseded = self._proc is not target_proc + # A later cancel has nothing to save, so tighten an in-flight save + # watchdog to the shorter cancel cap. + cancelling = cancel or self._cancel_requested + if superseded or not target_proc.is_alive(): + return + now = time.monotonic() + abs_timeout = _CANCEL_TIMEOUT_S if cancelling else _STOP_TIMEOUT_S + if complete_at is None and self._complete_seen.is_set(): + complete_at = now + if complete_at is not None and now - complete_at >= _STOP_GRACE_S: + reason = "worker still alive after save" + break + if now - started >= abs_timeout: + reason = "worker did not exit within the absolute timeout" + break + time.sleep(0.5) + + with self._lock: + superseded = self._proc is not target_proc + if superseded or not target_proc.is_alive(): + return + if complete_at is None: + # Backstop fired pre-completion: a save may still be in progress. + logger.warning( + "Stop watchdog: absolute timeout with no completion signal; " + "force-terminating a possibly-mid-save worker: %s", + reason, + ) + else: + logger.warning("Stop watchdog force-terminating stuck training worker: %s", reason) + # force_terminate can raise on a wedged child; finalize regardless. + try: + self.force_terminate(target_proc = target_proc) + except Exception: + logger.exception("Stop watchdog: force_terminate failed; finalizing anyway") + finally: + self._finalize_stopped_after_escalation( + target_proc = target_proc, watched_job_id = watched_job_id + ) + + def _finalize_stopped_after_escalation( + self, + target_proc: "Optional[mp.Process]" = None, + watched_job_id: Optional[str] = None, + ) -> None: + """Finalize parent state after a force-terminate so the UI leaves "Stopping..." + even if the worker is wedged in driver teardown; preserves output_dir so a saved + checkpoint is kept. No-ops if a new run already replaced the watched worker, so a + stale watchdog never marks a fresh run stopped or drops its handle. + + Supersession is checked on both the watched proc and job id: start_training sets + current_job_id before it installs the new _proc, so a stale watchdog entering that + startup window still sees the old (dead) handle and is caught by the job-id guard. + + The run's terminal DB state is recorded (create-if-needed + finish by captured id) + BEFORE _proc is dropped: a wedged worker still reports alive, so the pump never + reaches its own finalize and would bail on its _proc-is-None guard once the handle + is gone. While the handle is held is_training_active() stays true, so no new run can + start and current_job_id stays the watched run for the write. _proc is dropped last, + re-guarded on target_proc so a run that did replace the worker keeps its handle.""" + with self._lock: + if target_proc is not None and self._proc is not target_proc: + return # a new run replaced the worker; never touch its state + if watched_job_id is not None and self.current_job_id != watched_job_id: + return # a new run is already starting up; leave its state alone + run_id = self.current_job_id # == watched_job_id + self._progress.is_training = False + self._progress.status_message = "Training stopped." + # Create the row if a start-time create failed (no-op otherwise; skips when the pump + # is mid-create, in which case its create-then-finalize records the run instead). + self._ensure_db_run_created() + with self._lock: + claim = ( + bool(run_id) + and self.current_job_id == run_id + and self._db_run_created + and not self._run_finalized + ) + batch: list = [] + final_step = final_loss = duration = None + loss_history: list = [] + output_dir = self._output_dir + if claim: + self._run_finalized = True # claim this run's finalize + batch = list(self._metric_buffer) + del self._metric_buffer[: len(batch)] + final_step = self._progress.step + final_loss = self._progress.loss + if final_loss is not None and not math.isfinite(final_loss): + final_loss = None + duration = self._progress.elapsed_seconds + loss_history = list(self.loss_history) + if claim: + self._finish_stopped_run( + run_id, output_dir, batch, final_step, final_loss, duration, loss_history + ) + with self._lock: + if target_proc is None or self._proc is target_proc: + self._proc = None # drop only our handle, never a run that replaced it + + def _finish_stopped_run( + self, + run_id: str, + output_dir: Optional[str], + batch: list, + final_step: Optional[int], + final_loss: Optional[float], + duration: Optional[float], + loss_history: list, + ) -> None: + """Record a force-stopped run finished by its captured id, from state snapshotted + under the lock. insert_metrics_batch upserts and finish_run is an idempotent UPDATE, + so a concurrent pump finalize of the same run is harmless and a different current run + is never touched. The watchdog is the sole finalizer once _proc is dropped, so a + transient DB error (e.g. a SQLite lock) is retried a few times; on final failure the + finalize is unclaimed (only if the run is still current) so the row is not left + claimed-but-unfinalized.""" + for attempt in range(_DB_FINALIZE_RETRIES): + try: + from storage.studio_db import finish_run, insert_metrics_batch + from utils.downsample import downsample + + if batch: + insert_metrics_batch(run_id, batch) + sparkline = downsample(loss_history, 50) + finish_run( + id = run_id, + status = "stopped", + ended_at = datetime.now(timezone.utc).isoformat(), + final_step = final_step, + final_loss = final_loss, + duration_seconds = duration, + loss_sparkline = _json.dumps(sparkline), + output_dir = output_dir, + error_message = None, + ) + return + except Exception: + if attempt + 1 < _DB_FINALIZE_RETRIES: + time.sleep(_DB_FINALIZE_RETRY_S) + continue + logger.warning( + "Failed to finalize stopped run %s in DB after %d attempts", + run_id, + _DB_FINALIZE_RETRIES, + exc_info = True, + ) + with self._lock: + # Only if still current; a new run's finalize state is never touched. + if self.current_job_id == run_id: + self._run_finalized = False + + def force_terminate(self, target_proc: "Optional[mp.Process]" = None) -> None: + """Force-kill the training subprocess so state can be reset immediately. With + ``target_proc``, terminate only that handle and no-op if a new run has replaced + it, so the watchdog can never kill a fresh worker.""" + with self._lock: + proc = self._proc + if target_proc is not None and proc is not target_proc: + return # superseded by a new run; do not touch the new worker + if proc is not None and proc.is_alive(): + logger.info("Force-terminating training subprocess (pid=%s)", proc.pid) + proc.terminate() cancelled = self._cancel_requested output_dir = self._output_dir @@ -1468,6 +1696,8 @@ class TrainingBackend: "training cancelled", "training stopped", } + # Save is done by now; let the stop watchdog start its grace timer. + self._complete_seen.set() self._progress.is_training = False self._progress.is_completed = not stopped self._output_dir = event.get("output_dir") @@ -1532,90 +1762,135 @@ class TrainingBackend: self._finalize_run_in_db(**db_action_kwargs) def _ensure_db_run_created(self) -> None: - """Create the DB row if it doesn't exist yet. Called outside the lock.""" - if self._db_run_created or not self.current_job_id or not self._db_config: - return + """Create the DB row if it doesn't exist yet. An in-progress flag lets only one + caller create at a time, and ``_db_run_created`` is published only after + ``create_run`` commits, so a concurrent finalize never runs ``finish_run`` against a + not-yet-inserted row (a zero-row UPDATE that would leave the run stuck as running).""" + with self._lock: + if ( + self._db_run_created + or self._db_create_in_progress + or not self.current_job_id + or not self._db_config + ): + return + self._db_create_in_progress = True # only one caller creates + job_id = self.current_job_id + db_config = self._db_config + started_at = self._db_started_at or datetime.now(timezone.utc).isoformat() + total_steps = self._progress.total_steps or None + created = False try: from storage.studio_db import create_run dataset_name = ( - self._db_config.get("hf_dataset") - or next(iter(self._db_config.get("local_datasets") or []), None) - or _s3_dataset_name(self._db_config.get("s3_dataset")) + db_config.get("hf_dataset") + or next(iter(db_config.get("local_datasets") or []), None) + or _s3_dataset_name(db_config.get("s3_dataset")) or "unknown" ) create_run( - id = self.current_job_id, - model_name = self._db_config["model_name"], + id = job_id, + model_name = db_config["model_name"], dataset_name = dataset_name, - config_json = _json.dumps(self._db_config), - started_at = self._db_started_at or datetime.now(timezone.utc).isoformat(), - total_steps = self._progress.total_steps or None, + config_json = _json.dumps(db_config), + started_at = started_at, + total_steps = total_steps, ) - self._db_run_created = True + created = True except Exception: logger.warning("Failed to create DB run record for early failure", exc_info = True) + finally: + with self._lock: + # Publish the flags only if this is still the current run. A killed worker + # lets a new /start proceed mid-create, and these flags are backend-wide, so + # a stale create for the captured job must not satisfy the new run's DB state + # (the row was still created by id; the new run owns/creates its own row). + if self.current_job_id == job_id: + if created: + self._db_run_created = True # publish only after the insert commits + self._db_create_in_progress = False def _finalize_run_in_db( self, status: str, error_message: Optional[str] = None, output_dir: Optional[str] = None, + expected_job_id: Optional[str] = None, ) -> None: - """Flush remaining metrics and mark a run as finished in the DB.""" - if not self.current_job_id or not self._db_run_created or self._run_finalized: - return - self._flush_metrics_to_db() + """Flush remaining metrics and mark a run finished in the DB. Claims the finalize + under the lock so the watchdog and pump can't double-finalize, and no-ops when + ``expected_job_id`` no longer matches (a new run took over). The run id and final + progress are snapshotted under the lock and threaded through the flush/finish calls, + so a new run racing between this claim and the DB writes can't be flushed or marked + stopped under the old run's finalize.""" + with self._lock: + if expected_job_id is not None and self.current_job_id != expected_job_id: + return + if not self.current_job_id or not self._db_run_created or self._run_finalized: + return + self._run_finalized = True + run_id = self.current_job_id + final_step = self._progress.step + final_loss = self._progress.loss + if final_loss is not None and not math.isfinite(final_loss): + final_loss = None + duration = self._progress.elapsed_seconds + loss_history = list(self.loss_history) + self._flush_metrics_to_db(run_id = run_id) try: from storage.studio_db import finish_run from utils.downsample import downsample - sparkline = downsample(self.loss_history, 50) + sparkline = downsample(loss_history, 50) finish_run( - id = self.current_job_id, + id = run_id, status = status, ended_at = datetime.now(timezone.utc).isoformat(), - final_step = self._progress.step, - final_loss = self._progress.loss - if (self._progress.loss is not None and math.isfinite(self._progress.loss)) - else None, - duration_seconds = self._progress.elapsed_seconds, + final_step = final_step, + final_loss = final_loss, + duration_seconds = duration, loss_sparkline = _json.dumps(sparkline), output_dir = output_dir, error_message = error_message, ) - self._run_finalized = True except Exception: + with self._lock: + self._run_finalized = False # unclaim so a later flush can retry logger.warning("Failed to finalize run in DB (status=%s)", status, exc_info = True) - def _flush_metrics_to_db(self) -> None: - """Flush buffered metrics to the database and update live progress.""" - if not self._metric_buffer or not self.current_job_id or not self._db_run_created: - return - # Cap buffer to bound memory growth. - if len(self._metric_buffer) > 500: - logger.warning( - "Metric buffer exceeded 500 entries (%d) — trimming oldest", - len(self._metric_buffer), - ) - self._metric_buffer = self._metric_buffer[-500:] - # Snapshot before insert so metrics arriving during the write survive. - batch = list(self._metric_buffer) + def _flush_metrics_to_db(self, run_id: Optional[str] = None) -> None: + """Flush buffered metrics to the DB and update live progress. The target run id, + metric batch, and progress snapshot are all taken under the lock, so a concurrent + flush can't double-remove metrics and a racing new run can't redirect the write to + a different job. A finalizer passes ``run_id`` to pin the target to its captured run.""" + with self._lock: + target = run_id if run_id is not None else self.current_job_id + if not self._metric_buffer or not target or not self._db_run_created: + return + # Cap buffer to bound memory growth. + if len(self._metric_buffer) > 500: + logger.warning( + "Metric buffer exceeded 500 entries (%d) — trimming oldest", + len(self._metric_buffer), + ) + del self._metric_buffer[:-500] + # Claim the batch under the lock so a concurrent flush can't re-remove it. + batch = list(self._metric_buffer) + del self._metric_buffer[: len(batch)] + step = self._progress.step + loss = self._progress.loss + if loss is not None and not math.isfinite(loss): + loss = None + duration = self._progress.elapsed_seconds try: from storage.studio_db import insert_metrics_batch, update_run_progress - - insert_metrics_batch(self.current_job_id, batch) - del self._metric_buffer[: len(batch)] - update_run_progress( - id = self.current_job_id, - step = self._progress.step, - loss = self._progress.loss - if (self._progress.loss is not None and math.isfinite(self._progress.loss)) - else None, - duration_seconds = self._progress.elapsed_seconds, - ) + insert_metrics_batch(target, batch) + update_run_progress(id = target, step = step, loss = loss, duration_seconds = duration) except Exception: - # Leave buffer intact for retry on next flush + # Re-queue the claimed batch at the front so it retries on the next flush. + with self._lock: + self._metric_buffer[:0] = batch logger.warning("Failed to flush metrics to DB", exc_info = True) @staticmethod diff --git a/studio/backend/tests/test_training_stop_watchdog.py b/studio/backend/tests/test_training_stop_watchdog.py new file mode 100644 index 0000000000..457dfc8ea2 --- /dev/null +++ b/studio/backend/tests/test_training_stop_watchdog.py @@ -0,0 +1,824 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Stop-watchdog escalation for a stuck training stop. + +A save-stop signals the worker and waits for it to save and exit. On some platforms the +worker saves but then wedges in post-save GPU/driver teardown and never exits, leaving the +run stuck in "Stopping..." forever. These tests pin the bounded recovery: the watchdog +escalates to force_terminate() a short grace after "complete" (save done) or after an +absolute timeout (hang during save), and never force-kills a worker that exits cleanly. +Fakes only; no GPU, network, or subprocess. +""" + +from __future__ import annotations + +import contextlib +import logging +import queue +import sys +import threading +import time +import types as _types +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Stub the heavy module-level imports of core/training/training.py so it imports +# under CPU-only/no-network, then restore them (see the restore loop below). +_SAVED: dict = {} + + +def _stub(name, mod): + _SAVED[name] = sys.modules.get(name) + sys.modules[name] = mod + + +_lg = _types.ModuleType("loggers") +_lg.get_logger = lambda name: logging.getLogger(name) +_stub("loggers", _lg) +_stub("structlog", _types.ModuleType("structlog")) +_mpl = _types.ModuleType("matplotlib") +_plt = _types.ModuleType("matplotlib.pyplot") +_plt.Figure = type("Figure", (), {}) # referenced in a class-def annotation +_mpl.pyplot = _plt +_stub("matplotlib", _mpl) +_stub("matplotlib.pyplot", _plt) +_hw = _types.ModuleType("utils.hardware") +_hw.prepare_gpu_selection = lambda *a, **k: (None, None) +_stub("utils.hardware", _hw) +_npl = _types.ModuleType("utils.native_path_leases") +_npl.native_path_secret_removed_for_child_start = lambda: contextlib.nullcontext() +_npl.run_without_native_path_secret = lambda fn: fn +_stub("utils.native_path_leases", _npl) +_pth = _types.ModuleType("utils.paths") +_pth.outputs_root = lambda *a, **k: "/tmp/outputs" +_stub("utils.paths", _pth) + +# Whether core.training.training was already imported before this file ran; only +# evict it below if we were the one to create the (stub-bound) module instance. +_TRAINING_PRE_IMPORTED = "core.training.training" in sys.modules + +from core.training.training import TrainingBackend + +# Restore every stubbed module so this file never pollutes the shared session. +for _name in ( + "loggers", + "structlog", + "matplotlib", + "matplotlib.pyplot", + "utils.hardware", + "utils.native_path_leases", + "utils.paths", +): + _prev = _SAVED.get(_name) + if _prev is None: + sys.modules.pop(_name, None) + else: + sys.modules[_name] = _prev + +if not _TRAINING_PRE_IMPORTED: + sys.modules.pop("core.training.training", None) + sys.modules.pop("core.training", None) + +# The module globals hold the escalation timeouts and are the watchdog's own +# namespace; patch them here so tests run in well under a second. +_G = TrainingBackend._stop_watchdog_loop.__globals__ + + +class _FakeProc: + """A subprocess handle whose liveness and kill calls the test observes.""" + + def __init__(self, alive: bool = True): + self._alive = alive + self.pid = 4321 + self.terminated = False + self.killed = False + + def is_alive(self): + return self._alive + + def terminate(self): + self.terminated = True + + def kill(self): + self.killed = True + + def join(self, timeout = None): + pass + + +def _wait_until(predicate, timeout = 5.0): + deadline = time.time() + timeout + while time.time() < deadline: + if predicate(): + return True + time.sleep(0.01) + return predicate() + + +def _record_force_terminate(monkeypatch, b): + """Replace force_terminate + escalation finalize with recorders (no DB/OS).""" + calls: list = [] + monkeypatch.setattr(b, "force_terminate", lambda target_proc = None: calls.append("force")) + monkeypatch.setattr( + b, + "_finalize_stopped_after_escalation", + lambda target_proc = None, watched_job_id = None: calls.append("final"), + ) + return calls + + +# ---------------------------------------------------------------------------- +# (a) Escalate a short grace after "complete" (save done) if still alive. +# ---------------------------------------------------------------------------- + + +def test_watchdog_escalates_after_grace_once_complete_seen(monkeypatch): + monkeypatch.setitem(_G, "_STOP_GRACE_S", 0.05) + monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0) # ensure grace, not timeout, fires + b = TrainingBackend() + calls = _record_force_terminate(monkeypatch, b) + + proc = _FakeProc(alive = True) + b._proc = proc + b._complete_seen.set() # worker reported "complete" -> save is done + + b._start_stop_watchdog(cancel = False) + assert _wait_until( + lambda: calls == ["force", "final"] + ), "watchdog must force_terminate a worker still alive after the post-save grace" + b._stop_watchdog.join(timeout = 5) + + +# ---------------------------------------------------------------------------- +# (b) The absolute cap is a last-resort backstop, not a save killer. +# ---------------------------------------------------------------------------- + + +def test_watchdog_does_not_kill_save_still_saving_within_window(monkeypatch): + # save=True, no "complete" yet: a slow save in progress must not be force-killed + # inside the (long) absolute window. + monkeypatch.setitem(_G, "_STOP_GRACE_S", 100.0) + monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0) + b = TrainingBackend() + calls = _record_force_terminate(monkeypatch, b) + + proc = _FakeProc(alive = True) + b._proc = proc + b._start_stop_watchdog(cancel = False) + + time.sleep(0.3) + assert calls == [], "an in-progress save must not be killed within the absolute window" + assert b._stop_watchdog.is_alive() + + proc._alive = False + b._stop_watchdog.join(timeout = 5) + + +def test_watchdog_backstop_fires_for_save_after_absolute_timeout(monkeypatch): + # Past the long save=True cap with no completion: force-terminate as last resort. + monkeypatch.setitem(_G, "_STOP_GRACE_S", 100.0) # never trips (no complete) + monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 0.05) + b = TrainingBackend() + calls = _record_force_terminate(monkeypatch, b) + + b._proc = _FakeProc(alive = True) + b._start_stop_watchdog(cancel = False) + assert _wait_until( + lambda: calls == ["force", "final"] + ), "the absolute backstop must force_terminate a save that never completes" + b._stop_watchdog.join(timeout = 5) + + +def test_cancel_uses_shorter_absolute_timeout(monkeypatch): + # A cancel has nothing to save, so it escalates on the shorter cancel cap even before + # the long save cap elapses. + monkeypatch.setitem(_G, "_STOP_GRACE_S", 100.0) + monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0) # save cap would not fire + monkeypatch.setitem(_G, "_CANCEL_TIMEOUT_S", 0.05) + b = TrainingBackend() + calls = _record_force_terminate(monkeypatch, b) + + b._proc = _FakeProc(alive = True) + b._start_stop_watchdog(cancel = True) + assert _wait_until( + lambda: calls == ["force", "final"] + ), "a cancel must escalate on the shorter cancel timeout" + b._stop_watchdog.join(timeout = 5) + + +# ---------------------------------------------------------------------------- +# (c) No force-kill when the worker exits cleanly and promptly. +# ---------------------------------------------------------------------------- + + +def test_watchdog_no_op_on_clean_quick_exit(monkeypatch): + monkeypatch.setitem(_G, "_STOP_GRACE_S", 5.0) + monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 10.0) + b = TrainingBackend() + calls = _record_force_terminate(monkeypatch, b) + + proc = _FakeProc(alive = True) + b._proc = proc + b._complete_seen.set() # save done; worker is about to exit on its own + + b._start_stop_watchdog(cancel = False) + # Worker exits promptly, well before the grace period elapses. + time.sleep(0.1) + proc._alive = False + + b._stop_watchdog.join(timeout = 5) + assert not b._stop_watchdog.is_alive() + assert calls == [], "a clean quick exit must not trigger force_terminate" + + +def test_watchdog_no_op_when_worker_superseded(monkeypatch): + # A stale watchdog from a prior run must never kill a new run's worker: once + # self._proc is replaced, it exits silently. + monkeypatch.setitem(_G, "_STOP_GRACE_S", 0.05) + monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 0.05) + b = TrainingBackend() + calls = _record_force_terminate(monkeypatch, b) + + old_proc = _FakeProc(alive = True) + b._proc = old_proc + b._complete_seen.set() + b._start_stop_watchdog(cancel = False) + + # A new run takes over the handle before the grace elapses. + b._proc = _FakeProc(alive = True) + + b._stop_watchdog.join(timeout = 5) + assert calls == [], "watchdog must not force_terminate a superseded worker" + + +def test_new_run_gets_its_own_watchdog(monkeypatch): + # A stale watchdog sleeping on an old proc must not stop a new run's stop from + # creating its own watcher. + monkeypatch.setitem(_G, "_STOP_GRACE_S", 100.0) + monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0) + b = TrainingBackend() + _record_force_terminate(monkeypatch, b) + + old_proc = _FakeProc(alive = True) + b._proc = old_proc + b._start_stop_watchdog(cancel = False) + first_wd = b._stop_watchdog + + # New run: fresh worker replaces the handle; its stop must get a new watcher + # even though the old (superseded) watchdog is still alive. + new_proc = _FakeProc(alive = True) + b._proc = new_proc + b._start_stop_watchdog(cancel = False) + second_wd = b._stop_watchdog + + try: + assert first_wd.is_alive() + assert second_wd is not first_wd, "a new run must get its own watchdog" + assert b._stop_watchdog_proc is new_proc + finally: + old_proc._alive = False + new_proc._alive = False + first_wd.join(timeout = 5) + second_wd.join(timeout = 5) + + +def test_force_terminate_targets_only_captured_proc(): + # Superseded: force_terminate(target) must not touch a different current worker. + b = TrainingBackend() + old_proc = _FakeProc(alive = True) + new_proc = _FakeProc(alive = True) + b._proc = new_proc + b.force_terminate(target_proc = old_proc) + assert new_proc.terminated is False, "must not terminate the new run's worker" + assert old_proc.terminated is False, "must not terminate a handle that is not current" + + # Matching: the captured handle is the current worker, so it is terminated. + p = _FakeProc(alive = True) + b._proc = p + b.force_terminate(target_proc = p) + assert p.terminated is True + + +# ---------------------------------------------------------------------------- +# Post-escalation finalize leaves the parent ready for a new run. +# ---------------------------------------------------------------------------- + + +def test_finalize_runs_even_if_force_terminate_raises(monkeypatch): + # A wedged child can make force_terminate() raise; finalize must still run so the + # run does not stay stuck in "Stopping...". + monkeypatch.setitem(_G, "_STOP_GRACE_S", 0.05) + monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0) + b = TrainingBackend() + + def _boom(target_proc = None): + raise RuntimeError("kill() failed on wedged child") + + finalized: list = [] + monkeypatch.setattr(b, "force_terminate", _boom) + monkeypatch.setattr( + b, + "_finalize_stopped_after_escalation", + lambda target_proc = None, watched_job_id = None: finalized.append(True), + ) + + b._proc = _FakeProc(alive = True) + b._complete_seen.set() + b._start_stop_watchdog(cancel = False) + + assert _wait_until( + lambda: finalized == [True] + ), "finalize must run even when force_terminate raises" + b._stop_watchdog.join(timeout = 5) + + +def test_finalize_after_escalation_clears_state(monkeypatch): + # Even if the OS never reaps the wedged worker, the parent must report the run + # stopped so the UI leaves "Stopping..." and a new run can start. + b = TrainingBackend() + finstop: list = [] + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + + b._proc = _FakeProc(alive = True) # wedged: still reports alive + b._should_stop = True + b.current_job_id = "job_c" + b._db_run_created = True + b._progress.is_training = True + + b._finalize_stopped_after_escalation(watched_job_id = "job_c") + + assert b._proc is None, "the wedged handle must be dropped so is_training_active clears" + assert b._progress.is_training is False + assert b._progress.status_message == "Training stopped." + assert finstop and finstop[0][0] == "job_c", "the captured run must be finalized by id" + assert b.is_training_active() is False + + +def test_finalize_after_escalation_preserves_output_dir(monkeypatch): + # A save-stop that already emitted "complete" has the checkpoint dir; run history + # must record it even if the watchdog wins the finalize race against the pump. + b = TrainingBackend() + finstop: list = [] + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + + b._proc = _FakeProc(alive = True) + b._should_stop = True + b.current_job_id = "job_c" + b._db_run_created = True + b._output_dir = "/tmp/outputs/run-123" + + b._finalize_stopped_after_escalation(watched_job_id = "job_c") + + # _finish_stopped_run(run_id, output_dir, batch, final_step, final_loss, duration, loss_history) + assert finstop and finstop[0][0] == "job_c" + assert finstop[0][1] == "/tmp/outputs/run-123" + + +def test_stop_training_starts_watchdog_only_when_worker_alive(monkeypatch): + # No worker -> nothing to escalate; the watchdog must not spawn. + b = TrainingBackend() + b._proc = None + assert b.stop_training(save = True) is True + assert b._stop_watchdog is None + + +# ---------------------------------------------------------------------------- +# (d) A stale watchdog must never clobber a run that replaced its worker. +# ---------------------------------------------------------------------------- + + +def test_finalize_after_escalation_no_ops_when_superseded(monkeypatch): + # A /start can slip in while the watchdog force-terminates the old worker + # (is_training_active() is False once _should_stop is set and the old proc is dead). + # The escalation finalize must then leave the NEW run untouched, not drop its handle. + b = TrainingBackend() + finstop: list = [] + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + + old_proc = _FakeProc(alive = False) # force-terminated worker we were watching + new_proc = _FakeProc(alive = True) # a new run already took over + b._proc = new_proc + b.current_job_id = "job_new" + b._db_run_created = True + b._progress.is_training = True + + b._finalize_stopped_after_escalation(target_proc = old_proc) + + assert b._proc is new_proc, "must not drop the new run's handle" + assert b._progress.is_training is True, "must not mark the new run stopped" + assert finstop == [], "must not finalize the new run in the DB" + + +def test_finalize_after_escalation_runs_for_its_own_worker(monkeypatch): + # Common case: the watched worker is still current, so finalize proceeds and + # finalizes the captured run by id. + b = TrainingBackend() + finstop: list = [] + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + + proc = _FakeProc(alive = False) + b._proc = proc + b.current_job_id = "job_a" + b._db_run_created = True + b._progress.is_training = True + + b._finalize_stopped_after_escalation(target_proc = proc, watched_job_id = "job_a") + + assert b._proc is None + assert b._progress.is_training is False + assert finstop and finstop[0][0] == "job_a", "must finalize the captured run by id" + + +def test_finalize_after_escalation_no_ops_on_job_change_during_startup(monkeypatch): + # start_training updates current_job_id BEFORE it installs the new _proc, so a stale + # watchdog can enter while _proc is still the old (dead) handle. The job-id guard must + # catch this even though the proc-only guard would not. + b = TrainingBackend() + finstop: list = [] + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a)) + + old_proc = _FakeProc(alive = False) # old worker, dead; new _proc not installed yet + b._proc = old_proc # still the old handle (== target), so proc guard would pass + b.current_job_id = "job_new" # but the new run already claimed the job id + b._db_run_created = True + b._progress.is_training = True + + b._finalize_stopped_after_escalation(target_proc = old_proc, watched_job_id = "job_old") + + assert b._proc is old_proc, "must not drop the handle during a new run's startup" + assert b._progress.is_training is True, "must not mark the starting run stopped" + assert finstop == [], "must not finalize while a new run is starting up" + + +# ---------------------------------------------------------------------------- +# (e) A later cancel (save=False) tightens an in-flight save watchdog. +# ---------------------------------------------------------------------------- + + +def test_later_cancel_tightens_watchdog_timeout(monkeypatch): + monkeypatch.setitem(_G, "_STOP_GRACE_S", 100.0) # never trips (no complete) + monkeypatch.setitem(_G, "_STOP_TIMEOUT_S", 100.0) # save cap would not fire + monkeypatch.setitem(_G, "_CANCEL_TIMEOUT_S", 0.05) + b = TrainingBackend() + calls = _record_force_terminate(monkeypatch, b) + + b._proc = _FakeProc(alive = True) + b._start_stop_watchdog(cancel = False) # started as a save-stop with the long cap + time.sleep(0.15) + assert calls == [], "a save-stop must not escalate on the short cancel cap yet" + + # The user now cancels the in-flight stop: the watchdog must tighten its cap. + b._cancel_requested = True + assert _wait_until( + lambda: calls == ["force", "final"] + ), "a later cancel must tighten the watchdog to the shorter cancel cap" + b._stop_watchdog.join(timeout = 5) + + +# ---------------------------------------------------------------------------- +# (f) DB finalize/flush are safe when the watchdog and pump race (see Item 4). +# ---------------------------------------------------------------------------- + + +def _install_fake_db(monkeypatch): + """Stub storage.studio_db + utils.downsample so the real DB helpers run without + SQLite. Returns the recorder dict.""" + recs = {"created": [], "finished": [], "inserted": [], "insert_ids": [], "progress_ids": []} + fake_storage = _types.ModuleType("storage") + fake_db = _types.ModuleType("storage.studio_db") + fake_db.create_run = lambda **kw: recs["created"].append(kw) + fake_db.finish_run = lambda **kw: recs["finished"].append(kw) + fake_db.insert_metrics_batch = lambda job_id, batch: ( + recs["inserted"].extend(batch), + recs["insert_ids"].append(job_id), + ) + fake_db.update_run_progress = lambda **kw: recs["progress_ids"].append(kw.get("id")) + fake_storage.studio_db = fake_db + monkeypatch.setitem(sys.modules, "storage", fake_storage) + monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db) + fake_ds = _types.ModuleType("utils.downsample") + fake_ds.downsample = lambda seq, n: list(seq)[:n] + monkeypatch.setitem(sys.modules, "utils.downsample", fake_ds) + return recs + + +def test_finalize_run_in_db_single_winner_under_concurrency(monkeypatch): + # The watchdog and pump can both finalize; only one call may reach finish_run. + recs = _install_fake_db(monkeypatch) + b = TrainingBackend() + b.current_job_id = "job_x" + b._db_run_created = True + b._run_finalized = False + + start = threading.Barrier(8) + + def worker(): + start.wait() + b._finalize_run_in_db(status = "stopped") + + threads = [threading.Thread(target = worker) for _ in range(8)] + for t in threads: + t.start() + for t in threads: + t.join(timeout = 5) + + assert len(recs["finished"]) == 1, f"finalize must run once, got {len(recs['finished'])}" + assert b._run_finalized is True + + +def test_finalize_run_in_db_no_ops_on_job_mismatch(monkeypatch): + # A finalize captured for an old job must not finalize the run that replaced it. + recs = _install_fake_db(monkeypatch) + b = TrainingBackend() + b.current_job_id = "job_new" + b._db_run_created = True + b._run_finalized = False + + b._finalize_run_in_db(status = "stopped", expected_job_id = "job_old") + + assert recs["finished"] == [], "a superseded job id must not finalize the current run" + assert b._run_finalized is False + + +def test_concurrent_flush_claims_each_metric_once(monkeypatch): + # Concurrent flushes (pump periodic flush vs watchdog finalize flush) must not + # double-remove or drop buffered metrics. + recs = _install_fake_db(monkeypatch) + b = TrainingBackend() + b.current_job_id = "job_y" + b._db_run_created = True + b._metric_buffer[:] = [{"step": i} for i in range(200)] + + start = threading.Barrier(6) + + def worker(): + start.wait() + for _ in range(50): + b._flush_metrics_to_db() + + threads = [threading.Thread(target = worker) for _ in range(6)] + for t in threads: + t.start() + for t in threads: + t.join(timeout = 5) + b._flush_metrics_to_db() # drain any remainder + + steps = sorted(m["step"] for m in recs["inserted"]) + assert steps == list(range(200)), "each metric must be inserted exactly once" + assert b._metric_buffer == [], "the buffer must be fully drained" + + +def test_flush_pins_to_passed_run_id(monkeypatch): + # A finalizer flushes to the run it captured, even if a new /start has already + # changed current_job_id. + recs = _install_fake_db(monkeypatch) + b = TrainingBackend() + b.current_job_id = "job_new" # a new run is already live + b._db_run_created = True + b._metric_buffer[:] = [{"step": 1}, {"step": 2}] + + b._flush_metrics_to_db(run_id = "job_old") + + assert recs["insert_ids"] == ["job_old"], "metrics must go to the captured run, not the new one" + assert recs["progress_ids"] == ["job_old"] + + +def test_finalize_uses_snapshot_run_id_across_new_run(monkeypatch): + # If a new /start changes current_job_id after the finalize claim but before the DB + # writes, finish_run must still target the run captured under the lock. + recs = _install_fake_db(monkeypatch) + b = TrainingBackend() + b.current_job_id = "job_x" + b._db_run_created = True + b._run_finalized = False + + def hijack(run_id = None): + # Simulate a new run taking over during the flush (after the finalize claim). + b.current_job_id = "job_y" + + monkeypatch.setattr(b, "_flush_metrics_to_db", hijack) + + b._finalize_run_in_db(status = "stopped", expected_job_id = "job_x") + + assert [f["id"] for f in recs["finished"]] == [ + "job_x" + ], "finish_run must target the captured run, not the run that replaced it" + + +# ---------------------------------------------------------------------------- +# (g) DB row creation must not be published before the insert commits. +# ---------------------------------------------------------------------------- + + +def test_ensure_db_run_created_publishes_only_after_insert(monkeypatch): + # _db_run_created must stay False while create_run is in flight, so a concurrent + # finalize can't run finish_run (an UPDATE) against a not-yet-inserted row. + b = TrainingBackend() + b.current_job_id = "job_z" + b._db_config = {"model_name": "m"} + observed: dict = {} + + fake_storage = _types.ModuleType("storage") + fake_db = _types.ModuleType("storage.studio_db") + + def _create(**kw): + observed["flag_during_create"] = b._db_run_created + observed["in_progress_during_create"] = b._db_create_in_progress + + fake_db.create_run = _create + fake_storage.studio_db = fake_db + monkeypatch.setitem(sys.modules, "storage", fake_storage) + monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db) + + b._ensure_db_run_created() + + assert observed["flag_during_create"] is False, "flag must not be published before insert" + assert observed["in_progress_during_create"] is True + assert b._db_run_created is True, "flag must be published after a successful insert" + assert b._db_create_in_progress is False + + +def test_ensure_db_run_created_stays_unpublished_on_failure(monkeypatch): + # If create_run raises, neither flag stays set, so a later caller can retry. + b = TrainingBackend() + b.current_job_id = "job_z" + b._db_config = {"model_name": "m"} + + fake_storage = _types.ModuleType("storage") + fake_db = _types.ModuleType("storage.studio_db") + + def _boom_create(**kw): + raise RuntimeError("insert failed") + + fake_db.create_run = _boom_create + fake_storage.studio_db = fake_db + monkeypatch.setitem(sys.modules, "storage", fake_storage) + monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db) + + b._ensure_db_run_created() + + assert b._db_run_created is False, "a failed insert must not publish the row as created" + assert b._db_create_in_progress is False, "the in-progress flag must be cleared on failure" + + +def test_ensure_db_run_created_does_not_publish_for_a_new_run(monkeypatch): + # A killed worker lets a new /start proceed while the watchdog is still creating the old + # run's row. The stale create must not publish the backend-wide flags against the new + # current_job_id, or the new run would skip inserting its own row. + b = TrainingBackend() + b.current_job_id = "job_old" + b._db_config = {"model_name": "m"} + b._db_run_created = False + b._db_create_in_progress = False + + fake_storage = _types.ModuleType("storage") + fake_db = _types.ModuleType("storage.studio_db") + + def _create(**kw): + b.current_job_id = "job_new" # a new run takes over during the slow create + + fake_db.create_run = _create + fake_storage.studio_db = fake_db + monkeypatch.setitem(sys.modules, "storage", fake_storage) + monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db) + + b._ensure_db_run_created() + + assert b._db_run_created is False, "must not publish the created flag against the new run" + # The stale claim is left for start_training to reset, not satisfied for the new run. + assert b._db_create_in_progress is True, "must not clear the claim once the run is not current" + + +# ---------------------------------------------------------------------------- +# (h) The escalation finalizes the watched run by id (so it is never left running). +# ---------------------------------------------------------------------------- + + +def test_escalation_finalizes_watched_run_by_id_end_to_end(monkeypatch): + # Exercise the real _finish_stopped_run against a fake DB. The watched run is finalized + # by its captured id with its buffered metrics, so a new run that starts in the gap + # after the backend goes idle can never leave the stopped run recorded running. + recs = _install_fake_db(monkeypatch) + b = TrainingBackend() + b.current_job_id = "job_old" + b._db_run_created = True + b._proc = _FakeProc(alive = False) + b._progress.is_training = True + b._progress.step = 42 + b._metric_buffer[:] = [{"step": 41}, {"step": 42}] + + b._finalize_stopped_after_escalation(target_proc = b._proc, watched_job_id = "job_old") + + assert [f["id"] for f in recs["finished"]] == ["job_old"], "must finish the captured run by id" + assert recs["finished"][0]["status"] == "stopped" + assert recs["insert_ids"] == ["job_old"], "buffered metrics must land on the captured run" + assert b._metric_buffer == [], "the captured batch must be drained" + + +def test_escalation_defers_when_row_cannot_be_created_here(monkeypatch): + # If the row does not exist and cannot be created here (no db_config, or the pump is + # mid-create), the escalation must not claim _run_finalized or call _finish_stopped_run, + # so the pump's create-then-finalize records the run. Parent state still clears. + b = TrainingBackend() + called: list = [] + monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: called.append(a)) + + b._proc = _FakeProc(alive = False) + b.current_job_id = "job_q" + b._db_run_created = False # row not created yet + b._db_config = None # ... and cannot be created here + b._run_finalized = False + b._progress.is_training = True + + b._finalize_stopped_after_escalation(target_proc = b._proc, watched_job_id = "job_q") + + assert called == [], "must not finalize when the row can't be established here" + assert b._run_finalized is False, "must not claim the finalize the pump still owes" + assert b._progress.is_training is False, "parent state must still clear so the UI unsticks" + assert b._proc is None + + +def test_escalation_creates_row_then_finalizes_when_start_create_failed(monkeypatch): + # A wedged worker's pump can never finalize and would bail once _proc is dropped, so if + # the row was never created (start-time create failed) the escalation creates it and + # finalizes by id itself, recording the terminal state before dropping the handle. + recs = _install_fake_db(monkeypatch) + b = TrainingBackend() + b.current_job_id = "job_s" + b._db_config = {"model_name": "m"} # so _ensure_db_run_created can create the row + b._db_run_created = False # start-time create failed + b._proc = _FakeProc(alive = True) # wedged: still reports alive + b._should_stop = True + b._progress.is_training = True + + b._finalize_stopped_after_escalation(target_proc = b._proc, watched_job_id = "job_s") + + assert [c["id"] for c in recs["created"]] == ["job_s"], "must create the missing row" + assert [f["id"] for f in recs["finished"]] == ["job_s"], "must finish the created row by id" + assert b._proc is None, "handle dropped only after the terminal state is recorded" + assert b._db_run_created is True + + +def test_escalation_does_not_drop_a_new_runs_handle(monkeypatch): + # If a run replaces the worker while the finalize DB write is in flight, the final _proc + # drop must leave the new run's handle intact (re-guarded on target_proc). + b = TrainingBackend() + b.current_job_id = "job_old" + b._db_run_created = True + old_proc = _FakeProc(alive = False) + new_proc = _FakeProc(alive = True) + b._proc = old_proc + + def hijack(*a): + b._proc = new_proc # a new run takes over during the finalize + + monkeypatch.setattr(b, "_finish_stopped_run", hijack) + + b._finalize_stopped_after_escalation(target_proc = old_proc, watched_job_id = "job_old") + + assert b._proc is new_proc, "must not drop the handle a new run installed during finalize" + + +def _make_finish_raise(monkeypatch, calls): + fn = sys.modules["storage.studio_db"] + + def _boom(**kw): + calls.append(kw) + raise RuntimeError("database is locked") + + fn.finish_run = _boom + + +def test_finish_stopped_run_retries_then_unclaims_on_db_error(monkeypatch): + # The watchdog is the sole finalizer once _proc is dropped, so a transient DB error is + # retried a few times; on final failure the finalize is unclaimed (run still current). + monkeypatch.setitem(_G, "_DB_FINALIZE_RETRY_S", 0.0) + _install_fake_db(monkeypatch) + tries: list = [] + _make_finish_raise(monkeypatch, tries) + b = TrainingBackend() + b.current_job_id = "job_r" + b._run_finalized = True # the caller (escalation) already claimed + + b._finish_stopped_run("job_r", None, [{"step": 1}], 1, None, None, []) + + assert len(tries) == 3, "a transient DB error must be retried before giving up" + assert b._run_finalized is False, "a persistent DB error must unclaim the finalize" + + +def test_finish_stopped_run_error_leaves_new_run_untouched(monkeypatch): + # If the watched run was superseded, a DB error must not unclaim the new run's finalize. + monkeypatch.setitem(_G, "_DB_FINALIZE_RETRY_S", 0.0) + _install_fake_db(monkeypatch) + _make_finish_raise(monkeypatch, []) + b = TrainingBackend() + b.current_job_id = "job_new" # a new run is live + b._run_finalized = True # the new run's flag + + b._finish_stopped_run("job_old", None, [{"step": 1}], 1, None, None, []) + + assert b._run_finalized is True, "must not unclaim the new run's finalize" From 14d0e853fa5755542e318d8d3071023de6e1d3f4 Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Wed, 15 Jul 2026 18:09:45 +0800 Subject: [PATCH 08/17] fix(studio): recover MLX VLM image prompts (#7094) * fix(studio): recover mlx vlm image prompts * fix(studio): detect serialized vlm media items * studio: recover MLX VLM prompts when model_type only lives on _config _mlx_vlm_model_config only fell back to _config when config was entirely missing, so a model that exposes a config without a model_type (while _config carries it) skipped model-aware recovery. Prefer whichever of config / _config actually has a model_type. Adds a focused test. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Daniel Han Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../backend/core/inference/mlx_inference.py | 197 ++++++++++++++++-- .../tests/test_mlx_inference_backend.py | 112 ++++++++++ 2 files changed, 297 insertions(+), 12 deletions(-) diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 6287b184a6..163ade10c4 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -5,15 +5,120 @@ Drop-in replacement for InferenceBackend — same interface, uses mlx-lm/mlx-vlm instead of torch/transformers for model loading and generation. """ +import json import os import threading from typing import Optional, Generator +from core.inference.message_content import content_to_text from core.inference.runtime_context import runtime_context_length from loggers import get_logger logger = get_logger(__name__) +def _mlx_vlm_model_config(model): + """Return the loaded MLX model config and its type, preferring whichever of + config / _config actually carries a model_type.""" + + def _model_type(cfg): + return cfg.get("model_type") if isinstance(cfg, dict) else getattr(cfg, "model_type", None) + + configs = [ + cfg + for cfg in (getattr(model, "config", None), getattr(model, "_config", None)) + if cfg is not None + ] + for cfg in configs: + model_type = _model_type(cfg) + if model_type is not None: + return cfg, model_type + return (configs[0] if configs else None), None + + +def _render_registered_vlm_prompt(processor, model, messages, num_images): + """Render through mlx-vlm when it declares a formatter for this model.""" + from mlx_vlm import prompt_utils + + config, model_type = _mlx_vlm_model_config(model) + if config is None: + return None + if model_type not in getattr(prompt_utils, "MODEL_CONFIG", {}): + return None + + rendered = prompt_utils.apply_chat_template( + processor, + config, + messages, + add_generation_prompt = True, + num_images = num_images, + ) + if isinstance(rendered, str) and rendered.strip(): + return rendered + raise RuntimeError("mlx-vlm's registered renderer returned an empty prompt.") + + +def _count_vlm_images(content): + if isinstance(content, list): + return sum(_count_vlm_images(item) for item in content) + if not isinstance(content, dict): + return 0 + if str(content.get("type", "")).lower() in ("image", "image_url", "input_image"): + return 1 + return _count_vlm_images(content.get("content")) + + +def _vlm_media_reprs(content): + if isinstance(content, list): + values = ( + {str(content), json.dumps(content, ensure_ascii = False)} + if _count_vlm_images(content) + else set() + ) + for item in content: + values.update(_vlm_media_reprs(item)) + return values + if not isinstance(content, dict): + return set() + if str(content.get("type", "")).lower() in ("image", "image_url", "input_image"): + return {str(content), json.dumps(content, ensure_ascii = False)} + return _vlm_media_reprs(content.get("content")) + + +def _prompt_serializes_vlm_media(prompt, messages): + """Detect templates that embed the exact structured media object repr.""" + media_reprs = set() + for message in messages: + if isinstance(message, dict): + media_reprs.update(_vlm_media_reprs(message.get("content"))) + text_content = [ + content_to_text(message.get("content")) for message in messages if isinstance(message, dict) + ] + return any( + prompt.count(media_repr) > sum(content.count(media_repr) for content in text_content) + for media_repr in media_reprs + ) + + +def _vlm_prompt_issue(prompt, messages): + if not isinstance(prompt, str) or not prompt.strip(): + return "an empty prompt" + if _prompt_serializes_vlm_media(prompt, messages): + return "serialized structured image content" + return None + + +def _vlm_messages_have_tool_history(messages): + return any( + isinstance(message, dict) + and ( + message.get("role") == "tool" + or message.get("tool_calls") + or message.get("tool_call_id") + ) + for message in messages + ) + + def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps): """Map mlx stream stats onto the usage/timings shape llama-server emits.""" prompt_n = int(prompt_n or 0) @@ -422,9 +527,7 @@ class MLXInferenceBackend: {"type": "text", "text": content}, ] elif isinstance(content, list): - has_image = any( - p.get("type") == "image" for p in content if isinstance(p, dict) - ) + has_image = _count_vlm_images(content) > 0 if not has_image: content.insert(0, {"type": "image"}) break @@ -632,17 +735,87 @@ class MLXInferenceBackend: ): chat_target = getattr(self._processor, "tokenizer", self._processor) - prompt = apply_chat_template_for_generation( - chat_target, - messages, - tools = tools, - enable_thinking = enable_thinking, - reasoning_effort = reasoning_effort, - preserve_thinking = preserve_thinking, - ) - # mlx_vlm's stream_generate handles pixel_values (None for text-only) images = [image] if image is not None else None + attached_images = 0 if images is None else len(images) + structured_images = sum( + _count_vlm_images(message.get("content")) + for message in messages + if isinstance(message, dict) + ) + if structured_images != attached_images: + raise RuntimeError( + f"VLM conversation contains {structured_images} structured image " + f"item(s) for {attached_images} attached image(s)." + ) + prompt = None + has_tool_history = _vlm_messages_have_tool_history(messages) + prompt_error = None + try: + prompt = apply_chat_template_for_generation( + chat_target, + messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) + except Exception as exc: + if images is None or has_tool_history: + raise + prompt_error = exc + prompt_issue = ( + _vlm_prompt_issue(prompt, messages) if prompt_error is None else "a rendering error" + ) + if prompt_issue and has_tool_history: + raise RuntimeError( + f"VLM chat template returned {prompt_issue} and cannot be recovered " + "without dropping tool-call history." + ) from prompt_error + + if images is not None and prompt_issue: + if tools or any( + value is not None + for value in (enable_thinking, reasoning_effort, preserve_thinking) + ): + if prompt_error is not None: + raise prompt_error + raise RuntimeError( + f"VLM chat template returned {prompt_issue} and cannot be recovered " + "without dropping requested tools or reasoning controls." + ) + try: + recovered_prompt = _render_registered_vlm_prompt( + self._processor, + self._model, + messages, + len(images), + ) + except Exception as recovery_error: + if prompt_error is not None: + raise prompt_error + raise RuntimeError( + f"VLM chat template returned {prompt_issue}; model-aware " + f"recovery failed: {recovery_error}" + ) from recovery_error + if recovered_prompt is None: + if prompt_error is not None: + raise prompt_error + raise RuntimeError( + f"VLM chat template returned {prompt_issue}, and no registered " + "MLX VLM renderer was available for this model." + ) + recovered_issue = _vlm_prompt_issue(recovered_prompt, messages) + if recovered_issue: + if prompt_error is not None: + raise prompt_error + raise RuntimeError( + f"Model-aware VLM rendering returned {recovered_issue} for " + f"{attached_images} attached image(s)." + ) + prompt = recovered_prompt + elif prompt_issue: + raise RuntimeError(f"VLM chat template returned {prompt_issue}.") from prompt_error from core.inference.chat_template_helpers import detect_think_prefill diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index 55a3198a6b..29fbb45158 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -333,6 +333,118 @@ def test_mlx_generate_chat_response_accepts_template_kwargs(): ), f"{name!r} must default to None so existing callers stay valid" +def test_mlx_vlm_generation_selects_renderer_by_capability(monkeypatch): + from core.inference.mlx_inference import MLXInferenceBackend + + calls = {"generic": [], "model": [], "stream": []} + state = {"generic": "serialized", "model": " model-aware"} + prompt_utils = SimpleNamespace( + MODEL_CONFIG = {"deepseek_vl_v2": object()}, + apply_chat_template = lambda *_args, **kwargs: ( + calls["model"].append(kwargs) or state["model"] + ), + ) + mlx_vlm = types.ModuleType("mlx_vlm") + mlx_vlm.prompt_utils = prompt_utils + mlx_vlm.stream_generate = lambda *_args, **kwargs: ( + calls["stream"].append((_args, kwargs)) + or iter([SimpleNamespace(text = "ok", prompt_tokens = 3, generation_tokens = 1)]) + ) + monkeypatch.setitem(sys.modules, "mlx_vlm", mlx_vlm) + + def generic(_target, _messages, **kwargs): + calls["generic"].append(kwargs) + if isinstance(state["generic"], Exception): + raise state["generic"] + if state["generic"] == "serialized": + return f"User: {_messages[0]['content']}" + return state["generic"] + + monkeypatch.setattr( + "core.inference.chat_template_helpers.apply_chat_template_for_generation", + generic, + ) + backend = MLXInferenceBackend() + backend._model = SimpleNamespace(config = {"model_type": "deepseek_vl_v2"}) + backend._processor = SimpleNamespace(tokenizer = SimpleNamespace()) + args = ([{"role": "user", "content": [{"type": "image"}]}], object(), 0, 1, 0, 0, 1, 1, None) + tools = [{"function": {"name": "search"}}] + assert list(backend._generate_vlm(*args)) == ["ok"] + assert calls["model"][0]["num_images"] == 1 + assert calls["stream"][0][0][2] == " model-aware" + with pytest.raises(RuntimeError, match = "dropping requested tools"): + list(backend._generate_vlm(*args, tools = tools)) + with pytest.raises(RuntimeError, match = "dropping requested tools or reasoning"): + list(backend._generate_vlm(*args, enable_thinking = False)) + backend._processor = SimpleNamespace(chat_template = "template") + state["generic"] = " healthy generic" + assert list(backend._generate_vlm(*args, tools = tools, enable_thinking = False)) == ["ok"] + assert calls["generic"][-1]["enable_thinking"] is False + assert calls["stream"][-1][0][2] == " healthy generic" + state["generic"] = "generic prompt" + text_messages = [{"role": "user", "content": "hello"}] + assert list(backend._generate_vlm(*((text_messages, None) + args[2:]), tools = tools)) == ["ok"] + assert calls["generic"][-1]["tools"] == tools + assert calls["stream"][-1][0][2] == "generic prompt" + two_images = [{"role": "user", "content": [{"type": "image"}, {"type": "image"}]}] + with pytest.raises(RuntimeError, match = "2 structured image item"): + list(backend._generate_vlm(*((two_images,) + args[1:]), tools = tools)) + state["generic"] = "serialized" + tool_history = args[0] + [{"role": "assistant", "tool_calls": [{"id": "call-1"}]}] + with pytest.raises(RuntimeError, match = "tool-call history"): + list(backend._generate_vlm(*((tool_history,) + args[1:]), tools = tools)) + state["generic"] = ValueError("generic rendering failed") + state["model"] = f"User: {args[0][0]['content']}" + with pytest.raises(ValueError, match = "generic rendering failed"): + list(backend._generate_vlm(*args)) + + +def test_mlx_vlm_image_injection_reuses_media_aliases(monkeypatch): + from core.inference.mlx_inference import MLXInferenceBackend, _prompt_serializes_vlm_media + + media = [{"type": "image"}] + quoted = [{"role": "user", "content": media}, {"role": "user", "content": f"Explain {media}"}] + assert _prompt_serializes_vlm_media(f"\n{media[0]}", quoted[:1]) + assert not _prompt_serializes_vlm_media(f"\nExplain {media}", quoted) + assert _prompt_serializes_vlm_media(f"User: {media}\nExplain {media}", quoted) + quoted[1]["content"] = [{"type": "text", "text": f'Explain "this" {media}'}] + assert not _prompt_serializes_vlm_media(f'\nExplain "this" {media}', quoted) + json_media = [{"type": "image_url"}] + json_repr = '{"type": "image_url"}' + assert _prompt_serializes_vlm_media(f"\n{json_repr}", [{"content": json_media}]) + assert not _prompt_serializes_vlm_media( + f"\nExplain {json_repr}", + [{"content": json_media}, {"content": f"Explain {json_repr}"}], + ) + + backend = MLXInferenceBackend() + backend._model = object() + backend._is_vlm = True + captured = [] + backend._generate_vlm = lambda messages, *_args, **_kwargs: ( + captured.append(messages) or iter(()) + ) + messages = [{"role": "user", "content": [{"type": "image_url"}]}] + list(backend.generate_chat_response(messages, image = object())) + assert captured[0][0]["content"] == [{"type": "image_url"}] + + +def test_mlx_vlm_model_config_prefers_config_with_model_type(): + from core.inference.mlx_inference import _mlx_vlm_model_config + + # config present but missing model_type must fall back to _config + m = SimpleNamespace(config = {}, _config = {"model_type": "deepseek_vl_v2"}) + assert _mlx_vlm_model_config(m) == ({"model_type": "deepseek_vl_v2"}, "deepseek_vl_v2") + # an object config whose model_type is None also falls back + m = SimpleNamespace(config = SimpleNamespace(model_type = None), _config = {"model_type": "qwen2_vl"}) + assert _mlx_vlm_model_config(m)[1] == "qwen2_vl" + # a config that already carries a model_type is preferred and returned unchanged + assert _mlx_vlm_model_config(SimpleNamespace(config = {"model_type": "gemma3"})) == ( + {"model_type": "gemma3"}, + "gemma3", + ) + + def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): """Mac text path must route through apply_chat_template_for_generation so reasoning / tool kwargs reach the tokenizer.""" From 3eb3259f047ab830c4faf2086730209262e696b5 Mon Sep 17 00:00:00 2001 From: Kushida Date: Wed, 15 Jul 2026 14:37:03 +0300 Subject: [PATCH 09/17] fix(install): fail non-tauri installer errors (#7123) --- install.ps1 | 1 + tests/test_installer_skip_autostart.py | 23 ++++++++++++++++++++++- 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/install.ps1 b/install.ps1 index 100a3177ba..0b32d7cb6c 100644 --- a/install.ps1 +++ b/install.ps1 @@ -91,6 +91,7 @@ function Install-UnslothStudio { if ($TauriMode) { exit $Code } + throw $Message } # ── Parse flags ── diff --git a/tests/test_installer_skip_autostart.py b/tests/test_installer_skip_autostart.py index 8c283f4458..fcc199af44 100644 --- a/tests/test_installer_skip_autostart.py +++ b/tests/test_installer_skip_autostart.py @@ -1,4 +1,4 @@ -"""Regression tests for the installers' post-install autostart opt-out.""" +"""Regression tests for installer controls and process exits.""" from __future__ import annotations @@ -120,6 +120,27 @@ def test_windows_skip_autostart_bypasses_only_the_interactive_prompt(): ) +@pytest.mark.skipif(shutil.which("pwsh") is None, reason = "PowerShell is unavailable") +def test_windows_installer_invalid_package_fails(): + result = subprocess.run( + [ + "pwsh", + "-NoProfile", + "-NonInteractive", + "-File", + str(INSTALL_PS1), + "--package", + "bad!", + ], + capture_output = True, + text = True, + timeout = 30, + ) + + assert result.returncode != 0 + assert "package name contains invalid characters" in result.stdout + result.stderr + + def test_skip_autostart_is_documented_for_all_installers(): readme = README.read_text(encoding = "utf-8") assert "UNSLOTH_SKIP_AUTOSTART=1 sh" in readme From ee73bcb209dfe54818719e4b0b50ed560fa44639 Mon Sep 17 00:00:00 2001 From: Hyacinth-of-Security <144213008+lxcxjxhx@users.noreply.github.com> Date: Wed, 15 Jul 2026 19:39:38 +0800 Subject: [PATCH 10/17] Fix bare except clauses and remove duplicate MAX_FUSED_SIZE definition (#7138) * fix: replace bare except clauses and remove duplicate MAX_FUSED_SIZE definition * Also catch NameError in mllama RMSNorm patch/unpatch fallbacks If mllama exists but MllamaTextRMSNorm is missing, the module-level import fails so Unsloth_MllamaTextRMSNorm/MllamaTextRMSNorm stay undefined. The patch/unpatch module imports then succeed and reference the undefined name, raising NameError. Add NameError so these fallbacks stay no-ops as before. --------- Co-authored-by: lxcxjxhx Co-authored-by: Daniel Han <23090290+danielhanchen@users.noreply.github.com> --- unsloth/kernels/cross_entropy_loss.py | 3 --- unsloth/kernels/rms_layernorm.py | 6 +++--- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index 39446ec639..0d3bf08c35 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -285,9 +285,6 @@ _cross_entropy_backward = triton.heuristics( )(_cross_entropy_backward) -MAX_FUSED_SIZE = 65536 # 2**16 - - class Fast_CrossEntropyLoss(torch.autograd.Function): @staticmethod def forward( diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py index bc81c5533a..5fc867e0cb 100644 --- a/unsloth/kernels/rms_layernorm.py +++ b/unsloth/kernels/rms_layernorm.py @@ -270,7 +270,7 @@ try: return fast_rms_layernorm(self, X, gemma = False) -except: +except (ImportError, AttributeError): pass @@ -281,7 +281,7 @@ def patch_rms_layernorm(): try: import transformers.models.mllama.modeling_mllama transformers.models.mllama.modeling_mllama.MllamaTextRMSNorm = Unsloth_MllamaTextRMSNorm - except: + except (ImportError, AttributeError, NameError): pass return @@ -293,7 +293,7 @@ def unpatch_rms_layernorm(): try: import transformers.models.mllama.modeling_mllama transformers.models.mllama.modeling_mllama.MllamaTextRMSNorm = MllamaTextRMSNorm - except: + except (ImportError, AttributeError, NameError): pass return From d8094335b70ceb69bab74c5d59ac44906d825f34 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Wed, 15 Jul 2026 09:19:37 -0300 Subject: [PATCH 11/17] Studio: scope the seeded bootstrap password auto-fill to loopback clients (#7131) * Studio: scope the seeded bootstrap password auto-fill to loopback clients * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: block bootstrap injection through Cloudflare tunnels * Studio: require loopback host for bootstrap injection * Studio: add regression test for unparseable Host in bootstrap loopback gate * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Scope bootstrap auto-fill to a direct-loopback client (block proxy/tunnel headers and malformed Host) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Reject scope-id addresses in loopback check (fail closed on ::1%zone Host) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Reject malformed bracketed Host in loopback check (e.g. [::1]evil) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: danielhanchen <23090290+danielhanchen@users.noreply.github.com> --- studio/backend/main.py | 73 ++++++++++- .../tests/test_index_bootstrap_loopback.py | 123 ++++++++++++++++++ 2 files changed, 194 insertions(+), 2 deletions(-) create mode 100644 studio/backend/tests/test_index_bootstrap_loopback.py diff --git a/studio/backend/main.py b/studio/backend/main.py index 332f4e7d1f..8062b7b073 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -234,6 +234,7 @@ if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT: os.environ.setdefault("UNSLOTH_IS_PRESENT", "1") import hashlib +import ipaddress import mimetypes import re as _re import shutil @@ -1363,6 +1364,61 @@ def _canonical_origin(scheme: str, netloc: str) -> Optional[tuple[str, str, int] return (scheme, host, port) +def _is_loopback_ip(host: Optional[str]) -> bool: + """Return whether ``host`` is a loopback IP, including IPv4-mapped IPv6.""" + if not host or "%" in host: # a scope id (::1%eth0) is never a plain loopback + return False + try: + ip = ipaddress.ip_address(host) + except (TypeError, ValueError): + return False + mapped = getattr(ip, "ipv4_mapped", None) + return ip.is_loopback or (mapped is not None and mapped.is_loopback) + + +# A loopback peer carrying any of these is a proxy/tunnel relaying a remote +# client, so the peer is the proxy, not the caller: cloudflared sets +# cf-connecting-ip, reverse proxies set the rest (uvicorn only consumes +# x-forwarded-for, so the others survive to here). +_PROXIED_CLIENT_HEADERS = ( + "cf-connecting-ip", + "forwarded", + "x-forwarded-for", + "x-forwarded-host", + "x-real-ip", +) + + +def _host_header_is_loopback(host_header: Optional[str]) -> bool: + """Loopback/localhost check on the raw Host header. + + Reads the header directly so a malformed or absent Host cannot fall back to + ``request.url.hostname``'s (loopback) ASGI server address. + """ + if not host_header: + return False + host = host_header.strip() + if host.startswith("["): # [IPv6] or [IPv6]:port + end = host.find("]") + if end == -1 or (host[end + 1 :] and not host[end + 1 :].startswith(":")): + return False # unclosed bracket or junk after ] (e.g. [::1]evil) + host = host[1:end] + elif host.count(":") == 1: # host:port + host = host.split(":", 1)[0] + host = host.lower().rstrip(".") + return host == "localhost" or _is_loopback_ip(host) + + +def _is_local_bootstrap_request(request: Request) -> bool: + """Allow bootstrap injection only through a direct loopback authority.""" + client = request.client + if client is None or not _is_loopback_ip(client.host): + return False + if any(request.headers.get(h) is not None for h in _PROXIED_CLIENT_HEADERS): + return False + return _host_header_is_loopback(request.headers.get("host")) + + def _is_same_origin_request(request: Request) -> bool: """True when Origin is missing or matches request's scheme://host:port. @@ -1398,6 +1454,17 @@ def _is_same_origin_request(request: Request) -> bool: return origin_canon == self_canon +def _should_inject_bootstrap(request: Request) -> bool: + """Whether to embed the seeded bootstrap password in index.html.""" + if not _is_same_origin_request(request): + return False + if _IS_COLAB: + # Single-user notebook proxy: allow autofill, but never a public + # shareable tunnel (a Colab Cloudflare link sets cf-connecting-ip). + return request.headers.get("cf-connecting-ip") is None + return _is_local_bootstrap_request(request) + + def setup_frontend(app: FastAPI, build_path: Path): """Mount frontend static files (optional)""" if not build_path.exists(): @@ -1410,8 +1477,10 @@ def setup_frontend(app: FastAPI, build_path: Path): def _build_index_response(request: Request) -> Response: content = (build_path / "index.html").read_bytes() content = _strip_crossorigin(content) - # Bootstrap pw is same-origin only; Vary: Origin keeps caches honest. - if _is_same_origin_request(request): + # Bootstrap pw goes only to a same-origin, direct-loopback client (or + # Colab's single-user notebook proxy): a wildcard bind must not serve it + # in-page to a LAN or proxied peer. Vary: Origin keeps caches honest. + if _should_inject_bootstrap(request): content, nonce = _inject_bootstrap(content, app) else: nonce = None diff --git a/studio/backend/tests/test_index_bootstrap_loopback.py b/studio/backend/tests/test_index_bootstrap_loopback.py new file mode 100644 index 0000000000..87abace22c --- /dev/null +++ b/studio/backend/tests/test_index_bootstrap_loopback.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression coverage for bootstrap password exposure to remote clients.""" + +from types import SimpleNamespace + + +def _request( + client_host, + request_host = "127.0.0.1", + headers = None, +): + """Build a minimal request; ``None`` models an unresolved peer / absent Host.""" + client = None if client_host is None else SimpleNamespace(host = client_host, port = 0) + hdrs = {} + if request_host is not None: + hdrs["host"] = request_host + hdrs.update(headers or {}) + return SimpleNamespace(client = client, headers = hdrs, url = SimpleNamespace(hostname = request_host)) + + +def test_loopback_peers_are_local(): + from main import _is_local_bootstrap_request + cases = ( + ("127.0.0.1", "127.0.0.1"), + ("::1", "::1"), + ("::ffff:127.0.0.1", "::ffff:127.0.0.1"), + ("127.0.0.1", "localhost"), + ) + for peer, host in cases: + assert _is_local_bootstrap_request(_request(peer, host)) is True, (peer, host) + + +def test_non_loopback_peers_are_remote(): + from main import _is_local_bootstrap_request + + # ::1%eth0 is a scope-id'd address, which ipaddress treats as loopback on + # 3.9+; it must not count as a direct local peer. + for host in ("192.168.1.10", "::ffff:192.168.1.10", "::1%eth0"): + assert _is_local_bootstrap_request(_request(host)) is False, host + + +def test_absent_or_unparseable_peer_fails_safe(): + from main import _is_local_bootstrap_request + for host in (None, "localhost"): + assert _is_local_bootstrap_request(_request(host)) is False, host + + +def test_cloudflare_tunnel_clients_are_remote_despite_loopback_peer(): + from main import _is_local_bootstrap_request + for client_ip in ("203.0.113.7", ""): + request = _request("127.0.0.1", headers = {"cf-connecting-ip": client_ip}) + assert _is_local_bootstrap_request(request) is False, client_ip + + +def test_dns_rebinding_host_is_remote_despite_loopback_peer(): + from main import _is_local_bootstrap_request + for host in ("attacker.example", "192.168.1.10", None): + assert _is_local_bootstrap_request(_request("127.0.0.1", host)) is False, host + + +def test_unparseable_request_host_fails_safe(): + """A Host that makes ``request.url.hostname`` raise must fall to remote.""" + from main import _is_local_bootstrap_request + + class _RaisingURL: + @property + def hostname(self): + raise ValueError("malformed host") + + request = SimpleNamespace( + client = SimpleNamespace(host = "127.0.0.1", port = 0), headers = {}, url = _RaisingURL() + ) + assert _is_local_bootstrap_request(request) is False + + +def test_reverse_proxy_forwarded_headers_are_remote(): + """A loopback proxy relaying a remote client (non-Cloudflare headers) is remote.""" + from main import _is_local_bootstrap_request + for header in ("forwarded", "x-forwarded-for", "x-forwarded-host", "x-real-ip"): + request = _request("127.0.0.1", "localhost", headers = {header: "203.0.113.7"}) + assert _is_local_bootstrap_request(request) is False, header + + +def test_malformed_or_absent_host_is_remote(): + """A malformed/absent/scope-id Host must not fall back to the loopback server address.""" + from main import _is_local_bootstrap_request + + # incl. bracket smuggling: [::1]evil / unclosed [::1 must not reduce to ::1 + for host in ( + "e_vil", + "[malformed", + "", + None, + "[::1%25eth0]:8888", + "[::1]attacker", + "[::1]evil.com", + "[::1", + "[::1]x", + ): + assert _is_local_bootstrap_request(_request("127.0.0.1", host)) is False, host + + +def test_colab_allows_notebook_proxy_but_not_shareable_tunnel(monkeypatch): + """Colab autofills its single-user proxy, but not a public Cloudflare link.""" + import main + + monkeypatch.setattr(main, "_IS_COLAB", True) + # In-notebook proxy: same-origin, no tunnel header, injects off-loopback too. + assert main._should_inject_bootstrap(_request("10.0.0.2", "colab.proxy")) is True + # Shareable Cloudflare link marks visitors with cf-connecting-ip; withhold. + tunnel = _request("127.0.0.1", "localhost", headers = {"cf-connecting-ip": "203.0.113.7"}) + assert main._should_inject_bootstrap(tunnel) is False + + +def test_non_colab_gate_requires_local_client(monkeypatch): + """Outside Colab the gate injects only for a direct loopback client.""" + import main + + monkeypatch.setattr(main, "_IS_COLAB", False) + assert main._should_inject_bootstrap(_request("127.0.0.1", "localhost")) is True + assert main._should_inject_bootstrap(_request("192.168.1.10", "localhost")) is False From 8cfd1a21731994a7be0141ce329a8745a7ae16fa Mon Sep 17 00:00:00 2001 From: dylanschroers <60888108+dylanschroers@users.noreply.github.com> Date: Wed, 15 Jul 2026 08:25:03 -0400 Subject: [PATCH 12/17] fix: single-pass GGUF export for directly convertible outtypes in save.py (#7090) * Single-pass GGUF export for direct outtypes + parallel multi-quant save_to_gguf defaulted first_conversion to model_dtype before the block that picks the optimal base conversion, leaving that block dead since it landed (#3356). Every default export (fast_quantized -> q8_0) therefore ran two passes: convert HF -> 16-bit GGUF, then llama-quantize -> q8_0, writing a 2x-size intermediate that the cleanup step deletes again. - Route single-output exports whose type convert_hf_to_gguf.py emits directly (f32/f16/bf16/q8_0) through one conversion pass with no 16-bit intermediate. Measured on Qwen2.5-0.5B-Instruct (8-core CPU): bytes written 1525 MB -> 531 MB (2.9x less), peak extra disk 994 MB -> 0, wall time neutral on local NVMe (14.8s vs 15.4s). The dequantized q8_0 tensors are bit-identical to the two-pass output (max diff 0 over all 290 tensors, same quant-type table). On disk-capped runtimes (Kaggle 20 GB, Colab) the removed intermediate is the difference between an export that fits and one that dies - see the Kaggle error text this file already carries. imatrix runs keep the two-pass route since only llama-quantize can apply one; explicit first_conversion is still honored. - Run independent llama-quantize passes two at a time when several quant methods are requested (thread budget split between workers, outputs byte-identical, order preserved). Measured 1.38x wall-clock on q4_k_m+q5_k_m+q6_k. Sequential under UNSLOTH_ENABLE_LOGGING=1 to keep subprocess logs readable; kill switch UNSLOTH_PARALLEL_GGUF_QUANTS=0. Duplicate methods now quantize once. Co-Authored-By: Claude Fable 5 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Guard parallel GGUF quant on Kaggle and make multi-quant failures atomic Each llama-quantize pass loads the whole model into RAM, so running two at once on Kaggle can OOM a host that succeeded sequentially; skip the parallel path there. On a failed multi-quant export, stop launching queued passes and remove orphaned quant outputs so a failure leaves no partial GGUFs behind, keeping the 16-bit base for retry. Also accept 0/false/no/off/empty for UNSLOTH_PARALLEL_GGUF_QUANTS so a well-meant 'false' actually disables parallelism, and add tests/saving/test_gguf_single_pass_export.py to the CI saving bucket so the new tests run. * Preserve pre-existing outputs for canceled quant passes on failure The parallel cleanup unlinked every requested output name, so a failed rerun could delete a valid model..gguf left by an earlier successful export for a method whose pass was canceled and never ran this session. Skip canceled futures and only remove outputs from passes that actually executed. * Gate parallel GGUF quant on available memory and preserve prior outputs Skip the two-worker path when RAM cannot hold two full-model quantizations at once (and on Colab as well as Kaggle), so a multi-quant export that fit sequentially no longer OOMs. On failure, remove only outputs this run newly created, tracked against a pre-launch snapshot, so a rerun into an existing _gguf directory never deletes a valid artifact from an earlier export. --------- Co-authored-by: djs Co-authored-by: Claude Fable 5 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen --- .github/workflows/consolidated-tests-ci.yml | 2 + tests/saving/test_gguf_single_pass_export.py | 201 +++++++++++++ unsloth/save.py | 283 ++++++++++++------- 3 files changed, 389 insertions(+), 97 deletions(-) create mode 100644 tests/saving/test_gguf_single_pass_export.py diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 1bb4c2bb58..fa84471d36 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -272,6 +272,7 @@ jobs: tests/saving/test_export_api_surface.py \ tests/saving/test_export_dispatch.py \ tests/saving/test_imatrix_export.py \ + tests/saving/test_gguf_single_pass_export.py \ tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ tests/python/test_fast_language_model_text_only.py @@ -361,6 +362,7 @@ jobs: tests/saving/test_export_api_surface.py \ tests/saving/test_export_dispatch.py \ tests/saving/test_imatrix_export.py \ + tests/saving/test_gguf_single_pass_export.py \ tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ tests/python/test_fast_language_model_text_only.py \ diff --git a/tests/saving/test_gguf_single_pass_export.py b/tests/saving/test_gguf_single_pass_export.py new file mode 100644 index 0000000000..296fb0d8ca --- /dev/null +++ b/tests/saving/test_gguf_single_pass_export.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""CPU-only routing tests for the single-pass GGUF export and parallel quantization. + +With convert/quantize monkeypatched, verify save_to_gguf's pass planning: +- a single directly-convertible output type (f32/f16/bf16/q8_0) converts in ONE pass + with no llama-quantize step and no 16-bit intermediate, +- k-quants and imatrix runs keep the two-pass route, +- multiple quantize passes run through the bounded pool with request order preserved, +- quantize failures still raise the actionable RuntimeError. +""" + +from __future__ import annotations + +import contextlib +import os +import threading +import time + +import pytest + +import unsloth.save as save_mod + + +# -- _choose_first_conversion (pure planning logic) ---------------------------------------- + + +@pytest.mark.parametrize( + "methods, model_dtype, expected", + [ + (["q8_0"], "f16", "q8_0"), # default "fast_quantized" path: single pass + (["q8_0", "q8_0"], "bf16", "q8_0"), # duplicates collapse to a single pass + (["f32"], "f16", "f32"), # 16/32-bit outputs convert directly too + (["bf16"], "bf16", "bf16"), + (["q4_k_m"], "f16", "f16"), # k-quants need a 16-bit base + (["q4_k_m", "q8_0"], "bf16", "bf16"), # mixes need the shared base + (["q8_0", "f16"], "f16", "f16"), + ], +) +def test_choose_first_conversion(methods, model_dtype, expected): + assert save_mod._choose_first_conversion(methods, model_dtype) == expected + + +def test_choose_first_conversion_imatrix_forces_two_pass(): + # Only llama-quantize can apply an imatrix, so q8_0-only must keep the 16-bit base. + assert save_mod._choose_first_conversion(["q8_0"], "f16", has_imatrix = True) == "f16" + + +# -- save_to_gguf pass planning (mocked convert/quantize) ----------------------------------- + + +class _Harness: + """Monkeypatched convert/quantize recording calls and creating real files.""" + + def __init__( + self, + monkeypatch, + tmp_path, + quantize_delays = None, + quantize_error = None, + ): + self.tmp_path = tmp_path + self.convert_calls = [] + self.quantize_calls = [] + self.active = 0 + self.max_concurrency = 0 + self._lock = threading.Lock() + self._delays = quantize_delays or {} + self._error = quantize_error + + monkeypatch.setattr(save_mod, "check_llama_cpp", lambda: ("llama-quantize", "convert.py")) + monkeypatch.setattr( + save_mod, + "_download_convert_hf_to_gguf", + lambda: (str(tmp_path / "convert.py"), {"LlamaForCausalLM"}, set()), + ) + monkeypatch.setattr(save_mod, "use_local_gguf", contextlib.nullcontext) + monkeypatch.setattr(save_mod, "convert_to_gguf", self._convert) + monkeypatch.setattr(save_mod, "quantize_gguf", self._quantize) + + def _convert(self, **kwargs): + self.convert_calls.append(kwargs) + suffix = kwargs["quantization_type"] + if suffix == "None": + suffix = kwargs["model_dtype"] + out = self.tmp_path / f"{kwargs['model_name']}.{suffix.upper()}.gguf" + out.write_bytes(b"GGUF") + return [str(out)], False + + def _quantize( + self, + input_gguf, + output_gguf, + quant_type, + imatrix = None, + n_threads = None, + **kw, + ): + with self._lock: + self.active += 1 + self.max_concurrency = max(self.max_concurrency, self.active) + try: + if self._error is not None: + raise self._error + time.sleep(self._delays.get(quant_type, 0.02)) + self.quantize_calls.append({"quant_type": quant_type, "n_threads": n_threads}) + with open(output_gguf, "wb") as f: + f.write(b"GGUF") + return output_gguf + finally: + with self._lock: + self.active -= 1 + + +def _run(tmp_path, methods, **kwargs): + model_dir = tmp_path / "model_dir" + model_dir.mkdir(exist_ok = True) + return save_mod.save_to_gguf( + model_name = "testmodel", + model_type = "llama", + model_dtype = "float16", + model_directory = str(model_dir), + quantization_method = methods, + **kwargs, + ) + + +def test_q8_0_only_is_single_pass(monkeypatch, tmp_path): + h = _Harness(monkeypatch, tmp_path) + locations, want_full_precision, _ = _run(tmp_path, ["q8_0"]) + + assert len(h.convert_calls) == 1 + assert h.convert_calls[0]["quantization_type"] == "q8_0" + assert h.quantize_calls == [], "single-pass export must not launch llama-quantize" + assert want_full_precision is True, "the converted file IS the requested output" + assert len(locations) == 1 and locations[0].endswith("testmodel.Q8_0.gguf") + assert os.path.exists(locations[0]) + + +def test_fast_quantized_alias_is_single_pass(monkeypatch, tmp_path): + h = _Harness(monkeypatch, tmp_path) + _run(tmp_path, "fast_quantized") # the default of save_pretrained_gguf + assert h.convert_calls[0]["quantization_type"] == "q8_0" + assert h.quantize_calls == [] + + +def test_k_quant_keeps_two_pass(monkeypatch, tmp_path): + h = _Harness(monkeypatch, tmp_path) + locations, want_full_precision, _ = _run(tmp_path, ["q4_k_m"]) + + assert h.convert_calls[0]["quantization_type"] == "f16" + assert [c["quant_type"] for c in h.quantize_calls] == ["q4_k_m"] + assert want_full_precision is False + # The 16-bit intermediate must be cleaned up. + assert len(locations) == 1 and locations[0].endswith("testmodel.Q4_K_M.gguf") + + +def test_mixed_methods_share_16bit_base(monkeypatch, tmp_path): + h = _Harness(monkeypatch, tmp_path) + _run(tmp_path, ["q4_k_m", "q8_0"]) + assert h.convert_calls[0]["quantization_type"] == "f16" + assert sorted(c["quant_type"] for c in h.quantize_calls) == ["q4_k_m", "q8_0"] + + +def test_parallel_quants_preserve_request_order(monkeypatch, tmp_path): + # First method is the slowest: completion order != request order. + h = _Harness( + monkeypatch, tmp_path, quantize_delays = {"q4_k_m": 0.3, "q5_k_m": 0.05, "q6_k": 0.01} + ) + locations, _, _ = _run(tmp_path, ["q4_k_m", "q5_k_m", "q6_k"]) + + assert h.max_concurrency == 2, "quantize passes should overlap, bounded at 2" + quant_names = [os.path.basename(l) for l in locations if "F16" not in l] + assert quant_names == [ + "testmodel.Q6_K.gguf", # list is reversed by the cleanup block, as before + "testmodel.Q5_K_M.gguf", + "testmodel.Q4_K_M.gguf", + ] + assert all( + c["n_threads"] is not None for c in h.quantize_calls + ), "parallel workers must split the thread budget explicitly" + + +def test_parallel_quants_env_kill_switch(monkeypatch, tmp_path): + monkeypatch.setenv("UNSLOTH_PARALLEL_GGUF_QUANTS", "0") + h = _Harness(monkeypatch, tmp_path, quantize_delays = {"q4_k_m": 0.05, "q5_k_m": 0.05}) + _run(tmp_path, ["q4_k_m", "q5_k_m"]) + assert h.max_concurrency == 1 + + +def test_duplicate_methods_quantize_once(monkeypatch, tmp_path): + h = _Harness(monkeypatch, tmp_path) + _run(tmp_path, ["q4_k_m", "q4_k_m"]) + assert [c["quant_type"] for c in h.quantize_calls] == ["q4_k_m"] + + +def test_quantize_failure_raises_actionable_error(monkeypatch, tmp_path): + h = _Harness(monkeypatch, tmp_path, quantize_error = OSError("disk full")) + with pytest.raises(RuntimeError, match = "Quantization failed"): + _run(tmp_path, ["q4_k_m", "q5_k_m"]) diff --git a/unsloth/save.py b/unsloth/save.py index 0b408a1889..50a180f823 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -1734,6 +1734,33 @@ def get_executable(executables): return None +# Output types convert_hf_to_gguf.py can emit directly via --outtype. +_DIRECT_CONVERT_OUTTYPES = ("f32", "f16", "bf16", "q8_0") + + +def _choose_first_conversion( + quantization_methods, + model_dtype, + has_imatrix = False, +): + """Pick the dtype of the initial HF -> GGUF conversion. + + Single-pass fast path: when exactly one output type is requested and + convert_hf_to_gguf.py can emit it directly (f32/f16/bf16/q8_0), convert straight to + it - the llama-quantize pass and the 16-bit intermediate file are skipped entirely. + An imatrix forces the two-pass route since only llama-quantize can apply one. + + Every other case converts to the source dtype first, so each requested method is + quantized from weights identical to the checkpoint's. + """ + unique_methods = set(quantization_methods) + if len(unique_methods) == 1 and not has_imatrix: + only_method = next(iter(unique_methods)) + if only_method in _DIRECT_CONVERT_OUTTYPES: + return only_method + return model_dtype + + def save_to_gguf( model_name: str, model_type: str, @@ -1782,10 +1809,6 @@ def save_to_gguf( ) model_dtype = "f16" - # Check first_conversion as well - if first_conversion is None: - first_conversion = model_dtype - has_imatrix = imatrix is not None and str(imatrix) != "" if has_imatrix: # quantize_gguf gained the imatrix kwarg in a recent unsloth_zoo; fail fast (before the @@ -1834,32 +1857,12 @@ def save_to_gguf( first_conversion = "None" # No quantization for GPT-OSS # Only keep one conversion method since GPT-OSS doesn't quantize quantization_method = ["None"] - else: - if first_conversion is None: - # Check if q8_0 is the ONLY quantization method requested - if len(quantization_method) == 1 and quantization_method[0] == "q8_0": - first_conversion = "None" # Let llama-quantize do the direct conversion - else: - # For all other cases, choose the highest precision format - # that can be requantized to all requested formats - strength = 0 - for quant_method in quantization_method: - if quant_method == "f32": - strength = max(strength, 3) - elif quant_method == "f16": - strength = max(strength, 2) - elif quant_method == "bf16": - strength = max(strength, 1) - # Note: we don't set strength for q8_0 here since we handle it above - - if strength >= 3: - first_conversion = "f32" - elif strength >= 2: - first_conversion = "f16" - elif strength >= 1: - first_conversion = "bf16" - else: - first_conversion = "bf16" # requantizing from q8_0 disallowed in new llama.cpp default to bf16. + elif first_conversion is None: + first_conversion = _choose_first_conversion( + quantization_method, + model_dtype, + has_imatrix = has_imatrix, + ) # Check bfloat16 support again for first_conversion if first_conversion == "bf16" and not torch.cuda.is_bf16_supported(): @@ -1868,12 +1871,19 @@ def save_to_gguf( first_conversion_dtype = "" if first_conversion == "None" else first_conversion # Print conversion info + needs_quantize_pass = any(m != first_conversion for m in quantization_method) + if needs_quantize_pass: + second_step = f"[2] Converting GGUF {first_conversion_dtype} to {quantization_method} might take 10 minutes each." + total_line = "In total, you will have to wait at least 16 minutes." + else: + second_step = f"[2] Single-pass export: converting straight to {quantization_method} - no separate quantize step." + total_line = "In total, you will have to wait at least 6 minutes." print_info = ( f"==((====))== Unsloth: Conversion from HF to GGUF information\n" f" {chr(92)}{chr(92)} /| [0] Installing llama.cpp might take 3 minutes.\n" f"O^O/ {chr(92)}_/ {chr(92)} [1] Converting HF to GGUF {first_conversion_dtype} might take 3 minutes.\n" - f"{chr(92)} / [2] Converting GGUF {first_conversion_dtype} to {quantization_method} might take 10 minutes each.\n" - f' "-____-" In total, you will have to wait at least 16 minutes.\n' + f"{chr(92)} / {second_step}\n" + f' "-____-" {total_line}\n' ) print(print_info) @@ -1959,75 +1969,154 @@ def save_to_gguf( if not is_gpt_oss: base_gguf = initial_files[0] - quants_created = False - for quant_method in quantization_method: - if quant_method != first_conversion: + + # Deduplicate while keeping order; methods equal to the base conversion already + # exist on disk and need no quantize pass. + methods_to_quantize = [ + m for m in dict.fromkeys(quantization_method) if m != first_conversion + ] + + def _quantize_one(quant_method, n_threads = None): + output_location = os.path.join( + gguf_directory, f"{model_name}.{quant_method.upper()}.gguf" + ) + try: + if quant_method == "q2_k_l": + return _quantize_q2_k_l( + input_gguf = base_gguf, + output_gguf = output_location, + quantizer_location = quantizer_location, + n_threads = n_threads if n_threads is not None else n_cpus, + print_output = print_output, + imatrix = imatrix, + ) + else: + # Use unsloth-zoo's standard quantization for all other methods. Only pass + # imatrix when set so older unsloth_zoo (no imatrix kwarg) still works for + # plain quants; an imatrix that cannot be applied was rejected above. + quant_kwargs = dict( + input_gguf = base_gguf, + output_gguf = output_location, + quant_type = quant_method, + quantizer_location = quantizer_location, + print_output = print_output, + ) + if has_imatrix: + quant_kwargs["imatrix"] = imatrix + if n_threads is not None: + quant_kwargs["n_threads"] = n_threads + return quantize_gguf(**quant_kwargs) + except Exception as e: + if IS_KAGGLE_ENVIRONMENT: + raise RuntimeError( + f"Unsloth: Quantization failed for {output_location}\n" + "You are in a Kaggle environment, which might be the reason this is failing.\n" + "Kaggle only provides 20GB of disk space in the working directory.\n" + "Merging to 16bit for 7b models use 16GB of space.\n" + "This means using `model.{save_pretrained/push_to_hub}_merged` works, but\n" + "`model.{save_pretrained/push_to_hub}_gguf will use too much disk space.\n" + "You can try saving it to the `/tmp` directory for larger disk space.\n" + "I suggest you to save the 16bit model first, then use manual llama.cpp conversion.\n" + f"Error: {e}" + ) + else: + if IS_WINDOWS: + build_instructions = ( + f'cd "{LLAMA_CPP_DEFAULT_DIR}"\n' + f"cmake -S . -B build -DBUILD_SHARED_LIBS=OFF\n" + f"cmake --build build --config Release" + ) + else: + build_instructions = ( + f'cd "{LLAMA_CPP_DEFAULT_DIR}" && make clean && make all -j' + ) + + raise RuntimeError( + f"Unsloth: Quantization failed for {output_location}\n" + "You might have to compile llama.cpp yourself, then run this again.\n" + "You do not need to close this Python program. Run the following commands in a new terminal:\n" + f'git clone --recursive https://github.com/ggerganov/llama.cpp "{LLAMA_CPP_DEFAULT_DIR}"\n' + f"{build_instructions}\n" + "Once that's done, redo the quantization.\n" + f"Error: {e}" + ) + + # Outputs already on disk pre-date this run; never delete them on a failure. + preexisting_outputs = { + m + for m in methods_to_quantize + if os.path.exists(os.path.join(gguf_directory, f"{model_name}.{m.upper()}.gguf")) + } + # Each llama-quantize pass loads the whole base GGUF into RAM, so only run two at + # once when the host has headroom for two copies, else a multi-quant export that + # fit sequentially could OOM. + try: + base_bytes = sum( + os.path.getsize(f) + for f in initial_files + if "-mmproj" not in os.path.basename(f).lower() + ) + mem_ok = psutil.virtual_memory().available >= int(2.5 * base_bytes) + except Exception: + mem_ok = False + # Independent llama-quantize runs on the same base GGUF can overlap. Kept at 2 + # workers; run sequentially when streaming logs (UNSLOTH_ENABLE_LOGGING), on + # Kaggle/Colab, when RAM is tight, or when the kill switch (0/false/no/off/empty) + # is set. + _parallel_flag = os.environ.get("UNSLOTH_PARALLEL_GGUF_QUANTS", "1").strip().lower() + parallel_quants = ( + len(methods_to_quantize) > 1 + and not print_output + and not IS_KAGGLE_ENVIRONMENT + and not IS_COLAB_ENVIRONMENT + and mem_ok + and _parallel_flag not in ("0", "false", "no", "off", "") + ) + if parallel_quants: + max_workers = min(2, len(methods_to_quantize)) + # Split the thread budget so total threads match the sequential run. + per_worker_threads = max(1, n_cpus // max_workers) + print( + f"Unsloth: [2] Converting GGUF {first_conversion_dtype} into " + f"{methods_to_quantize}, {max_workers} at a time. This might take 10 minutes each..." + ) + from concurrent.futures import ThreadPoolExecutor, wait, FIRST_EXCEPTION + + quantized_files = [None] * len(methods_to_quantize) + with ThreadPoolExecutor(max_workers = max_workers) as pool: + future_to_idx = { + pool.submit(_quantize_one, method, per_worker_threads): i + for i, method in enumerate(methods_to_quantize) + } + done, pending = wait(future_to_idx, return_when = FIRST_EXCEPTION) + # Do not start queued passes after a failure (avoid filling the disk). + for fut in pending: + fut.cancel() + first_exc = next((f.exception() for f in done if f.exception() is not None), None) + if first_exc is not None: + # Remove only outputs this run newly created; a file that pre-dated the + # run (or a canceled pass that never wrote) is left intact, so a rerun + # never deletes a prior artifact. Base kept for retry. + wait(future_to_idx) + for method in methods_to_quantize: + if method in preexisting_outputs: + continue + Path( + os.path.join(gguf_directory, f"{model_name}.{method.upper()}.gguf") + ).unlink(missing_ok = True) + raise first_exc + for fut, i in future_to_idx.items(): + quantized_files[i] = fut.result() + else: + quantized_files = [] + for quant_method in methods_to_quantize: print( f"Unsloth: [2] Converting GGUF {first_conversion_dtype} into {quant_method}. This might take 10 minutes..." ) - output_location = os.path.join( - gguf_directory, f"{model_name}.{quant_method.upper()}.gguf" - ) - try: - if quant_method == "q2_k_l": - quantized_file = _quantize_q2_k_l( - input_gguf = base_gguf, - output_gguf = output_location, - quantizer_location = quantizer_location, - n_threads = n_cpus, - print_output = print_output, - imatrix = imatrix, - ) - else: - # Use unsloth-zoo's standard quantization for all other methods. Only pass - # imatrix when set so older unsloth_zoo (no imatrix kwarg) still works for - # plain quants; an imatrix that cannot be applied was rejected above. - quant_kwargs = dict( - input_gguf = base_gguf, - output_gguf = output_location, - quant_type = quant_method, - quantizer_location = quantizer_location, - print_output = print_output, - ) - if has_imatrix: - quant_kwargs["imatrix"] = imatrix - quantized_file = quantize_gguf(**quant_kwargs) - all_saved_locations.append(quantized_file) - quants_created = True - except Exception as e: - if IS_KAGGLE_ENVIRONMENT: - raise RuntimeError( - f"Unsloth: Quantization failed for {output_location}\n" - "You are in a Kaggle environment, which might be the reason this is failing.\n" - "Kaggle only provides 20GB of disk space in the working directory.\n" - "Merging to 16bit for 7b models use 16GB of space.\n" - "This means using `model.{save_pretrained/push_to_hub}_merged` works, but\n" - "`model.{save_pretrained/push_to_hub}_gguf will use too much disk space.\n" - "You can try saving it to the `/tmp` directory for larger disk space.\n" - "I suggest you to save the 16bit model first, then use manual llama.cpp conversion.\n" - f"Error: {e}" - ) - else: - if IS_WINDOWS: - build_instructions = ( - f'cd "{LLAMA_CPP_DEFAULT_DIR}"\n' - f"cmake -S . -B build -DBUILD_SHARED_LIBS=OFF\n" - f"cmake --build build --config Release" - ) - else: - build_instructions = ( - f'cd "{LLAMA_CPP_DEFAULT_DIR}" && make clean && make all -j' - ) + quantized_files.append(_quantize_one(quant_method)) - raise RuntimeError( - f"Unsloth: Quantization failed for {output_location}\n" - "You might have to compile llama.cpp yourself, then run this again.\n" - "You do not need to close this Python program. Run the following commands in a new terminal:\n" - f'git clone --recursive https://github.com/ggerganov/llama.cpp "{LLAMA_CPP_DEFAULT_DIR}"\n' - f"{build_instructions}\n" - "Once that's done, redo the quantization.\n" - f"Error: {e}" - ) + all_saved_locations.extend(quantized_files) + quants_created = len(quantized_files) > 0 print("Unsloth: Model files cleanup...") want_full_precision = first_conversion in quantization_method if quants_created: From 815f24297069c9c90ab617cb74bb191b93e56d0f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 15 Jul 2026 05:25:26 -0700 Subject: [PATCH 13/17] Studio: offer the latest transformers release for brand-new architectures (#7056) * Studio: offer the latest transformers release for brand-new architectures When a model's config.json model_type is absent from every installed transformers overlay (base 4.57.x and the .venv_t5_530/550/510 sidecars), Studio now checks, unauthenticated and cached, whether the newest transformers ships it: - utils/transformers_latest.py fetches the latest release version from https://pypi.org/pypi/transformers/json and the CONFIG_MAPPING_NAMES sources for that tag and for main from raw.githubusercontent.com (never api.github.com), parsing them with the same AST extractor the static router uses (no code execution, no trust_remote_code). Results are cached in memory and in a JSON snapshot under studio_root()/cache with a one day ttl; fetches are bounded to 5s with one retry and a failure backoff, and offline mode or the new kill switch UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS=1 short-circuits to None. - POST /api/inference/validate gains requires_transformers_upgrade plus a transformers_upgrade payload (model_type, pypi_version, supported_in_pypi, supported_in_main) so the frontend can raise the install consent dialog before /load, mirroring the existing remote-code consent flow. The check fires only when the model_type is unknown to all installed overlays and the hardcoded tier tables. - POST /api/inference/install-latest-transformers provisions a new persistent .venv_t5_latest sidecar after user consent, pinned to the exact PyPI version (re-verified server-side) with the same --target/--no-deps recipe as the fixed sidecars. A JSON pin marker inside the dir records the installed package set, so restarts revalidate it and routing resolves the new highest-ranked tier automatically. A dependency preflight (compat_plan) compares the release's requires_dist against the running env: unsatisfied tokenizers/safetensors floors are shadow-installed as exact pins into the sidecar, anything else unsatisfied blocks the install with a clear message. Routing for every already-supported model_type is unchanged: the hardcoded lists and the 530/550/510 static resolver run first, the new tier only participates once its venv exists, and the probe order gains the latest sidecar only when provisioned. Verified against live PyPI and GitHub (transformers 5.13.0: 674 model_types, 26 absent from all installed overlays, e.g. cosmos3_omni; 4 dev-only on main) and with a real sidecar install plus restart persistence. 64 new tests; the existing 200-test transformers_version suite passes unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Latest-transformers check: fetch outside the lock, serialize installs Release the module lock during the network refresh so a slow fetch cannot stall other threads in the ASGI pool; concurrent callers during a fetch get None (the graceful fallthrough) via an in-flight flag instead of stacking fetches. Serialize install_latest_transformers with an in-progress flag so concurrent consents cannot race the sidecar delete and recreate; the loser gets a structured already-in-progress refusal. * Latest-transformers check: LoRA bases, pin-gated mapping, live reverify Run the upgrade check over the [adapter, base] target set so a LoRA whose base model is a brand-new architecture surfaces the prompt (the worker activates transformers for the base, not the adapter). Gate the latest overlay's mapping lookup on a valid pin marker, matching activation and the probe order, so a partial or manual .venv_t5_latest dir cannot be routed to and then refused at activation. Re-verify the requested version against a live PyPI snapshot at install time, falling back to the cached one on fetch failure, so a release published inside the cache TTL is not silently missed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Latest-transformers check: nested config types and latest-tier vision probe Collect every model_type in the config (top level plus each nested sub-config) and signal on the first one missing from all installed overlays, so a supported wrapper carrying a brand-new backbone still surfaces the upgrade prompt; wrappers instantiate sub-configs through CONFIG_MAPPING and would fail on the nested type. Route the vision capability subprocess through the pinned latest sidecar when the model resolves to the latest tier, so latest-only VLMs are not misclassified as text-only; every other tier keeps the 5.5 sidecar used today. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Latest tier: nested routing, vision probe after raw miss, safe upgrades Route by every model_type in the config: a nested sub-config type can raise the tier (wrappers instantiate sub-configs through CONFIG_MAPPING), so a supported wrapper with a latest-only backbone routes to latest once installed instead of staying on default. An unknown nested type never vetoes; the primary type keeps its previous semantics. The collector is shared with the upgrade checker. Vision detection: when the raw heuristics say False for a model that routes to the latest tier, run the AutoConfig subprocess under the pinned latest sidecar instead of trusting heuristics built from older transformers. Provisioning: stage-and-swap. Build the new sidecar in .venv_t5_latest.staging and swap it in only when the install and pin marker are complete, so a failed upgrade never destroys a previously working sidecar; restore the old dir if the final swap fails. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Latest-transformers checker, vision subprocess, and cache fixes Require the latest release to support every missing model_type (the primary included) before prompting; a nested-only match cannot make the model loadable, so no install is offered for it. The vision-check subprocess now unions the active sidecar's own registry mappings into the inlined parent-process detection sets, so architectures only the sidecar knows classify correctly. A successful sidecar install clears the tier probe cache, the latest tier's model_type mapping, and the vision-detection cache so the new venv takes effect without a restart. Tests for all three. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Aggregate upgrade support flags and keep install off /v1 The upgrade signal now reports supported_in_pypi only when the latest release covers every missing model_type; a mix with a main-only nested type surfaces as dev-only so no PyPI install is offered that would still fail at load. The consented install endpoint moves to studio_router so it is not reachable through the OpenAI-compatible /v1 mount. Tests for both. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Honor the latest-transformers kill switch in routing With UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS set after the sidecar was provisioned, the latest tier still joined mapping and probe routing because only the pin was checked. Both admission points now also check the kill switch, so operators can roll back a problematic sidecar without deleting files. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Repair the latest sidecar through stage-and-swap The lazy repair path installed into the live .venv_t5_latest, which _ensure_venv_dir wipes first, so a failed repair deleted the pinned sidecar and its marker. Both the consented install and the repair now share one stage-and-swap helper: the incomplete-but-pinned dir survives any failure and a later attempt can still repair it. * Tighten comments * Remove the staging dir when a latest-sidecar install fails A pip failure inside _ensure_venv_dir returns False without raising, so the except cleanup never ran and the partial .venv_t5_latest.staging leaked until a later attempt. Also note on the validate response fields that frontend consumption ships in the follow-up PR. * Add the transformers-upgrade consent dialog to the frontend When /validate reports requires_transformers_upgrade, every explicit load path (chat runtime and the compare composer) now pauses on a consent dialog modeled on the remote-code one: it names the model_type and the latest PyPI transformers version, and on Accept calls /api/inference/install-latest-transformers itself, shows an installing state, and resumes the original load automatically on success. Errors surface in the dialog with a retry; Cancel aborts the load like the trust dialog's deny path. Architectures shipped only on transformers main get a dev-only notice with no install button. Background auto-load skips upgrade-requiring candidates instead of prompting, mirroring the trust_remote_code rule. The dialog mounts once in the root layout and runs before the security dialogs, since no load can proceed without the runtime. * Route a non-installable new architecture to the custom-code consent as a last resort When the upgrade dialog has no installable PyPI release (the architecture is only on transformers main, which Studio never installs), the dialog now says so explicitly, and when the model also declares custom (auto_map) code it offers Continue with custom code: resolving the paused load into the existing trust_remote_code consent gate instead of hard-aborting. Models with no custom code keep the Cancel-only notice. The backend returns no upgrade signal at all for architectures unknown to both PyPI and main, so those still route straight to the unchanged security gate. * Force a 16-bit load for models on the latest-transformers sidecar Live validation with Zyphra/ZAYA1-8B (model_type zaya, shipped by transformers 5.13.1 but unknown to every installed tier) surfaced a generation crash when the consented sidecar load kept the default bnb 4-bit quantization: transformers' grouped-MoE kernels feed the packed uint8 expert weights straight into torch._grouped_mm, and generation dies (plain 16-bit works). New latest_tier_active_for() mirrors the sidecar activation's tier resolution and never raises; the inference worker flips load_in_4bit off when it reports true, and the load route applies the same flip so the pre-load VRAM guard and the worker command agree. Fixed tiers are untouched. With the guard, ZAYA1-8B loads and generates correctly in Studio chat. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Offer the custom-code fallback when a latest-sidecar install fails * Fail remote mapping fetches wholesale and mirror the 16-bit flip in validate A transient fetch or parse failure of one auto-mapping file no longer caches a partial latest-release map for the TTL (a real 404 on pre-5.10 tags is still tolerated), and validate_model now applies the same latest-sidecar 16-bit sizing flip as /load before the training guard so the two agree. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten comments in the latest-transformers changes * Resolve remote LoRA bases, fold nested tiers, and guard the sidecar swap latest_tier_active_for now resolves a remote adapter's base model the same way worker pre-activation does (and returns early without a sidecar pin), a hardcoded fast-path tier is raised when a nested sub-config's model_type needs a higher sidecar, and the install route refuses to swap .venv_t5_latest while training runs on it and unloads a latest-tier chat model first. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Gate the sidecar install on worker liveness and size installable upgrades 16-bit The install route now refuses while any training or export runs (tier re-resolution without the load token is unreliable for gated repos), holds the inference lifecycle gate across the unload and the swap so no load can interleave, and passes the model name to unload_model. validate_model runs the upgrade check before the training guard and sizes an installable upgrade as 16-bit, matching what /load and the worker will force after the consented install. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Close the sidecar install races and honor the kill switch over cached mappings Training starts and mutating export routes now refuse while a transformers install is in progress (shared is_install_in_progress flag), the chat unload and idle export-worker teardown moved into a before_swap hook that runs only once the staged install succeeded, and _config_model_types checks the kill switch before returning a cached latest mapping. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Reserve the sidecar swap before the gate wait and abort it on failed teardown The install-in-progress flag moved into a shared sidecar swap reservation in transformers_version, taken by the install route before awaiting the inference lifecycle gate (so training and export starts see it for the whole window) and by the lazy .venv_t5_latest repair path. The before_swap hook now raises when the chat unload or export teardown reports failure, leaving the previous sidecar untouched. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Back the sidecar swap reservation with a cross-process lock file The lazy repair runs inside worker subprocesses, where a module-level flag is invisible to the parent's route checks. The reservation now also creates a lock file next to .venv_t5_latest (O_EXCL, owner-only removal, stale after two hours for crashed owners), so is_install_in_progress sees a repair from any Studio process. * Hand the swap reservation to the installer thread and harden pre-swap teardown A cancelled install request no longer releases the reservation while the installer thread is still staging (the thread owns and releases it, shielded from cancellation). The route refuses while another inference request is generating, export teardown runs before the chat unload and is judged by worker liveness rather than the cleanup return value, and a live inference worker with no active model (failed load residue) is shut down before the swap. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep the lifecycle gate with the installer and recheck the swap at spawn time The gate moved into the shielded install task so a cancelled POST cannot release the guard /load honors while the installer still runs, cached latest probe results are ignored while the kill switch is set, and the training and export subprocess spawns recheck the sidecar swap reservation right before spawning (the route-level guards are one-shot and validation can outlast an install's start). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Close the spawn-registration windows against the sidecar install Training marks the spawn in progress before its reservation recheck and is_training_active honors the flag, so the install route sees a start that has passed proc.start() but not yet recorded _proc. Export load-checkpoint rechecks the reservation after setting _export_active and before tearing down the old worker, so losing the race keeps the loaded checkpoint instead of surfacing a 500. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Refine the install-window interleavings around worker teardown The inference busy count is rechecked under the lifecycle gate (streams start by taking that gate, so nothing slips past a held gate), the training handshake moved ahead of the VRAM-freeing before_spawn hook so a lost race leaves chat/export intact, the export spawn-time check is op-aware (inside an active op the install is the side that aborts), and the Xet-stall respawn waits out a transient reservation instead of stranding the run. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Track the install's server-side unload and guard export ops against the swap The upgrade dialog store records when its install actually ran (the server unloads the active chat model before swapping), and the load flow then marks the previous model as unloaded so a later cancelled gate still triggers rollback; the custom-code fallback leaves the flag unset. _run_export gained the same reservation handshake as load_checkpoint so an install cannot block behind an hours-long export op instead of returning 409. * Tighten comments in the install-guard and upgrade-consent changes * Surface install-race refusals cleanly and roll back after a failed swap unload /load refuses while the sidecar swap is reserved so a load cannot succeed and immediately be unloaded by the pre-swap teardown, worker starts that lose the install race raise a typed SidecarSwapInProgress mapped to 409 instead of a 500, the install response reports model_unloaded even on a structured failure so the client can restore its state, and the compare flow tracks the server-side unload like the primary load path and clears a stale checkpoint on abort. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Type the export install races, scope the lock release, and keep the unload signal Export load-checkpoint and export ops raise SidecarSwapInProgress (mapped to 409 in every export route) instead of a 400-shaped failure, the export spawn check distinguishes repair reservations (always refused) from install ones (op-aware), the swap lock release only unlinks a lock this process wrote so a stale-superseded owner cannot drop the new owner's live lock, and the frontend unload signal survives a superseding consent via read-and-clear consumption instead of a reset. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Finalize a stalled run when the respawn loses the install race and latch the unload signal The Xet-stall respawn timeout now finalizes the run as a failure instead of raising into the pump's broad finalization catch (which stranded it in a training state with no worker), and a successful install retry ORs the model_unloaded signal with the latched value so a failed-after-unload first attempt still triggers rollback. * Recheck the swap under the load gate and latch the unload before resolver checks /load rechecks the sidecar reservation after acquiring the lifecycle gate (an install can reserve while the load queues on it), and the dialog store latches model_unloaded as soon as the install response arrives, before any resolver-identity guard, so a superseded consent's unload still reaches whichever load consumes the signal next. * Report cleared-state unload failures, guard queued installs, and fold name tiers A failed chat unload that still cleared the orchestrator's model state now reports model_unloaded so the client rolls back, the installer aborts with a 409 when a model load completed while it waited on the lifecycle gate, and the fixed-tier name fast path consults the config mapping when a latest sidecar is pinned so an accepted upgrade routes to the sidecar it installed (no I/O added to the unpinned path). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Report cleared-state unload failures and harden the spawn handshake flag The failed-unload branch in before_swap now detects that the orchestrator cleared its model state and reports model_unloaded before aborting (the earlier commit claimed this fix but a scripting error dropped the edit), the installer's queued-load check compares a load generation counter so a same-model reload is caught, and both training spawn sites wrap everything after the handshake in a guard that resets _spawn_in_progress on any exception so a failed start cannot wedge is_training_active. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Bump the load generation when the load is published, not at load start A start-time bump is already visible when the installer snapshots mid-load, so a same-model reload completing after the snapshot looked unchanged and could be unloaded by the swap. The counter now increments alongside the active_model_name publish. * Self-heal a broken pinned sidecar, guard lazy repairs, and refresh stale retries A valid pin whose transformers source dir vanished now triggers the repair from the routing path (with a five minute backoff after failures) instead of silently routing latest-only models to older tiers, the lazy repair refuses while parent-visible chat/training/export workers are active since it has no teardown of its own, and a version-mismatch install failure carries the superseding release so the dialog's Retry re-requests a version that can succeed. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Flip latest-tier loads to 16-bit outside chat and protect export state Training and export workers now apply the same latest-sidecar 16-bit flip as the chat worker so a brand-new grouped-MoE architecture cannot reach bnb 4-bit through those paths, the latest-tier vision override returns None on an inconclusive probe so a transient failure is not cached as not-vision, and the install route refuses while an idle export checkpoint is loaded rather than discard it with no rollback signal on a failed swap. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address parallel-review findings on the sidecar guards and install checks The training route sizes latest-tier jobs 16-bit before GPU selection, the inference subprocess spawn rechecks the swap reservation like training and export (covering the OpenAI auto-switch path) with the typed error mapped to a retryable 409, compat_plan blocks the install when dependency metadata cannot be fetched instead of proceeding unverified, snapshot model-type lists must contain only strings, and pin-marker package specs are validated against the sidecar's own package set before ever reaching pip. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Parent-only repairs, live-owner locks, remote-base activation, pre-teardown recheck Lazy sidecar repairs now refuse inside worker children (whose empty backend singletons cannot see live siblings) and run only in the parent where the active-worker guard is real, swap-lock staleness requires the owner pid to be dead so a slow live install is never superseded, both activation entry points resolve a remote adapter's base model like the inference worker and latest_tier_active_for already do, and load_model rechecks the reservation before tearing down the old worker so losing the race keeps the current model loaded. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Check workers under the repair reservation and keep state on refused swaps The lazy repair now reserves first and checks workers under the reservation (worker starts set their active markers before rechecking, so every interleaving aborts one side), with export ops and in-flight inference loads counted as active. The inference pre-teardown and spawn guards refuse only repair reservations since an install shares the load's lifecycle gate and aborts via its queued-load snapshot, a SidecarSwapInProgress raised before teardown no longer clears the live model mirrors, and an export spawn abort after teardown clears current_checkpoint so the page cannot claim a loaded checkpoint with no worker. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Repair a present-but-incomplete latest sidecar from routing The routing self-heal only fired when the pinned sidecar's transformers/ dir was missing. A sidecar that kept transformers/ but lost another pinned package still routed models to the latest tier, and workers refuse parent-only repairs, so every load failed until a manual reinstall. Routing now validates the full pin (via _venv_dir_is_valid) and repairs any incomplete sidecar under the same swap reservation and 5-minute backoff. * Treat an unrepaired latest sidecar as unavailable in routing When the pinned sidecar is incomplete and the lazy repair fails (offline, pip failure, workers active) or is inside the backoff window, routing returned the source dir anyway, sending models to a tier whose worker activation is known to fail. Return None instead so models an older tier supports keep loading there until a repair succeeds, matching the behavior when the sidecar dir is missing entirely. * Harden sidecar swap and repair against crash, survivor, and 16-bit paths Reclaim a swap lock as soon as its recorded owner PID is dead instead of waiting out the two-hour cutoff, so a crash mid-install no longer wedges /load, training, export, and repair for hours. A lock whose PID cannot be read yet still uses the long cutoff so the create-before-write window is never mistaken for dead. Probe process liveness with OpenProcess on Windows: os.kill(pid, 0) there is CTRL_C_EVENT (a real Ctrl+C via GenerateConsoleCtrlEvent), not a harmless check, and psutil is not always present. Return whether _shutdown_subprocess actually killed the worker and keep the live handle when it survives terminate/kill (an uninterruptible CUDA syscall can outlive SIGKILL). The pre-swap liveness guard now trusts that result, so the destructive .venv_t5_latest rename cannot proceed while a live worker still holds sidecar modules. Recover a sidecar stranded at .old when a swap's activation rename and its rollback both fail: reading the pin restores it when no swap holds the reservation, so latest-tier models are not permanently broken. Resolve the latest tier in the parent for export loads and for explicitly 16-bit training runs, not only 4-bit ones: tier resolution self-heals an incomplete sidecar, and repairs are parent-only, so those paths could not recover before. Sidecar integrity and quantization are independent. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Revert the parent-side latest-tier repair probe on training and export loads The probe ran before the route freed VRAM, so a resident chat or export worker made _workers_active_for_repair() refuse the parent-only repair; the route then tore that worker down and spawned a child that also cannot repair, so an incomplete sidecar still failed to load. Repairing correctly requires running the repair between the worker teardown and the child spawn, decoupled from VRAM sizing, which is a larger change tracked separately. Restore the prior behavior so these paths match the reviewed form and do not partially attempt a repair that cannot complete while workers are resident. * Honor failed worker shutdowns on load and revalidate the cached latest mapping The fresh-load paths spawned a new worker straight after _shutdown_subprocess without checking its result, so a worker that outlived terminate/kill (a wedged CUDA syscall) had its handle overwritten by the replacement while it still held GPU memory, and is_worker_alive/the pre-swap guard could no longer see it. Both the inference load and the export checkpoint load now abort when the old worker did not exit, so the load can be retried once it does. _config_model_types returned a cached latest mapping without re-checking the sidecar, so a sidecar deleted or broken in-process after its first parse was never re-validated: routing kept sending latest-only models to the stale latest tier while activation failed. The cached latest mapping is now dropped and re-resolved (self-healing) when the sidecar is no longer intact. * Drop cached latest mapping when the pin is gone; keep 4-bit for custom-code fallback _latest_sidecar_intact now returns False when the pin marker itself is gone, not just when a pinned package is missing. Otherwise a cached latest mapping outlived a deleted pin: _config_model_types kept returning it, so routing sent latest-only models to a tier whose worker activation then failed (no pinned version) until restart. It now drops the cache and re-resolves to no latest tier. The _overlay_transformers_dir caller already gates on a present pin, so it is unaffected. validate_model forced 16-bit sizing whenever a PyPI upgrade was merely offered, even for a model that can fall back to its own auto_map code. /load loads such a model 4-bit without the install, and the install route refuses while training is active, so 16-bit sizing here returned a VRAM 409 for the only viable 4-bit path. The offered-upgrade flip is now gated on the absence of a custom-code fallback; an already-active latest sidecar still always sizes 16-bit. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/export/orchestrator.py | 90 +- studio/backend/core/export/worker.py | 11 + studio/backend/core/inference/orchestrator.py | 87 +- studio/backend/core/inference/worker.py | 12 + studio/backend/core/training/training.py | 312 +++-- studio/backend/core/training/worker.py | 15 +- studio/backend/models/inference.py | 63 + studio/backend/routes/export.py | 35 + studio/backend/routes/inference.py | 319 ++++- studio/backend/routes/training.py | 41 +- .../tests/test_chat_load_during_training.py | 10 +- .../tests/test_orchestrator_unload_cancel.py | 29 + .../test_shutdown_preserves_live_worker.py | 145 +++ .../backend/tests/test_transformers_latest.py | 1099 +++++++++++++++++ .../tests/test_transformers_version.py | 646 ++++++++++ studio/backend/utils/models/model_config.py | 42 +- studio/backend/utils/transformers_latest.py | 607 +++++++++ studio/backend/utils/transformers_version.py | 812 +++++++++++- studio/frontend/src/app/routes/__root.tsx | 2 + .../src/features/chat/api/chat-adapter.ts | 5 + .../chat/hooks/use-chat-model-runtime.ts | 32 + .../src/features/chat/shared-composer.tsx | 39 + .../frontend/src/features/chat/types/api.ts | 6 + .../api/transformers-upgrade-api.ts | 32 + .../transformers-upgrade-dialog.tsx | 167 +++ .../hooks/use-transformers-upgrade-consent.ts | 28 + .../features/transformers-upgrade/index.ts | 8 + .../transformers-upgrade-dialog-store.ts | 139 +++ .../features/transformers-upgrade/types.ts | 16 + 29 files changed, 4641 insertions(+), 208 deletions(-) create mode 100644 studio/backend/tests/test_shutdown_preserves_live_worker.py create mode 100644 studio/backend/tests/test_transformers_latest.py create mode 100644 studio/backend/utils/transformers_latest.py create mode 100644 studio/frontend/src/features/transformers-upgrade/api/transformers-upgrade-api.ts create mode 100644 studio/frontend/src/features/transformers-upgrade/components/transformers-upgrade-dialog.tsx create mode 100644 studio/frontend/src/features/transformers-upgrade/hooks/use-transformers-upgrade-consent.ts create mode 100644 studio/frontend/src/features/transformers-upgrade/index.ts create mode 100644 studio/frontend/src/features/transformers-upgrade/stores/transformers-upgrade-dialog-store.ts create mode 100644 studio/frontend/src/features/transformers-upgrade/types.ts diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 671ef363f5..31bbbdc748 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -132,6 +132,11 @@ class ExportOrchestrator: """True while an export / load / cleanup command is running.""" return self._export_active + def is_worker_alive(self) -> bool: + """True while the persistent export subprocess is running (op or idle).""" + proc = self._proc + return proc is not None and proc.is_alive() + def was_cancelled(self) -> bool: """True if the in-flight (or most recent) run was cancelled by the user.""" return self._cancel_requested @@ -204,6 +209,23 @@ class ExportOrchestrator: def _spawn_subprocess(self, config: dict) -> None: """Spawn a new export subprocess.""" + # Last-resort recheck for spawns outside an active op. Inside an op, _export_active is set and + # load_checkpoint already rechecked, so a reservation here is an install about to observe + # is_export_active() and abort; raising would kill this export for an install that never proceeds. + from utils.transformers_version import sidecar_swap_in_progress + + from utils.transformers_version import sidecar_swap_kind + + _swap_kind = sidecar_swap_kind() + # Inside an active op an INSTALL reservation is about to abort on the + # is_export_active check, but a lazy REPAIR has no such check and can be + # rebuilding the sidecar right now, so it must always refuse the spawn. + if _swap_kind == "repair" or (_swap_kind is not None and not self._export_active): + from utils.transformers_version import SidecarSwapInProgress + raise SidecarSwapInProgress( + "A transformers installation is replacing the latest sidecar; " + "retry when it completes." + ) from utils.native_path_leases import ( native_path_secret_removed_for_child_start, run_without_native_path_secret, @@ -231,11 +253,17 @@ class ExportOrchestrator: adopt_pid(self._proc.pid) # bind to parent lifetime (Windows job / sweep) logger.info("Export subprocess started (pid=%s)", self._proc.pid) - def _shutdown_subprocess(self, timeout: float = 10.0) -> None: - """Gracefully shut down the export subprocess.""" + def _shutdown_subprocess(self, timeout: float = 10.0) -> bool: + """Gracefully shut down the export subprocess. + + Returns True only once the worker is confirmed dead. If it survives + terminate/kill (e.g. wedged in an uninterruptible CUDA syscall that outlives + SIGKILL) the live handle is KEPT, not nulled, so is_worker_alive() and the + pre-swap liveness guard can still observe the survivor instead of a cleared + handle and refuse the destructive sidecar swap.""" if self._proc is None or not self._proc.is_alive(): self._proc = None - return + return True self._drain_queue() @@ -265,10 +293,20 @@ class ExportOrchestrator: except Exception: pass + if self._proc is not None and self._proc.is_alive(): + # Survived SIGKILL (uninterruptible syscall): keep the handle so callers + # and the pre-swap guard see a live worker rather than a nulled one. + logger.error( + "Export subprocess still alive after terminate/kill; " + "preserving its handle for the pre-swap liveness check" + ) + return False + self._proc = None self._cmd_queue = None self._resp_queue = None logger.info("Export subprocess shut down") + return True def _cleanup(self): """atexit handler.""" @@ -409,14 +447,44 @@ class ExportOrchestrator: self._export_active = True op_success, op_message = False, "" try: + # Handshake with the sidecar install route: _export_active is set above, so either this + # recheck refuses BEFORE tearing down the old worker (keeping the loaded checkpoint), or + # the install sees is_export_active() and 409s. The spawn-time recheck stays as a last resort. + from utils.transformers_version import sidecar_swap_in_progress + + if sidecar_swap_in_progress(): + from utils.transformers_version import SidecarSwapInProgress + op_message = ( + "A transformers installation is replacing the latest " + "sidecar; retry when it completes." + ) + raise SidecarSwapInProgress(op_message) # Always kill any existing subprocess and spawn fresh. if self._ensure_subprocess_alive(): - self._shutdown_subprocess() + if self._shutdown_subprocess() is False: + # Survivor still holds GPU memory (a wedged CUDA syscall outliving + # SIGKILL); its handle is kept so is_worker_alive() and the pre-swap + # guard still see it. Do not spawn a second worker over it -- fail so + # the load can retry once it exits. + op_message = ( + "The current export worker did not exit and still holds GPU " + "memory; not starting a new checkpoint load over it. Retry shortly." + ) + return False, op_message elif self._proc is not None: self._shutdown_subprocess(timeout = 2) logger.info("Spawning fresh export subprocess for '%s'", checkpoint_path) - self._spawn_subprocess(sub_config) + try: + self._spawn_subprocess(sub_config) + except Exception: + # The old worker is already gone; a stale current_checkpoint + # would make the Export page claim a loaded checkpoint that + # the next op then fails on with "no subprocess running". + self.current_checkpoint = None + self.is_vision = False + self.is_peft = False + raise try: resp = self._wait_response("loaded") @@ -560,6 +628,18 @@ class ExportOrchestrator: self._export_active = True op_success, op_message, op_output_path = False, "", None try: + # Handshake with the sidecar install route (see load_checkpoint): _export_active is set + # above, so this recheck refuses before the command is sent, or the install sees the active + # op and 409s. Without it, an install would block in cleanup_memory behind a long export op. + from utils.transformers_version import sidecar_swap_in_progress + + if sidecar_swap_in_progress(): + from utils.transformers_version import SidecarSwapInProgress + op_message = ( + "A transformers installation is replacing the latest " + "sidecar; retry when it completes." + ) + raise SidecarSwapInProgress(op_message) cmd = {"type": "export", "export_type": export_type, **params} try: self._send_cmd(cmd) diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index 7828116236..08993a9a08 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -236,6 +236,17 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: checkpoint_path = cmd["checkpoint_path"] max_seq_length = cmd.get("max_seq_length", 2048) load_in_4bit = cmd.get("load_in_4bit", True) + # Latest-sidecar checkpoints load 16-bit here too: bnb 4-bit feeds quantized + # expert weights into unvalidated paths (same flip as the chat worker). + if load_in_4bit: + from utils.transformers_version import latest_tier_active_for + if latest_tier_active_for(checkpoint_path, cmd.get("hf_token")): + load_in_4bit = False + logger.info( + "Latest-transformers sidecar active for %s - forcing a 16-bit " + "export load (4-bit is disabled for brand-new architectures)", + checkpoint_path, + ) trust_remote_code = cmd.get("trust_remote_code", False) # Auto-enable trust_remote_code for NemotronH/Nano models. diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 2b1ceca75a..6d0b13ced9 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -174,6 +174,21 @@ class InferenceOrchestrator: def _spawn_subprocess(self, config: dict) -> None: """Spawn a new inference subprocess.""" + # Same recheck as the training/export spawns, REPAIR reservations only: a + # repair swaps without holding the lifecycle gate this load's caller owns, + # while an install cannot swap until this gate is released (and then its + # queued-load snapshot aborts it), so tolerating installs here lets the + # load win instead of failing both sides. Also covers the OpenAI + # auto-switch path, which enters _load_model_impl without route guards. + from utils.transformers_version import ( + SidecarSwapInProgress, + sidecar_swap_kind, + ) + + if sidecar_swap_kind() == "repair": + raise SidecarSwapInProgress( + "A transformers repair is replacing the latest sidecar; retry when it completes." + ) from utils.native_path_leases import ( native_path_secret_removed_for_child_start, run_without_native_path_secret, @@ -210,12 +225,24 @@ class InferenceOrchestrator: if self._cancel_event is not None: self._cancel_event.set() - def _shutdown_subprocess(self, timeout: float = 10.0) -> None: - """Gracefully shut down the inference subprocess.""" + def is_worker_alive(self) -> bool: + """True while the inference subprocess is running, even with no model + active (a failed load can leave a live worker holding sidecar modules).""" + proc = self._proc + return proc is not None and proc.is_alive() + + def _shutdown_subprocess(self, timeout: float = 10.0) -> bool: + """Gracefully shut down the inference subprocess. + + Returns True only once the worker is confirmed dead. If it survives + terminate/kill (e.g. wedged in an uninterruptible CUDA syscall that outlives + SIGKILL) the live handle is KEPT, not nulled, so is_worker_alive() and the + pre-swap liveness guard can still observe the survivor instead of a cleared + handle and refuse the destructive sidecar swap.""" self._stop_dispatcher() # before killing subprocess if self._proc is None or not self._proc.is_alive(): self._proc = None - return + return True # 1. Cancel any ongoing generation first (instant via mp.Event) self._cancel_generation() @@ -252,12 +279,22 @@ class InferenceOrchestrator: except Exception: pass + if self._proc is not None and self._proc.is_alive(): + # Survived SIGKILL (uninterruptible syscall): keep the handle so callers + # and the pre-swap guard see a live worker rather than a nulled one. + logger.error( + "Inference subprocess still alive after terminate/kill; " + "preserving its handle for the pre-swap liveness check" + ) + return False + self._proc = None self._cmd_queue = None self._resp_queue = None self._cancel_event = None self._drain_event = None logger.info("Inference subprocess shut down") + return True def _cleanup(self): """atexit handler.""" @@ -882,6 +919,13 @@ class InferenceOrchestrator: # Public API — same interface as InferenceBackend # ------------------------------------------------------------------ + # Monotonic count of PUBLISHED loads; lets the install route detect a load + # (including a same-model reload) that completed while it waited on the gate. + # Bumped when the load result is published, not at load start: a start-time + # bump is already visible when the installer snapshots mid-load, so the + # completed reload would look unchanged and get unloaded by the swap. + load_generation: int = 0 + def load_model( self, config, # ModelConfig @@ -935,13 +979,36 @@ class InferenceOrchestrator: sub_config["resolved_gpu_ids"] = resolved_gpu_ids sub_config["gpu_selection"] = gpu_selection + # Recheck the sidecar reservation BEFORE tearing the old worker down, + # for REPAIRS only: an install holds this same lifecycle gate, so it + # cannot swap while this load runs, and its queued-load snapshot + # aborts it after this load publishes -- the load wins cleanly. + # Raising here (repair) keeps the current model loaded. + from utils.transformers_version import ( + SidecarSwapInProgress, + sidecar_swap_kind, + ) + + if sidecar_swap_kind() == "repair": + raise SidecarSwapInProgress( + "A transformers repair is replacing the latest sidecar; " + "retry when it completes." + ) + # Always kill the existing subprocess and spawn fresh: reusing one # after unsloth patches torch internals breaks getsource on reload. if self._ensure_subprocess_alive(): self._cancel_generation() time.sleep(0.3) - self._shutdown_subprocess() - + if self._shutdown_subprocess() is False: + # The worker survived terminate/kill (e.g. a wedged CUDA syscall that + # outlives SIGKILL). Its handle is kept, so is_worker_alive() and the + # pre-swap guard still see it; do not spawn a second worker over one + # still holding GPU memory. Fail so the load can retry once it exits. + raise RuntimeError( + "The current inference worker did not exit and still holds GPU " + "memory; not starting a new model over it. Retry shortly." + ) elif self._proc is not None: self._shutdown_subprocess(timeout = 2) @@ -1030,6 +1097,7 @@ class InferenceOrchestrator: return False model_info = resp.get("model_info", {}) self.active_model_name = model_info.get("identifier", model_name) + self.load_generation += 1 # A load always spawns a fresh subprocess holding only this model, so # mirror that. A lingering stale name would pass unload_model's "not in # self.models" guard, and the worker's absent-name fallback would unload @@ -1061,8 +1129,15 @@ class InferenceOrchestrator: self.models.clear() raise Exception(error) - except Exception: + except Exception as exc: self.loading_models.discard(model_name) + from utils.transformers_version import SidecarSwapInProgress + + if isinstance(exc, SidecarSwapInProgress) and self._ensure_subprocess_alive(): + # Raised before the old worker was torn down: the previous model + # is still live, so keep the mirrors (clearing them would let the + # installer treat the worker as inactive and kill it unreported). + raise self.active_model_name = None self.models.clear() raise diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index d56353ee56..e4628dcea8 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -291,6 +291,18 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: hf_token = _clean_token(config.get("hf_token")) load_in_4bit = _resolve_lora_4bit(mc, config.get("load_in_4bit", True)) + # Latest-transformers sidecar models load 16-bit: bnb 4-bit feeds quantized + # expert weights into unvalidated paths (e.g. grouped-MoE torch._grouped_mm). + if load_in_4bit: + from utils.transformers_version import latest_tier_active_for + if latest_tier_active_for(config["model_name"], hf_token): + load_in_4bit = False + logger.info( + "Latest-transformers sidecar active for %s - forcing a 16-bit " + "load (4-bit is disabled for brand-new architectures)", + config["model_name"], + ) + trust_remote_code = config.get("trust_remote_code", False) if not trust_remote_code and _needs_nemotron_trust(config["model_name"], hf_token = hf_token): trust_remote_code = True diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 6b32ec873d..406c780e81 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -881,93 +881,125 @@ class TrainingBackend: else: defer_auto_selection = True - # Synchronous validation passed -> free VRAM (export + chat) now, before - # auto-selection and the spawn, so placement sees the freed memory. - if before_spawn is not None: - try: - before_spawn() - except Exception: - logger.warning("before_spawn hook failed; continuing", exc_info = True) + # Handshake with the sidecar install route: mark the spawn in progress BEFORE rechecking + # the reservation, so either this recheck aborts, or the install's is_training_active() + # sees this flag (or the recorded proc) and refuses. + from utils.transformers_version import sidecar_swap_in_progress - if defer_auto_selection: - resolved_gpu_ids, gpu_selection = prepare_gpu_selection(None, **gpu_selection_kwargs) - config["resolved_gpu_ids"] = resolved_gpu_ids - config["gpu_selection"] = gpu_selection - - from .worker import run_training_process + self._spawn_in_progress = True + if sidecar_swap_in_progress(): + self._spawn_in_progress = False + from utils.transformers_version import SidecarSwapInProgress + raise SidecarSwapInProgress( + "A transformers installation is replacing the latest sidecar; " + "retry when it completes." + ) + # Any exception between the handshake above and the flag reset below would + # otherwise leave _spawn_in_progress latched, wedging is_training_active + # (and the install route) until restart. try: - with native_path_secret_removed_for_child_start(): - event_queue = _CTX.Queue() - stop_queue = _CTX.Queue() + # Synchronous validation passed -> free VRAM (export + chat) now, before + # auto-selection and the spawn, so placement sees the freed memory. Runs AFTER the handshake + # so a lost race to an install can't tear down chat/export for a training run that never spawns. + if before_spawn is not None: + try: + before_spawn() + except Exception: + logger.warning("before_spawn hook failed; continuing", exc_info = True) - proc = _CTX.Process( - target = run_without_native_path_secret, - args = (run_training_process,), - kwargs = { - "event_queue": event_queue, - "stop_queue": stop_queue, - "config": config, - }, - daemon = True, - ) - proc.start() - from utils.process_lifetime import adopt_pid + if defer_auto_selection: + try: + resolved_gpu_ids, gpu_selection = prepare_gpu_selection( + None, **gpu_selection_kwargs + ) + except Exception: + # Flag is already set; a failed GPU selection must not leave is_training_active stuck True. + self._spawn_in_progress = False + raise + config["resolved_gpu_ids"] = resolved_gpu_ids + config["gpu_selection"] = gpu_selection + + from .worker import run_training_process + + try: + with native_path_secret_removed_for_child_start(): + event_queue = _CTX.Queue() + stop_queue = _CTX.Queue() + + proc = _CTX.Process( + target = run_without_native_path_secret, + args = (run_training_process,), + kwargs = { + "event_queue": event_queue, + "stop_queue": stop_queue, + "config": config, + }, + daemon = True, + ) + proc.start() + from utils.process_lifetime import adopt_pid + + adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep) + except Exception: + logger.error("Failed to start training subprocess", exc_info = True) + self._spawn_in_progress = False + return False + + logger.info("Training subprocess started (pid=%s)", proc.pid) + + # Reset state (old pump thread dead, proc.start() succeeded). + self.current_job_id = job_id + self._should_stop = False + self._cancel_requested = False + self._complete_seen.clear() + self._progress = TrainingProgress( + is_training = True, status_message = "Initializing training..." + ) + self.loss_history.clear() + self.lr_history.clear() + self.step_history.clear() + self.grad_norm_history.clear() + self.grad_norm_step_history.clear() + self.eval_loss_history.clear() + self.eval_step_history.clear() + self.eval_enabled = False + self._output_dir = None + self._metric_buffer.clear() + self._run_finalized = False + self._db_run_created = False + self._db_create_in_progress = False # a stale watchdog create can't block this run + self._db_total_steps_set = False + self._db_config = _sanitize_db_config(config) + self._db_started_at = datetime.now(timezone.utc).isoformat() + # Start each job Xet-first; keep config so a stall can respawn over HTTP. + self._last_full_config = config + self._in_model_load = False + self._xet_fallback_used = False + self._needs_xet_respawn = False + + # Create the DB run row before the pump can consume events, so it appears + # in history during model loading and a fast terminal worker can't race the + # pump into a duplicate create/finalize. From here the pump only finalizes. + self._ensure_db_run_created() + + # Assign handles and start the pump together under the lock so a concurrent + # poll can't see a live _proc with no pump and spawn a duplicate. + new_pump = threading.Thread(target = self._pump_loop, daemon = True) + with self._lock: + self._pump_running = False + self._event_queue = event_queue + self._stop_queue = stop_queue + self._proc = proc + self._pump_thread = new_pump + new_pump.start() + self._spawn_in_progress = False + + return True - adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep) except Exception: - logger.error("Failed to start training subprocess", exc_info = True) - return False - - logger.info("Training subprocess started (pid=%s)", proc.pid) - - # Reset state (old pump thread dead, proc.start() succeeded). - self.current_job_id = job_id - self._should_stop = False - self._cancel_requested = False - self._complete_seen.clear() - self._progress = TrainingProgress( - is_training = True, status_message = "Initializing training..." - ) - self.loss_history.clear() - self.lr_history.clear() - self.step_history.clear() - self.grad_norm_history.clear() - self.grad_norm_step_history.clear() - self.eval_loss_history.clear() - self.eval_step_history.clear() - self.eval_enabled = False - self._output_dir = None - self._metric_buffer.clear() - self._run_finalized = False - self._db_run_created = False - self._db_create_in_progress = False # a stale watchdog create can't block this run - self._db_total_steps_set = False - self._db_config = _sanitize_db_config(config) - self._db_started_at = datetime.now(timezone.utc).isoformat() - # Start each job Xet-first; keep config so a stall can respawn over HTTP. - self._last_full_config = config - self._in_model_load = False - self._xet_fallback_used = False - self._needs_xet_respawn = False - - # Create the DB run row before the pump can consume events, so it appears - # in history during model loading and a fast terminal worker can't race the - # pump into a duplicate create/finalize. From here the pump only finalizes. - self._ensure_db_run_created() - - # Assign handles and start the pump together under the lock so a concurrent - # poll can't see a live _proc with no pump and spawn a duplicate. - new_pump = threading.Thread(target = self._pump_loop, daemon = True) - with self._lock: - self._pump_running = False - self._event_queue = event_queue - self._stop_queue = stop_queue - self._proc = proc - self._pump_thread = new_pump - new_pump.start() - - return True + self._spawn_in_progress = False + raise def stop_training(self, save: bool = True) -> bool: """Send stop signal to the training subprocess.""" @@ -1266,50 +1298,84 @@ class TrainingBackend: from .worker import run_training_process - try: - with native_path_secret_removed_for_child_start(): - event_queue = _CTX.Queue() - stop_queue = _CTX.Queue() - new_proc = _CTX.Process( - target = run_without_native_path_secret, - args = (run_training_process,), - kwargs = { - "event_queue": event_queue, - "stop_queue": stop_queue, - "config": config, - }, - daemon = True, - ) - new_proc.start() - from utils.process_lifetime import adopt_pid + # This run is active, so an install request 409s rather than proceeds: a reservation seen here + # is transient (an aborting install or short lazy repair). Wait it out instead of stranding the + # stalled run; only a wedged reservation fails the respawn. + from utils.transformers_version import sidecar_swap_in_progress - adopt_pid(new_proc.pid) # bind to parent lifetime (Windows job / sweep) - except Exception: - logger.error("Failed to respawn training subprocess", exc_info = True) - with self._lock: - # No replacement pump will run; clear the flag so a later run can't - # inherit a stale _pump_running=True and spawn a duplicate. - self._pump_running = False - self._progress.is_training = False - self._progress.error = "Failed to recover stalled model download" - self._ensure_db_run_created() - self._finalize_run_in_db( - status = "error", - error_message = "Failed to recover stalled model download", + self._spawn_in_progress = True + _swap_wait_deadline = time.time() + 120 + while sidecar_swap_in_progress() and time.time() < _swap_wait_deadline: + time.sleep(1) + if sidecar_swap_in_progress(): + # Raising here would land in the pump's broad finalization catch and + # strand the run in a training state with no worker: finalize it as a + # failure explicitly instead. + self._spawn_in_progress = False + msg = ( + "A transformers installation is replacing the latest sidecar; " + "cannot respawn the training worker." ) + logger.error(msg) + with self._lock: + self._progress.is_training = False + self._progress.error = msg + self._ensure_db_run_created() + self._finalize_run_in_db(status = "error", error_message = msg) return - logger.info("Training subprocess respawned with Xet disabled (pid=%s)", new_proc.pid) - new_pump = threading.Thread(target = self._pump_loop, daemon = True) - with self._lock: - self._in_model_load = False - self._event_queue = event_queue - self._stop_queue = stop_queue - self._proc = new_proc - self._pump_thread = new_pump - # Start under the lock so _ensure_pump_alive can never observe the - # new pump as a not-yet-started (dead) thread and spawn a duplicate. - new_pump.start() + # Reset the handshake flag on any unexpected failure past this point, so a + # crashed respawn cannot wedge is_training_active until restart. + try: + try: + with native_path_secret_removed_for_child_start(): + event_queue = _CTX.Queue() + stop_queue = _CTX.Queue() + new_proc = _CTX.Process( + target = run_without_native_path_secret, + args = (run_training_process,), + kwargs = { + "event_queue": event_queue, + "stop_queue": stop_queue, + "config": config, + }, + daemon = True, + ) + new_proc.start() + from utils.process_lifetime import adopt_pid + + adopt_pid(new_proc.pid) # bind to parent lifetime (Windows job / sweep) + except Exception: + logger.error("Failed to respawn training subprocess", exc_info = True) + self._spawn_in_progress = False + with self._lock: + # No replacement pump will run; clear the flag so a later run can't + # inherit a stale _pump_running=True and spawn a duplicate. + self._pump_running = False + self._progress.is_training = False + self._progress.error = "Failed to recover stalled model download" + self._ensure_db_run_created() + self._finalize_run_in_db( + status = "error", + error_message = "Failed to recover stalled model download", + ) + return + + logger.info("Training subprocess respawned with Xet disabled (pid=%s)", new_proc.pid) + new_pump = threading.Thread(target = self._pump_loop, daemon = True) + with self._lock: + self._in_model_load = False + self._event_queue = event_queue + self._stop_queue = stop_queue + self._proc = new_proc + self._spawn_in_progress = False + self._pump_thread = new_pump + # Start under the lock so _ensure_pump_alive can never observe the + # new pump as a not-yet-started (dead) thread and spawn a duplicate. + new_pump.start() + except Exception: + self._spawn_in_progress = False + raise def _ensure_pump_alive(self) -> bool: """Restart the event pump if it crashed, even after the worker exited. @@ -1342,6 +1408,10 @@ class TrainingBackend: def is_training_active(self) -> bool: """Check if training is currently active.""" + # A spawn past its sidecar-swap recheck counts as active even before _proc is recorded, + # so an install cannot slip in mid-spawn. + if getattr(self, "_spawn_in_progress", False): + return True # Self-heal a crashed pump first: a dead pump must never leave the worker # training invisibly behind a frozen UI. Cheap enough for per-second polls. self._ensure_pump_alive() diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 130e6ece64..1c84b8268f 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -3019,11 +3019,24 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> ), xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1", ) + # Latest-sidecar models load 16-bit here too: bnb 4-bit feeds quantized + # expert weights into unvalidated paths (same flip as the chat worker). + _train_load_in_4bit = config["load_in_4bit"] + if _train_load_in_4bit: + from utils.transformers_version import latest_tier_active_for + if latest_tier_active_for(model_name, hf_token): + _train_load_in_4bit = False + logger.info( + "Latest-transformers sidecar active for %s - forcing a 16-bit " + "training load (4-bit is disabled for brand-new architectures)", + model_name, + ) + try: success = trainer.load_model( model_name = model_name, max_seq_length = config["max_seq_length"], - load_in_4bit = config["load_in_4bit"], + load_in_4bit = _train_load_in_4bit, full_finetuning = not use_lora, hf_token = hf_token, is_dataset_image = config.get("is_dataset_image", False), diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 53441bdb3a..2bdc00dce1 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -140,6 +140,27 @@ class ValidateModelRequest(BaseModel): ) +class TransformersUpgradeInfo(BaseModel): + """A model architecture no installed transformers ships, but a newer release does.""" + + model_type: str = Field( + ..., description = "config.json model_type unknown to every installed transformers" + ) + pypi_version: Optional[str] = Field( + None, description = "Latest transformers release on PyPI at check time" + ) + supported_in_pypi: bool = Field( + False, + description = "True if the latest PyPI release ships this model_type; Studio can " + "install it into a persistent sidecar after user consent.", + ) + supported_in_main: bool = Field( + False, + description = "True if transformers GitHub main ships this model_type (dev-only; " + "not installable through Studio yet).", + ) + + class ValidateModelResponse(BaseModel): """Result of model validation. @@ -167,6 +188,48 @@ class ValidateModelResponse(BaseModel): description = "Native training context length, read from the GGUF header when the file " "is already downloaded locally; None for non-GGUF, gated, or not-yet-downloaded models.", ) + # Additive fields; the consuming consent dialog ships in a follow-up frontend PR. + requires_transformers_upgrade: bool = Field( + False, + description = "True when the model's architecture is unknown to every installed " + "transformers but a newer transformers ships it; the UI should offer the " + "install-latest-transformers consent dialog (or the dev-only notice).", + ) + transformers_upgrade: Optional[TransformersUpgradeInfo] = Field( + None, + description = "Details for the transformers-upgrade dialog; set only when " + "requires_transformers_upgrade is true.", + ) + + +class InstallLatestTransformersRequest(BaseModel): + """Consented request to install the latest transformers release into a sidecar.""" + + version: str = Field( + ..., + min_length = 1, + max_length = 64, + description = "Exact transformers version to install; must match the current " + "latest PyPI release reported by /validate.", + ) + + +class InstallLatestTransformersResponse(BaseModel): + """Result of the consented latest-transformers sidecar install.""" + + success: bool = Field(..., description = "Whether the sidecar was provisioned") + version: str = Field(..., description = "The requested transformers version") + message: str = Field(..., description = "Human-readable result") + model_unloaded: bool = Field( + False, + description = "Whether the active chat model was unloaded before the swap " + "(reported even on failure, so the client can restore its state)", + ) + latest_version: Optional[str] = Field( + None, + description = "On a version-mismatch failure: the release that superseded " + "the requested one, so the client can retry with it", + ) class GenerateRequest(BaseModel): diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index a7fd7cbec7..d44e2ac021 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -51,7 +51,17 @@ def _ensure_export_supported() -> None: Keeps the backend authoritative even if a client bypasses the UI gate. Read-only endpoints (scan/status/logs) are intentionally NOT gated so the Export page can still render the reason. + Also refuses (409) while a latest-transformers install is swapping .venv_t5_latest: an + export worker spawned mid-swap could activate a half-replaced sidecar. """ + from utils.transformers_latest import is_install_in_progress + + if is_install_in_progress(): + raise HTTPException( + status_code = 409, + detail = "A transformers installation is in progress. Retry when it completes.", + ) + from utils.hardware import export_capability cap = export_capability() @@ -97,6 +107,11 @@ async def load_checkpoint( except HTTPException: raise except Exception as e: + from utils.transformers_version import SidecarSwapInProgress + + if isinstance(e, SidecarSwapInProgress): + # Expected loss of the race against a sidecar install: retryable 409. + raise HTTPException(status_code = 409, detail = str(e)) logger.error(f"Error loading checkpoint: {e}", exc_info = True) raise HTTPException( status_code = 500, @@ -308,6 +323,11 @@ async def export_merged_model( except HTTPException: raise except Exception as e: + from utils.transformers_version import SidecarSwapInProgress + + if isinstance(e, SidecarSwapInProgress): + # Expected loss of the race against a sidecar install: retryable 409. + raise HTTPException(status_code = 409, detail = str(e)) logger.error(f"Error exporting merged model: {e}", exc_info = True) raise HTTPException( status_code = 500, @@ -347,6 +367,11 @@ async def export_base_model( except HTTPException: raise except Exception as e: + from utils.transformers_version import SidecarSwapInProgress + + if isinstance(e, SidecarSwapInProgress): + # Expected loss of the race against a sidecar install: retryable 409. + raise HTTPException(status_code = 409, detail = str(e)) logger.error(f"Error exporting base model: {e}", exc_info = True) raise HTTPException( status_code = 500, @@ -388,6 +413,11 @@ async def export_gguf( except HTTPException: raise except Exception as e: + from utils.transformers_version import SidecarSwapInProgress + + if isinstance(e, SidecarSwapInProgress): + # Expected loss of the race against a sidecar install: retryable 409. + raise HTTPException(status_code = 409, detail = str(e)) logger.error(f"Error exporting GGUF model: {e}", exc_info = True) raise HTTPException( status_code = 500, @@ -428,6 +458,11 @@ async def export_lora_adapter( except HTTPException: raise except Exception as e: + from utils.transformers_version import SidecarSwapInProgress + + if isinstance(e, SidecarSwapInProgress): + # Expected loss of the race against a sidecar install: retryable 409. + raise HTTPException(status_code = 409, detail = str(e)) logger.error(f"Error exporting LoRA adapter: {e}", exc_info = True) raise HTTPException( status_code = 500, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 8141ff073a..073cd64400 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1693,6 +1693,9 @@ from models.inference import ( CompletionUsage, ValidateModelRequest, ValidateModelResponse, + TransformersUpgradeInfo, + InstallLatestTransformersRequest, + InstallLatestTransformersResponse, TextContentPart, ImageContentPart, ImageUrl, @@ -3836,11 +3839,25 @@ async def load_model( GGUF models load via llama-server (llama.cpp) instead of Unsloth. """ + # A sidecar install that has reserved the swap must not lose to a load that + # then gets unloaded by the pre-swap teardown. Rechecked under the gate: an + # install can reserve while this request queues on the gate, so the pre-gate + # check alone is only a fast path. + from core.inference.llama_keepwarm import inference_lifecycle_gate + from utils.transformers_version import sidecar_swap_in_progress + + _swap_409 = HTTPException( + status_code = 409, + detail = "A transformers installation is in progress. Retry when it completes.", + ) + if sidecar_swap_in_progress(): + raise _swap_409 # Hold the lifecycle gate across the load so idle auto-unload can't unload the # model mid-load. Auto-switch calls _load_model_impl directly since it already # holds this gate. - from core.inference.llama_keepwarm import inference_lifecycle_gate async with inference_lifecycle_gate(): + if sidecar_swap_in_progress(): + raise _swap_409 return await _load_model_impl(request, fastapi_request, current_subject) @@ -4037,6 +4054,17 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre f"Resolved load_in_4bit={effective_load_in_4bit} for '{model_log_label}' " f"from adapter_config.json / base model (requested {request.load_in_4bit})" ) + # Latest-sidecar models load 16-bit (worker refuses bnb 4-bit); size the guard + # to match. Off-loop: tier resolution reads configs. + if effective_load_in_4bit and not config.is_gguf: + from utils.transformers_version import latest_tier_active_for + if await asyncio.to_thread(latest_tier_active_for, config.identifier, request.hf_token): + effective_load_in_4bit = False + logger.info( + f"Latest-transformers sidecar active for '{model_log_label}' - " + "sizing and loading in 16-bit (4-bit is disabled for brand-new " + "architectures)" + ) # Refuse a load that would OOM active training, before the unload step below # frees the resident model. Off-loop: guard does sync nvidia-smi / HF work. @@ -4470,6 +4498,11 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre logger.warning("GGUF runtime missing while loading '%s': %s", model_log_label, e) raise HTTPException(status_code = 400, detail = str(e)) except Exception as e: + from utils.transformers_version import SidecarSwapInProgress + + if isinstance(e, SidecarSwapInProgress): + # Lost the spawn-time race to a sidecar install/repair: retryable 409. + raise HTTPException(status_code = 409, detail = str(e)) # Friendlier message for models Unsloth cannot load. if native_grant_backed: redacted_msg = redact_native_paths(str(e)) @@ -4598,16 +4631,6 @@ async def validate_model( detail = "gpu_ids is not supported for GGUF models yet.", ) effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit) - # Off-loop: guard does sync nvidia-smi / HF work. - await asyncio.to_thread( - _guard_chat_load_against_training, - config, - model_identifier = model_identifier, - hf_token = request.hf_token, - load_in_4bit = effective_load_in_4bit, - max_seq_length = request.max_seq_length, - requested_gpu_ids = effective_gpu_ids, - ) # Both checks cover the [adapter, base] set (matching the scan route and workers): # either repo can ship auto_map code or a poisoned pickle. @@ -4624,16 +4647,69 @@ async def validate_model( security_targets = list(dict.fromkeys(security_targets)) is_gguf = getattr(config, "is_gguf", False) - # A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a - # mixed repo are inert for this load, so gating on them is a false positive. Only - # run the remote-code/security preflight for non-GGUF loads. + # Does a newer transformers ship this model_type? Static overlay first, cached + # PyPI/main snapshot only for unknown types. Never fails validation; run before + # the training guard so an installable upgrade sizes as 16-bit. + transformers_upgrade: Optional[TransformersUpgradeInfo] = None + if not is_gguf: + from utils.transformers_latest import check_upgrade_for_model + + # Cover [adapter, base]: the worker activates transformers for the base model. + for _target in security_targets: + _upgrade = await asyncio.to_thread( + check_upgrade_for_model, _target, request.hf_token + ) + if _upgrade is not None: + transformers_upgrade = TransformersUpgradeInfo(**_upgrade) + break + + # Whether the model can load on the CURRENT transformers through its own remote + # code (auto_map, or the YAML trust default). Computed before the 16-bit flip + # because a model with this fallback still loads 4-bit without the offered install, + # exactly as /load does. requires_trust_remote_code = False - requires_security_review = False if not is_gguf: requires_trust_remote_code = any( _requires_trust_remote_code_for_model(_t, request.hf_token) for _t in security_targets ) + + # Mirror /load's latest-sidecar 16-bit flip so the guard sizes it the same way. An + # ALREADY-ACTIVE latest sidecar always forces 16-bit (the worker will). A merely + # OFFERED (not yet installed) upgrade forces 16-bit only when the model has NO + # custom-code fallback: with auto_map it still loads 4-bit on the current + # transformers (as /load does without a successful install), and the install route + # refuses while training is active, so sizing 16-bit here would 409 the only viable + # 4-bit path. /load re-sizes 16-bit after a successful install and re-guards there. + if effective_load_in_4bit and not is_gguf: + from utils.transformers_version import latest_tier_active_for + _install_only_upgrade = ( + transformers_upgrade is not None + and transformers_upgrade.supported_in_pypi + and transformers_upgrade.pypi_version + and not requires_trust_remote_code + ) + if _install_only_upgrade or await asyncio.to_thread( + latest_tier_active_for, config.identifier, request.hf_token + ): + effective_load_in_4bit = False + # Off-loop: guard does sync nvidia-smi / HF work. + await asyncio.to_thread( + _guard_chat_load_against_training, + config, + model_identifier = model_identifier, + hf_token = request.hf_token, + load_in_4bit = effective_load_in_4bit, + max_seq_length = request.max_seq_length, + requested_gpu_ids = effective_gpu_ids, + ) + + # A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a + # mixed repo are inert for this load, so gating on them is a false positive. Only + # run the security preflight for non-GGUF loads (requires_trust_remote_code was + # already resolved above for the sizing flip). + requires_security_review = False + if not is_gguf: requires_security_review = any( _requires_security_review_for_model(_t, request.hf_token) for _t in security_targets ) @@ -4676,6 +4752,8 @@ async def validate_model( requires_trust_remote_code = requires_trust_remote_code, requires_security_review = requires_security_review, context_length = context_length, + requires_transformers_upgrade = transformers_upgrade is not None, + transformers_upgrade = transformers_upgrade, ) except HTTPException: @@ -4720,6 +4798,217 @@ async def validate_model( ) +# studio_router only: admin action, kept off the OpenAI-compatible /v1 mount. +@studio_router.post( + "/install-latest-transformers", response_model = InstallLatestTransformersResponse +) +async def install_latest_transformers_route( + request: InstallLatestTransformersRequest, current_subject: str = Depends(get_current_subject) +): + """ + Consented install of the latest transformers release into the persistent + .venv_t5_latest sidecar. + + Called after the user confirms the transformers-upgrade dialog raised by /validate + (requires_transformers_upgrade). The requested version must match the current latest + PyPI release (re-verified server-side); the sidecar then participates in routing on + this and every future start. A pip install runs off-loop, so this can take a minute. + """ + from utils.transformers_latest import install_latest_transformers + from utils.transformers_version import end_sidecar_swap, try_begin_sidecar_swap + + # The install stage-and-swaps .venv_t5_latest in place; a live worker would + # lazy-import from the new version mid-run, mixing incompatible modules. Gate on + # worker LIVENESS not tier (no HF token here, so tier re-resolution is unreliable + # for gated repos): training and export are refused, the chat model unloaded. + # Reserve the swap FIRST, before any await: training/export starts check this + # reservation, so raising it after the gate wait would let a worker slip in. + if not try_begin_sidecar_swap(): + raise HTTPException( + status_code = 409, + detail = "A transformers installation is already in progress.", + ) + # Until the installer thread takes over, this coroutine owns the reservation + # and must release it on any early exit (the 409 refusals below). + owns_reservation = True + try: + from core.export import get_export_backend + from core.training import get_training_backend + + if get_training_backend().is_training_active(): + raise HTTPException( + status_code = 409, + detail = ( + "A training run is active. Wait for it to finish before " + "installing a new transformers version." + ), + ) + _export = get_export_backend() + if _export.is_export_active(): + raise HTTPException( + status_code = 409, + detail = ( + "An export is running. Wait for it to finish before " + "installing a new transformers version." + ), + ) + # A loaded (idle) export checkpoint would be torn down by the pre-swap + # cleanup; if the swap then failed, that state would be silently lost + # with no rollback signal. Make the user unload it deliberately first. + if getattr(_export, "current_checkpoint", None): + raise HTTPException( + status_code = 409, + detail = ( + "An export checkpoint is loaded. Unload it from the Export " + "page before installing a new transformers version." + ), + ) + # In-flight streams passed the middleware already, so the lifecycle gate can't + # protect them and the swap's unload would kill them mid-stream; mirror the + # auto-switch busy check. This route is not middleware-counted and pending + # requests stay blocked in the middleware, so neither is subtracted here. + from core.inference.llama_keepwarm import ( + inference_lifecycle_gate, + note_model_unloaded, + other_inference_request_count, + ) + + if other_inference_request_count(current_request_counted = False, include_pending = False) > 0: + raise HTTPException( + status_code = 409, + detail = ( + "Another inference request is in progress. Wait for it to " + "finish before installing a new transformers version." + ), + ) + + # Hold the lifecycle gate /load holds so no HF worker can start (or be mid-load + # with active_model_name unset) while the sidecar is swapped. Teardown runs via + # before_swap, only once the staged install succeeded: a failed pip/compat check + # must not leave the user with their model gone. GGUF stays loaded (llama-server + # never imports transformers). + backend = get_inference_backend() + export_backend = get_export_backend() + + unloaded_chat = {"v": False} + + def _unload_before_swap() -> None: + # Runs on the install thread, inside the gate held by _gated_install. Any + # failure raises so the previous sidecar stays untouched (a worker that did + # not tear down cleanly may still lazy-import from it). Export teardown runs + # FIRST so its failure aborts while the chat model is still loaded; + # cleanup_memory shuts the subprocess down even when its command fails, so + # judge by worker liveness, not its return value. + export_backend.cleanup_memory() + export_alive = getattr(export_backend, "is_worker_alive", None) + if callable(export_alive) and export_alive(): + raise RuntimeError("Export worker still alive before the transformers swap") + active = getattr(backend, "active_model_name", None) + if active: + if not backend.unload_model(active): + # A failed unload still clears the orchestrator's model state, + # so the model is gone from the parent's view even though the + # swap aborts: report it so the client rolls back instead of + # pointing at an unloaded model. + if getattr(backend, "active_model_name", None) != active: + unloaded_chat["v"] = True + note_model_unloaded() + raise RuntimeError(f"Could not unload '{active}' before the transformers swap") + note_model_unloaded() + unloaded_chat["v"] = True + logger.info( + "Unloaded '%s' before swapping in transformers %s", + active, + request.version, + ) + # A failed load can leave a live worker with no active model that + # still holds sidecar modules (and blocks the rename on Windows). + worker_alive = getattr(backend, "is_worker_alive", None) + if callable(worker_alive) and worker_alive(): + # _shutdown_subprocess keeps the handle when the worker outlives SIGKILL, + # so both its False result and the liveness recheck catch a survivor + # rather than the recheck being fooled by a nulled handle. + stopped = backend._shutdown_subprocess() + if not stopped or worker_alive(): + raise RuntimeError("Inference worker still alive before the transformers swap") + + def _run_install() -> dict: + # Owns the reservation from here: releasing in the thread, not the route, + # keeps it held if the request is cancelled while the install still stages. + try: + return install_latest_transformers(request.version, _unload_before_swap, True) + finally: + end_sidecar_swap() + + # Snapshot before waiting on the gate: a /load already holding it can + # complete meanwhile (including a same-model reload with new settings), + # and the installer must not unload a model whose successful LoadResponse + # the client is about to render. The generation counter catches reloads + # the name alone would miss. + active_before_gate = ( + getattr(backend, "active_model_name", None), + getattr(backend, "load_generation", 0), + ) + + async def _gated_install() -> dict: + # Held by THIS task, not the request coroutine: a cancelled POST unwinding an + # `async with` here would drop the only guard /load honors mid-install. + async with inference_lifecycle_gate(): + _active_now = ( + getattr(backend, "active_model_name", None), + getattr(backend, "load_generation", 0), + ) + if _active_now != active_before_gate: + end_sidecar_swap() + raise HTTPException( + status_code = 409, + detail = ( + "A model load completed while the install was waiting. " + "Retry the install." + ), + ) + # Recheck under the gate: new streams bump their in-flight count while + # holding it, so once held nothing slips past (the pre-gate check is only + # a fast path and can be outlasted by a wait on a long /load). + if ( + other_inference_request_count( + current_request_counted = False, include_pending = False + ) + > 0 + ): + end_sidecar_swap() + raise HTTPException( + status_code = 409, + detail = ( + "Another inference request is in progress. Wait for " + "it to finish before installing a new transformers " + "version." + ), + ) + return await asyncio.to_thread(_run_install) + + install_task = asyncio.ensure_future(_gated_install()) + owns_reservation = False + # shield: a cancelled request stops waiting, but the installer runs to + # completion (holding the gate) instead of being torn down mid-swap. + result = await asyncio.shield(install_task) + finally: + if owns_reservation: + end_sidecar_swap() + if not result["success"]: + if result.get("latest_version"): + # Structured failure so the dialog can update to the newer release + # and offer a retry that can actually succeed. + return InstallLatestTransformersResponse(**result, model_unloaded = unloaded_chat["v"]) + if unloaded_chat["v"]: + # The chat model is already gone even though the swap failed; return a + # structured failure (not a bare 400) so the client can restore its + # model state instead of pointing at an unloaded model. + return InstallLatestTransformersResponse(**result, model_unloaded = True) + raise HTTPException(status_code = 400, detail = result["message"]) + return InstallLatestTransformersResponse(**result, model_unloaded = unloaded_chat["v"]) + + @router.post("/unload", response_model = UnloadResponse) async def unload_model(request: UnloadRequest, current_subject: str = Depends(get_current_subject)): """ diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 1da1c4f425..5e633f4896 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -146,6 +146,16 @@ async def start_training( # No in-process ensure_transformers_version(): the subprocess # (worker.py) activates the correct version before importing ML libs. + # A consented latest-transformers install stage-and-swaps .venv_t5_latest; + # a worker spawned mid-swap could activate a half-replaced sidecar. + from utils.transformers_latest import is_install_in_progress + + if is_install_in_progress(): + raise HTTPException( + status_code = 409, + detail = ("A transformers installation is in progress. Retry when it completes."), + ) + backend = get_training_backend() # S3 dataset loading needs the optional boto3 dependency. Reject early @@ -341,6 +351,24 @@ async def start_training( "s3_config": request.s3_config.model_dump() if request.s3_config else None, } + # Latest-sidecar models size and train 16-bit (same flip as chat load): + # 4-bit is disabled for brand-new architectures, so VRAM coexistence + # checks must not underestimate against a load the worker will refuse. + if training_kwargs["load_in_4bit"]: + from utils.transformers_version import latest_tier_active_for + if await asyncio.to_thread( + latest_tier_active_for, + training_kwargs["model_name"], + training_kwargs["hf_token"] or None, + ): + training_kwargs["load_in_4bit"] = False + logger.info( + "Latest-transformers sidecar active for %s - sizing and " + "training in 16-bit (4-bit is disabled for brand-new " + "architectures)", + training_kwargs["model_name"], + ) + # Training page has no trust_remote_code toggle, so honor the YAML default # -- but only for genuine first-party (unsloth/nvidia) Hub repos, never a # local path or a name merely starting with "unsloth/". @@ -426,9 +454,16 @@ async def start_training( logger.warning("Chat/training VRAM coordination failed; proceeding: %s", e) # The hook runs only once start guards pass -> VRAM freed iff training starts. - success = backend.start_training( - job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs - ) + from utils.transformers_version import SidecarSwapInProgress + + try: + success = backend.start_training( + job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs + ) + except SidecarSwapInProgress as exc: + # Expected loss of the race against a sidecar install: a retryable + # 409 matching the route-entry guard, not an internal error. + raise HTTPException(status_code = 409, detail = str(exc)) if not success: progress_error = backend.trainer.training_progress.error diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index 487c3c7ce0..63dba8579c 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -651,11 +651,19 @@ class TestLoadModelGuardIntegration(unittest.TestCase): inf._shutdown_subprocess = MagicMock() llama = SimpleNamespace(is_loaded = False, model_identifier = None, hf_variant = None) llama.unload_model = MagicMock() - cfg = SimpleNamespace(is_gguf = False, is_lora = False, path = None, base_model = None) + cfg = SimpleNamespace( + is_gguf = False, + is_lora = False, + path = None, + base_model = None, + identifier = "unsloth/Qwen3-1.7B", + ) request = LoadRequest(model_path = "unsloth/Qwen3-1.7B") info = {"required_gb": 40.0, "usable_gb": 5.0, "needed_gb": 50.0, "mode": "auto"} with ( + # Pin the latest-sidecar tier check so the guard path stays offline. + patch("utils.transformers_version.latest_tier_active_for", return_value = False), patch.object(self.route, "validate_extra_args", return_value = None), patch.object( self.route, diff --git a/studio/backend/tests/test_orchestrator_unload_cancel.py b/studio/backend/tests/test_orchestrator_unload_cancel.py index fe3c6d5a0d..fb80b6d061 100644 --- a/studio/backend/tests/test_orchestrator_unload_cancel.py +++ b/studio/backend/tests/test_orchestrator_unload_cancel.py @@ -874,6 +874,35 @@ def test_load_model_aborts_when_cancelled_before_spawn(monkeypatch): assert o.models == {} +def test_load_model_aborts_when_old_worker_survives_shutdown(monkeypatch): + # A wedged worker that outlives terminate/kill makes _shutdown_subprocess return + # False. load_model must not spawn a second worker over it (double GPU allocation + + # the survivor's handle is lost); it aborts so the load can retry once it exits. + import types + + from utils import transformers_version as tv + + o = _bare_orchestrator() + o.active_model_name = "old" + o.models = {"old": {}} + o.loading_models = set() + monkeypatch.setattr(tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda *a, **k: ([0], "sel")) + monkeypatch.setattr(orch_mod.time, "sleep", lambda *_a, **_k: None) + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_cancel_generation", lambda: None) + monkeypatch.setattr(o, "_shutdown_subprocess", lambda *a, **k: False) # survivor + monkeypatch.setattr( + o, "_spawn_subprocess", lambda cfg: pytest.fail("must not spawn over a live survivor") + ) + + with pytest.raises(RuntimeError, match = "did not exit"): + o.load_model(types.SimpleNamespace(identifier = "new", gguf_variant = None)) + # The except path cleared the loading marker and mirrors. + assert "new" not in o.loading_models + assert o.active_model_name is None + + def test_load_model_proceeds_when_not_cancelled(monkeypatch): # Guard against a false abort: an uncancelled load keeps its marker and spawns. o = _bare_orchestrator() diff --git a/studio/backend/tests/test_shutdown_preserves_live_worker.py b/studio/backend/tests/test_shutdown_preserves_live_worker.py new file mode 100644 index 0000000000..faf273411c --- /dev/null +++ b/studio/backend/tests/test_shutdown_preserves_live_worker.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: AGPL-3.0-only +"""_shutdown_subprocess returns whether the worker actually died, and preserves the +live handle when it survives terminate/kill. + +A GPU worker wedged in an uninterruptible CUDA syscall can outlive SIGKILL. If shutdown +nulled its handle anyway, is_worker_alive() would report False and the pre-swap liveness +guard would let the destructive .venv_t5_latest rename proceed while a live worker still +holds sidecar transformers modules (breaking the rename on Windows). The methods must keep +the handle and return False so callers can refuse the swap. +""" + +import pytest + +from core.export.orchestrator import ExportOrchestrator +from core.inference.orchestrator import InferenceOrchestrator + + +class _FakeProc: + """A subprocess handle that dies only on the requested step (or never).""" + + def __init__(self, dies_on = None): + self._alive = True + self._dies_on = dies_on # None | "join" | "terminate" | "kill" + + def is_alive(self): + return self._alive + + def join(self, timeout = None): + if self._dies_on == "join": + self._alive = False + + def terminate(self): + if self._dies_on == "terminate": + self._alive = False + + def kill(self): + if self._dies_on == "kill": + self._alive = False + + +def _bare_inference(): + o = InferenceOrchestrator.__new__(InferenceOrchestrator) + o._stop_dispatcher = lambda: None + o._cancel_generation = lambda: None + o._drain_queue = lambda: [] + + class _Q: + def put(self, *a, **k): + pass + + o._cmd_queue = _Q() + o._resp_queue = _Q() + o._cancel_event = None + o._drain_event = None + return o + + +def _bare_export(): + o = ExportOrchestrator.__new__(ExportOrchestrator) + o._drain_queue = lambda: [] + + class _Q: + def put(self, *a, **k): + pass + + o._cmd_queue = _Q() + o._resp_queue = _Q() + return o + + +@pytest.fixture(autouse = True) +def _no_sleep(monkeypatch): + # _shutdown_subprocess sleeps 0.5s after cancelling; keep the tests instant. + import core.inference.orchestrator as inf_mod + monkeypatch.setattr(inf_mod.time, "sleep", lambda *_a, **_k: None) + + +class TestInferenceShutdownReturn: + def test_worker_that_dies_returns_true_and_clears_handle(self): + o = _bare_inference() + o._proc = _FakeProc(dies_on = "terminate") + assert o._shutdown_subprocess(timeout = 0.01) is True + assert o._proc is None + assert o.is_worker_alive() is False + + def test_survivor_returns_false_and_keeps_handle(self): + o = _bare_inference() + o._proc = _FakeProc(dies_on = None) # outlives terminate AND kill + assert o._shutdown_subprocess(timeout = 0.01) is False + assert o._proc is not None + # is_worker_alive stays truthful, so the pre-swap guard can refuse the swap. + assert o.is_worker_alive() is True + + def test_already_dead_returns_true(self): + o = _bare_inference() + o._proc = _FakeProc(dies_on = "join") + o._proc._alive = False + assert o._shutdown_subprocess(timeout = 0.01) is True + assert o._proc is None + + +class TestExportShutdownReturn: + def test_worker_that_dies_returns_true_and_clears_handle(self): + o = _bare_export() + o._proc = _FakeProc(dies_on = "terminate") + assert o._shutdown_subprocess(timeout = 0.01) is True + assert o._proc is None + assert o.is_worker_alive() is False + + def test_survivor_returns_false_and_keeps_handle(self): + o = _bare_export() + o._proc = _FakeProc(dies_on = None) + assert o._shutdown_subprocess(timeout = 0.01) is False + assert o._proc is not None + assert o.is_worker_alive() is True + + +class TestSpawnPathsHonorFailedShutdown: + """A fresh-load path must not spawn a second worker over one that outlived + terminate/kill: the survivor still holds GPU memory and its handle would be lost.""" + + def test_export_load_checkpoint_aborts_when_worker_survives(self, monkeypatch): + import threading + + import utils.transformers_version as tv + + o = ExportOrchestrator.__new__(ExportOrchestrator) + o._lock = threading.RLock() + o._proc = _FakeProc(dies_on = None) # survivor + o.clear_logs = lambda: None + o._cancel_requested = False + o._active_op_kind = None + o._export_active = False + o._ensure_subprocess_alive = lambda: True + o._shutdown_subprocess = lambda *a, **k: False + o._spawn_subprocess = lambda cfg: pytest.fail("must not spawn over a live survivor") + o._record_op_finished = lambda *a, **k: None + monkeypatch.setattr(tv, "sidecar_swap_in_progress", lambda: False) + + ok, msg = o.load_checkpoint(checkpoint_path = "ckpt") + + assert ok is False + assert "did not exit" in msg + # The finally cleared the op flags even though we returned early. + assert o._export_active is False diff --git a/studio/backend/tests/test_transformers_latest.py b/studio/backend/tests/test_transformers_latest.py new file mode 100644 index 0000000000..20616dccba --- /dev/null +++ b/studio/backend/tests/test_transformers_latest.py @@ -0,0 +1,1099 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the latest-transformers support check and the consented sidecar install.""" + +import ast +import json +import os +import textwrap +import time +import pytest +from pathlib import Path + + +# The backend uses "from utils..." imports; ensure the backend dir is on sys.path. +import sys + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Stub the custom logger before importing the modules under test. +import types as _types + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +import utils.transformers_latest as tl +import utils.transformers_version as tv +from utils.transformers_latest import ( + check_upgrade_for_model, + install_latest_transformers, + latest_transformers_supports, + _fetch_remote_model_types, + _model_types_from_config, +) +from utils.transformers_version import ( + _config_mapping_cache, + _config_json_cache, + _higher_tier, + _is_valid_version_string, + _model_types_from_source, + _tier_from_config_mapping, + _venv_t5_latest_packages, + activate_transformers_for_subprocess, + ensure_latest_transformers_venv, + get_transformers_tier, + latest_venv_pinned_version, +) + + +# A CONFIG_MAPPING_NAMES source exercising every construct the AST extractor supports. +_MAPPING_SOURCE = """ +from collections import OrderedDict +CONFIG_MAPPING_NAMES = OrderedDict( + [ + ("llama", "LlamaConfig"), + ("gemma4", "Gemma4Config"), + ], + **{"qwen3_moe": "Qwen3MoeConfig"}, +) +CONFIG_MAPPING_NAMES.update({"brandnew_arch": "BrandNewConfig"}) +""" + +_MAIN_ONLY_SOURCE = """ +CONFIG_MAPPING_NAMES = { + "llama": "LlamaConfig", + "gemma4": "Gemma4Config", + "qwen3_moe": "Qwen3MoeConfig", + "brandnew_arch": "BrandNewConfig", + "dev_only_arch": "DevOnlyConfig", +} +""" + + +class _FakeResponse: + def __init__(self, body: bytes): + self._body = body + + def read(self): + return self._body + + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + +def _fake_urlopen_factory(counter: dict): + """urlopen stub serving the PyPI JSON and both refs' mapping sources.""" + + def _fake_urlopen(req, timeout = None): + url = req.full_url if hasattr(req, "full_url") else str(req) + counter[url] = counter.get(url, 0) + 1 + counter["__total__"] = counter.get("__total__", 0) + 1 + if url == tl._PYPI_JSON_URL: + return _FakeResponse(json.dumps({"info": {"version": "5.13.0"}}).encode()) + if "/v5.13.0/" in url and url.endswith("auto_mappings.py"): + return _FakeResponse(_MAPPING_SOURCE.encode()) + if "/v5.13.0/" in url and url.endswith("configuration_auto.py"): + return _FakeResponse(b"CONFIG_MAPPING_NAMES = {}\n") + if "/main/" in url and url.endswith("auto_mappings.py"): + return _FakeResponse(_MAIN_ONLY_SOURCE.encode()) + if "/main/" in url and url.endswith("configuration_auto.py"): + return _FakeResponse(b"CONFIG_MAPPING_NAMES = {}\n") + raise AssertionError(f"unexpected URL fetched: {url}") + + return _fake_urlopen + + +@pytest.fixture(autouse = True) +def _isolated_caches(tmp_path: Path, monkeypatch): + """Fresh in-memory + on-disk caches per test; no accidental real studio_root writes.""" + tl.clear_caches() + monkeypatch.setattr(tl, "_cache_file", lambda: tmp_path / "transformers_latest_check.json") + # The sidecar swap reservation writes a lock file next to the venv dir; + # point it at tmp so tests never touch the real studio root. + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest")) + monkeypatch.delenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", raising = False) + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + yield + tl.clear_caches() + + +def _no_network(monkeypatch, exc = None): + """Fail every urlopen and return a counter; tests assert n == 0 to prove no fetch + happened (check_upgrade_for_model swallows exceptions, so a raising stub alone + cannot prove the negative).""" + calls = {"n": 0} + + def _raise(*args, **kwargs): + calls["n"] += 1 + raise (exc or OSError("network fetch attempted")) + + monkeypatch.setattr("urllib.request.urlopen", _raise) + return calls + + +# --- AST extraction shared with the static router --- + + +class TestModelTypesFromSource: + def test_ordereddict_update_and_unpacking(self): + keys = _model_types_from_source(_MAPPING_SOURCE) + assert keys == {"llama", "gemma4", "qwen3_moe", "brandnew_arch"} + + def test_plain_dict_literal(self): + keys = _model_types_from_source(_MAIN_ONLY_SOURCE) + assert "dev_only_arch" in keys and "llama" in keys + + def test_syntax_error_raises_for_caller_to_handle(self): + with pytest.raises(SyntaxError): + _model_types_from_source("def broken(:\n") + + +class TestFetchRemoteModelTypes: + def test_merges_both_auto_files(self, monkeypatch): + counter = {} + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory(counter)) + keys = _fetch_remote_model_types("v5.13.0") + assert keys is not None and "brandnew_arch" in keys + + def test_all_fetches_failing_returns_none(self, monkeypatch): + _no_network(monkeypatch, exc = OSError("no route")) + assert _fetch_remote_model_types("main") is None + + def test_empty_mapping_treated_as_failure(self, monkeypatch): + monkeypatch.setattr( + "urllib.request.urlopen", + lambda req, timeout = None: _FakeResponse(b"CONFIG_MAPPING_NAMES = {}\n"), + ) + assert _fetch_remote_model_types("main") is None + + def test_transient_failure_of_one_file_fails_whole_lookup(self, monkeypatch): + # One file times out: the partial map must not be returned and cached. + def _fake(req, timeout = None): + url = req.full_url if hasattr(req, "full_url") else str(req) + if url.endswith("configuration_auto.py"): + return _FakeResponse(_MAPPING_SOURCE.encode()) + raise OSError("timed out") + + monkeypatch.setattr("urllib.request.urlopen", _fake) + assert _fetch_remote_model_types("main") is None + + def test_missing_auto_mappings_404_still_succeeds(self, monkeypatch): + # Pre-5.10 tags have no auto_mappings.py; a 404 must not fail the lookup. + import urllib.error + + def _fake(req, timeout = None): + url = req.full_url if hasattr(req, "full_url") else str(req) + if url.endswith("configuration_auto.py"): + return _FakeResponse(_MAPPING_SOURCE.encode()) + raise urllib.error.HTTPError(url, 404, "Not Found", None, None) + + monkeypatch.setattr("urllib.request.urlopen", _fake) + keys = _fetch_remote_model_types("v5.9.0") + assert keys is not None and "brandnew_arch" in keys + + def test_unparseable_file_fails_whole_lookup(self, monkeypatch): + def _fake(req, timeout = None): + url = req.full_url if hasattr(req, "full_url") else str(req) + if url.endswith("configuration_auto.py"): + return _FakeResponse(_MAPPING_SOURCE.encode()) + return _FakeResponse(b"def broken(:\n") + + monkeypatch.setattr("urllib.request.urlopen", _fake) + assert _fetch_remote_model_types("main") is None + + +# --- latest_transformers_supports: snapshot, cache, offline, kill switch --- + + +class TestLatestTransformersSupports: + def test_supported_in_pypi(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + result = latest_transformers_supports("brandnew_arch") + assert result == { + "pypi_version": "5.13.0", + "supported_in_pypi": True, + "supported_in_main": True, + } + + def test_dev_only_arch_reported_main_only(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + result = latest_transformers_supports("dev_only_arch") + assert result["supported_in_pypi"] is False + assert result["supported_in_main"] is True + + def test_unknown_everywhere(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + result = latest_transformers_supports("no_such_arch") + assert result["supported_in_pypi"] is False and result["supported_in_main"] is False + + def test_network_failure_returns_none(self, monkeypatch): + _no_network(monkeypatch, exc = OSError("down")) + assert latest_transformers_supports("brandnew_arch") is None + + def test_offline_returns_none_without_fetch(self, monkeypatch): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + calls = _no_network(monkeypatch) + assert latest_transformers_supports("brandnew_arch") is None + assert calls["n"] == 0 + + def test_kill_switch_returns_none_without_fetch(self, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1") + calls = _no_network(monkeypatch) + assert latest_transformers_supports("brandnew_arch") is None + assert calls["n"] == 0 + + def test_memory_cache_hit_avoids_refetch(self, monkeypatch): + counter = {} + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory(counter)) + latest_transformers_supports("brandnew_arch") + first_total = counter["__total__"] + latest_transformers_supports("some_other_arch") + assert counter["__total__"] == first_total + + def test_disk_cache_survives_restart(self, monkeypatch): + counter = {} + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory(counter)) + latest_transformers_supports("brandnew_arch") + # Simulate a restart: memory gone, disk snapshot stays, network unavailable. + tl.clear_caches() + _no_network(monkeypatch) + result = latest_transformers_supports("brandnew_arch") + assert result is not None and result["supported_in_pypi"] is True + + def test_expired_snapshot_refetches(self, monkeypatch): + counter = {} + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory(counter)) + latest_transformers_supports("brandnew_arch") + stale = dict(tl._memory_snapshot, fetched_at = time.time() - tl._CACHE_TTL_SECONDS - 1) + tl.clear_caches() + tl._save_snapshot_file(stale) + first_total = counter["__total__"] + latest_transformers_supports("brandnew_arch") + assert counter["__total__"] > first_total + + def test_corrupt_disk_cache_ignored(self, monkeypatch, tmp_path: Path): + counter = {} + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory(counter)) + tl._cache_file().write_text("{not json", encoding = "utf-8") + result = latest_transformers_supports("brandnew_arch") + assert result is not None and counter["__total__"] > 0 + + def test_failure_backoff_skips_immediate_retry(self, monkeypatch): + calls = {"n": 0} + + def _fail(*args, **kwargs): + calls["n"] += 1 + raise OSError("down") + + monkeypatch.setattr("urllib.request.urlopen", _fail) + assert latest_transformers_supports("brandnew_arch") is None + first = calls["n"] + assert latest_transformers_supports("brandnew_arch") is None + assert calls["n"] == first # backed off, no second network attempt + + +# --- check_upgrade_for_model: the tier hook --- + + +def _local_model(tmp_path: Path, model_type: str) -> str: + d = tmp_path / f"model_{model_type}" + d.mkdir() + (d / "config.json").write_text(json.dumps({"model_type": model_type})) + return str(d) + + +_FAKE_OVERLAYS = { + "default": frozenset({"llama", "bert", "gpt2"}), + "530": frozenset({"qwen3_moe", "qwen3_next"}), + "550": frozenset({"gemma4"}), + "510": frozenset({"gemma4_unified"}), + "latest": frozenset(), +} + + +def _fake_overlays(monkeypatch, overlays = None): + overlays = overlays or _FAKE_OVERLAYS + fake = lambda tier: overlays.get(tier, frozenset()) + monkeypatch.setattr(tv, "_config_model_types", fake) + monkeypatch.setattr(tl, "_config_model_types", fake) + + +class TestCheckUpgradeForModel: + def test_unknown_type_supported_in_pypi_signals(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + result = check_upgrade_for_model(_local_model(tmp_path, "brandnew_arch")) + assert result == { + "model_type": "brandnew_arch", + "pypi_version": "5.13.0", + "supported_in_pypi": True, + "supported_in_main": True, + } + + def test_dev_only_type_signals_main_only(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + result = check_upgrade_for_model(_local_model(tmp_path, "dev_only_arch")) + assert result["supported_in_pypi"] is False and result["supported_in_main"] is True + + def test_unknown_everywhere_falls_through(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + assert check_upgrade_for_model(_local_model(tmp_path, "no_such_arch")) is None + + def test_offline_falls_through_without_fetch(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + calls = _no_network(monkeypatch) + assert check_upgrade_for_model(_local_model(tmp_path, "brandnew_arch")) is None + assert calls["n"] == 0 + + def test_network_failure_falls_through(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + _no_network(monkeypatch, exc = OSError("down")) + assert check_upgrade_for_model(_local_model(tmp_path, "brandnew_arch")) is None + + def test_known_default_type_never_fetches(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + calls = _no_network(monkeypatch) + assert check_upgrade_for_model(_local_model(tmp_path, "llama")) is None + assert calls["n"] == 0 + + def test_known_sidecar_type_never_fetches(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + calls = _no_network(monkeypatch) + assert check_upgrade_for_model(_local_model(tmp_path, "gemma4_unified")) is None + assert calls["n"] == 0 + + def test_hardcoded_tier_type_never_fetches_even_without_overlays( + self, tmp_path: Path, monkeypatch + ): + # Sidecar overlays unreadable, but the hardcoded tables route it. + _fake_overlays( + monkeypatch, + {"default": frozenset({"llama"})}, + ) + calls = _no_network(monkeypatch) + assert check_upgrade_for_model(_local_model(tmp_path, "qwen3_5_moe")) is None + assert calls["n"] == 0 + + def test_unreadable_default_overlay_bails_out(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch, {"default": frozenset()}) + calls = _no_network(monkeypatch) + assert check_upgrade_for_model(_local_model(tmp_path, "brandnew_arch")) is None + assert calls["n"] == 0 + + def test_no_model_type_falls_through(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + _no_network(monkeypatch) + d = tmp_path / "no_type" + d.mkdir() + (d / "config.json").write_text(json.dumps({"architectures": ["Whatever"]})) + assert check_upgrade_for_model(str(d)) is None + + def test_nested_model_type_is_used(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + d = tmp_path / "nested" + d.mkdir() + (d / "config.json").write_text(json.dumps({"text_config": {"model_type": "brandnew_arch"}})) + result = check_upgrade_for_model(str(d)) + assert result is not None and result["model_type"] == "brandnew_arch" + + def test_never_raises_on_internal_error(self, monkeypatch): + monkeypatch.setattr( + tl, "_load_config_json", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom")) + ) + assert check_upgrade_for_model("some/model") is None + + +class TestNestedModelTypeExtraction: + def test_top_level_wins(self): + assert _model_types_from_config( + {"model_type": "a", "text_config": {"model_type": "b"}} + ) == ["a", "b"] + + def test_nested_fallback(self): + assert _model_types_from_config({"llm_config": {"model_type": "b"}}) == ["b"] + + def test_missing_returns_none(self): + assert _model_types_from_config({}) == [] + + +# --- Routing parity: overlay-shipped model_types route as before, never remote-check --- + + +class TestRoutingParity: + def test_all_overlay_types_route_identically_and_never_check(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + calls = _no_network(monkeypatch) + expected_tier = { + "llama": "default", + "bert": "default", + "gpt2": "default", + "qwen3_moe": "530", + "qwen3_next": "530", + "gemma4": "550", + "gemma4_unified": "510", + } + for model_type, tier in expected_tier.items(): + cfg = {"model_type": model_type} + assert _tier_from_config_mapping(cfg) == tier, model_type + assert check_upgrade_for_model(_local_model(tmp_path, model_type)) is None + assert calls["n"] == 0 + + def test_real_installed_mappings_route_without_checker(self, monkeypatch, tmp_path: Path): + """Parity over the REAL installed overlays (base + any provisioned sidecar): + every shipped model_type resolves statically, so the remote checker never + fires and routing is byte-identical with the feature enabled.""" + _no_network(monkeypatch) + seen = 0 + for tier in ("default", "530", "550", "510"): + types = tv._config_model_types(tier) + if not types: + continue # overlay not provisioned in this environment + for model_type in types: + assert _tier_from_config_mapping({"model_type": model_type}) is not None + seen += 1 + if seen == 0: + pytest.skip("no transformers overlay available in this environment") + + def test_get_tier_unchanged_by_kill_switch(self, tmp_path: Path, monkeypatch): + _fake_overlays(monkeypatch) + _no_network(monkeypatch) + path = _local_model(tmp_path, "no_such_arch") + _config_json_cache.clear() + tier_default = get_transformers_tier(path, probe = False) + monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1") + _config_json_cache.clear() + assert get_transformers_tier(path, probe = False) == tier_default == "default" + + +# --- .venv_t5_latest provisioning and routing participation --- + + +class TestLatestVenvProvisioning: + def test_version_string_validation(self): + assert _is_valid_version_string("5.13.0") + assert _is_valid_version_string("5.14.0rc1") + assert not _is_valid_version_string("5.13.0; rm -rf /") + assert not _is_valid_version_string("git+https://evil") + assert not _is_valid_version_string("") + + def test_packages_pin_exact_version(self): + pkgs = _venv_t5_latest_packages("5.13.0") + assert pkgs[0] == "transformers==5.13.0" + assert any(p.startswith("huggingface_hub==") for p in pkgs) + + def test_ensure_latest_writes_pin_and_invalidates_cache(self, tmp_path: Path, monkeypatch): + venv_dir = tmp_path / ".venv_t5_latest" + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + recorded = {} + + def _fake_ensure(dir_, packages, label): + recorded["dir"] = dir_ + recorded["packages"] = packages + Path(dir_).mkdir(parents = True, exist_ok = True) + return True + + monkeypatch.setattr(tv, "_ensure_venv_dir", _fake_ensure) + _config_mapping_cache["latest"] = frozenset({"stale"}) + assert ensure_latest_transformers_venv("5.13.0") is True + # Stage-and-swap: pip installs into staging, the live dir is the swap result. + assert recorded["dir"] == str(venv_dir) + ".staging" + assert "transformers==5.13.0" in recorded["packages"] + assert venv_dir.is_dir() + assert not Path(str(venv_dir) + ".staging").exists() + assert latest_venv_pinned_version() == "5.13.0" + assert "latest" not in _config_mapping_cache + + def test_ensure_latest_upgrade_failure_keeps_old_sidecar(self, tmp_path: Path, monkeypatch): + venv_dir = tmp_path / ".venv_t5_latest" + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + venv_dir.mkdir(parents = True) + (venv_dir / tv._LATEST_PIN_MARKER).write_text( + json.dumps({"version": "5.12.0", "packages": ["transformers==5.12.0"]}) + ) + (venv_dir / "transformers").mkdir() + monkeypatch.setattr(tv, "_venv_dir_is_valid", lambda *a, **k: True) + # Install fails mid-flight: the previous sidecar and pin survive. + monkeypatch.setattr(tv, "_ensure_venv_dir", lambda *a, **k: False) + assert ensure_latest_transformers_venv("5.13.0") is False + assert latest_venv_pinned_version() == "5.12.0" + assert (venv_dir / "transformers").is_dir() + assert not Path(str(venv_dir) + ".staging").exists() + + def test_ensure_latest_rejects_bad_version(self, tmp_path: Path, monkeypatch): + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + monkeypatch.setattr( + tv, + "_ensure_venv_dir", + lambda *a: (_ for _ in ()).throw(AssertionError("must not install")), + ) + assert ensure_latest_transformers_venv("5.13.0 && curl evil") is False + + def test_ensure_latest_offline_refuses(self, tmp_path: Path, monkeypatch): + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setattr( + tv, + "_ensure_venv_dir", + lambda *a: (_ for _ in ()).throw(AssertionError("must not install")), + ) + assert ensure_latest_transformers_venv("5.13.0") is False + + def test_unpinned_sidecar_never_installs(self, tmp_path: Path, monkeypatch): + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + monkeypatch.setattr( + tv, + "_ensure_venv_dir", + lambda *a: (_ for _ in ()).throw(AssertionError("must not install")), + ) + assert tv._ensure_venv_t5_latest_exists() is False + + def test_pinned_sidecar_repairs_with_same_version(self, tmp_path: Path, monkeypatch): + venv_dir = tmp_path / ".venv_t5_latest" + venv_dir.mkdir() + (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + monkeypatch.setattr(tv, "_venv_dir_is_valid", lambda *a: False) + recorded = {} + + def _fake_ensure(dir_, packages, label): + recorded["dir"] = dir_ + recorded["packages"] = packages + Path(dir_).mkdir(parents = True, exist_ok = True) + return True + + monkeypatch.setattr(tv, "_ensure_venv_dir", _fake_ensure) + assert tv._ensure_venv_t5_latest_exists() is True + # Repair also stage-and-swaps, never installing into the live dir. + assert recorded["dir"] == str(venv_dir) + ".staging" + assert "transformers==5.13.0" in recorded["packages"] + assert latest_venv_pinned_version() == "5.13.0" + + +class TestLatestTierRouting: + def test_latest_outranks_510(self): + assert _higher_tier("latest", "510") == "latest" + assert _higher_tier("510", "latest") == "latest" + + def test_tier_from_mapping_prefers_lowest_but_reaches_latest(self, monkeypatch): + overlays = dict(_FAKE_OVERLAYS) + overlays["latest"] = frozenset({"brandnew_arch"}) + _fake_overlays(monkeypatch, overlays) + assert _tier_from_config_mapping({"model_type": "brandnew_arch"}) == "latest" + # Anything a lower tier ships stays on the lower tier. + assert _tier_from_config_mapping({"model_type": "qwen3_moe"}) == "530" + + def test_overlay_dir_for_latest(self, tmp_path: Path, monkeypatch): + venv_dir = tmp_path / ".venv_t5_latest" + (venv_dir / "transformers").mkdir(parents = True) + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + # Unpinned dir is ignored: activation refuses an unpinned sidecar. + assert tv._overlay_transformers_dir("latest") is None + (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") + assert tv._overlay_transformers_dir("latest") == str(venv_dir / "transformers") + + def test_probe_order_excludes_unprovisioned_latest(self, tmp_path: Path, monkeypatch): + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + assert tv._probe_tier_order() == tv._PROBE_TIER_ORDER + + def test_probe_order_includes_provisioned_latest(self, tmp_path: Path, monkeypatch): + venv_dir = tmp_path / ".venv_t5_latest" + venv_dir.mkdir() + (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + assert tv._probe_tier_order() == tv._PROBE_TIER_ORDER + ("latest",) + + def test_activation_prepends_latest_dir(self, tmp_path: Path, monkeypatch): + venv_dir = tmp_path / ".venv_t5_latest" + venv_dir.mkdir() + (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + monkeypatch.setattr(tv, "get_transformers_tier", lambda *a, **k: "latest") + monkeypatch.setattr(tv, "_ensure_venv_t5_latest_exists", lambda: True) + old_sys_path = list(sys.path) + old_pp = os.environ.get("PYTHONPATH") + try: + activate_transformers_for_subprocess("some/brand-new-model") + assert sys.path[0] == str(venv_dir) + assert os.environ["PYTHONPATH"].split(os.pathsep)[0] == str(venv_dir) + finally: + sys.path[:] = old_sys_path + if old_pp is None: + os.environ.pop("PYTHONPATH", None) + else: + os.environ["PYTHONPATH"] = old_pp + + def test_activation_raises_when_latest_missing(self, tmp_path: Path, monkeypatch): + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / ".venv_t5_latest")) + monkeypatch.setattr(tv, "get_transformers_tier", lambda *a, **k: "latest") + with pytest.raises(RuntimeError, match = "venv_t5_latest"): + activate_transformers_for_subprocess("some/brand-new-model") + + +# --- install_latest_transformers: the consent endpoint helper --- + + +class TestInstallLatestTransformers: + def test_success_path(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + monkeypatch.setattr(tl, "compat_plan", lambda v: ((), [])) + recorded = {} + + def _fake_ensure( + version, + extra_packages = (), + before_swap = None, + ): + recorded["args"] = (version, extra_packages) + return True + + monkeypatch.setattr(tl, "ensure_latest_transformers_venv", _fake_ensure) + monkeypatch.setattr(tl, "latest_venv_pinned_version", lambda: "5.13.0") + result = install_latest_transformers("5.13.0") + assert result["success"] is True and result["version"] == "5.13.0" + assert recorded["args"] == ("5.13.0", ()) + + def test_version_mismatch_rejected(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + monkeypatch.setattr( + tl, + "ensure_latest_transformers_venv", + lambda v, extra_packages = (): (_ for _ in ()).throw(AssertionError("must not install")), + ) + result = install_latest_transformers("4.99.0") + assert result["success"] is False and "not the latest" in result["message"] + + def test_offline_rejected(self, monkeypatch): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + _no_network(monkeypatch) + result = install_latest_transformers("5.13.0") + assert result["success"] is False and "offline" in result["message"].lower() + + def test_kill_switch_rejected(self, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1") + _no_network(monkeypatch) + result = install_latest_transformers("5.13.0") + assert result["success"] is False + + def test_install_failure_reported(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + monkeypatch.setattr(tl, "compat_plan", lambda v: ((), [])) + monkeypatch.setattr( + tl, + "ensure_latest_transformers_venv", + lambda v, extra_packages = (), before_swap = None: False, + ) + result = install_latest_transformers("5.13.0") + assert result["success"] is False and "failed" in result["message"] + + def test_blocked_by_incompatible_deps(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + monkeypatch.setattr(tl, "compat_plan", lambda v: ((), ["numpy>=99.0"])) + monkeypatch.setattr( + tl, + "ensure_latest_transformers_venv", + lambda v, extra_packages = (): (_ for _ in ()).throw(AssertionError("must not install")), + ) + result = install_latest_transformers("5.13.0") + assert result["success"] is False and "numpy>=99.0" in result["message"] + + def test_compat_shadows_passed_to_installer(self, monkeypatch): + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + monkeypatch.setattr(tl, "compat_plan", lambda v: (("tokenizers==0.23.0",), [])) + recorded = {} + + def _fake_ensure( + version, + extra_packages = (), + before_swap = None, + ): + recorded["extras"] = extra_packages + return True + + monkeypatch.setattr(tl, "ensure_latest_transformers_venv", _fake_ensure) + monkeypatch.setattr(tl, "latest_venv_pinned_version", lambda: "5.13.0") + result = install_latest_transformers("5.13.0") + assert result["success"] is True + assert recorded["extras"] == ("tokenizers==0.23.0",) + + +class TestCompatPlan: + def _patch_env(self, monkeypatch, requires, installed): + monkeypatch.setattr(tl, "_fetch_requires_dist", lambda v: requires) + + def _ver(name): + from importlib.metadata import PackageNotFoundError + + key = name.lower().replace("_", "-") + if key not in installed: + raise PackageNotFoundError(name) + return installed[key] + + monkeypatch.setattr("importlib.metadata.version", _ver) + + def test_satisfied_env_needs_nothing(self, monkeypatch): + self._patch_env( + monkeypatch, + ["tokenizers<=0.23.0,>=0.22.0", "safetensors>=0.8.0", "numpy>=1.17"], + {"tokenizers": "0.22.2", "safetensors": "0.8.0", "numpy": "2.4.4"}, + ) + extras, blockers = tl.compat_plan("5.13.0") + assert extras == () and blockers == [] + + def test_unsatisfied_shadowable_dep_pinned(self, monkeypatch): + self._patch_env( + monkeypatch, + ["tokenizers>=0.24.0"], + {"tokenizers": "0.22.2"}, + ) + monkeypatch.setattr(tl, "_resolve_exact_version", lambda name, spec: "0.24.1") + extras, blockers = tl.compat_plan("5.99.0") + assert extras == ("tokenizers==0.24.1",) and blockers == [] + + def test_unsatisfied_non_shadowable_dep_blocks(self, monkeypatch): + self._patch_env(monkeypatch, ["numpy>=99.0"], {"numpy": "2.4.4"}) + extras, blockers = tl.compat_plan("5.99.0") + assert extras == () and blockers == ["numpy>=99.0"] + + def test_cli_only_dep_ignored(self, monkeypatch): + self._patch_env(monkeypatch, ["typer"], {}) + extras, blockers = tl.compat_plan("5.13.0") + assert extras == () and blockers == [] + + def test_sidecar_provided_hub_checked_against_recipe_pin(self, monkeypatch): + self._patch_env(monkeypatch, ["huggingface-hub<2.0,>=1.5.0"], {"huggingface-hub": "0.36.2"}) + extras, blockers = tl.compat_plan("5.13.0") + assert extras == () and blockers == [] # 1.8.0 sidecar pin satisfies it + + def test_sidecar_provided_hub_out_of_range_blocks(self, monkeypatch): + self._patch_env(monkeypatch, ["huggingface-hub>=2.1"], {"huggingface-hub": "0.36.2"}) + extras, blockers = tl.compat_plan("5.99.0") + assert blockers == ["huggingface-hub>=2.1"] + + def test_unfetchable_requires_dist_blocks_install(self, monkeypatch): + # Proceeding unverified could pin a sidecar whose imports crash workers. + monkeypatch.setattr(tl, "_fetch_requires_dist", lambda v: None) + extras, blockers = tl.compat_plan("5.13.0") + assert extras == () and len(blockers) == 1 and "retry" in blockers[0] + + def test_extra_marker_requirements_skipped(self, monkeypatch): + self._patch_env( + monkeypatch, + ['torch>=99.0; extra == "torch"', 'pytest; python_version < "3.0"'], + {}, + ) + extras, blockers = tl.compat_plan("5.13.0") + assert extras == () and blockers == [] + + +def test_get_snapshot_dedupes_concurrent_fetch(monkeypatch): + """While one thread is fetching, other callers return None instead of stacking fetches.""" + with tl._lock: + tl._is_fetching = True + calls = {"n": 0} + + def boom(): + calls["n"] += 1 + raise AssertionError("must not fetch while another fetch is in flight") + + monkeypatch.setattr(tl, "_refresh_snapshot", boom) + assert tl._get_snapshot() is None + assert calls["n"] == 0 + tl.clear_caches() + + +def test_install_serialized(): + """A second install call while one is in progress gets a structured refusal.""" + from utils.transformers_version import try_begin_sidecar_swap + + assert try_begin_sidecar_swap() is True + out = tl.install_latest_transformers("5.13.0") + assert out["success"] is False + assert "already in progress" in out["message"] + tl.clear_caches() + + +def test_install_in_progress_reflects_reservation(): + """is_install_in_progress mirrors the shared sidecar swap reservation, so a + lazy repair (which takes the same reservation) also blocks worker starts.""" + from utils.transformers_version import end_sidecar_swap, try_begin_sidecar_swap + + assert tl.is_install_in_progress() is False + assert try_begin_sidecar_swap() is True + try: + assert tl.is_install_in_progress() is True + finally: + end_sidecar_swap() + assert tl.is_install_in_progress() is False + + +def test_upgrade_check_sees_nested_model_types(monkeypatch): + """A supported wrapper with a brand-new nested backbone must still signal.""" + cfg = { + "model_type": "llava", # in every installed overlay + "text_config": {"model_type": "zz_brand_new_llm"}, + } + monkeypatch.setattr(tl, "_load_config_json", lambda *a, **k: cfg) + monkeypatch.setattr( + tl, + "latest_transformers_supports", + lambda mt: { + "pypi_version": "5.13.0", + "supported_in_pypi": mt == "zz_brand_new_llm", + "supported_in_main": mt == "zz_brand_new_llm", + }, + ) + out = tl.check_upgrade_for_model("some-org/wrapped-new-backbone") + assert out is not None + assert out["model_type"] == "zz_brand_new_llm" + + +def test_upgrade_check_ignores_nested_known_types(monkeypatch): + """All nested types known to installed overlays -> no signal, no remote call.""" + cfg = { + "model_type": "llava", + "text_config": {"model_type": "llama"}, + "vision_config": {"model_type": "clip_vision_model"}, + } + monkeypatch.setattr(tl, "_load_config_json", lambda *a, **k: cfg) + calls = [] + monkeypatch.setattr(tl, "latest_transformers_supports", lambda mt: calls.append(mt) or None) + assert tl.check_upgrade_for_model("some-org/normal-vlm") is None + assert calls == [] + + +def test_upgrade_check_requires_primary_supported(monkeypatch): + """Latest supporting only a nested type must not prompt: routing still + cannot load the primary, so the install would not fix the model.""" + cfg = { + "model_type": "zz_new_wrapper", + "text_config": {"model_type": "zz_new_llm"}, + } + monkeypatch.setattr(tl, "_load_config_json", lambda *a, **k: cfg) + monkeypatch.setattr( + tl, + "latest_transformers_supports", + lambda mt: { + "pypi_version": "5.13.0", + "supported_in_pypi": mt == "zz_new_llm", + "supported_in_main": mt == "zz_new_llm", + }, + ) + assert tl.check_upgrade_for_model("some-org/half-supported") is None + + +def test_upgrade_check_requires_every_missing_type(monkeypatch): + """Primary supported but a nested backbone missing from latest -> no prompt + (CONFIG_MAPPING would still fail on the sub-config); all supported -> signal + carries the primary type.""" + cfg = { + "model_type": "zz_new_wrapper", + "text_config": {"model_type": "zz_new_llm"}, + } + monkeypatch.setattr(tl, "_load_config_json", lambda *a, **k: cfg) + monkeypatch.setattr( + tl, + "latest_transformers_supports", + lambda mt: { + "pypi_version": "5.13.0", + "supported_in_pypi": mt == "zz_new_wrapper", + "supported_in_main": mt == "zz_new_wrapper", + }, + ) + assert tl.check_upgrade_for_model("some-org/half-supported") is None + + monkeypatch.setattr( + tl, + "latest_transformers_supports", + lambda mt: { + "pypi_version": "5.13.0", + "supported_in_pypi": True, + "supported_in_main": True, + }, + ) + out = tl.check_upgrade_for_model("some-org/fully-supported") + assert out is not None and out["model_type"] == "zz_new_wrapper" + + +def test_install_success_invalidates_capability_caches(monkeypatch): + """A successful install must drop tier probes, the latest mapping, and the + vision-detection cache so the new sidecar takes effect without a restart.""" + from utils.models import model_config as mc + + monkeypatch.setattr("urllib.request.urlopen", _fake_urlopen_factory({})) + monkeypatch.setattr(tl, "compat_plan", lambda v: ((), [])) + monkeypatch.setattr( + tl, "ensure_latest_transformers_venv", lambda v, extra_packages = (), before_swap = None: True + ) + monkeypatch.setattr(tl, "latest_venv_pinned_version", lambda: "5.13.0") + + tv._probe_tier_cache["stale/model"] = "default" + tv._config_mapping_cache["latest"] = frozenset({"stale_type"}) + tv._config_mapping_cache["default"] = frozenset({"llama"}) + mc._vision_detection_cache[("stale/model", None, False)] = False + + result = install_latest_transformers("5.13.0") + assert result["success"] is True + assert tv._probe_tier_cache == {} + assert "latest" not in tv._config_mapping_cache + assert tv._config_mapping_cache.get("default") == frozenset({"llama"}) # untouched + assert mc._vision_detection_cache == {} + + tv._probe_tier_cache.clear() + tv._config_mapping_cache.clear() + tl.clear_caches() + + +def test_vision_subprocess_unions_sidecar_registry(): + """The embedded vision-check script must extend the inlined parent sets with + the ACTIVE sidecar's registry so sidecar-only architectures classify.""" + from utils.models import model_config as mc + + script = mc._VISION_CHECK_SCRIPT + ast.parse(script) + stub_registry = { + "MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES": { + "zz_sidecar_vlm": "ZzSidecarForConditionalGeneration" + }, + } + ns = {} + # Exec only the registry-union block against a stubbed sidecar registry. + body = script.split("from transformers import AutoConfig", 1)[1] + body = body.split("kwargs = {", 1)[0] + helpers = script.split("sys.path.insert(0, backend_dir)", 1)[1] + helpers = helpers.split("try:", 1)[0] + exec(helpers, ns) + + class _FakeMa: + MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES = stub_registry[ + "MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES" + ] + + import sys as _sys + import types as _types + + fake_pkg = _types.ModuleType("transformers.models.auto") + fake_pkg.modeling_auto = _FakeMa + saved = { + k: _sys.modules.get(k) + for k in ("transformers.models.auto", "transformers.models.auto.modeling_auto") + } + _sys.modules["transformers.models.auto"] = fake_pkg + _sys.modules["transformers.models.auto.modeling_auto"] = _FakeMa + try: + exec(textwrap.dedent(body), ns) + finally: + for k, v in saved.items(): + if v is None: + _sys.modules.pop(k, None) + else: + _sys.modules[k] = v + + assert "zz_sidecar_vlm" in ns["_VLM_MODEL_TYPES"] + assert "ZzSidecarForConditionalGeneration" in ns["_VLM_CLASS_NAMES"] + + class _Cfg: + architectures = ["ZzSidecarForConditionalGeneration"] + model_type = "zz_sidecar_vlm" + + assert ns["_is_vlm"](_Cfg()) is True + + +def test_upgrade_check_mixed_pypi_main_reports_dev_only(monkeypatch): + """Primary in the PyPI release but a nested type only on main: no install + may be offered (CONFIG_MAPPING would fail on the nested sub-config), so the + aggregate must read as main-only.""" + cfg = { + "model_type": "zz_new_wrapper", + "text_config": {"model_type": "zz_new_llm"}, + } + monkeypatch.setattr(tl, "_load_config_json", lambda *a, **k: cfg) + monkeypatch.setattr( + tl, + "latest_transformers_supports", + lambda mt: { + "pypi_version": "5.13.0", + "supported_in_pypi": mt == "zz_new_wrapper", + "supported_in_main": True, + }, + ) + out = tl.check_upgrade_for_model("some-org/mixed-support") + assert out is not None + assert out["model_type"] == "zz_new_wrapper" + assert out["supported_in_pypi"] is False # no install offered + assert out["supported_in_main"] is True + + +def test_install_endpoint_not_mounted_on_v1(): + """The consented pip-install endpoint is a Studio admin action; it must live + on studio_router (kept off the OpenAI-compatible /v1 mount), not router.""" + from routes import inference as ri + + path = "/install-latest-transformers" + assert path in [r.path for r in ri.studio_router.routes] + assert path not in [r.path for r in ri.router.routes] + + +def test_kill_switch_removes_provisioned_latest_from_routing(tmp_path, monkeypatch): + """UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS must roll back a provisioned latest + sidecar: no overlay mapping, no probe participation, no file deletion needed.""" + venv_dir = tmp_path / ".venv_t5_latest" + (venv_dir / "transformers").mkdir(parents = True) + (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + + assert tv._overlay_transformers_dir("latest") == str(venv_dir / "transformers") + assert tv._probe_tier_order() == tv._PROBE_TIER_ORDER + ("latest",) + + monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1") + tv._config_mapping_cache.pop("latest", None) + assert tv._overlay_transformers_dir("latest") is None + assert tv._probe_tier_order() == tv._PROBE_TIER_ORDER + tv._config_mapping_cache.pop("latest", None) + + +def test_repair_failure_preserves_pin_and_live_dir(tmp_path, monkeypatch): + """A failed lazy repair must not delete the incomplete-but-pinned live + sidecar: the pin survives so a later attempt can still repair it.""" + venv_dir = tmp_path / ".venv_t5_latest" + venv_dir.mkdir() + (venv_dir / tv._LATEST_PIN_MARKER).write_text("5.13.0") + (venv_dir / "partial_file").write_text("x") + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + monkeypatch.setattr(tv, "_venv_dir_is_valid", lambda *a: False) + monkeypatch.setattr(tv, "_ensure_venv_dir", lambda *a, **k: False) + + from utils.transformers_version import latest_venv_pinned_version + + assert tv._ensure_venv_t5_latest_exists() is False + assert venv_dir.is_dir() + assert (venv_dir / "partial_file").exists() + assert latest_venv_pinned_version() == "5.13.0" + assert not (tmp_path / ".venv_t5_latest.staging").exists() + + +def test_failed_staging_install_removes_staging_dir(tmp_path, monkeypatch): + """A pip failure inside _ensure_venv_dir returns False without raising, so + the except cleanup never runs; the partial staging dir must still go.""" + venv_dir = tmp_path / ".venv_t5_latest" + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(venv_dir)) + + def _fake_ensure(dir_, packages, label): + Path(dir_).mkdir(parents = True, exist_ok = True) + (Path(dir_) / "partial").write_text("x") + return False + + monkeypatch.setattr(tv, "_ensure_venv_dir", _fake_ensure) + assert ensure_latest_transformers_venv("5.13.0") is False + assert not Path(str(venv_dir) + ".staging").exists() diff --git a/studio/backend/tests/test_transformers_version.py b/studio/backend/tests/test_transformers_version.py index b9b5abb9e5..a6e6803a5c 100644 --- a/studio/backend/tests/test_transformers_version.py +++ b/studio/backend/tests/test_transformers_version.py @@ -2550,3 +2550,649 @@ class TestHfEndpointUnreachable: t0 = time.time() result = hf_endpoint_unreachable(timeout = 2) assert result is True and (time.time() - t0) < 6.0 + + +class TestLatestTierActiveFor: + """latest_tier_active_for: the 16-bit guard for the consented latest sidecar.""" + + @staticmethod + def _pin( + monkeypatch, + tv, + version = "5.13.1", + ): + monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: version) + monkeypatch.setattr(tv, "_remote_lora_base", lambda name, hf_token = None: None) + + def test_true_when_tier_latest(self, monkeypatch): + import utils.transformers_version as tv + + self._pin(monkeypatch, tv) + monkeypatch.setattr(tv, "get_transformers_tier", lambda *a, **k: "latest") + assert tv.latest_tier_active_for("Zyphra/ZAYA1-8B") is True + + def test_false_for_fixed_tiers(self, monkeypatch): + import utils.transformers_version as tv + self._pin(monkeypatch, tv) + for tier in ("default", "530", "550", "510"): + monkeypatch.setattr(tv, "get_transformers_tier", lambda *a, _t = tier, **k: _t) + assert tv.latest_tier_active_for("some/model") is False + + def test_false_without_pin_and_no_resolution(self, monkeypatch): + """No sidecar pin returns False before any tier or network resolution.""" + import utils.transformers_version as tv + + def _boom(*a, **k): + raise AssertionError("must not resolve without a pin") + + monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: None) + monkeypatch.setattr(tv, "_remote_lora_base", _boom) + monkeypatch.setattr(tv, "get_transformers_tier", _boom) + assert tv.latest_tier_active_for("Zyphra/ZAYA1-8B") is False + + def test_never_raises(self, monkeypatch): + import utils.transformers_version as tv + + def _boom(*a, **k): + raise RuntimeError("tier resolution exploded") + + self._pin(monkeypatch, tv) + monkeypatch.setattr(tv, "get_transformers_tier", _boom) + assert tv.latest_tier_active_for("some/model") is False + + def test_remote_lora_base_is_resolved(self, monkeypatch): + """A remote adapter is judged by its base model, like worker activation.""" + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: "5.13.1") + monkeypatch.setattr(tv, "_remote_lora_base", lambda name, hf_token = None: "Zyphra/ZAYA1-8B") + tiers = {"Zyphra/ZAYA1-8B": "latest"} + monkeypatch.setattr( + tv, "get_transformers_tier", lambda name, *a, **k: tiers.get(name, "default") + ) + assert tv.latest_tier_active_for("someuser/zaya-lora") is True + + def test_local_checkpoint_config_upgrades(self, monkeypatch, tmp_path): + """An adapter dir with its own config.json merges tiers like activation does.""" + import utils.transformers_version as tv + + adapter = tmp_path / "ckpt" + adapter.mkdir() + (adapter / "adapter_config.json").write_text("{}") + (adapter / "adapter_model.safetensors").write_text("x") + (adapter / "config.json").write_text("{}") + self._pin(monkeypatch, tv) + monkeypatch.setattr(tv, "_resolve_base_model", lambda name: "base/model") + tiers = {"base/model": "default", str(adapter): "latest"} + monkeypatch.setattr( + tv, "get_transformers_tier", lambda name, *a, **k: tiers.get(name, "default") + ) + assert tv.latest_tier_active_for(str(adapter)) is True + + +class TestLatestTierForces16Bit: + """The inference worker and load route refuse bnb 4-bit on the latest sidecar.""" + + def _read(self, rel): + backend_dir = Path(__file__).resolve().parent.parent + return (backend_dir / rel).read_text() + + def test_worker_guard_present(self): + src = self._read("core/inference/worker.py") + assert "latest_tier_active_for" in src, ( + "core/inference/worker.py must force load_in_4bit=False when " + "latest_tier_active_for(model) is true: transformers' grouped-MoE " + "kernels crash on bnb-quantized expert weights for brand-new " + "architectures." + ) + + def test_route_guard_present(self): + src = self._read("routes/inference.py") + assert "latest_tier_active_for" in src, ( + "routes/inference.py must size the VRAM guard with the same 16-bit " + "flip the worker applies for latest-sidecar models." + ) + + def test_validate_route_mirrors_16bit_flip(self): + # Without the same flip, /validate sizes 4-bit and /load then 409s. + src = self._read("routes/inference.py") + body = src.split("async def validate_model", 1)[1].split("\nasync def ", 1)[0] + assert "latest_tier_active_for" in body, ( + "validate_model must apply the latest-sidecar 16-bit flip before " + "_guard_chat_load_against_training so /validate and /load agree." + ) + # First-time loads have no pin yet, so an installable upgrade must also size 16-bit. + assert body.index("check_upgrade_for_model") < body.index( + "_guard_chat_load_against_training" + ), "the upgrade check must run before the training guard" + assert ( + "supported_in_pypi" in body.split("_guard_chat_load_against_training")[0] + ), "an installable upgrade must force 16-bit sizing for the guard" + + def test_validate_offered_upgrade_preserves_custom_code_4bit(self): + # A merely-offered (not installed) upgrade must NOT force 16-bit sizing when the + # model has a custom-code (auto_map) fallback: /load loads it 4-bit without the + # install, and the install route refuses during active training, so 16-bit sizing + # here would 409 the only viable 4-bit path. + src = self._read("routes/inference.py") + body = src.split("async def validate_model", 1)[1].split("\nasync def ", 1)[0] + flip = body.split("Mirror /load's latest-sidecar 16-bit flip", 1)[1].split( + "_guard_chat_load_against_training", 1 + )[0] + assert "not requires_trust_remote_code" in flip, ( + "the offered-upgrade 16-bit flip must be gated on the absence of a custom-code " + "fallback so /validate does not 409 a 4-bit load /load would allow" + ) + # requires_trust_remote_code must be resolved before the flip consumes it. + assert body.index("requires_trust_remote_code = any(") < body.index( + "not requires_trust_remote_code" + ) + + def test_install_route_guards_active_latest_workers(self): + # Stage-and-swap replaces .venv_t5_latest in place, so a live worker on the + # old sidecar would lazy-import files from the new version. + src = self._read("routes/inference.py") + body = src.split("async def install_latest_transformers_route", 1)[1].split( + "\nasync def ", 1 + )[0] + assert ( + "is_training_active" in body + and "is_export_active" in body + and "inference_lifecycle_gate" in body + ), ( + "install_latest_transformers_route must refuse while training or export " + "runs, and hold the lifecycle gate while unloading the chat model and " + "swapping the sidecar." + ) + # The unload (via before_swap so failed installs keep the model), the export-worker + # teardown, and the install must all sit INSIDE the gate so no /load interleaves. + assert "unload_model(active)" in body + assert "cleanup_memory()" in body + # Export teardown precedes the chat unload so its failure aborts with the model still loaded. + assert body.index("cleanup_memory()") < body.index("unload_model(active)") + assert "install_latest_transformers(" in body and "_unload_before_swap" in body + # The gate must be owned by the shielded task, not the request coroutine: a cancelled + # POST unwinding an async-with would release the only guard /load honors mid-install. + gated_task = body.split("async def _gated_install", 1)[1] + assert "inference_lifecycle_gate():" in gated_task + assert "asyncio.to_thread(_run_install)" in gated_task + # The reservation must be taken BEFORE the (awaitable) gate wait, or a + # training/export start could slip in while this request queues on the gate. + assert body.index("try_begin_sidecar_swap()") < body.index( + "inference_lifecycle_gate():" + ), "the swap reservation must be raised before waiting on the lifecycle gate" + # A failed teardown must abort the swap (raise), not fall through to it. + assert body.count("raise RuntimeError") >= 3, ( + "export, chat-unload, and idle-worker teardown failures must raise so " + "the staged install never swaps under a live worker" + ) + # The installer thread owns (and releases) the reservation, shielded from + # request cancellation, so a cancelled POST cannot unlock a live swap. + assert "asyncio.shield" in body and "end_sidecar_swap()" in body + # In-flight generation streams predate the gate; the route refuses rather than kill them + # via the before_swap unload. The count is rechecked UNDER the gate, since a wait on a + # long /load outlasts the pre-gate fast path and streams take this same gate. + assert "other_inference_request_count" in body + gated_task = body.split("async def _gated_install", 1)[1] + assert "other_inference_request_count" in gated_task + + def test_start_routes_refuse_during_install(self): + # A worker spawned mid-swap could activate a half-replaced sidecar. + training = self._read("routes/training.py") + start = training.split("async def start_training", 1)[1].split("\nasync def ", 1)[0] + assert ( + "is_install_in_progress" in start + ), "training /start must refuse while a transformers install is in progress" + export = self._read("routes/export.py") + helper = export.split("def _ensure_export_supported", 1)[1].split("\ndef ", 1)[0] + assert ( + "is_install_in_progress" in helper + ), "mutating export routes must refuse while a transformers install is in progress" + + def test_spawn_sites_recheck_reservation(self): + # The route-level guards are one-shot; validation between them and the + # actual spawn can outlast an install's start, so the spawn itself rechecks. + training = self._read("core/training/training.py") + assert ( + training.count("sidecar_swap_in_progress()") >= 2 + ), "both training spawn sites must recheck the sidecar swap reservation" + export = self._read("core/export/orchestrator.py") + spawn = export.split("def _spawn_subprocess", 1)[1].split("\n def ", 1)[0] + assert ( + "sidecar_swap_kind()" in spawn + ), "the export subprocess spawn must recheck the sidecar swap reservation" + # Training marks the spawn active BEFORE its recheck, so either side sees the other: + # is_training_active covers the window between proc.start() and the _proc assignment. + assert training.index("self._spawn_in_progress = True") < training.index( + "if sidecar_swap_in_progress():" + ) + active = training.split("def is_training_active", 1)[1].split("\n def ", 1)[0] + assert "_spawn_in_progress" in active + # Export load-checkpoint refuses BEFORE tearing down the old worker, so a + # lost race against an install keeps the loaded checkpoint (no bare 500). + loadck = export.split("def load_checkpoint", 1)[1].split("\n def ", 1)[0] + assert loadck.index("sidecar_swap_in_progress()") < loadck.index("_shutdown_subprocess()") + # The training handshake precedes the VRAM-freeing before_spawn hook, so + # losing the race never tears down chat/export for a run that won't spawn. + assert training.index("self._spawn_in_progress = True") < training.index("before_spawn()") + # The spawn-time export check is op-aware for installs (the install side + # aborts on is_export_active) but always refuses for repairs, which have + # no such abort and can be rebuilding the sidecar right now. + assert ( + '_swap_kind == "repair" or (_swap_kind is not None and not self._export_active)' + in spawn + ) + + +class TestSidecarSwapReservation: + """The lazy repair takes the same reservation the install route and worker starts use.""" + + def _repair_setup(self, monkeypatch, tmp_path): + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest")) + monkeypatch.setattr( + tv, + "_latest_pin_data", + lambda: { + "version": "5.99.0", + "packages": ["transformers==5.99.0"], + }, + ) + monkeypatch.setattr(tv, "_venv_dir_is_valid", lambda d, p: False) + monkeypatch.setattr(tv, "_env_offline", lambda: False) + return tv + + def test_repair_holds_reservation_during_swap(self, monkeypatch, tmp_path): + tv = self._repair_setup(monkeypatch, tmp_path) + seen = {} + + def _fake_swap( + version, + packages, + before_swap = None, + ): + seen["active_during_swap"] = tv.sidecar_swap_in_progress() + return True + + monkeypatch.setattr(tv, "_stage_and_swap_latest_venv", _fake_swap) + assert tv._ensure_venv_t5_latest_exists() is True + assert seen["active_during_swap"] is True + assert tv.sidecar_swap_in_progress() is False + + def test_foreign_process_lock_file_visible(self, monkeypatch, tmp_path): + """A repair in a LIVE worker subprocess is seen (via the lock file) by this + process, and its lock is never broken while the owner is alive.""" + import os + import time + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest")) + lock = tv._swap_lock_path() + lock.parent.mkdir(parents = True, exist_ok = True) + # A live owner (this process): visible and never reclaimed, even once aged past + # the cutoff -- a slow but live pip install must keep its lock. + lock.write_text('{"pid": %d}' % os.getpid()) + assert tv.sidecar_swap_in_progress() is True + assert tv.try_begin_sidecar_swap() is False + old_ts = time.time() - 3 * 60 * 60 + os.utime(lock, (old_ts, old_ts)) + assert tv.sidecar_swap_in_progress() is True + assert tv.try_begin_sidecar_swap() is False + + def test_dead_owner_lock_reclaimed_promptly(self, monkeypatch, tmp_path): + """A fresh lock whose recorded owner is dead is reclaimed at once, not after the + long cutoff: a crash mid-install must not wedge loads/training/export for hours.""" + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest")) + lock = tv._swap_lock_path() + lock.parent.mkdir(parents = True, exist_ok = True) + # 999999 is not a live PID: a fresh dead-owner lock is immediately stale. + lock.write_text('{"pid": 999999, "kind": "install"}') + assert tv._pid_alive(999999) is False + assert tv.sidecar_swap_in_progress() is False + assert tv.try_begin_sidecar_swap() is True + try: + assert lock.is_file() + finally: + tv.end_sidecar_swap() + assert not lock.exists() + + def test_unreadable_pid_lock_uses_age_cutoff(self, monkeypatch, tmp_path): + """A lock with no readable owner PID (mid create-before-write, or corrupt) is not + reclaimed while fresh -- only after the long cutoff -- so a lock a live owner just + created is not stolen before its PID lands.""" + import os + import time + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest")) + lock = tv._swap_lock_path() + lock.parent.mkdir(parents = True, exist_ok = True) + lock.write_text("") # created but metadata not yet written + assert tv.sidecar_swap_in_progress() is True + old_ts = time.time() - (tv._SWAP_LOCK_STALE_SECS + 60) + os.utime(lock, (old_ts, old_ts)) + assert tv.sidecar_swap_in_progress() is False + + def test_repair_refused_while_install_holds_reservation(self, monkeypatch, tmp_path): + tv = self._repair_setup(monkeypatch, tmp_path) + + def _must_not_run(*a, **k): + raise AssertionError("repair must not swap while an install is in progress") + + monkeypatch.setattr(tv, "_stage_and_swap_latest_venv", _must_not_run) + assert tv.try_begin_sidecar_swap() is True + try: + assert tv._ensure_venv_t5_latest_exists() is False + finally: + tv.end_sidecar_swap() + + +class TestRecoverStrandedSidecar: + """A swap whose activation rename AND rollback both fail strands the previous sidecar + at .old with no live dir (its pin marker went with it). Reading the pin self-heals it, + but never while a swap legitimately holds the reservation.""" + + def _setup(self, monkeypatch, tmp_path): + import utils.transformers_version as tv + + live = str(tmp_path / "venv_t5_latest") + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", live) + # Stranded state: live gone, previous sidecar (with its marker) sits at .old. + retired = Path(live + ".old") + retired.mkdir(parents = True) + (retired / tv._LATEST_PIN_MARKER).write_text( + '{"version": "5.99.0", "packages": ["transformers==5.99.0"]}' + ) + return tv, Path(live), retired + + def test_stranded_old_recovered_on_pin_read(self, monkeypatch, tmp_path): + tv, live, retired = self._setup(monkeypatch, tmp_path) + data = tv._latest_pin_data() + assert live.is_dir() + assert not retired.exists() + assert data is not None and data["version"] == "5.99.0" + + def test_stranded_recovery_skipped_during_swap(self, monkeypatch, tmp_path): + tv, live, retired = self._setup(monkeypatch, tmp_path) + assert tv.try_begin_sidecar_swap() is True + try: + # A swap holds the reservation and may be mid-rename; do not race it. + assert tv._latest_pin_data() is None + assert not live.exists() + assert retired.is_dir() + finally: + tv.end_sidecar_swap() + # Once the swap is done, the next pin read recovers the stranded sidecar. + assert tv._latest_pin_data() is not None + assert live.is_dir() + + +class TestCachedLatestMappingRevalidated: + """A cached 'latest' mapping is dropped and re-resolved when the sidecar since broke + in-process, so routing self-heals instead of trusting a mapping parsed from a sidecar + that no longer exists (which would keep routing latest-only models to a broken tier).""" + + def test_broken_sidecar_drops_cached_latest_mapping(self, monkeypatch): + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "_config_mapping_cache", {"latest": frozenset({"brandnew"})}) + monkeypatch.setattr(tv, "_latest_sidecar_intact", lambda: False) + seen = {"n": 0} + + def _fake_overlay(tier): + seen["n"] += 1 + return None # broken/unavailable -> empty, uncached + + monkeypatch.setattr(tv, "_overlay_transformers_dir", _fake_overlay) + assert tv._config_model_types("latest") == frozenset() + assert seen["n"] == 1 # re-resolved, not served from the stale cache + assert "latest" not in tv._config_mapping_cache + + def test_intact_sidecar_serves_cached_latest_mapping(self, monkeypatch): + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "_config_mapping_cache", {"latest": frozenset({"brandnew"})}) + monkeypatch.setattr(tv, "_latest_sidecar_intact", lambda: True) + monkeypatch.setattr( + tv, + "_overlay_transformers_dir", + lambda tier: pytest.fail("intact sidecar must serve the cache without re-resolving"), + ) + assert tv._config_model_types("latest") == frozenset({"brandnew"}) + + def test_non_latest_cache_not_revalidated(self, monkeypatch): + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "_config_mapping_cache", {"530": frozenset({"gemma3"})}) + monkeypatch.setattr( + tv, + "_latest_sidecar_intact", + lambda: pytest.fail("non-latest tiers must not pay the sidecar-intact check"), + ) + assert tv._config_model_types("530") == frozenset({"gemma3"}) + + def test_deleted_pin_drops_cached_latest_mapping(self, monkeypatch, tmp_path): + # A pin marker deleted after the mapping was cached makes _latest_pin_data None; + # the cache must be dropped (not trusted), so routing re-resolves to no latest tier + # rather than routing to a latest tier that then fails worker activation. + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest")) + monkeypatch.setattr(tv, "_latest_tier_disabled", lambda: False) + monkeypatch.setattr(tv, "_config_mapping_cache", {"latest": frozenset({"brandnew"})}) + # No pin marker on disk -> _latest_pin_data() is None -> not intact. + assert tv._latest_sidecar_intact() is False + assert tv._config_model_types("latest") == frozenset() + assert "latest" not in tv._config_mapping_cache + + +class TestOverlayRepairsIncompleteSidecar: + """Routing self-heals a pinned latest sidecar that is present but incomplete, + not only one whose transformers/ dir vanished: workers refuse parent-only + repairs, so a sidecar missing a pinned package would fail every load.""" + + def _setup(self, monkeypatch, tmp_path, valid): + import utils.transformers_version as tv + + live = tmp_path / "venv_t5_latest" + (live / "transformers").mkdir(parents = True) + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(live)) + monkeypatch.setattr(tv, "_latest_tier_disabled", lambda: False) + monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: "5.99.0") + monkeypatch.setattr( + tv, + "_latest_pin_data", + lambda: {"version": "5.99.0", "packages": ["transformers==5.99.0", "tiktoken"]}, + ) + monkeypatch.setattr(tv, "_venv_dir_is_valid", lambda d, p: valid) + monkeypatch.setattr(tv, "_latest_repair_failed_at", 0.0) + return tv + + def test_incomplete_sidecar_triggers_repair(self, monkeypatch, tmp_path): + tv = self._setup(monkeypatch, tmp_path, valid = False) + called = {"n": 0} + + def _fake_repair(): + called["n"] += 1 + return True + + monkeypatch.setattr(tv, "_ensure_venv_t5_latest_exists", _fake_repair) + src = tv._overlay_transformers_dir("latest") + assert called["n"] == 1 + assert src == str(tmp_path / "venv_t5_latest" / "transformers") + + def test_intact_sidecar_skips_repair(self, monkeypatch, tmp_path): + tv = self._setup(monkeypatch, tmp_path, valid = True) + + def _must_not_run(): + raise AssertionError("intact sidecar must not trigger a repair") + + monkeypatch.setattr(tv, "_ensure_venv_t5_latest_exists", _must_not_run) + assert tv._overlay_transformers_dir("latest") == str( + tmp_path / "venv_t5_latest" / "transformers" + ) + + def test_failed_repair_backs_off(self, monkeypatch, tmp_path): + tv = self._setup(monkeypatch, tmp_path, valid = False) + called = {"n": 0} + + def _fake_repair(): + called["n"] += 1 + return False + + monkeypatch.setattr(tv, "_ensure_venv_t5_latest_exists", _fake_repair) + # A failed repair must not route through the broken sidecar, neither on + # the failing attempt nor while the backoff suppresses the next attempt. + assert tv._overlay_transformers_dir("latest") is None + assert tv._overlay_transformers_dir("latest") is None + assert called["n"] == 1 + + +class TestStageAndSwapBeforeSwap: + """before_swap fires only when the staged install succeeded and the swap is next.""" + + def _setup(self, monkeypatch, tmp_path, build_ok): + import utils.transformers_version as tv + + live = tmp_path / "venv_latest" + monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(live)) + + def _fake_build(target, packages, label): + if build_ok: + Path(target).mkdir(parents = True, exist_ok = True) + return build_ok + + monkeypatch.setattr(tv, "_ensure_venv_dir", _fake_build) + return tv, live + + def test_called_after_successful_staging(self, monkeypatch, tmp_path): + tv, live = self._setup(monkeypatch, tmp_path, build_ok = True) + calls = [] + assert tv._stage_and_swap_latest_venv( + "5.99.0", ("transformers==5.99.0",), before_swap = lambda: calls.append(1) + ) + assert calls == [1] and live.is_dir() + + def test_not_called_when_staging_fails(self, monkeypatch, tmp_path): + tv, live = self._setup(monkeypatch, tmp_path, build_ok = False) + calls = [] + assert not tv._stage_and_swap_latest_venv( + "5.99.0", ("transformers==5.99.0",), before_swap = lambda: calls.append(1) + ) + assert calls == [] and not live.exists() + + def test_failure_in_before_swap_keeps_previous_sidecar(self, monkeypatch, tmp_path): + tv, live = self._setup(monkeypatch, tmp_path, build_ok = True) + live.mkdir() + (live / "sentinel").write_text("old") + + def _boom(): + raise RuntimeError("worker teardown failed") + + assert not tv._stage_and_swap_latest_venv( + "5.99.0", ("transformers==5.99.0",), before_swap = _boom + ) + assert (live / "sentinel").read_text() == "old" + + +class TestKillSwitchBeatsMappingCache: + def test_cached_latest_probe_ignored_when_disabled(self, monkeypatch): + import utils.transformers_version as tv + + key = tv._probe_cache_key("some/model") + monkeypatch.setitem(tv._probe_tier_cache, key, "latest") + monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1") + # With the switch set, the cached latest entry must not short-circuit; + # the probe re-resolves against the non-latest order (stub it to 530). + monkeypatch.setattr(tv, "_probe_tier_venvs", lambda: {}) + monkeypatch.setattr(tv, "_probe_tier_order", lambda: ()) + assert tv._probe_tier("some/model", None, "test") != "latest" + # Cached non-latest entries and the unset switch still short-circuit. + monkeypatch.delenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS") + assert tv._probe_tier("some/model", None, "test") == "latest" + + def test_cached_latest_mapping_ignored_when_disabled(self, monkeypatch): + import utils.transformers_version as tv + + monkeypatch.setitem(tv._config_mapping_cache, "latest", frozenset({"brandnew"})) + # The cache is trusted only when the sidecar is intact; hold it intact so this + # test isolates the kill switch, not the sidecar-revalidation path. + monkeypatch.setattr(tv, "_latest_sidecar_intact", lambda: True) + monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1") + assert tv._config_model_types("latest") == frozenset() + monkeypatch.delenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS") + assert tv._config_model_types("latest") == frozenset({"brandnew"}) + + +class TestRaiseTierForNested: + """_raise_tier_for_nested: a wrapper's nested model_type can raise a fast-path tier.""" + + def _patch_types(self, monkeypatch, per_tier): + import utils.transformers_version as tv + monkeypatch.setattr( + tv, "_config_model_types", lambda tier: frozenset(per_tier.get(tier, ())) + ) + + def test_nested_latest_only_type_raises(self, monkeypatch): + import utils.transformers_version as tv + + self._patch_types(monkeypatch, {"550": {"gemma4"}, "latest": {"gemma4", "brandnew_arch"}}) + cfg = {"model_type": "gemma4", "text_config": {"model_type": "brandnew_arch"}} + assert tv._raise_tier_for_nested(cfg, "550") == "latest" + + def test_never_lowers_a_fast_path_tier(self, monkeypatch): + import utils.transformers_version as tv + + # Mapping alone would say 530, but the fast path (e.g. a name override) said 550. + self._patch_types(monkeypatch, {"530": {"qwen3_5"}, "550": {"qwen3_5"}}) + assert tv._raise_tier_for_nested({"model_type": "qwen3_5"}, "550") == "550" + + def test_no_config_keeps_tier(self): + import utils.transformers_version as tv + assert tv._raise_tier_for_nested(None, "550") == "550" + + def test_unknown_nested_type_never_vetoes(self, monkeypatch): + import utils.transformers_version as tv + + # A nested type unknown everywhere (not even latest) keeps the fast path. + self._patch_types(monkeypatch, {"550": {"gemma4"}, "latest": {"gemma4"}}) + cfg = {"model_type": "gemma4", "text_config": {"model_type": "unreleased"}} + assert tv._raise_tier_for_nested(cfg, "550") == "550" + + def test_name_fast_path_folds_when_latest_pinned(self, monkeypatch): + """A fixed-tier name match with a latest-only model_type routes to latest + once the sidecar is pinned; without a pin the name tier stands (no I/O).""" + import utils.transformers_version as tv + + self._patch_types(monkeypatch, {"550": {"gemma4"}, "latest": {"brandnew_arch"}}) + monkeypatch.setattr(tv, "_tier_from_name", lambda name: ("550", "gemma-4")) + monkeypatch.setattr( + tv, "_load_config_json", lambda name, tok = None: {"model_type": "brandnew_arch"} + ) + monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: "5.99.0") + assert tv.get_transformers_tier("org/gemma-4-new", probe = False) == "latest" + monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: None) + monkeypatch.setattr( + tv, + "_load_config_json", + lambda name, tok = None: (_ for _ in ()).throw(AssertionError("no I/O without a pin")), + ) + assert tv.get_transformers_tier("org/gemma-4-new", probe = False) == "550" + + def test_fast_path_folds_nested_tier(self, monkeypatch, tmp_path): + """End to end: a local wrapper config on a fixed fast path routes to latest + when its nested type only exists in the installed latest sidecar.""" + import utils.transformers_version as tv + + ckpt = tmp_path / "wrapper" + ckpt.mkdir() + (ckpt / "config.json").write_text( + json.dumps({"model_type": "gemma4", "text_config": {"model_type": "brandnew_arch"}}) + ) + self._patch_types(monkeypatch, {"550": {"gemma4"}, "latest": {"gemma4", "brandnew_arch"}}) + monkeypatch.setattr(tv, "_config_needs_510", lambda cfg: False) + monkeypatch.setattr(tv, "_config_needs_550", lambda cfg: True) + assert tv.get_transformers_tier(str(ckpt), probe = False) == "latest" diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 281ca24281..284bbb5745 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -698,6 +698,25 @@ if backend_dir not in sys.path: try: from transformers import AutoConfig + # Union the ACTIVE sidecar's registry into the inlined parent-process sets + # so architectures only the sidecar knows still classify correctly. + try: + from transformers.models.auto import modeling_auto as _ma + for _attr in ("MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES", + "MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES"): + _d = dict(getattr(_ma, _attr, None) or {}) + _VLM_MODEL_TYPES |= set(_d) + _VLM_CLASS_NAMES |= set(_d.values()) + for _attr in ("MODEL_FOR_CTC_MAPPING_NAMES", + "MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES", + "MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES", + "MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING_NAMES", + "MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING_NAMES", + "MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES"): + _AUDIO_ONLY_MODEL_TYPES |= set(dict(getattr(_ma, _attr, None) or {})) + except Exception: + pass + # Capability detection never executes model repo code. kwargs = {"trust_remote_code": False} if token: @@ -727,13 +746,23 @@ def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None) """ token_arg = hf_token or "" + # Latest-only architectures need the latest sidecar for AutoConfig; + # other tiers keep the 5.5 sidecar. + sidecar_dir = _VENV_T5_DIR + try: + from utils.transformers_version import _VENV_T5_LATEST_DIR, get_transformers_tier + if get_transformers_tier(model_name, hf_token, probe = False) == "latest": + sidecar_dir = _VENV_T5_LATEST_DIR + except Exception: + pass + try: result = subprocess.run( [ sys.executable, "-c", _VISION_CHECK_SCRIPT, - _VENV_T5_DIR, + sidecar_dir, _BACKEND_DIR, model_name, token_arg, @@ -876,6 +905,17 @@ def _is_vision_model_uncached( model_name, hf_token = hf_token, local_files_only = local_files_only ) if raw is not None: + if raw is False and not local_files_only: + # Raw heuristics predate latest-only architectures; on the latest tier, + # trust that sidecar's AutoConfig probe over the heuristic False. An + # inconclusive probe (sidecar mid-repair, timeout) is transient: return + # None so the heuristic False is not cached and the model is re-probed. + try: + from utils.transformers_version import get_transformers_tier + if get_transformers_tier(model_name, hf_token, probe = False) == "latest": + return _is_vision_model_subprocess(model_name, hf_token = hf_token) + except Exception: + pass return raw # Raw read failed transiently: fall back to AutoConfig (remote code DISABLED), via a diff --git a/studio/backend/utils/transformers_latest.py b/studio/backend/utils/transformers_latest.py new file mode 100644 index 0000000000..40c8f729a5 --- /dev/null +++ b/studio/backend/utils/transformers_latest.py @@ -0,0 +1,607 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Latest-transformers support check for brand-new model architectures. + +When a model's ``model_type`` is absent from every installed transformers overlay +(base 4.57.x plus the .venv_t5_530/550/510 sidecars and, if provisioned, .venv_t5_latest), +Studio cannot load it today. This module answers, without authentication, code execution, +or trust_remote_code: + + 1. Does the LATEST transformers release on PyPI ship this ``model_type``? + 2. Does transformers ``main`` on GitHub ship it (dev-only, not yet installable)? + +Sources (all unauthenticated; raw.githubusercontent.com is not API rate-limited and +api.github.com is deliberately never used): + - https://pypi.org/pypi/transformers/json -> latest release version + - https://raw.githubusercontent.com/huggingface/transformers/{ref}/src/transformers/ + models/auto/configuration_auto.py + auto_mappings.py -> CONFIG_MAPPING_NAMES + +The fetched sources are parsed with the same AST extractor the static router uses +(:func:`utils.transformers_version._model_types_from_source`), so the remote answer is +computed exactly like the local overlay answer. + +Results are cached in memory and in a small JSON snapshot under ``studio_root()/cache`` +(ttl ~1 day) so repeated tier resolutions never re-fetch; failures are backed off in +memory. Every fetch is bounded (<=5s, one retry), so a hung network cannot block model +loading. Fully offline-safe: offline env vars or the kill switch +``UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS=1`` make every check return None (current +behavior preserved). + +The consented install path (:func:`install_latest_transformers`) provisions the +persistent ``.venv_t5_latest`` sidecar via +:func:`utils.transformers_version.ensure_latest_transformers_venv`. +""" + +import json +import os +import threading +import time +from pathlib import Path + +from loggers import get_logger +from utils.paths.storage_roots import studio_root as _studio_root +from utils.transformers_version import ( + _env_offline, + _load_config_json, + _model_types_from_source, + _tier_from_config_mapping, + _config_model_types, + _NESTED_CONFIG_KEYS, + _TIER_RANK, + _model_types_from_config, + _TRANSFORMERS_510_MODEL_TYPES, + _TRANSFORMERS_530_MODEL_TYPES, + _TRANSFORMERS_550_MODEL_TYPES, + ensure_latest_transformers_venv, + latest_venv_pinned_version, +) + +logger = get_logger(__name__) + +_PYPI_JSON_URL = "https://pypi.org/pypi/transformers/json" +_RAW_URL = ( + "https://raw.githubusercontent.com/huggingface/transformers/{ref}" + "/src/transformers/models/auto/{name}" +) +_AUTO_FILES = ("configuration_auto.py", "auto_mappings.py") + +_FETCH_TIMEOUT_SECONDS = 5.0 +_FETCH_RETRIES = 1 +_CACHE_TTL_SECONDS = 24 * 60 * 60 +_FAILURE_BACKOFF_SECONDS = 300 + +_CACHE_FILE_NAME = "transformers_latest_check.json" +_SNAPSHOT_SCHEMA = 1 + +# Snapshot: {"schema", "fetched_at", "pypi_version", "pypi_model_types", "main_model_types"}. +# Install-in-progress state lives in utils.transformers_version (the sidecar swap reservation). +_lock = threading.Lock() +_memory_snapshot: dict | None = None +_last_failure_at: float = 0.0 +_is_fetching: bool = False + +_TRUE_VALUES = {"1", "true", "yes", "on"} + + +def _disabled() -> bool: + """True if the operator disabled the latest-transformers check entirely.""" + return ( + os.environ.get("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "").strip().lower() in _TRUE_VALUES + ) + + +def _cache_file() -> Path: + return _studio_root() / "cache" / _CACHE_FILE_NAME + + +# Sentinel for HTTP 404 (absent at ref), distinct from transient failures. +_FETCH_MISSING = "__unsloth_fetch_missing__" + + +def _fetch_text(url: str) -> str | None: + """GET *url* with a bounded timeout and one retry; None on any failure. + + Returns ``_FETCH_MISSING`` (without retrying) on HTTP 404 so callers can tell + "absent at this ref" apart from "network flaked". + """ + import urllib.error + import urllib.request + + for attempt in range(1 + _FETCH_RETRIES): + try: + req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-studio"}) + with urllib.request.urlopen(req, timeout = _FETCH_TIMEOUT_SECONDS) as resp: + return resp.read().decode("utf-8", "replace") + except urllib.error.HTTPError as exc: + if exc.code == 404: + return _FETCH_MISSING + logger.debug("Fetch failed (attempt %d) for %s: %s", attempt + 1, url, exc) + except Exception as exc: + logger.debug("Fetch failed (attempt %d) for %s: %s", attempt + 1, url, exc) + return None + + +def _fetch_latest_pypi_version() -> str | None: + """Latest transformers release version from PyPI's unauthenticated JSON API.""" + body = _fetch_text(_PYPI_JSON_URL) + if body is None or body == _FETCH_MISSING: + return None + try: + version = json.loads(body).get("info", {}).get("version") + except Exception as exc: + logger.debug("Could not parse PyPI JSON: %s", exc) + return None + return version if isinstance(version, str) and version else None + + +def _fetch_remote_model_types(ref: str) -> frozenset[str] | None: + """CONFIG_MAPPING_NAMES keys at *ref* (a release tag like ``v5.12.0`` or ``main``). + + Fetches configuration_auto.py plus auto_mappings.py (the 5.10+ split) from + raw.githubusercontent.com and parses them with the shared AST extractor. A file + that 404s (auto_mappings.py on pre-5.10 tags) is skipped, but a transient fetch + or parse failure of EITHER file fails the whole lookup: most model types live in + auto_mappings.py on current releases, so a partial map cached for the TTL would + make /validate skip the upgrade prompt for architectures the release does ship. + An empty result is likewise a failure so it is never cached as "supports nothing". + """ + keys: set[str] = set() + fetched_any = False + for name in _AUTO_FILES: + source = _fetch_text(_RAW_URL.format(ref = ref, name = name)) + if source is None: + return None + if source == _FETCH_MISSING: + continue + fetched_any = True + try: + keys |= _model_types_from_source(source) + except Exception as exc: + logger.debug("Could not parse %s at %s: %s", name, ref, exc) + return None + if not fetched_any or not keys: + return None + return frozenset(keys) + + +def _load_snapshot_file() -> dict | None: + """Persisted snapshot from disk, or None (missing/corrupt/old schema).""" + try: + with open(_cache_file(), encoding = "utf-8") as f: + data = json.load(f) + except Exception: + return None + if not isinstance(data, dict) or data.get("schema") != _SNAPSHOT_SCHEMA: + return None + if not isinstance(data.get("fetched_at"), (int, float)): + return None + if not isinstance(data.get("pypi_version"), str): + return None + for key in ("pypi_model_types", "main_model_types"): + value = data.get(key) + if not isinstance(value, list) or not all(isinstance(v, str) for v in value): + return None + return data + + +def _save_snapshot_file(snapshot: dict) -> None: + """Atomic best-effort write (tmp + os.replace, Windows-safe); failures only log.""" + path = _cache_file() + tmp = path.with_name(path.name + ".tmp") + try: + path.parent.mkdir(parents = True, exist_ok = True) + tmp.write_text(json.dumps(snapshot), encoding = "utf-8") + os.replace(tmp, path) + except Exception as exc: + logger.debug("Could not persist %s: %s", path, exc) + try: + tmp.unlink(missing_ok = True) + except Exception: + pass + + +def _snapshot_is_fresh(snapshot: dict | None) -> bool: + return ( + snapshot is not None + and (time.time() - float(snapshot.get("fetched_at", 0))) < _CACHE_TTL_SECONDS + ) + + +def _refresh_snapshot() -> dict | None: + """Fetch a fresh snapshot from PyPI + raw.githubusercontent.com; None on failure. + + The PyPI version and its tagged mapping are required; the ``main`` mapping is + best-effort (recorded as an empty list plus ``main_checked=False`` when unavailable, + so a dev-only architecture is reported as "unknown" rather than "unsupported"). + """ + version = _fetch_latest_pypi_version() + if version is None: + return None + pypi_types = _fetch_remote_model_types(f"v{version}") + if pypi_types is None: + return None + main_types = _fetch_remote_model_types("main") + return { + "schema": _SNAPSHOT_SCHEMA, + "fetched_at": time.time(), + "pypi_version": version, + "pypi_model_types": sorted(pypi_types), + "main_model_types": sorted(main_types) if main_types is not None else [], + "main_checked": main_types is not None, + } + + +def _get_snapshot() -> dict | None: + """Current support snapshot: memory -> disk -> network, with TTL and failure backoff. + + The network refresh runs outside the lock so a slow fetch cannot stall other + threads in the ASGI pool; _is_fetching deduplicates concurrent refreshes + (losers return None, the graceful fallthrough, rather than waiting). + """ + global _memory_snapshot, _last_failure_at, _is_fetching + with _lock: + if _snapshot_is_fresh(_memory_snapshot): + return _memory_snapshot + disk = _load_snapshot_file() + if _snapshot_is_fresh(disk): + _memory_snapshot = disk + return disk + if _disabled() or _env_offline(): + return None + if time.time() - _last_failure_at < _FAILURE_BACKOFF_SECONDS: + return None + if _is_fetching: + return None + _is_fetching = True + fresh = None + try: + fresh = _refresh_snapshot() + finally: + with _lock: + _is_fetching = False + if fresh is None: + _last_failure_at = time.time() + else: + _memory_snapshot = fresh + if fresh is None: + # A stale positive could offer a version PyPI no longer serves; be strict. + return None + _save_snapshot_file(fresh) + return fresh + + +def clear_caches() -> None: + """Test helper: drop the in-memory snapshot, failure backoff, and busy flags.""" + global _memory_snapshot, _last_failure_at, _is_fetching + with _lock: + _memory_snapshot = None + _last_failure_at = 0.0 + _is_fetching = False + from utils.transformers_version import end_sidecar_swap + + end_sidecar_swap() + + +def latest_transformers_supports(model_type: str) -> dict | None: + """Whether the newest transformers (PyPI release and/or GitHub main) ships *model_type*. + + Returns ``{"pypi_version": str, "supported_in_pypi": bool, "supported_in_main": bool}`` + or None when the answer is unavailable (offline, kill switch, network failure) — the + caller must then fall through to current behavior. Cached (memory + JSON snapshot on + disk, ttl ~1 day) so repeated tier resolutions never re-fetch. + """ + if not isinstance(model_type, str) or not model_type: + return None + if _disabled() or _env_offline(): + return None + snapshot = _get_snapshot() + if snapshot is None: + return None + return { + "pypi_version": snapshot["pypi_version"], + "supported_in_pypi": model_type in set(snapshot["pypi_model_types"]), + "supported_in_main": model_type in set(snapshot["main_model_types"]), + } + + +# model_types the hardcoded tier tables already route; never remote-check these. +def _hardcoded_model_types() -> frozenset[str]: + return frozenset( + _TRANSFORMERS_530_MODEL_TYPES + | _TRANSFORMERS_550_MODEL_TYPES + | _TRANSFORMERS_510_MODEL_TYPES + ) + + +def check_upgrade_for_model(model_name: str, hf_token: str | None = None) -> dict | None: + """Upgrade signal for *model_name*, or None when current routing already handles it. + + The tier hook for the pre-load ``/validate`` path: fires ONLY when the model's + ``model_type`` is absent from every installed overlay (and from the hardcoded tier + tables), i.e. exactly when today's load would fail with an unrecognized-architecture + error. Returns ``{"model_type", "pypi_version", "supported_in_pypi", + "supported_in_main"}`` when the newest transformers knows the type, else None. + + Never raises; every network touch is bounded and cached. Offline or with the + ``UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS`` kill switch it returns None immediately. + """ + try: + if _disabled() or _env_offline(): + return None + cfg = _load_config_json(model_name, hf_token) + if not isinstance(cfg, dict): + return None + candidates = _model_types_from_config(cfg) + if not candidates: + return None + # Without a readable base mapping every type looks brand new; bail out. + if not _config_model_types("default"): + return None + hardcoded = _hardcoded_model_types() + missing = [ + candidate + for candidate in candidates + if candidate not in hardcoded + and not any(candidate in _config_model_types(tier) for tier in _TIER_RANK) + ] + if not missing: + return None + # Latest must load EVERY missing type (wrappers build nested sub-configs + # through CONFIG_MAPPING) or the load still fails. + supports = [latest_transformers_supports(candidate) for candidate in missing] + if any( + s is None or not (s["supported_in_pypi"] or s["supported_in_main"]) for s in supports + ): + return None + # Offer the PyPI install only if the release ships every missing type; a + # main-only type in the mix surfaces as dev-only. + model_type = missing[0] + supported_in_pypi = all(s["supported_in_pypi"] for s in supports) + supported_in_main = all(s["supported_in_pypi"] or s["supported_in_main"] for s in supports) + logger.info( + "Model %s has model_type=%s unknown to every installed transformers " + "(latest PyPI %s: %s, main: %s)", + model_name, + model_type, + supports[0]["pypi_version"], + "supported" if supported_in_pypi else "unsupported", + "supported" if supported_in_main else "unsupported", + ) + return { + "model_type": model_type, + "pypi_version": supports[0]["pypi_version"], + "supported_in_pypi": supported_in_pypi, + "supported_in_main": supported_in_main, + } + except Exception as exc: + logger.debug("Latest-transformers check failed for '%s': %s", model_name, exc) + return None + + +# --- Dependency compatibility preflight ------------------------------------------------------ +# Sidecars install transformers --no-deps atop the base env. Before installing, compare +# requires_dist: unsatisfied shadowable deps become exact --target pins, anything else blocks. + +# Safe to shadow inside the sidecar dir (pure wheels, no torch coupling). +_SHADOWABLE_DEPS = frozenset({"tokenizers", "safetensors"}) +# Provided by the sidecar recipe; checked against its pin, not the base env. +_SIDECAR_PROVIDED = {"huggingface-hub": "1.8.0", "hf-xet": "1.4.2"} +# CLI-only; never imported at runtime in Studio's workers. +_IGNORED_DEPS = frozenset({"typer"}) + + +def _canonical_dep_name(name: str) -> str: + return name.lower().replace("_", "-") + + +def _fetch_requires_dist(version: str) -> list[str] | None: + """Core (marker-free, non-extra) requires_dist of transformers *version* from PyPI.""" + body = _fetch_text(f"https://pypi.org/pypi/transformers/{version}/json") + if body is None or body == _FETCH_MISSING: + return None + try: + reqs = json.loads(body).get("info", {}).get("requires_dist") + except Exception: + return None + if not isinstance(reqs, list): + return None + return [r for r in reqs if isinstance(r, str)] + + +def _resolve_exact_version(name: str, specifier) -> str | None: + """Newest PyPI release of *name* satisfying *specifier* (exact pin for the shadow).""" + body = _fetch_text(f"https://pypi.org/pypi/{name}/json") + if body is None or body == _FETCH_MISSING: + return None + try: + from packaging.version import InvalidVersion, Version + + releases = json.loads(body).get("releases", {}) + best = None + for candidate in releases: + try: + parsed = Version(candidate) + except InvalidVersion: + continue + if parsed.is_prerelease or not specifier.contains(candidate): + continue + if best is None or parsed > Version(best): + best = candidate + return best + except Exception as exc: + logger.debug("Could not resolve an exact %s version: %s", name, exc) + return None + + +def compat_plan(version: str) -> tuple[tuple[str, ...], list[str]]: + """(extra exact pins to shadow-install, blocking requirement strings) for *version*. + + Compares the release's core requires_dist against the running base env (the env the + workers overlay the sidecar onto). A requirement the base env satisfies needs nothing; + an unsatisfied shadowable dep becomes an exact pin inside the sidecar; any other + unsatisfied requirement is a blocker. An unavailable requires_dist BLOCKS the + install: proceeding unverified could pin a sidecar whose imports then crash the + workers, and the caller just reached PyPI for the version check so a retry is cheap. + """ + reqs = _fetch_requires_dist(version) + if reqs is None: + return (), ["dependency metadata for this release (could not be fetched from PyPI; retry)"] + try: + from importlib.metadata import PackageNotFoundError + from importlib.metadata import version as _installed_version + from packaging.requirements import InvalidRequirement, Requirement + except Exception: + return (), [] + extras: list[str] = [] + blockers: list[str] = [] + for raw in reqs: + try: + req = Requirement(raw) + except InvalidRequirement: + continue + if req.extras or (req.marker is not None and not req.marker.evaluate()): + continue + name = _canonical_dep_name(req.name) + if name in _IGNORED_DEPS: + continue + if name in _SIDECAR_PROVIDED: + if not req.specifier.contains(_SIDECAR_PROVIDED[name], prereleases = True): + blockers.append(raw) + continue + try: + installed = _installed_version(req.name) + except PackageNotFoundError: + installed = None + if installed is not None and req.specifier.contains(installed, prereleases = True): + continue + if name in _SHADOWABLE_DEPS: + exact = _resolve_exact_version(name, req.specifier) + if exact is None: + blockers.append(raw) + else: + extras.append(f"{name}=={exact}") + else: + blockers.append(raw) + return tuple(extras), blockers + + +def is_install_in_progress() -> bool: + """True while a latest-transformers install or lazy repair holds the sidecar swap + reservation. Training and export starts check this so a fresh worker never + activates the sidecar mid-swap.""" + from utils.transformers_version import sidecar_swap_in_progress + return sidecar_swap_in_progress() + + +def install_latest_transformers( + version: str, + before_swap = None, + reserved: bool = False, +) -> dict: + """Consented install of the latest transformers sidecar; returns a structured result. + + Guards: the requested *version* must match the current PyPI latest from the (cached) + snapshot, so a client cannot pin an arbitrary package version through this endpoint. + On success ``.venv_t5_latest`` is provisioned and pinned; routing then resolves the + new tier automatically on this and every future start. *before_swap* is forwarded + to the stage-and-swap: it runs only after the staged install succeeded, right + before the live sidecar is replaced. *reserved* means the caller already holds the + sidecar swap reservation (the install route takes it before waiting on the + inference lifecycle gate, so worker starts see it for the whole window). + """ + from utils.transformers_version import end_sidecar_swap, try_begin_sidecar_swap + + if not reserved and not try_begin_sidecar_swap(): + return { + "success": False, + "version": version, + "message": "A transformers installation is already in progress.", + } + try: + return _install_latest_transformers_locked(version, before_swap = before_swap) + finally: + if not reserved: + end_sidecar_swap() + + +def _install_latest_transformers_locked(version: str, before_swap = None) -> dict: + """Body of install_latest_transformers; runs with the in-progress flag held.""" + if _disabled(): + return { + "success": False, + "version": version, + "message": "Latest-transformers installs are disabled " + "(UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS).", + } + if _env_offline(): + return { + "success": False, + "version": version, + "message": "Cannot install: Studio is in offline mode.", + } + # Re-verify against a LIVE snapshot (a release may land inside the cache TTL); + # fall back to the cached one on fetch failure. + global _memory_snapshot + snapshot = _refresh_snapshot() + if snapshot is not None: + with _lock: + _memory_snapshot = snapshot + _save_snapshot_file(snapshot) + else: + snapshot = _get_snapshot() + if snapshot is None: + return { + "success": False, + "version": version, + "message": "Could not verify the latest transformers release on PyPI.", + } + if version != snapshot["pypi_version"]: + return { + "success": False, + "version": version, + "message": f"Requested version {version!r} is not the latest transformers " + f"release ({snapshot['pypi_version']}).", + # Lets the consent dialog retry with the release that superseded the + # one /validate saw, instead of re-sending the stale version forever. + "latest_version": snapshot["pypi_version"], + } + extra_packages, blockers = compat_plan(version) + if blockers: + return { + "success": False, + "version": version, + "message": "Cannot install transformers " + f"{version}: this environment does not satisfy {', '.join(blockers)}. " + "A Studio update is required first.", + } + if not ensure_latest_transformers_venv(version, extra_packages, before_swap = before_swap): + return { + "success": False, + "version": version, + "message": f"Installing transformers {version} failed; see the Studio logs.", + } + _invalidate_capability_caches() + return { + "success": True, + "version": version, + "message": f"Installed transformers {version} into the latest sidecar " + f"(pinned: {latest_venv_pinned_version()}).", + } + + +def _invalidate_capability_caches(): + """Drop caches computed before the new sidecar existed: tier probes and the + latest tier's model_type mapping (stale on upgrade) plus vision detection + (a raw-heuristic False may now defer to the sidecar AutoConfig probe).""" + try: + from utils import transformers_version as tv + tv._probe_tier_cache.clear() + tv._config_mapping_cache.pop("latest", None) + except Exception: + pass + try: + from utils.models import model_config as mc + mc._vision_detection_cache.clear() + except Exception: + pass diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index a69673f081..9f9f8aa3de 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -35,9 +35,12 @@ import json import structlog from loggers import get_logger import os +import re import shutil import subprocess import sys +import threading +import time from pathlib import Path from utils.native_path_leases import child_env_without_native_path_secret @@ -235,8 +238,12 @@ _VENV_T5_DIR = _VENV_T5_550_DIR # reuses the workspace torch (torch-agnostic). _VENV_LLMCOMPRESSOR_DIR = str(_studio_root() / ".venv_llmcompressor") -# Tier precedence: higher rank wins in _higher_tier. -_TIER_RANK = {"default": 0, "530": 1, "550": 2, "510": 3} +# User-consented "latest transformers" sidecar (utils/transformers_latest.py); pinned version in a marker file. +_VENV_T5_LATEST_DIR = str(_studio_root() / ".venv_t5_latest") +_LATEST_PIN_MARKER = ".unsloth_pinned_transformers" + +# Tier precedence: higher rank wins in _higher_tier. "latest" outranks every fixed tier. +_TIER_RANK = {"default": 0, "530": 1, "550": 2, "510": 3, "latest": 4} def _higher_tier(a: str, b: str) -> str: @@ -254,20 +261,40 @@ def activate_transformers_for_subprocess(model_name: str, hf_token: str | None = ``hf_token`` is forwarded to tier detection so a gated/private model whose only 5.x signal is an authenticated config/tokenizer reaches the right sidecar, not the default. """ - # Pre-resolve only LoRA adapters; full checkpoints go to get_transformers_tier so their - # local config.json drives the tier (a full checkpoint with a private/offline - # _name_or_path must not resolve to an unreachable HF id and skip its own config). + # Pre-resolve LoRA adapters (local dir or remote adapter repo); full checkpoints + # go to get_transformers_tier so their local config.json drives the tier (a full + # checkpoint with a private/offline _name_or_path must not resolve to an + # unreachable HF id and skip its own config). Remote adapters activate for their + # BASE model, matching latest_tier_active_for and the inference worker. if _is_lora_adapter_dir(Path(model_name)): resolved = _resolve_base_model(model_name) else: - resolved = model_name + resolved = _remote_lora_base(model_name, hf_token = hf_token) or model_name tier = get_transformers_tier(resolved, hf_token) if model_name != resolved and _safe_is_file(Path(model_name) / "config.json"): # Gate on a real local config.json: a checkpoint carries config the base may not # surface, but path names alone must not upgrade a plain adapter. tier = _higher_tier(tier, get_transformers_tier(model_name, hf_token)) - if tier == "510": + if tier == "latest": + pinned = latest_venv_pinned_version() + if pinned is None or not _ensure_venv_t5_latest_exists(): + raise RuntimeError( + f"Cannot activate the latest-transformers sidecar: " + f".venv_t5_latest missing or unpinned at {_VENV_T5_LATEST_DIR}" + ) + if _VENV_T5_LATEST_DIR not in sys.path: + sys.path.insert(0, _VENV_T5_LATEST_DIR) + logger.info( + "Prepended transformers %s venv to sys.path from %s " + "(path only; the loaded version is confirmed later by " + "'Subprocess loaded transformers ...' on first import)", + pinned, + _VENV_T5_LATEST_DIR, + ) + _pp = os.environ.get("PYTHONPATH", "") + os.environ["PYTHONPATH"] = _VENV_T5_LATEST_DIR + (os.pathsep + _pp if _pp else "") + elif tier == "510": if not _ensure_venv_t5_510_exists(): raise RuntimeError( f"Cannot activate transformers {TRANSFORMERS_510_VERSION}: " @@ -322,6 +349,34 @@ def activate_transformers_for_subprocess(model_name: str, hf_token: str | None = logger.info("Using default transformers (4.57.x) for %s", model_name) +def latest_tier_active_for(model_name: str, hf_token: str | None = None) -> bool: + """True when *model_name* routes to the consented latest-transformers sidecar. + + Mirrors the inference worker's pre-activation resolution (local adapter dir, + then a remote adapter's Hub adapter_config.json). ``latest`` only wins when + the sidecar exists with a valid pin, i.e. exactly the loads that will import + the newest release. Never raises: any resolution failure returns False so + callers treat the model as a known tier. + """ + try: + # No consented sidecar pin means nothing routes to latest; return before + # any resolution so the common case costs no config or network reads. + if latest_venv_pinned_version() is None: + return False + if _is_lora_adapter_dir(Path(model_name)): + resolved = _resolve_base_model(model_name) + else: + # A remote LoRA activates the sidecar for its BASE model; sizing and the + # worker's 4-bit guard must see that base too, not the adapter repo. + resolved = _remote_lora_base(model_name, hf_token = hf_token) or model_name + tier = get_transformers_tier(resolved, hf_token) + if model_name != resolved and _safe_is_file(Path(model_name) / "config.json"): + tier = _higher_tier(tier, get_transformers_tier(model_name, hf_token)) + return tier == "latest" + except Exception: + return False + + def _has_adapter_weights(path: Path) -> bool: """True if *path* holds LoRA adapter weight files (``adapter_model.*``).""" try: @@ -881,17 +936,85 @@ def _cached_config_json(model_name: str, hf_token: str | None) -> dict | None: _config_mapping_cache: dict[str, frozenset[str]] = {} +def _latest_tier_disabled() -> bool: + """Kill switch shared with utils.transformers_latest: lets operators roll + back a provisioned latest sidecar without deleting files.""" + return os.environ.get("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "").strip().lower() in ( + "1", + "true", + "yes", + "on", + ) + + +# Failed lazy repairs back off so a broken sidecar can't turn every routing +# call into a pip install attempt. +_latest_repair_failed_at: float = 0.0 +_LATEST_REPAIR_BACKOFF_SECS = 5 * 60 + + +def _latest_sidecar_intact() -> bool: + """The pinned latest sidecar exists with its transformers dir and every pinned + package. False when the pin itself is gone: a cached 'latest' mapping must then be + dropped (routing re-resolves to no latest tier), not trusted, and a sidecar that kept + transformers/ but lost a pinned package must self-heal rather than route models to a + latest tier that fails activation in workers, which refuse parent-only repairs. + + (_overlay_transformers_dir only calls this after gating on a present pin, so the + pin-missing case here is the cache-revalidation caller whose pin was deleted after + the mapping was first cached.)""" + pin = _latest_pin_data() + if pin is None: + return False + return _venv_dir_is_valid(_VENV_T5_LATEST_DIR, tuple(pin["packages"])) + + def _overlay_transformers_dir(tier: str) -> str | None: """transformers source dir for a tier, located without importing it.""" + global _latest_repair_failed_at if tier != "default": - root = {"530": _VENV_T5_530_DIR, "550": _VENV_T5_550_DIR, "510": _VENV_T5_510_DIR}.get(tier) + # latest requires a valid pin and the kill switch off. + if tier == "latest" and (_latest_tier_disabled() or latest_venv_pinned_version() is None): + return None + root = { + "530": _VENV_T5_530_DIR, + "550": _VENV_T5_550_DIR, + "510": _VENV_T5_510_DIR, + "latest": _VENV_T5_LATEST_DIR, + }.get(tier) src = os.path.join(root, "transformers") if root else None + if src and tier == "latest" and not _latest_sidecar_intact(): + # A valid pin whose sidecar vanished or lost a pinned package (partial + # deletion, disk issue, interrupted external edits) must self-heal, or + # latest-only models either silently route to older tiers or reach a + # worker that cannot repair, failing every load until a manual + # reinstall. Repair under the swap reservation; back off after a + # failure so routing calls don't hammer pip. + repaired = False + if time.time() - _latest_repair_failed_at >= _LATEST_REPAIR_BACKOFF_SECS: + if _ensure_venv_t5_latest_exists(): + _latest_repair_failed_at = 0.0 + repaired = True + else: + _latest_repair_failed_at = time.time() + if not repaired: + # Still broken: treat the overlay as unavailable rather than route + # models to a tier whose worker activation is known to fail. Models + # an older tier supports keep loading there until a repair succeeds, + # matching the behavior when the sidecar dir is missing entirely. + return None return src if src and _safe_is_dir(Path(src)) else None # default: the base 4.x transformers. find_spec resolves to a 5.x sidecar if one # is already on sys.path, so skip any .venv_t5_* / llmcompressor overlay dir. sidecars = tuple( os.path.abspath(d) + os.sep - for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR, _VENV_T5_510_DIR, _VENV_LLMCOMPRESSOR_DIR) + for d in ( + _VENV_T5_530_DIR, + _VENV_T5_550_DIR, + _VENV_T5_510_DIR, + _VENV_T5_LATEST_DIR, + _VENV_LLMCOMPRESSOR_DIR, + ) ) candidates = [] try: @@ -930,11 +1053,47 @@ def _mapping_first_keys(value: ast.AST) -> set[str]: return {n.value for n in nodes if isinstance(n, ast.Constant) and isinstance(n.value, str)} +def _model_types_from_source(source: str) -> set[str]: + """model_type keys of CONFIG_MAPPING_NAMES in *source* (AST only, no execution). + + Handles the direct ``CONFIG_MAPPING_NAMES = ...`` binding (dict literal or + OrderedDict/dict call over 2-tuple lists and **{...} unpacking) and any + ``CONFIG_MAPPING_NAMES.update({...})`` mutation. Shared by the on-disk overlay + reader below and the remote latest-release checker (utils/transformers_latest.py). + """ + keys: set[str] = set() + tree = ast.parse(source) + for node in ast.walk(tree): + if isinstance(node, ast.Assign) and any( + isinstance(t, ast.Name) and t.id == "CONFIG_MAPPING_NAMES" for t in node.targets + ): + keys |= _mapping_first_keys(node.value) + elif isinstance(node, ast.Expr) and isinstance(node.value, ast.Call): + fn = node.value.func + if ( + isinstance(fn, ast.Attribute) + and fn.attr == "update" + and isinstance(fn.value, ast.Name) + and fn.value.id == "CONFIG_MAPPING_NAMES" + ): + keys |= _mapping_first_keys(node.value) + return keys + + def _config_model_types(tier: str) -> frozenset[str]: """model_type keys in a tier's CONFIG_MAPPING_NAMES (5.10 moved it to auto_mappings.py).""" + # Kill switch beats the cache: a stale mapping must not keep routing latest-only models until restart. + if tier == "latest" and _latest_tier_disabled(): + return frozenset() cached = _config_mapping_cache.get(tier) if cached is not None: - return cached + # A cached 'latest' mapping can outlive the sidecar it was parsed from: if the + # pinned sidecar was since deleted or lost a package in this process, drop the + # cache so routing re-resolves through _overlay_transformers_dir (which self-heals) + # instead of routing latest-only models to a broken tier until restart. + if tier != "latest" or _latest_sidecar_intact(): + return cached + _config_mapping_cache.pop("latest", None) tdir = _overlay_transformers_dir(tier) if tdir is None: return frozenset() # overlay not provisioned yet; do not cache so a later call re-reads @@ -944,22 +1103,7 @@ def _config_model_types(tier: str) -> frozenset[str]: if not _safe_is_file(path): continue try: - tree = ast.parse(path.read_text(encoding = "utf-8")) - for node in ast.walk(tree): - # direct binding, or a CONFIG_MAPPING_NAMES.update({...}) mutation - if isinstance(node, ast.Assign) and any( - isinstance(t, ast.Name) and t.id == "CONFIG_MAPPING_NAMES" for t in node.targets - ): - keys |= _mapping_first_keys(node.value) - elif isinstance(node, ast.Expr) and isinstance(node.value, ast.Call): - fn = node.value.func - if ( - isinstance(fn, ast.Attribute) - and fn.attr == "update" - and isinstance(fn.value, ast.Name) - and fn.value.id == "CONFIG_MAPPING_NAMES" - ): - keys |= _mapping_first_keys(node.value) + keys |= _model_types_from_source(path.read_text(encoding = "utf-8")) except Exception: continue result = frozenset(keys) @@ -967,23 +1111,73 @@ def _config_model_types(tier: str) -> frozenset[str]: return result -def _tier_from_config_mapping(cfg: dict) -> str | None: - """Lowest tier whose transformers ships cfg's model_type, or None if unknown.""" - model_type = cfg.get("model_type") - if not isinstance(model_type, str): - for key in _NESTED_CONFIG_KEYS: - sub = cfg.get(key) - if isinstance(sub, dict) and isinstance(sub.get("model_type"), str): - model_type = sub["model_type"] - break - if not isinstance(model_type, str): - return None +def _model_types_from_config(cfg: dict) -> list[str]: + """All model_types in the config: the primary (top-level, else first nested) + first, then every other nested sub-config. Wrappers instantiate sub-configs + through CONFIG_MAPPING, so nested types matter for routing too.""" + seen: list[str] = [] + + def add(value): + if isinstance(value, str) and value and value not in seen: + seen.append(value) + + add(cfg.get("model_type")) + for key in _NESTED_CONFIG_KEYS: + sub = cfg.get(key) + if isinstance(sub, dict): + add(sub.get("model_type")) + for value in cfg.values(): + if isinstance(value, dict): + add(value.get("model_type")) + return seen + + +def _lowest_tier_for(model_type: str) -> str | None: for tier in sorted(_TIER_RANK, key = _TIER_RANK.get): if model_type in _config_model_types(tier): return tier return None +def _tier_from_config_mapping(cfg: dict) -> str | None: + """Lowest tier able to load every model_type in cfg, or None when the + primary type is unknown everywhere. A nested type can raise the tier (its + sub-config is built through CONFIG_MAPPING); an unknown nested type never + vetoes, since no installed tier could load it either way (the latest + checker handles surfacing the install prompt for it).""" + types = _model_types_from_config(cfg) + if not types: + return None + best = _lowest_tier_for(types[0]) + if best is None: + return None + for model_type in types[1:]: + tier = _lowest_tier_for(model_type) + if tier is not None and _TIER_RANK[tier] > _TIER_RANK[best]: + best = tier + return best + + +def _raise_tier_for_nested(cfg: dict | None, tier: str) -> str: + """Raise *tier* when the mapping resolver needs a higher one for *cfg*. + + A wrapper's top-level model_type can match a hardcoded fast path while a + nested text/vision config's type only exists in a newer sidecar (e.g. the + installed latest); its sub-config is built through CONFIG_MAPPING, so the + fast-path tier would fail to load it. Raise-only: never lowers a fast-path + match, so name overrides (Qwen3.6) keep their tier. Never raises an + exception: a resolution failure keeps the fast-path tier.""" + if not isinstance(cfg, dict): + return tier + try: + mapped = _tier_from_config_mapping(cfg) + if mapped is not None and _TIER_RANK.get(mapped, 0) > _TIER_RANK.get(tier, 0): + return mapped + except Exception: + pass + return tier + + # --- AutoConfig probe: general tier resolution for ambiguous models ---------- # When the cheap signals only say "needs some 5.x", parse config.json with the built-in # parser in each candidate sidecar (lowest first) instead of guessing. Generalizes beyond @@ -1039,9 +1233,19 @@ def _probe_tier_venvs(): "530": (_VENV_T5_530_DIR, _ensure_venv_t5_530_exists), "550": (_VENV_T5_550_DIR, _ensure_venv_t5_550_exists), "510": (_VENV_T5_510_DIR, _ensure_venv_t5_510_exists), + "latest": (_VENV_T5_LATEST_DIR, _ensure_venv_t5_latest_exists), } +def _probe_tier_order() -> tuple[str, ...]: + """Sidecar probe order. The consented "latest" sidecar joins only once it is + provisioned (pin marker present): an absent optional tier must not flip the probe's + skipped-tier bookkeeping, keeping pre-latest behavior byte-identical.""" + if not _latest_tier_disabled() and latest_venv_pinned_version() is not None: + return _PROBE_TIER_ORDER + ("latest",) + return _PROBE_TIER_ORDER + + def _probe_autoconfig(target_dir: str, model_name: str, hf_token: str | None) -> bool | None: """Parse config.json with the built-in parser inside *target_dir*'s sidecar. True = parses, False = parse/version failure (escalate), None = transient @@ -1119,7 +1323,7 @@ def _probe_tier( stays on the default. Cached per _probe_cache_key (process lifetime). No Hub sha is resolved: that would import huggingface_hub before the sidecar is on sys.path. """ - if os.environ.get("UNSLOTH_DISABLE_TIER_PROBE", "").lower() in ("1", "true", "yes"): + if os.environ.get("UNSLOTH_DISABLE_TIER_PROBE", "").lower() in ("1", "true", "yes", "on"): return floor key = _probe_cache_key(model_name) # Key by probe mode: the default-first path can return 'default', which must not be @@ -1127,7 +1331,10 @@ def _probe_tier( if include_default or floor != "530": key = f"{key}\0floor={floor}:def={int(include_default)}" if key in _probe_tier_cache: - return _probe_tier_cache[key] + cached = _probe_tier_cache[key] + # Kill switch beats the cache (like _config_model_types): a stale 'latest' probe must not keep activating it. + if cached != "latest" or not _latest_tier_disabled(): + return cached def _cache(tier: str, *, skipped: bool) -> str: # Do not pin a result that depended on a skipped lower tier: once that sidecar is @@ -1137,7 +1344,8 @@ def _probe_tier( return tier venvs = _probe_tier_venvs() - order = (("default",) + _PROBE_TIER_ORDER) if include_default else _PROBE_TIER_ORDER + sidecar_order = _probe_tier_order() + order = (("default",) + sidecar_order) if include_default else sidecar_order probed_count = 0 skipped_any = False for tier in order: @@ -1264,17 +1472,21 @@ def get_transformers_tier( cfg = _load_config_json(model_name, hf_token) if cfg is not None: if _config_needs_510(cfg): + tier = _raise_tier_for_nested(cfg, "510") logger.info( - "Transformers tier 510 selected for %s (local config.json check)", + "Transformers tier %s selected for %s (local config.json check)", + tier, model_name, ) - return "510" + return tier if _config_needs_550(cfg): + tier = _raise_tier_for_nested(cfg, "550") logger.info( - "Transformers tier 550 selected for %s (local config.json check)", + "Transformers tier %s selected for %s (local config.json check)", + tier, model_name, ) - return "550" + return tier if _config_needs_530(cfg): # Qwen3.6 reuses Qwen3.5 config ids but needs 5.5 by name. Only a real # Hub id (or the folder basename) may override 530, so a stale local @@ -1287,17 +1499,20 @@ def get_transformers_tier( ) override = _higher_tier_name_override(hint_src) if override is not None: + override = _raise_tier_for_nested(cfg, override) logger.info( "Transformers tier %s selected for %s (name overrides 530 config)", override, model_name, ) return override + tier = _raise_tier_for_nested(cfg, "530") logger.info( - "Transformers tier 530 selected for %s (local config.json check)", + "Transformers tier %s selected for %s (local config.json check)", + tier, model_name, ) - return "530" + return tier # Unknown arch: resolve the base id from config. A resolved local dir # recurses (config check); a Hub id uses name rules only (no network). resolved = _resolve_base_model(model_name) @@ -1359,6 +1574,13 @@ def get_transformers_tier( result = _tier_from_name(model_name) if result is not None: tier, match = result + # With a consented latest sidecar pinned, a name that matches a fixed + # tier can still carry a latest-only model_type (e.g. a newer variant + # reusing a family name); consult the config so an accepted upgrade + # actually routes to the sidecar it installed. Costs a config read only + # in the pinned case, keeping the pre-latest path I/O-free. + if latest_venv_pinned_version() is not None: + tier = _raise_tier_for_nested(_load_config_json(model_name, hf_token), tier) logger.info( "Transformers tier %s selected for %s (substring match: %s)", tier, @@ -1369,11 +1591,13 @@ def get_transformers_tier( # --- Slow config fallbacks (network for HF IDs; authenticated with hf_token) -------- if _check_config_needs_510(model_name, hf_token): - logger.info("Transformers tier 510 selected for %s (config.json check)", model_name) - return "510" + tier = _raise_tier_for_nested(_load_config_json(model_name, hf_token), "510") + logger.info("Transformers tier %s selected for %s (config.json check)", tier, model_name) + return tier if _check_config_needs_550(model_name, hf_token): - logger.info("Transformers tier 550 selected for %s (config.json check)", model_name) - return "550" + tier = _raise_tier_for_nested(_load_config_json(model_name, hf_token), "550") + logger.info("Transformers tier %s selected for %s (config.json check)", tier, model_name) + return tier if _check_config_needs_530(model_name, hf_token): # Qwen3.6 reuses Qwen3.5 config ids but needs 5.5 by name; honor a real Hub-id name # hint from _name_or_path before selecting 530. @@ -1383,14 +1607,16 @@ def get_transformers_tier( base if isinstance(base, str) and base != model_name else None ) if override is not None: + override = _raise_tier_for_nested(remote_cfg, override) logger.info( "Transformers tier %s selected for %s (name overrides 530 config)", override, model_name, ) return override - logger.info("Transformers tier 530 selected for %s (config.json check)", model_name) - return "530" + tier = _raise_tier_for_nested(remote_cfg, "530") + logger.info("Transformers tier %s selected for %s (config.json check)", tier, model_name) + return tier # _load_config_json (not the cache-only reader) so a config served from the hub # cache during a transient outage still feeds the mapping resolver. remote_cfg = _load_config_json(model_name, hf_token) @@ -1657,6 +1883,471 @@ def _ensure_venv_t5_exists() -> bool: return _ensure_venv_t5_550_exists() +# --- User-consented "latest transformers" sidecar (.venv_t5_latest) -------------------------- +# Provisioned via ensure_latest_transformers_venv() after the user confirms the upgrade popup +# (utils/transformers_latest.py); pinned in a marker file so restarts revalidate and routing auto-picks it. + +# PEP 440-ish release strings only (guards the pip install spec against injection). +_LATEST_VERSION_RE = r"[0-9]+(\.[0-9]+)*((a|b|rc)[0-9]+)?(\.post[0-9]+)?(\.dev[0-9]+)?" + + +def _is_valid_version_string(version: str) -> bool: + import re + return isinstance(version, str) and re.fullmatch(_LATEST_VERSION_RE, version) is not None + + +# Only the sidecar recipe's own packages, as plain (optionally ==pinned) specs, may +# come from the on-disk pin marker; anything else (URLs, extras, options) is rebuilt. +_PIN_SPEC_RE = re.compile(r"^[A-Za-z0-9_.-]+(==[A-Za-z0-9_.+-]+)?$") +_PIN_ALLOWED_NAMES = frozenset( + { + "transformers", + "huggingface_hub", + "huggingface-hub", + "hf_xet", + "hf-xet", + "tiktoken", + "tokenizers", + "safetensors", + } +) + + +def _is_safe_pin_spec(spec: str) -> bool: + if not _PIN_SPEC_RE.match(spec): + return False + name = spec.split("==", 1)[0].lower().replace("_", "-") + return name in {n.replace("_", "-") for n in _PIN_ALLOWED_NAMES} + + +def _recover_stranded_latest_sidecar() -> None: + """Restore a sidecar stranded at ``.old`` by a swap whose activation rename AND its + rollback both failed (e.g. a lingering worker file handle on Windows blocked both). + + That double failure leaves no live dir and the pin marker gone with it, so the + sidecar reads as unprovisioned and never self-heals. Recover only when no live dir + exists and no swap is in flight: the reservation is held throughout the swap, so the + transient live-absent window of a legitimate swap never triggers a restore.""" + live = Path(_VENV_T5_LATEST_DIR) + retired = Path(_VENV_T5_LATEST_DIR + ".old") + try: + if live.exists() or not retired.is_dir() or sidecar_swap_in_progress(): + return + os.rename(retired, live) + logger.info("Recovered .venv_t5_latest from a stranded .old after a failed swap") + except OSError: + pass + + +def _latest_pin_data() -> dict | None: + """Parsed pin marker: {"version": str, "packages": [specs...]}, or None. + + The marker is JSON; a plain version string (older/simpler writers) is tolerated and + expanded with the default package set. + """ + _recover_stranded_latest_sidecar() + marker = Path(_VENV_T5_LATEST_DIR) / _LATEST_PIN_MARKER + try: + if not marker.is_file(): + return None + raw = marker.read_text(encoding = "utf-8").strip() + except Exception: + return None + try: + data = json.loads(raw) + except ValueError: + data = raw + if isinstance(data, str): + if not _is_valid_version_string(data): + return None + return {"version": data, "packages": list(_venv_t5_latest_packages(data))} + if not isinstance(data, dict): + return None + version = data.get("version") + if not _is_valid_version_string(version): + return None + packages = data.get("packages") + if not ( + isinstance(packages, list) + and packages + and all(isinstance(p, str) and _is_safe_pin_spec(p) for p in packages) + ): + # Malformed or unexpected specs (the pin is user-writable on disk) never + # reach pip: rebuild the canonical set for the pinned version instead. + packages = list(_venv_t5_latest_packages(version)) + return {"version": version, "packages": packages} + + +def latest_venv_pinned_version() -> str | None: + """Exact transformers version pinned in .venv_t5_latest's marker, or None if the + sidecar was never provisioned (or the marker is unreadable/invalid).""" + data = _latest_pin_data() + return data["version"] if data else None + + +def _venv_t5_latest_packages(version: str, extra_packages: tuple[str, ...] = ()) -> tuple[str, ...]: + """Package set for the latest sidecar; mirrors the fixed .venv_t5_* sidecars. + *extra_packages* carries dep-compat shadows (e.g. a newer tokenizers) computed by + utils.transformers_latest before install.""" + return ( + f"transformers=={version}", + "huggingface_hub==1.8.0", + "hf_xet==1.4.2", + "tiktoken", + ) + tuple(extra_packages) + + +# Single reservation for ANY .venv_t5_latest replacement (consented install or lazy repair), +# checked by training/export starts so no worker spawns mid-swap. Backed by a lock FILE (not just +# this flag) so a lazy repair running in a worker subprocess stays visible to the parent's route +# checks; the in-process flag marks ownership (only the owner unlinks the file). +_sidecar_swap_lock = threading.Lock() +_sidecar_swap_active = False +_sidecar_swap_token: str | None = None +_sidecar_swap_kind: str | None = None +# An install is minutes; a lock this old is a crashed owner, not a live swap. +_SWAP_LOCK_STALE_SECS = 2 * 60 * 60 + + +def _swap_lock_path() -> Path: + return Path(_VENV_T5_LATEST_DIR + ".swaplock") + + +def _pid_alive(pid) -> bool: + if not isinstance(pid, int) or pid <= 0: + return False + try: + import psutil + return psutil.pid_exists(pid) + except Exception: + pass + if os.name == "nt": + # os.kill(pid, 0) is NOT a POSIX signal-0 liveness probe on Windows: signal 0 + # is CTRL_C_EVENT, so CPython routes it through GenerateConsoleCtrlEvent (a real + # Ctrl+C to that console group) rather than a harmless check. Probe via OpenProcess. + try: + import ctypes + from ctypes import wintypes + + kernel32 = ctypes.WinDLL("kernel32", use_last_error = True) + kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] + kernel32.OpenProcess.restype = wintypes.HANDLE + # PROCESS_QUERY_LIMITED_INFORMATION: minimal right, granted across integrity levels. + handle = kernel32.OpenProcess(0x1000, False, pid) + if handle: + kernel32.CloseHandle(handle) + return True + # ERROR_ACCESS_DENIED means the process exists but we may not query it. + return ctypes.get_last_error() == 5 + except Exception: + return False + try: + os.kill(pid, 0) + return True + except OSError: + return False + except Exception: + return False + + +def _swap_lock_is_stale(path: Path) -> bool: + """Stale when the recorded owner is provably dead: a crashed installer is reclaimed + at once, not after the long cutoff, so `/load`, training, export, and repair are not + wedged for hours after a crash. A live but slow pip install keeps its lock (its PID + is alive), so breaking it and racing two swaps on the same staging dirs stays + impossible. Only a lock whose PID can't be read (mid-write or corrupt) falls back to + the age cutoff, so the create-before-metadata-write window is never mistaken for dead.""" + try: + age = time.time() - path.stat().st_mtime + except OSError: + return False + data = _read_swap_lock(path) or {} + pid = data.get("pid") + if not isinstance(pid, int) or pid <= 0: + return age > _SWAP_LOCK_STALE_SECS + return not _pid_alive(pid) + + +class SidecarSwapInProgress(RuntimeError): + """A worker start lost the race to a .venv_t5_latest install/repair; retryable.""" + + +def _read_swap_lock(path: Path) -> dict | None: + try: + data = json.loads(path.read_text(encoding = "utf-8")) + return data if isinstance(data, dict) else {} + except FileNotFoundError: + return None + except OSError: + return {} + except Exception: + return {} + + +def try_begin_sidecar_swap(kind: str = "install") -> bool: + """Reserve the sidecar swap window; False when one is already reserved + (in this process or, via the lock file, in any worker subprocess). + *kind* is "install" (consented route) or "repair" (lazy venv repair).""" + global _sidecar_swap_active, _sidecar_swap_token, _sidecar_swap_kind + with _sidecar_swap_lock: + if _sidecar_swap_active: + return False + token = f"{os.getpid()}-{time.time_ns()}" + path = _swap_lock_path() + try: + path.parent.mkdir(parents = True, exist_ok = True) + except OSError: + pass + for attempt in range(2): + try: + fd = os.open(str(path), os.O_CREAT | os.O_EXCL | os.O_WRONLY) + break + except FileExistsError: + if attempt or not _swap_lock_is_stale(path): + return False + try: + path.unlink() + except OSError: + return False + except OSError: + # Lock file not creatable (odd filesystem): fall back to the process-local reservation. + fd = None + break + if fd is not None: + try: + with os.fdopen(fd, "w") as f: + f.write( + json.dumps( + {"pid": os.getpid(), "at": time.time(), "token": token, "kind": kind} + ) + ) + except OSError: + pass + _sidecar_swap_active = True + _sidecar_swap_token = token + _sidecar_swap_kind = kind + return True + + +def end_sidecar_swap() -> None: + """Release the reservation taken by :func:`try_begin_sidecar_swap`.""" + global _sidecar_swap_active, _sidecar_swap_token, _sidecar_swap_kind + with _sidecar_swap_lock: + if _sidecar_swap_active: + # Only the file WE wrote is removed: if this reservation was declared + # stale and superseded, unlinking blindly would drop the new owner's + # live lock and unguard its in-flight swap. + path = _swap_lock_path() + data = _read_swap_lock(path) + if data is not None and data.get("token", _sidecar_swap_token) == _sidecar_swap_token: + try: + path.unlink() + except OSError: + pass + _sidecar_swap_active = False + _sidecar_swap_token = None + _sidecar_swap_kind = None + + +def sidecar_swap_in_progress() -> bool: + """True while a .venv_t5_latest install or repair holds the reservation, + in this process or any other Studio process (lock file).""" + return sidecar_swap_kind() is not None + + +def sidecar_swap_kind() -> str | None: + """The active reservation's kind ("install" / "repair"), or None when idle. + Lets guards that rely on the install route's own abort-on-active-worker + checks keep refusing for repairs, which have no such checks.""" + with _sidecar_swap_lock: + if _sidecar_swap_active: + return _sidecar_swap_kind or "install" + path = _swap_lock_path() + try: + if not path.is_file() or _swap_lock_is_stale(path): + return None + except OSError: + return None + data = _read_swap_lock(path) or {} + kind = data.get("kind") + return kind if kind in ("install", "repair") else "install" + + +def _stage_and_swap_latest_venv( + version: str, + packages: tuple[str, ...], + before_swap = None, +) -> bool: + """Stage-and-swap: build the new sidecar next to the live one and swap only + once complete, so a failed install or marker write never destroys a + previously working .venv_t5_latest or its pin. Shared by the consented + install and the lazy repair path. *before_swap* (optional callable) runs + after the staging build succeeds and immediately before the live dir is + replaced, so callers can tear down workers only when the swap is certain; + if it raises, the previous sidecar is left untouched.""" + staging = _VENV_T5_LATEST_DIR + ".staging" + retired = _VENV_T5_LATEST_DIR + ".old" + shutil.rmtree(staging, ignore_errors = True) + try: + if not _ensure_venv_dir(staging, packages, f"transformers {version} (latest)"): + # No exception, so the except cleanup below never runs; drop the partial dir. + shutil.rmtree(staging, ignore_errors = True) + return False + (Path(staging) / _LATEST_PIN_MARKER).write_text( + json.dumps({"version": version, "packages": list(packages)}), encoding = "utf-8" + ) + if before_swap is not None: + before_swap() + shutil.rmtree(retired, ignore_errors = True) + if os.path.isdir(_VENV_T5_LATEST_DIR): + os.rename(_VENV_T5_LATEST_DIR, retired) + try: + os.rename(staging, _VENV_T5_LATEST_DIR) + except OSError: + # Restore the previous sidecar if the final swap fails. + if not os.path.isdir(_VENV_T5_LATEST_DIR) and os.path.isdir(retired): + os.rename(retired, _VENV_T5_LATEST_DIR) + raise + except Exception as exc: + logger.error("Could not provision transformers %s into .venv_t5_latest: %s", version, exc) + shutil.rmtree(staging, ignore_errors = True) + return False + shutil.rmtree(retired, ignore_errors = True) + # CONFIG_MAPPING_NAMES may have changed: drop the cached key set. + _config_mapping_cache.pop("latest", None) + logger.info("Provisioned .venv_t5_latest with transformers %s", version) + return True + + +def _workers_active_for_repair() -> bool: + """Best-effort: any parent-visible chat/training/export worker alive. Never + raises; unavailable backends (worker subprocess, early startup) count idle.""" + try: + from core.training import get_training_backend + if get_training_backend().is_training_active(): + return True + except Exception: + pass + try: + from core.export import get_export_backend + + _export = get_export_backend() + if _export.is_export_active(): + return True + _alive = getattr(_export, "is_worker_alive", None) + if callable(_alive) and _alive(): + return True + except Exception: + pass + try: + from core.inference import get_inference_backend + + backend = get_inference_backend() + if getattr(backend, "active_model_name", None): + return True + # An in-flight load counts too: its worker spawns moments later. + if getattr(backend, "loading_models", None): + return True + _alive = getattr(backend, "is_worker_alive", None) + if callable(_alive) and _alive(): + return True + except Exception: + pass + return False + + +def _ensure_venv_t5_latest_exists() -> bool: + """Ensure .venv_t5_latest/ holds its pinned transformers version. + + Never installs without a pin: an unprovisioned sidecar (no marker) returns False so + routing and probing behave exactly as before the feature existed. With a pin present + it repairs a broken dir the same way the fixed sidecars do. + """ + pin = _latest_pin_data() + if pin is None: + return False + version = pin["version"] + packages = tuple(pin["packages"]) + if _venv_dir_is_valid(_VENV_T5_LATEST_DIR, packages): + return True + if _env_offline(): + logger.warning( + ".venv_t5_latest (transformers %s) is incomplete and offline mode is set; " + "cannot repair it.", + version, + ) + return False + # Repairs are a parent-process action: a worker child's backend singletons are + # empty, so it cannot see live siblings that may still lazy-import from the + # sidecar. Fail activation in the child instead; the parent's routing + # self-heal (guarded below) performs the actual repair. + try: + import multiprocessing as _mp + if _mp.parent_process() is not None: + logger.warning( + ".venv_t5_latest is incomplete; repairs run in the parent process. " + "Retry after the parent repairs the sidecar." + ) + return False + except Exception: + pass + # Same stage-and-swap as the install, under the same reservation so training/export starts + # (which check sidecar_swap_in_progress) wait out a lazy repair; a failed repair keeps the pin. + if not try_begin_sidecar_swap(kind = "repair"): + logger.warning( + "Cannot repair .venv_t5_latest: another sidecar install or repair is in progress." + ) + return False + try: + # Worker check UNDER the reservation (the install route quiesces workers; + # a repair has none): worker starts set their active markers BEFORE + # rechecking the reservation, so either this check sees them and aborts, + # or their recheck sees this reservation and aborts -- no interleaving + # lets a worker spawn against a mid-swap sidecar. + if _workers_active_for_repair(): + logger.warning( + "Cannot repair .venv_t5_latest: active chat/training/export workers " + "may be importing from it. Retry when they are idle." + ) + return False + return _stage_and_swap_latest_venv(version, packages) + finally: + end_sidecar_swap() + + +def ensure_latest_transformers_venv( + version: str, + extra_packages: tuple[str, ...] = (), + before_swap = None, +) -> bool: + """Provision .venv_t5_latest/ pinned to *version* (user-consented install path). + + Reuses the same --target/--no-deps installer as the fixed sidecars, then writes the pin + marker (version + full package set) so the venv persists across restarts and + :func:`latest_venv_pinned_version` / routing pick it up automatically. + *extra_packages* carries dep-compat shadows (see utils.transformers_latest). + Returns True on success. + """ + if not _is_valid_version_string(version): + logger.error("Refusing to install invalid transformers version %r", version) + return False + if _env_offline(): + logger.warning( + "Cannot install transformers %s: HF/transformers offline mode is set.", version + ) + return False + packages = _venv_t5_latest_packages(version, extra_packages) + pin = _latest_pin_data() + if ( + pin is not None + and pin["version"] == version + and tuple(pin["packages"]) == packages + and _venv_dir_is_valid(_VENV_T5_LATEST_DIR, packages) + ): + return True + return _stage_and_swap_latest_venv(version, packages, before_swap = before_swap) + + # --- llm-compressor-main shadow (FP8/FP4 export of newer-transformers models) --------------------- # Exact, reproducible pins (bump deliberately in review). Full 40-char SHA validated to FP8-quantize # Qwen3.5 / Gemma-4 / Llama. @@ -1819,7 +2510,7 @@ def _activate_venv(venv_dir: str, label: str) -> None: def _deactivate_5x() -> None: """Remove all .venv_t5_*/ dirs from sys.path, purge stale modules, reimport.""" - for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR, _VENV_T5_510_DIR): + for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR, _VENV_T5_510_DIR, _VENV_T5_LATEST_DIR): while d in sys.path: sys.path.remove(d) logger.info("Removed venv_t5 dirs from sys.path") @@ -1853,14 +2544,25 @@ def ensure_transformers_version(model_name: str) -> None: if _is_lora_adapter_dir(Path(model_name)): resolved = _resolve_base_model(model_name) else: - resolved = model_name + # A remote adapter's tier is its BASE model's (see activation above). + resolved = _remote_lora_base(model_name) or model_name tier = get_transformers_tier(resolved) if model_name != resolved and _safe_is_file(Path(model_name) / "config.json"): # Gate on a real local config.json: a checkpoint carries config the base may not # surface, but path names alone must not upgrade a plain adapter. tier = _higher_tier(tier, get_transformers_tier(model_name)) - if tier == "510": + if tier == "latest": + pinned = latest_venv_pinned_version() + if pinned is None: + raise RuntimeError( + f"Cannot activate the latest-transformers sidecar: " + f"no pin marker at {_VENV_T5_LATEST_DIR}" + ) + target_version = pinned + venv_dir = _VENV_T5_LATEST_DIR + ensure_fn = _ensure_venv_t5_latest_exists + elif tier == "510": target_version = TRANSFORMERS_510_VERSION venv_dir = _VENV_T5_510_DIR ensure_fn = _ensure_venv_t5_510_exists diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index 6f68917224..ba56ce7525 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -16,6 +16,7 @@ import { type ChatSearch, } from "@/features/chat"; import { RemoteCodeConsentDialog } from "@/features/security"; +import { TransformersUpgradeDialog } from "@/features/transformers-upgrade"; import { useTrainingUnloadGuard } from "@/features/training"; import { useExportRuntimeLifecycle } from "@/features/export"; import { hasAuthToken } from "@/features/auth"; @@ -230,6 +231,7 @@ function RootLayout() { {!isAuthFlowRoute && } + {hideNavbar ? (
}> diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 9ed11deccc..12c1c3b385 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1454,6 +1454,11 @@ async function autoLoadSmallestModel(): Promise<{ blockedByTrustRemoteCode = true; return false; } + // Never install packages from a background load; explicit loads raise the upgrade dialog. + if (validation.requires_transformers_upgrade) { + hadNonTrustFailure = true; + return false; + } return true; } diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index a659b7f83e..4646f67cb4 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -4,6 +4,10 @@ import { createElement, useCallback, useRef, useState } from "react"; import { toast } from "@/lib/toast"; import { confirmRemoteCodeIfNeeded } from "@/features/security"; +import { + confirmTransformersUpgradeIfNeeded, + useTransformersUpgradeDialogStore, +} from "@/features/transformers-upgrade"; import { consumeNativePathToken } from "@/features/native-intents/api"; import { notifyNative, @@ -245,6 +249,10 @@ function getTrustRemoteCodeRequiredMessage(modelName: string): string { return `${modelName} was not loaded because its custom code was not approved. Load it again to review the code and approve it.`; } +function getTransformersUpgradeRequiredMessage(modelName: string): string { + return `${modelName} was not loaded because it needs a newer transformers release that was not installed. Load it again to install it.`; +} + /** * Reconcile the chat runtime store against `/api/inference/status`: refresh the * models/loras catalogs and either re-pin the active checkpoint or clear the @@ -626,6 +634,30 @@ export function useChatModelRuntime() { is_lora: isLora, gguf_variant: ggufVariant ?? null, }); + // Upgrade consent runs before the security dialogs; Accept installs and the load continues. + if (validation.requires_transformers_upgrade) { + const upgraded = await confirmTransformersUpgradeIfNeeded({ + modelName: modelId, + upgrade: validation.transformers_upgrade, + // No installable release: custom-code models may fall back to the trust_remote_code gate below. + trustRemoteCodeFallback: validation.requires_trust_remote_code, + }); + // The install unloads the previous model before the swap (even when + // the swap then fails), so any exit after this point must roll back. + // False for the custom-code fallback, which resolves without installing. + if ( + useTransformersUpgradeDialogStore + .getState() + .consumeServerUnloadedChat() + && currentCheckpoint + ) { + previousWasUnloaded = true; + } + if (!upgraded) { + throw new Error(getTransformersUpgradeRequiredMessage(displayName)); + } + } + if (abortCtrl.signal.aborted) throw new Error("Cancelled"); // Open the consent dialog when the model needs custom-code consent or has a // flagged unsafe file. Fires even when trustRemoteCode is preset on, since the // worker requires a matching fingerprint that only the dialog produces. diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index a50af46e85..47a0720dac 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -67,6 +67,10 @@ import { KnowledgeBaseComposerButton } from "@/features/rag/components/knowledge import { NewProjectDialog } from "./components/new-project-dialog"; import { useChatProjects } from "./hooks/use-chat-projects"; import { confirmRemoteCodeIfNeeded } from "@/features/security"; +import { + confirmTransformersUpgradeIfNeeded, + useTransformersUpgradeDialogStore, +} from "@/features/transformers-upgrade"; import { loadModel, validateModel } from "./api/chat-api"; import { parseExternalModelId, @@ -929,6 +933,10 @@ export function SharedComposer({ return parts[parts.length - 1] || id; } + // Set when an accepted transformers install unloaded the active model + // server-side; a later failure must then clear the stale checkpoint. + let upgradeUnloadedActive = false; + // Helper: load a model and update store checkpoint async function ensureModelLoaded( sel: CompareModelSelection, @@ -955,6 +963,31 @@ export function SharedComposer({ trust_remote_code: loadTrustRemoteCode, chat_template_override: effectiveChatTemplateOverride, }); + // Upgrade dialog first (mirrors the primary load path). + if (validation.requires_transformers_upgrade) { + const upgraded = await confirmTransformersUpgradeIfNeeded({ + modelName: sel.id, + upgrade: validation.transformers_upgrade, + // No installable release: custom-code models may fall back to the trust_remote_code gate below. + trustRemoteCodeFallback: validation.requires_trust_remote_code, + }); + // The install unloads the active model before the swap (even when the + // swap then fails); if a later gate cancels or the load fails, the UI + // must stop pointing at that unloaded model. + if ( + useTransformersUpgradeDialogStore + .getState() + .consumeServerUnloadedChat() + && currentStore.params.checkpoint + ) { + upgradeUnloadedActive = true; + } + if (!upgraded) { + throw new Error( + `${modelDisplayName(sel.id)} needs a newer transformers release to load.`, + ); + } + } if ( validation.requires_trust_remote_code || validation.requires_security_review @@ -990,6 +1023,7 @@ export function SharedComposer({ tensor_parallel: currentStore.tensorParallel, }); saveSpeculativeType(specSettings.speculativeType); + upgradeUnloadedActive = false; const store = useChatRuntimeStore.getState(); store.setCheckpoint( resp.model, @@ -1097,6 +1131,11 @@ export function SharedComposer({ toast.success("Compare complete", { id: toastId, duration: 2000 }); } catch (err) { compareStepSucceededRef.current = false; + // The install already unloaded the previously active model; drop the + // checkpoint so the UI does not keep pointing at an unloaded model. + if (upgradeUnloadedActive) { + useChatRuntimeStore.getState().clearCheckpoint(); + } toast.error("Compare failed", { id: toastId, description: err instanceof Error ? err.message : "Unknown error", diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 954e88e86b..1a9a967263 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -1,6 +1,8 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +import type { TransformersUpgradeInfo } from "@/features/transformers-upgrade"; + export interface BackendModelDetails { id: string; name?: string | null; @@ -78,6 +80,10 @@ export interface ValidateModelResponse { requires_security_review?: boolean; /** Native context length from the local GGUF header; null until downloaded. */ context_length?: number | null; + /** Architecture only shipped by a newer transformers; UI pauses on the upgrade dialog. */ + requires_transformers_upgrade?: boolean; + /** Set only when requires_transformers_upgrade. */ + transformers_upgrade?: TransformersUpgradeInfo | null; } export interface GgufVariantDetail { diff --git a/studio/frontend/src/features/transformers-upgrade/api/transformers-upgrade-api.ts b/studio/frontend/src/features/transformers-upgrade/api/transformers-upgrade-api.ts new file mode 100644 index 0000000000..2df9f712ff --- /dev/null +++ b/studio/frontend/src/features/transformers-upgrade/api/transformers-upgrade-api.ts @@ -0,0 +1,32 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { authFetch } from "@/features/auth"; +import { readFastApiError } from "@/lib/format-fastapi-error"; + +interface InstallLatestTransformersResponse { + success: boolean; + version: string; + message: string; + /** The server unloaded the active chat model before the swap (set even on a + * structured failure, so callers can restore their model state). */ + model_unloaded?: boolean; + /** On a version-mismatch failure: the release that superseded the requested + * one, so Retry can use it. */ + latest_version?: string | null; +} + +/** Consented install of the latest transformers into the sidecar; synchronous, can take minutes. */ +export async function installLatestTransformers( + version: string, +): Promise { + const response = await authFetch("/api/inference/install-latest-transformers", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ version }), + }); + if (!response.ok) { + throw new Error(await readFastApiError(response)); + } + return (await response.json()) as InstallLatestTransformersResponse; +} diff --git a/studio/frontend/src/features/transformers-upgrade/components/transformers-upgrade-dialog.tsx b/studio/frontend/src/features/transformers-upgrade/components/transformers-upgrade-dialog.tsx new file mode 100644 index 0000000000..98c590e5e4 --- /dev/null +++ b/studio/frontend/src/features/transformers-upgrade/components/transformers-upgrade-dialog.tsx @@ -0,0 +1,167 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; +import { Spinner } from "@/components/ui/spinner"; +import { cn } from "@/lib/utils"; +import { PackageIcon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { useTransformersUpgradeDialogStore } from "../stores/transformers-upgrade-dialog-store"; + +function modelDisplayName(modelName: string | null): string { + if (!modelName) return "This model"; + return modelName.split("/").pop() || modelName; +} + +/** Root-mounted consent dialog for models needing a newer transformers; + * Install runs the sidecar install and resumes the paused load on success. */ +export function TransformersUpgradeDialog() { + const open = useTransformersUpgradeDialogStore((s) => s.open); + const modelName = useTransformersUpgradeDialogStore((s) => s.modelName); + const upgrade = useTransformersUpgradeDialogStore((s) => s.upgrade); + const phase = useTransformersUpgradeDialogStore((s) => s.phase); + const errorMessage = useTransformersUpgradeDialogStore((s) => s.errorMessage); + const trustRemoteCodeFallback = useTransformersUpgradeDialogStore( + (s) => s.trustRemoteCodeFallback, + ); + const install = useTransformersUpgradeDialogStore((s) => s.install); + const resolve = useTransformersUpgradeDialogStore((s) => s.resolve); + + const displayName = modelDisplayName(modelName); + const modelType = upgrade?.model_type ?? "unknown"; + const version = upgrade?.pypi_version ?? null; + // Only released PyPI versions are installable; dev (main) builds are never offered. + const installable = Boolean(upgrade?.supported_in_pypi && version); + const devOnly = !installable && Boolean(upgrade?.supported_in_main); + const installing = phase === "installing"; + + return ( + { + // Escape/overlay dismiss must not abandon an in-flight install. + if (!next && !installing) resolve(false); + }} + > + + +
+
+ +
+
+
+ New model architecture + + + {displayName} + {" "} + uses the{" "} + {modelType}{" "} + architecture, which your installed transformers does not + support yet.{" "} + {installable ? ( + <> + Install transformers{" "} + + {version} + {" "} + from PyPI to load it. The install runs once and can take + a minute; loading continues automatically afterwards. + + ) : devOnly ? ( + <> + Even the latest transformers release on PyPI does not + support it yet: the architecture is only available on the + transformers development branch (main), and Studio does + not install development builds. Support arrives with the + next transformers release on PyPI. + + ) : ( + <> + No released transformers version supports it yet, so it + cannot be loaded. + + )} + {!installable && trustRemoteCodeFallback ? ( + <> + {" "} + This model also ships its own modeling code; you can + continue and review enabling that custom code instead. + + ) : null} + +
+ + {phase === "error" && errorMessage ? ( +

+ {errorMessage} +

+ ) : null} + + {installing ? ( +

+ + Installing transformers {version}... This can take a minute. +

+ ) : null} +
+
+ + + + Cancel + {installable ? ( + <> + {phase === "error" && trustRemoteCodeFallback ? ( + // Install failed but the model ships custom code: offer the + // caller's trust_remote_code gate instead of forcing a retry. + resolve(true)} + > + Continue with custom code + + ) : null} + { + // Keep the dialog open; the store closes it on success. + event.preventDefault(); + void install(); + }} + > + {installing ? ( + <> + + Installing... + + ) : phase === "error" ? ( + "Retry install" + ) : ( + `Install transformers ${version}` + )} + + + ) : trustRemoteCodeFallback ? ( + // No installable release but the model ships custom code: continue + // into the caller's trust_remote_code gate as the last resort. + resolve(true)}> + Continue with custom code + + ) : null} + + + + ); +} diff --git a/studio/frontend/src/features/transformers-upgrade/hooks/use-transformers-upgrade-consent.ts b/studio/frontend/src/features/transformers-upgrade/hooks/use-transformers-upgrade-consent.ts new file mode 100644 index 0000000000..7d79d08d9c --- /dev/null +++ b/studio/frontend/src/features/transformers-upgrade/hooks/use-transformers-upgrade-consent.ts @@ -0,0 +1,28 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { useTransformersUpgradeDialogStore } from "../stores/transformers-upgrade-dialog-store"; +import type { TransformersUpgradeInfo } from "../types"; + +interface ConfirmArgs { + modelName: string; + /** validate's transformers_upgrade payload; null/undefined skips the dialog. */ + upgrade: TransformersUpgradeInfo | null | undefined; + /** When no release is installable, offer continuing into the caller's custom-code gate. */ + trustRemoteCodeFallback?: boolean; +} + +/** Pause a load needing a newer transformers on the consent dialog and run the install. + * Resolves true when the load can continue; false on cancel or not-installable with no fallback. */ +export async function confirmTransformersUpgradeIfNeeded({ + modelName, + upgrade, + trustRemoteCodeFallback, +}: ConfirmArgs): Promise { + if (!upgrade) return true; + return useTransformersUpgradeDialogStore + .getState() + .requestConsent(modelName, upgrade, { + trustRemoteCodeFallback: Boolean(trustRemoteCodeFallback), + }); +} diff --git a/studio/frontend/src/features/transformers-upgrade/index.ts b/studio/frontend/src/features/transformers-upgrade/index.ts new file mode 100644 index 0000000000..4318fd7830 --- /dev/null +++ b/studio/frontend/src/features/transformers-upgrade/index.ts @@ -0,0 +1,8 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +export { TransformersUpgradeDialog } from "./components/transformers-upgrade-dialog"; +export { confirmTransformersUpgradeIfNeeded } from "./hooks/use-transformers-upgrade-consent"; +export { installLatestTransformers } from "./api/transformers-upgrade-api"; +export { useTransformersUpgradeDialogStore } from "./stores/transformers-upgrade-dialog-store"; +export type { TransformersUpgradeInfo } from "./types"; diff --git a/studio/frontend/src/features/transformers-upgrade/stores/transformers-upgrade-dialog-store.ts b/studio/frontend/src/features/transformers-upgrade/stores/transformers-upgrade-dialog-store.ts new file mode 100644 index 0000000000..9e307fb1a1 --- /dev/null +++ b/studio/frontend/src/features/transformers-upgrade/stores/transformers-upgrade-dialog-store.ts @@ -0,0 +1,139 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { create } from "zustand"; +import { installLatestTransformers } from "../api/transformers-upgrade-api"; +import type { TransformersUpgradeInfo, TransformersUpgradePhase } from "../types"; + +type Resolver = (installed: boolean) => void; + +// One in-flight consent; a new request resolves any prior pending one as declined. +let pendingResolver: Resolver | null = null; + +interface TransformersUpgradeDialogStore { + open: boolean; + modelName: string | null; + upgrade: TransformersUpgradeInfo | null; + phase: TransformersUpgradePhase; + errorMessage: string | null; + /** Model ships custom code; without a PyPI install the load may fall back to trust_remote_code. */ + trustRemoteCodeFallback: boolean; + /** True once this consent's install completed. The install unloads the previous + * model before swapping, so the caller must treat it as already unloaded; the + * custom-code fallback resolves true without installing and leaves it loaded. */ + installRan: boolean; + /** True when the server unloaded the active chat model during this consent, + * including a swap that failed AFTER the unload: callers must then treat + * their previous model as gone and roll back on any later cancel. */ + serverUnloadedChat: boolean; + /** Read-and-clear serverUnloadedChat: each waiter consumes the signal once, + * so a superseding consent can neither erase it before the old waiter reads + * it nor leak it into an unrelated later load. */ + consumeServerUnloadedChat: () => boolean; + /** Open the dialog for a paused load; resolves true on install success or custom-code fallback. */ + requestConsent: ( + modelName: string, + upgrade: TransformersUpgradeInfo, + options?: { trustRemoteCodeFallback?: boolean }, + ) => Promise; + /** Accept/Retry: run the install; on success resolve(true) and close. */ + install: () => Promise; + resolve: (installed: boolean) => void; +} + +export const useTransformersUpgradeDialogStore = + create()((set, get) => ({ + open: false, + modelName: null, + upgrade: null, + phase: "consent", + errorMessage: null, + trustRemoteCodeFallback: false, + installRan: false, + serverUnloadedChat: false, + requestConsent: (modelName, upgrade, options) => + new Promise((resolve) => { + pendingResolver?.(false); + pendingResolver = resolve; + set({ + open: true, + modelName, + upgrade, + phase: "consent", + errorMessage: null, + trustRemoteCodeFallback: Boolean(options?.trustRemoteCodeFallback), + installRan: false, + }); + }), + consumeServerUnloadedChat: () => { + const value = get().serverUnloadedChat; + if (value) set({ serverUnloadedChat: false }); + return value; + }, + install: async () => { + const { upgrade, phase } = get(); + const version = upgrade?.pypi_version; + if (!version || phase === "installing") return; + const requestResolver = pendingResolver; + set({ phase: "installing", errorMessage: null }); + let result: Awaited>; + try { + result = await installLatestTransformers(version); + // Latch the server-side unload IMMEDIATELY, before any resolver-identity + // guard: even a superseded consent's install may have unloaded the chat + // model, and the signal must survive for whichever load consumes it next. + if (result.model_unloaded) { + set({ serverUnloadedChat: true }); + } + } catch (error) { + // Ignore the failure if a newer request superseded this consent. + if (pendingResolver === requestResolver) { + set({ + phase: "error", + errorMessage: + error instanceof Error && error.message + ? error.message + : "Failed to install transformers.", + }); + } + return; + } + if (pendingResolver === requestResolver) { + if (result.success) { + // serverUnloadedChat was latched above (and is never reset here): a + // retry after a failed-after-unload attempt reports false because the + // model is already gone, and a superseded install may have set it too. + set({ installRan: true }); + get().resolve(true); + return; + } + // Structured failure: the swap failed but may have already unloaded the + // chat model; record that so a later cancel still rolls the caller back. + // A version mismatch also carries the superseding release, so Retry + // re-requests a version that can actually succeed. + const { upgrade } = get(); + set({ + phase: "error", + errorMessage: result.message || "Failed to install transformers.", + serverUnloadedChat: + get().serverUnloadedChat || Boolean(result.model_unloaded), + ...(result.latest_version && upgrade + ? { upgrade: { ...upgrade, pypi_version: result.latest_version } } + : {}), + }); + } + }, + resolve: (installed) => { + const resolver = pendingResolver; + pendingResolver = null; + set({ + open: false, + modelName: null, + upgrade: null, + phase: "consent", + errorMessage: null, + trustRemoteCodeFallback: false, + }); + resolver?.(installed); + }, + })); diff --git a/studio/frontend/src/features/transformers-upgrade/types.ts b/studio/frontend/src/features/transformers-upgrade/types.ts new file mode 100644 index 0000000000..cf6bbe3d2f --- /dev/null +++ b/studio/frontend/src/features/transformers-upgrade/types.ts @@ -0,0 +1,16 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +/** Wire shape of `transformers_upgrade` from /api/inference/validate. */ +export interface TransformersUpgradeInfo { + /** config.json model_type unknown to installed transformers. */ + model_type: string; + /** Latest transformers release on PyPI at check time. */ + pypi_version?: string | null; + /** Latest PyPI release ships this model_type (installable after consent). */ + supported_in_pypi?: boolean; + /** Only transformers main ships it (dev-only; not installable). */ + supported_in_main?: boolean; +} + +export type TransformersUpgradePhase = "consent" | "installing" | "error"; From 1bf3509fea0cfa4de8de56a87cf215d7fd1d55cf Mon Sep 17 00:00:00 2001 From: Wasim Yousef Said Date: Wed, 15 Jul 2026 15:06:03 +0200 Subject: [PATCH 14/17] Fix agent workspace isolation and Hermes one-shot resume (#7103) * Fix coding agent workspace and resume handling * Handle attached Hermes flags and OpenClaw paths * Add Codex model metadata catalog * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix Codex reasoning summary metadata * Preserve Hermes hook approval on resumed one-shots --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- pyproject.toml | 1 + unsloth_cli/codex_fallback_prompt.md | 275 +++++++++++++++++++++++++++ unsloth_cli/commands/start.py | 152 ++++++++++++++- unsloth_cli/tests/test_start.py | 271 +++++++++++++++++++++++++- 4 files changed, 686 insertions(+), 13 deletions(-) create mode 100644 unsloth_cli/codex_fallback_prompt.md diff --git a/pyproject.toml b/pyproject.toml index 2b79121c82..917247c216 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,6 +42,7 @@ version = {attr = "unsloth.models._utils.__version__"} include-package-data = true [tool.setuptools.package-data] +unsloth_cli = ["codex_fallback_prompt.md"] studio = [ "*.sh", "*.ps1", diff --git a/unsloth_cli/codex_fallback_prompt.md b/unsloth_cli/codex_fallback_prompt.md new file mode 100644 index 0000000000..5d5fd6c9f8 --- /dev/null +++ b/unsloth_cli/codex_fallback_prompt.md @@ -0,0 +1,275 @@ +You are a coding agent running in the Codex CLI, a terminal-based coding assistant. Codex CLI is an open source project led by OpenAI. You are expected to be precise, safe, and helpful. + +Your capabilities: + +- Receive user prompts and other context provided by the harness, such as files in the workspace. +- Communicate with the user by streaming thinking & responses, and by making & updating plans. +- Emit function calls to run terminal commands and apply patches. Depending on how this specific run is configured, you can request that these function calls be escalated to the user for approval before running. More on this in the "Sandbox and approvals" section. + +Within this context, Codex refers to the open-source agentic coding interface (not the old Codex language model built by OpenAI). + +# How you work + +## Personality + +Your default personality and tone is concise, direct, and friendly. You communicate efficiently, always keeping the user clearly informed about ongoing actions without unnecessary detail. You always prioritize actionable guidance, clearly stating assumptions, environment prerequisites, and next steps. Unless explicitly asked, you avoid excessively verbose explanations about your work. + +# AGENTS.md spec +- Repos often contain AGENTS.md files. These files can appear anywhere within the repository. +- These files are a way for humans to give you (the agent) instructions or tips for working within the container. +- Some examples might be: coding conventions, info about how code is organized, or instructions for how to run or test code. +- Instructions in AGENTS.md files: + - The scope of an AGENTS.md file is the entire directory tree rooted at the folder that contains it. + - For every file you touch in the final patch, you must obey instructions in any AGENTS.md file whose scope includes that file. + - Instructions about code style, structure, naming, etc. apply only to code within the AGENTS.md file's scope, unless the file states otherwise. + - More-deeply-nested AGENTS.md files take precedence in the case of conflicting instructions. + - Direct system/developer/user instructions (as part of a prompt) take precedence over AGENTS.md instructions. +- The contents of the AGENTS.md file at the root of the repo and any directories from the CWD up to the root are included with the developer message and don't need to be re-read. When working in a subdirectory of CWD, or a directory outside the CWD, check for any AGENTS.md files that may be applicable. + +## Responsiveness + +### Preamble messages + +Before making tool calls, send a brief preamble to the user explaining what you’re about to do. When sending preamble messages, follow these principles and examples: + +- **Logically group related actions**: if you’re about to run several related commands, describe them together in one preamble rather than sending a separate note for each. +- **Keep it concise**: be no more than 1-2 sentences, focused on immediate, tangible next steps. (8–12 words for quick updates). +- **Build on prior context**: if this is not your first tool call, use the preamble message to connect the dots with what’s been done so far and create a sense of momentum and clarity for the user to understand your next actions. +- **Keep your tone light, friendly and curious**: add small touches of personality in preambles feel collaborative and engaging. +- **Exception**: Avoid adding a preamble for every trivial read (e.g., `cat` a single file) unless it’s part of a larger grouped action. + +**Examples:** + +- “I’ve explored the repo; now checking the API route definitions.” +- “Next, I’ll patch the config and update the related tests.” +- “I’m about to scaffold the CLI commands and helper functions.” +- “Ok cool, so I’ve wrapped my head around the repo. Now digging into the API routes.” +- “Config’s looking tidy. Next up is patching helpers to keep things in sync.” +- “Finished poking at the DB gateway. I will now chase down error handling.” +- “Alright, build pipeline order is interesting. Checking how it reports failures.” +- “Spotted a clever caching util; now hunting where it gets used.” + +## Planning + +You have access to an `update_plan` tool which tracks steps and progress and renders them to the user. Using the tool helps demonstrate that you've understood the task and convey how you're approaching it. Plans can help to make complex, ambiguous, or multi-phase work clearer and more collaborative for the user. A good plan should break the task into meaningful, logically ordered steps that are easy to verify as you go. + +Note that plans are not for padding out simple work with filler steps or stating the obvious. The content of your plan should not involve doing anything that you aren't capable of doing (i.e. don't try to test things that you can't test). Do not use plans for simple or single-step queries that you can just do or answer immediately. + +Do not repeat the full contents of the plan after an `update_plan` call — the harness already displays it. Instead, summarize the change made and highlight any important context or next step. + +Before running a command, consider whether or not you have completed the previous step, and make sure to mark it as completed before moving on to the next step. It may be the case that you complete all steps in your plan after a single pass of implementation. If this is the case, you can simply mark all the planned steps as completed. Sometimes, you may need to change plans in the middle of a task: call `update_plan` with the updated plan and make sure to provide an `explanation` of the rationale when doing so. + +Use a plan when: + +- The task is non-trivial and will require multiple actions over a long time horizon. +- There are logical phases or dependencies where sequencing matters. +- The work has ambiguity that benefits from outlining high-level goals. +- You want intermediate checkpoints for feedback and validation. +- When the user asked you to do more than one thing in a single prompt +- The user has asked you to use the plan tool (aka "TODOs") +- You generate additional steps while working, and plan to do them before yielding to the user + +### Examples + +**High-quality plans** + +Example 1: + +1. Add CLI entry with file args +2. Parse Markdown via CommonMark library +3. Apply semantic HTML template +4. Handle code blocks, images, links +5. Add error handling for invalid files + +Example 2: + +1. Define CSS variables for colors +2. Add toggle with localStorage state +3. Refactor components to use variables +4. Verify all views for readability +5. Add smooth theme-change transition + +Example 3: + +1. Set up Node.js + WebSocket server +2. Add join/leave broadcast events +3. Implement messaging with timestamps +4. Add usernames + mention highlighting +5. Persist messages in lightweight DB +6. Add typing indicators + unread count + +**Low-quality plans** + +Example 1: + +1. Create CLI tool +2. Add Markdown parser +3. Convert to HTML + +Example 2: + +1. Add dark mode toggle +2. Save preference +3. Make styles look good + +Example 3: + +1. Create single-file HTML game +2. Run quick sanity check +3. Summarize usage instructions + +If you need to write a plan, only write high quality plans, not low quality ones. + +## Task execution + +You are a coding agent. Please keep going until the query is completely resolved, before ending your turn and yielding back to the user. Only terminate your turn when you are sure that the problem is solved. Autonomously resolve the query to the best of your ability, using the tools available to you, before coming back to the user. Do NOT guess or make up an answer. + +You MUST adhere to the following criteria when solving queries: + +- Working on the repo(s) in the current environment is allowed, even if they are proprietary. +- Analyzing code for vulnerabilities is allowed. +- Showing user code and tool call details is allowed. +- Use the `apply_patch` tool to edit files (NEVER try `applypatch` or `apply-patch`, only `apply_patch`): {"command":["apply_patch","*** Begin Patch\\n*** Update File: path/to/file.py\\n@@ def example():\\n- pass\\n+ return 123\\n*** End Patch"]} + +If completing the user's task requires writing or modifying files, your code and final answer should follow these coding guidelines, though user instructions (i.e. AGENTS.md) may override these guidelines: + +- Fix the problem at the root cause rather than applying surface-level patches, when possible. +- Avoid unneeded complexity in your solution. +- Do not attempt to fix unrelated bugs or broken tests. It is not your responsibility to fix them. (You may mention them to the user in your final message though.) +- Update documentation as necessary. +- Keep changes consistent with the style of the existing codebase. Changes should be minimal and focused on the task. +- Use `git log` and `git blame` to search the history of the codebase if additional context is required. +- NEVER add copyright or license headers unless specifically requested. +- Do not waste tokens by re-reading files after calling `apply_patch` on them. The tool call will fail if it didn't work. The same goes for making folders, deleting folders, etc. +- Do not `git commit` your changes or create new git branches unless explicitly requested. +- Do not add inline comments within code unless explicitly requested. +- Do not use one-letter variable names unless explicitly requested. +- NEVER output inline citations like "【F:README.md†L5-L14】" in your outputs. The CLI is not able to render these so they will just be broken in the UI. Instead, if you output valid filepaths, users will be able to click on them to open the files in their editor. + +## Validating your work + +If the codebase has tests or the ability to build or run, consider using them to verify that your work is complete. + +When testing, your philosophy should be to start as specific as possible to the code you changed so that you can catch issues efficiently, then make your way to broader tests as you build confidence. If there's no test for the code you changed, and if the adjacent patterns in the codebases show that there's a logical place for you to add a test, you may do so. However, do not add tests to codebases with no tests. + +Similarly, once you're confident in correctness, you can suggest or use formatting commands to ensure that your code is well formatted. If there are issues you can iterate up to 3 times to get formatting right, but if you still can't manage it's better to save the user time and present them a correct solution where you call out the formatting in your final message. If the codebase does not have a formatter configured, do not add one. + +For all of testing, running, building, and formatting, do not attempt to fix unrelated bugs. It is not your responsibility to fix them. (You may mention them to the user in your final message though.) + +Be mindful of whether to run validation commands proactively. In the absence of behavioral guidance: + +- When running in the non-interactive approval mode **never**, proactively run tests, lint and do whatever you need to ensure you've completed the task. +- When working in interactive approval modes like **untrusted**, or **on-request**, hold off on running tests or lint commands until the user is ready for you to finalize your output, because these commands take time to run and slow down iteration. Instead suggest what you want to do next, and let the user confirm first. +- When working on test-related tasks, such as adding tests, fixing tests, or reproducing a bug to verify behavior, you may proactively run tests regardless of approval mode. Use your judgement to decide whether this is a test-related task. + +## Ambition vs. precision + +For tasks that have no prior context (i.e. the user is starting something brand new), you should feel free to be ambitious and demonstrate creativity with your implementation. + +If you're operating in an existing codebase, you should make sure you do exactly what the user asks with surgical precision. Treat the surrounding codebase with respect, and don't overstep (i.e. changing filenames or variables unnecessarily). You should balance being sufficiently ambitious and proactive when completing tasks of this nature. + +You should use judicious initiative to decide on the right level of detail and complexity to deliver based on the user's needs. This means showing good judgment that you're capable of doing the right extras without gold-plating. This might be demonstrated by high-value, creative touches when scope of the task is vague; while being surgical and targeted when scope is tightly specified. + +## Sharing progress updates + +For especially longer tasks that you work on (i.e. requiring many tool calls, or a plan with multiple steps), you should provide progress updates back to the user at reasonable intervals. These updates should be structured as a concise sentence or two (no more than 8-10 words long) recapping progress so far in plain language: this update demonstrates your understanding of what needs to be done, progress so far (i.e. files explores, subtasks complete), and where you're going next. + +Before doing large chunks of work that may incur latency as experienced by the user (i.e. writing a new file), you should send a concise message to the user with an update indicating what you're about to do to ensure they know what you're spending time on. Don't start editing or writing large files before informing the user what you are doing and why. + +The messages you send before tool calls should describe what is immediately about to be done next in very concise language. If there was previous work done, this preamble message should also include a note about the work done so far to bring the user along. + +## Presenting your work and final message + +Your final message should read naturally, like an update from a concise teammate. For casual conversation, brainstorming tasks, or quick questions from the user, respond in a friendly, conversational tone. You should ask questions, suggest ideas, and adapt to the user’s style. If you've finished a large amount of work, when describing what you've done to the user, you should follow the final answer formatting guidelines to communicate substantive changes. You don't need to add structured formatting for one-word answers, greetings, or purely conversational exchanges. + +You can skip heavy formatting for single, simple actions or confirmations. In these cases, respond in plain sentences with any relevant next step or quick option. Reserve multi-section structured responses for results that need grouping or explanation. + +The user is working on the same computer as you, and has access to your work. As such there's no need to show the full contents of large files you have already written unless the user explicitly asks for them. Similarly, if you've created or modified files using `apply_patch`, there's no need to tell users to "save the file" or "copy the code into a file"—just reference the file path. + +If there's something that you think you could help with as a logical next step, concisely ask the user if they want you to do so. Good examples of this are running tests, committing changes, or building out the next logical component. If there’s something that you couldn't do (even with approval) but that the user might want to do (such as verifying changes by running the app), include those instructions succinctly. + +Brevity is very important as a default. You should be very concise (i.e. no more than 10 lines), but can relax this requirement for tasks where additional detail and comprehensiveness is important for the user's understanding. + +### Final answer structure and style guidelines + +You are producing plain text that will later be styled by the CLI. Follow these rules exactly. Formatting should make results easy to scan, but not feel mechanical. Use judgment to decide how much structure adds value. + +**Section Headers** + +- Use only when they improve clarity — they are not mandatory for every answer. +- Choose descriptive names that fit the content +- Keep headers short (1–3 words) and in `**Title Case**`. Always start headers with `**` and end with `**` +- Leave no blank line before the first bullet under a header. +- Section headers should only be used where they genuinely improve scanability; avoid fragmenting the answer. + +**Bullets** + +- Use `-` followed by a space for every bullet. +- Merge related points when possible; avoid a bullet for every trivial detail. +- Keep bullets to one line unless breaking for clarity is unavoidable. +- Group into short lists (4–6 bullets) ordered by importance. +- Use consistent keyword phrasing and formatting across sections. + +**Monospace** + +- Wrap all commands, file paths, env vars, and code identifiers in backticks (`` `...` ``). +- Apply to inline examples and to bullet keywords if the keyword itself is a literal file/command. +- Never mix monospace and bold markers; choose one based on whether it’s a keyword (`**`) or inline code/path (`` ` ``). + +**File References** +When referencing files in your response, make sure to include the relevant start line and always follow the below rules: + * Use inline code to make file paths clickable. + * Each reference should have a stand alone path. Even if it's the same file. + * Accepted: absolute, workspace‑relative, a/ or b/ diff prefixes, or bare filename/suffix. + * Line/column (1‑based, optional): :line[:column] or #Lline[Ccolumn] (column defaults to 1). + * Do not use URIs like file://, vscode://, or https://. + * Do not provide range of lines + * Examples: src/app.ts, src/app.ts:42, b/server/index.js#L10, C:\repo\project\main.rs:12:5 + +**Structure** + +- Place related bullets together; don’t mix unrelated concepts in the same section. +- Order sections from general → specific → supporting info. +- For subsections (e.g., “Binaries” under “Rust Workspace”), introduce with a bolded keyword bullet, then list items under it. +- Match structure to complexity: + - Multi-part or detailed results → use clear headers and grouped bullets. + - Simple results → minimal headers, possibly just a short list or paragraph. + +**Tone** + +- Keep the voice collaborative and natural, like a coding partner handing off work. +- Be concise and factual — no filler or conversational commentary and avoid unnecessary repetition +- Use present tense and active voice (e.g., “Runs tests” not “This will run tests”). +- Keep descriptions self-contained; don’t refer to “above” or “below”. +- Use parallel structure in lists for consistency. + +**Don’t** + +- Don’t use literal words “bold” or “monospace” in the content. +- Don’t nest bullets or create deep hierarchies. +- Don’t output ANSI escape codes directly — the CLI renderer applies them. +- Don’t cram unrelated keywords into a single bullet; split for clarity. +- Don’t let keyword lists run long — wrap or reformat for scanability. + +Generally, ensure your final answers adapt their shape and depth to the request. For example, answers to code explanations should have a precise, structured explanation with code references that answer the question directly. For tasks with a simple implementation, lead with the outcome and supplement only with what’s needed for clarity. Larger changes can be presented as a logical walkthrough of your approach, grouping related steps, explaining rationale where it adds value, and highlighting next actions to accelerate the user. Your answers should provide the right level of detail while being easily scannable. + +For casual greetings, acknowledgements, or other one-off conversational messages that are not delivering substantive information or structured results, respond naturally without section headers or bullet formatting. + +# Tool Guidelines + +## Shell commands + +When using the shell, you must adhere to the following guidelines: + +- When searching for text or files, prefer using `rg` or `rg --files` respectively because `rg` is much faster than alternatives like `grep`. (If the `rg` command is not found, then use alternatives.) +- Do not use python scripts to attempt to output larger chunks of a file. + +## `update_plan` + +A tool named `update_plan` is available to you. You can use it to keep an up‑to‑date, step‑by‑step plan for the task. + +To create a new plan, call `update_plan` with a short list of 1‑sentence steps (no more than 5-7 words each) with a `status` for each step (`pending`, `in_progress`, or `completed`). + +When steps have been completed, use `update_plan` to mark each finished step as `completed` and the next step you are working on as `in_progress`. There should always be exactly one `in_progress` step until everything is done. You can mark multiple items as complete in a single `update_plan` call. + +If all steps are complete, ensure you call `update_plan` to mark all steps as `completed`. diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index 48c0aca34b..fa39fdf761 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -170,6 +170,43 @@ def _hermes_install_hint() -> str: return _HERMES_WINDOWS_INSTALL_HINT if os.name == "nt" else _HERMES_POSIX_INSTALL_HINT +def _hermes_resume_oneshot_args(args: list[str]) -> list[str]: + """Route resumed one-shot prompts through Hermes' session-aware chat command.""" + has_resume = any( + arg in ("--resume", "-r", "--continue", "-c") + or arg.startswith(("--resume=", "--continue=")) + or (len(arg) > 2 and arg.startswith(("-r", "-c"))) + for arg in args + ) + if not has_resume: + return args + + rewritten = list(args) + for index, arg in enumerate(rewritten): + if arg in ("-z", "--oneshot"): + rewritten[index] = "-q" + elif len(arg) > 2 and arg.startswith("-z"): + # argparse accepts attached short-option values (`-zPROMPT` and + # `-z=PROMPT`); preserve the value byte-for-byte when switching to -q. + rewritten[index] = f"-q{arg[2:]}" + elif arg.startswith("--oneshot="): + rewritten[index] = f"--query={arg.partition('=')[2]}" + else: + continue + if any(item == "--usage-file" or item.startswith("--usage-file=") for item in args): + raise typer.BadParameter( + "Hermes cannot resume a one-shot session with --usage-file; remove that option." + ) + prefix = ["chat", "-Q"] + if "--yolo" not in rewritten: + prefix.append("--yolo") + if "--accept-hooks" not in rewritten: + prefix.append("--accept-hooks") + rewritten = prefix + rewritten + return rewritten + return args + + class LoadOptions(NamedTuple): """Model-load knobs forwarded to /api/inference/load when --model triggers a load.""" @@ -840,6 +877,60 @@ def _merge_codex_config(existing: str, base: str) -> str: ) +# Keep custom-model behavior aligned with Codex's own unknown-model fallback. This +# Apache-2.0 prompt is copied from openai/codex rust-v0.144.0 models-manager/prompt.md. +_CODEX_FALLBACK_PROMPT = Path(__file__).parent.parent / "codex_fallback_prompt.md" +_CODEX_MODEL_CATALOG_MIN_VERSION = (0, 110, 0) + + +def _codex_supports_model_catalog() -> bool: + executable = shutil.which("codex") + if executable is None: + # A --no-launch recipe may be copied to another machine; assume a current Codex. + return True + try: + output = subprocess.check_output( + [executable, "--version"], text = True, timeout = 10, stderr = subprocess.DEVNULL + ) + except Exception: + return False + match = re.search(r"(\d+)\.(\d+)\.(\d+)", output) + return bool(match) and tuple(int(part) for part in match.groups()) >= ( + _CODEX_MODEL_CATALOG_MIN_VERSION + ) + + +def _codex_model_catalog(model: dict) -> dict: + """Return conservative metadata for a Studio model unknown to Codex's built-in catalog.""" + model_id = model["id"] + window = model.get("context_length") or model.get("max_context_length") + entry = { + "slug": model_id, + "display_name": model_id, + "description": "Model served by Unsloth Studio", + "supported_reasoning_levels": [], + "shell_type": "default", + "visibility": "none", + "supported_in_api": True, + "priority": 99, + "availability_nux": None, + "upgrade": None, + "base_instructions": _CODEX_FALLBACK_PROMPT.read_text(encoding = "utf-8"), + "supports_reasoning_summaries": False, + "supports_reasoning_summary_parameter": False, + "support_verbosity": False, + "default_verbosity": None, + "apply_patch_tool_type": None, + "truncation_policy": {"mode": "bytes", "limit": 10_000}, + "supports_parallel_tool_calls": False, + "experimental_supported_tools": [], + } + if window: + entry["context_window"] = int(window) + entry["max_context_window"] = int(window) + return {"models": [entry]} + + def write_codex_config(base: str, model: dict, home: Path) -> None: home.mkdir(parents = True, exist_ok = True) @@ -857,6 +948,16 @@ def write_codex_config(base: str, model: dict, home: Path) -> None: f'model_provider = "{_CODEX_PROFILE}"\n' f"model = {json.dumps(model['id'])}\n" ) + if _codex_supports_model_catalog() and _CODEX_FALLBACK_PROMPT.is_file(): + catalog = home / "model-catalog.json" + catalog_text = json.dumps(_codex_model_catalog(model), indent = 2) + "\n" + if not catalog.exists() or catalog.read_text(encoding = "utf-8") != catalog_text: + catalog.write_text(catalog_text, encoding = "utf-8") + typer.echo(f"Updated {catalog}") + # Resolve relative to the profile file. This also survives WSL launching a Windows + # Codex binary, where a Linux absolute path inside TOML would not be usable. + profile_text += f"model_catalog_json = {json.dumps(catalog.name)}\n" + window = model.get("context_length") or model.get("max_context_length") if window: profile_text += f"model_context_window = {int(window)}\n" @@ -875,6 +976,16 @@ def _wsl_windows_executable(command: list) -> Optional[str]: return None +def _wsl_windows_path(path: Path) -> str: + try: + translated = subprocess.check_output(["wslpath", "-w", str(path)], text = True).strip() + except (OSError, subprocess.CalledProcessError) as exc: + _fail(f"Could not translate WSL path {path}: {exc}") + if not translated: + _fail(f"Could not translate WSL path {path}") + return translated + + def _looks_like_path(value: str) -> bool: # A var only wants the WSLENV /p flag if its value is a filesystem path: an # absolute POSIX path (/...), a UNC path (\\...), or a drive-qualified Windows @@ -1184,6 +1295,7 @@ def write_openclaw_config( model: dict, path: Path, yolo: bool = False, + workspace_path: Optional[str] = None, ) -> None: config = _read_json_object(path) if config is None: @@ -1208,8 +1320,23 @@ def write_openclaw_config( "models": [provider_model], } # Pin a default model, else OpenClaw drops into its setup agent ("no models available"). - defaults = _subdict(_subdict(config, "agents"), "defaults") + agents = _subdict(config, "agents") + defaults = _subdict(agents, "defaults") _subdict(defaults, "model")["primary"] = f"unsloth/{model['id']}" + # OPENCLAW_STATE_DIR does not relocate the workspace. Keep it beside the managed + # config so ephemeral launches avoid ~/.openclaw and persisted sessions retain it. + workspace = path.parent / "workspace" + workspace.mkdir(parents = True, exist_ok = True, mode = 0o700) + defaults["workspace"] = workspace_path or str(workspace) + # Per-agent paths override agents.defaults.workspace and OPENCLAW_STATE_DIR. This + # config is itself an isolated Unsloth copy, so remove stale explicit paths and let + # OpenClaw resolve every listed agent beneath the managed defaults/state directory. + agent_list = agents.get("list") + if isinstance(agent_list, list): + for agent_config in agent_list: + if isinstance(agent_config, dict): + agent_config.pop("workspace", None) + agent_config.pop("agentDir", None) # Unauthenticated loopback gateway: without auth.mode=none the client won't open # the websocket. The daemon must still be started separately (`openclaw gateway`). gateway = _subdict(config, "gateway") @@ -1339,9 +1466,11 @@ def write_opencode_config( tools = ("edit", "bash", "webfetch") if yolo: # OpenCode has no --yolo flag; auto-approve is the config `permission` block - # (singular). Allow the prompting tools so tool calls don't block on the TUI. This - # rides inline (OPENCODE_CONFIG_CONTENT) so --yolo works even over a project config. + # (singular). Allow the prompting tools and paths outside the launch directory so + # tool calls don't block on the TUI. This rides inline (OPENCODE_CONFIG_CONTENT) so + # --yolo works even over a project config. session_permission = {t: "allow" for t in tools} + session_permission["external_directory"] = {"*": "allow"} config["permission"] = dict(session_permission) else: # Undo only what --yolo wrote: our yolo sets an explicit per-tool "allow" for these @@ -1358,6 +1487,8 @@ def write_opencode_config( for tool in tools: if permission.get(tool) == "allow": permission[tool] = "ask" + if permission.get("external_directory") == {"*": "allow"}: + permission["external_directory"] = {"*": "ask"} if json.dumps(config, sort_keys = True) != before: _write_private_json(path, config) typer.echo(f"Updated {path}") @@ -1629,8 +1760,18 @@ def openclaw( ) with _session_config("openclaw", launch, persist = persist) as cfg: config_path = cfg / "openclaw.json" + workspace_path = None + if _wsl_windows_executable(command): + workspace_path = _wsl_windows_path(cfg / "workspace") # key lives in the config, not the env; --yolo writes the exec policy here too. - write_openclaw_config(base, key, entry, config_path, yolo = yolo) + write_openclaw_config( + base, + key, + entry, + config_path, + yolo = yolo, + workspace_path = workspace_path, + ) # Scope both config and state so OpenClaw never touches the user's ~/.openclaw. env = {"OPENCLAW_CONFIG_PATH": str(config_path), "OPENCLAW_STATE_DIR": str(cfg)} _run(base, entry, env, command, launch = launch, install_hint = install_hint) @@ -1729,6 +1870,8 @@ def hermes( persist: bool = _PERSIST_OPTION, ): """Point Hermes (Nous Research) at the running Studio server and start it.""" + native_args = [*_yolo_command_flags("hermes", yolo), *ctx.args] + command = ["hermes", *_hermes_resume_oneshot_args(native_args)] base, key, entry = _connect( api_key, model, @@ -1736,7 +1879,6 @@ def hermes( serve = serve, launch = launch, ) - command = ["hermes", *_yolo_command_flags("hermes", yolo), *ctx.args] install_hint = _hermes_install_hint() with _session_config("hermes", launch, persist = persist) as home: # HERMES_HOME relocates hermes' whole home dir (config.yaml, sessions, state) diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index 065972b275..7e76465144 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -294,17 +294,58 @@ def test_merge_codex_config_keeps_user_oss_provider(): assert _parse_toml(merged)["oss_provider"] == "ollama" -def test_write_codex_config_profile(tmp_path): +def test_write_codex_config_profile(tmp_path, monkeypatch): + monkeypatch.setattr(start, "_codex_supports_model_catalog", lambda: True) start.write_codex_config(BASE, MODEL, tmp_path) profile = _parse_toml((tmp_path / "unsloth_api.config.toml").read_text()) assert profile["oss_provider"] == "unsloth_api" assert profile["model_provider"] == "unsloth_api" assert profile["model"] == MODEL["id"] assert profile["model_context_window"] == 131072 + + catalog_path = Path(profile["model_catalog_json"]) + assert catalog_path == Path("model-catalog.json") + catalog = json.loads((tmp_path / catalog_path).read_text()) + assert catalog["models"][0]["slug"] == MODEL["id"] + assert catalog["models"][0]["context_window"] == 131072 + assert catalog["models"][0]["max_context_window"] == 131072 + assert catalog["models"][0]["supports_reasoning_summary_parameter"] is False + assert catalog["models"][0]["supports_parallel_tool_calls"] is False + + assert catalog["models"][0]["base_instructions"] == start._CODEX_FALLBACK_PROMPT.read_text() config = _parse_toml((tmp_path / "config.toml").read_text()) assert config["model_providers"]["unsloth_api"]["env_key"] == "UNSLOTH_STUDIO_AUTH_TOKEN" +def test_write_codex_config_catalog_without_context_length(tmp_path, monkeypatch): + monkeypatch.setattr(start, "_codex_supports_model_catalog", lambda: True) + start.write_codex_config(BASE, {"id": "unsloth/no-window"}, tmp_path) + profile = _parse_toml((tmp_path / "unsloth_api.config.toml").read_text()) + catalog = json.loads((tmp_path / profile["model_catalog_json"]).read_text()) + entry = catalog["models"][0] + assert entry["slug"] == "unsloth/no-window" + assert "context_window" not in entry + assert "max_context_window" not in entry + + +@pytest.mark.parametrize( + ("version", "expected"), + [("codex-cli 0.109.0", False), ("codex-cli 0.110.0", True), ("codex-cli 0.144.4", True)], +) +def test_codex_model_catalog_version_gate(monkeypatch, version, expected): + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex") + monkeypatch.setattr(start.subprocess, "check_output", lambda *args, **kwargs: version) + assert start._codex_supports_model_catalog() is expected + + +def test_write_codex_config_omits_catalog_for_old_codex(tmp_path, monkeypatch): + monkeypatch.setattr(start, "_codex_supports_model_catalog", lambda: False) + start.write_codex_config(BASE, MODEL, tmp_path) + profile = _parse_toml((tmp_path / "unsloth_api.config.toml").read_text()) + assert "model_catalog_json" not in profile + assert not (tmp_path / "model-catalog.json").exists() + + @pytest.fixture() def fake_studio(tmp_path, monkeypatch): calls = [] @@ -742,7 +783,12 @@ def test_opencode_inline_config_beats_project_config(fake_studio): assert result.exit_code == 0, result.output inline = _opencode_inline_config(result.output) assert inline["model"] == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}" - assert inline["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"} + assert inline["permission"] == { + "edit": "allow", + "bash": "allow", + "webfetch": "allow", + "external_directory": {"*": "allow"}, + } assert "sk-unsloth" not in result.output # key stays in the private file, not the env @@ -1611,12 +1657,50 @@ def test_write_openclaw_config_fresh(tmp_path): ] # The default model must be pinned or OpenClaw has nothing active. assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" + assert config["agents"]["defaults"]["workspace"] == str(tmp_path / "workspace") + assert (tmp_path / "workspace").is_dir() assert config["gateway"]["mode"] == "local" assert config["gateway"]["auth"]["mode"] == "none" # unauth loopback gateway if os.name != "nt": # the file holds an API key assert path.stat().st_mode & 0o777 == 0o600 +def test_write_openclaw_config_clears_per_agent_path_overrides(tmp_path): + path = tmp_path / "openclaw.json" + path.write_text( + json.dumps( + { + "agents": { + "defaults": {"workspace": "/old/default"}, + "list": [ + { + "id": "main", + "default": True, + "workspace": "/old/main-workspace", + "agentDir": "/old/main-agent", + "model": "keep/me", + }, + { + "id": "reviewer", + "workspace": "/old/reviewer-workspace", + "agentDir": "/old/reviewer-agent", + }, + ], + } + } + ) + ) + + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path) + + agents = json.loads(path.read_text())["agents"] + assert agents["defaults"]["workspace"] == str(tmp_path / "workspace") + assert agents["list"] == [ + {"id": "main", "default": True, "model": "keep/me"}, + {"id": "reviewer"}, + ] + + def test_write_openclaw_config_preserves_and_idempotent(tmp_path): path = tmp_path / "openclaw.json" path.write_text( @@ -1660,11 +1744,32 @@ def test_connect_openclaw_no_launch(fake_studio, tmp_path): config = json.loads(config_path.read_text()) assert config["models"]["providers"]["unsloth"]["apiKey"] == "sk-unsloth-feedfacefeedface" assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" + assert config["agents"]["defaults"]["workspace"] == str( + tmp_path / "agents" / "openclaw" / "workspace" + ) assert _launch_command(result.output) == ["openclaw", "tui", "--local"] # OpenAI /v1/chat/completions works on either backend — no GGUF gate. assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) +@pytest.mark.skipif(os.name == "nt", reason = "WSL scenario") +def test_connect_openclaw_wsl_windows_shim_translates_workspace(fake_studio, tmp_path, monkeypatch): + windows_workspace = r"\\wsl.localhost\Ubuntu\tmp\openclaw\workspace" + monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") + monkeypatch.setattr( + start.shutil, "which", lambda _: "/mnt/c/Users/x/AppData/Roaming/npm/openclaw" + ) + monkeypatch.setattr(start.subprocess, "check_output", lambda *args, **kwargs: windows_workspace) + + result = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch"]) + + assert result.exit_code == 0, result.output + config_path = tmp_path / "agents" / "openclaw" / "openclaw.json" + config = json.loads(config_path.read_text()) + assert config["agents"]["defaults"]["workspace"] == windows_workspace + assert (config_path.parent / "workspace").is_dir() + + def test_connect_openclaw_no_launch_keeps_explicit_subcommand(fake_studio): result = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch", "crestodian"]) assert result.exit_code == 0, result.output @@ -2090,7 +2195,12 @@ def test_yolo_opencode_writes_permission_block(fake_studio, tmp_path): result = CliRunner().invoke(start.start_app, ["opencode", "--yolo", "--no-launch"]) assert result.exit_code == 0, result.output config = json.loads((tmp_path / "agents" / "opencode" / "opencode.json").read_text()) - assert config["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"} + assert config["permission"] == { + "edit": "allow", + "bash": "allow", + "webfetch": "allow", + "external_directory": {"*": "allow"}, + } def test_no_yolo_opencode_has_no_permission_block(fake_studio, tmp_path): @@ -2112,6 +2222,7 @@ def test_no_yolo_opencode_flips_prior_yolo_allow_to_ask(fake_studio, tmp_path): "edit": "allow", "bash": "allow", "webfetch": "allow", + "external_directory": {"*": "allow"}, } plain = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) assert plain.exit_code == 0, plain.output @@ -2119,6 +2230,7 @@ def test_no_yolo_opencode_flips_prior_yolo_allow_to_ask(fake_studio, tmp_path): "edit": "ask", "bash": "ask", "webfetch": "ask", + "external_directory": {"*": "ask"}, } @@ -2152,7 +2264,12 @@ def test_write_opencode_config_yolo_unit(tmp_path): path = tmp_path / "opencode.json" start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True) config = json.loads(path.read_text()) - assert config["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"} + assert config["permission"] == { + "edit": "allow", + "bash": "allow", + "webfetch": "allow", + "external_directory": {"*": "allow"}, + } def test_write_openclaw_config_yolo_unit(tmp_path): @@ -2180,7 +2297,12 @@ def test_no_launch_rerun_clears_stale_opencode_yolo_permissions(fake_studio, tmp config = json.loads(config_path.read_text()) # The yolo allow policy is replaced by a prompting one, not deleted (which would # revert to OpenCode's permissive "allow" default). - assert config["permission"] == {"edit": "ask", "bash": "ask", "webfetch": "ask"} + assert config["permission"] == { + "edit": "ask", + "bash": "ask", + "webfetch": "ask", + "external_directory": {"*": "ask"}, + } # The session provider survives the cleanup. assert start._OPENCODE_PROVIDER in config["provider"] @@ -2218,7 +2340,12 @@ def test_write_opencode_config_yolo_then_plain_unit(tmp_path): start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = False) config = json.loads(path.read_text()) # A plain rerun replaces the yolo allow policy with a prompting one. - assert config["permission"] == {"edit": "ask", "bash": "ask", "webfetch": "ask"} + assert config["permission"] == { + "edit": "ask", + "bash": "ask", + "webfetch": "ask", + "external_directory": {"*": "ask"}, + } def test_openclaw_non_yolo_keeps_runtime_approvals(tmp_path): @@ -2670,8 +2797,7 @@ def test_default_launch_has_no_resume_token(fake_studio, monkeypatch): def test_resume_persist_only_agents_have_no_resume_token(fake_studio, monkeypatch): - # openclaw/hermes persist their session dir but have no non-interactive resume - # selector, so --persist must not append a token; their own picker resumes. + # Persistence alone must not select a session. for agent in ("openclaw", "hermes"): monkeypatch.setattr(start.shutil, "which", lambda _, a = agent: f"/usr/local/bin/{a}") captured = _capture_launch(monkeypatch, [agent, "--persist"]) @@ -2679,6 +2805,135 @@ def test_resume_persist_only_agents_have_no_resume_token(fake_studio, monkeypatc assert "--continue" not in captured["command"] +@pytest.mark.parametrize( + ("args", "expected"), + [ + ( + ["--resume", "session-id", "-z", "follow up"], + [ + "chat", + "-Q", + "--yolo", + "--accept-hooks", + "--resume", + "session-id", + "-q", + "follow up", + ], + ), + ( + ["-rsession-id", "-zfollow up"], + ["chat", "-Q", "--yolo", "--accept-hooks", "-rsession-id", "-qfollow up"], + ), + ( + ["-c=project", "-z=follow up"], + ["chat", "-Q", "--yolo", "--accept-hooks", "-c=project", "-q=follow up"], + ), + ( + ["-r", "session-id", "--oneshot=follow up"], + [ + "chat", + "-Q", + "--yolo", + "--accept-hooks", + "-r", + "session-id", + "--query=follow up", + ], + ), + ( + ["--continue", "project", "--oneshot", "follow up"], + [ + "chat", + "-Q", + "--yolo", + "--accept-hooks", + "--continue", + "project", + "-q", + "follow up", + ], + ), + ( + ["--yolo", "--resume", "session-id", "-z", "follow up"], + [ + "chat", + "-Q", + "--accept-hooks", + "--yolo", + "--resume", + "session-id", + "-q", + "follow up", + ], + ), + ( + ["--accept-hooks", "--resume", "session-id", "-z", "follow up"], + [ + "chat", + "-Q", + "--yolo", + "--accept-hooks", + "--resume", + "session-id", + "-q", + "follow up", + ], + ), + ( + ["--resume", "chat", "-z", "follow up"], + [ + "chat", + "-Q", + "--yolo", + "--accept-hooks", + "--resume", + "chat", + "-q", + "follow up", + ], + ), + (["--resume", "session-id"], ["--resume", "session-id"]), + (["-z", "new session"], ["-z", "new session"]), + ], +) +def test_hermes_resume_oneshot_args(args, expected): + assert start._hermes_resume_oneshot_args(args) == expected + + +def test_hermes_resume_oneshot_uses_session_aware_chat(fake_studio, monkeypatch): + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/hermes") + captured = _capture_launch( + monkeypatch, + ["hermes", "--persist", "--resume", "session-id", "-z", "follow up"], + ) + assert captured["command"][1:] == [ + "chat", + "-Q", + "--yolo", + "--accept-hooks", + "--resume", + "session-id", + "-q", + "follow up", + ] + + +@pytest.mark.parametrize("usage_arg", ["--usage-file", "--usage-file=usage.json"]) +def test_hermes_resume_oneshot_rejects_usage_file(monkeypatch, usage_arg): + monkeypatch.setattr( + start, + "_connect", + lambda *args, **kwargs: pytest.fail("argument validation must run before connect"), + ) + argv = ["hermes", "--resume", "session-id", "-z", "follow up", usage_arg] + if usage_arg == "--usage-file": + argv.append("usage.json") + result = CliRunner().invoke(start.start_app, argv) + assert result.exit_code == 2 + assert "cannot resume a one-shot session with --usage-file" in result.output + + def test_native_resume_flag_passes_through_unchanged(fake_studio, monkeypatch): # The persistence flag is --persist, NOT --resume, so an agent's own # `--resume ` (e.g. `unsloth start claude --resume `) still flows From e1e38419dfb661d26235794da903d4b7b73b26da Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Wed, 15 Jul 2026 06:07:21 -0700 Subject: [PATCH 15/17] Studio: permission levels for chat tool calls (Ask, Approve for me, Off, Full access) (#7079) * Studio: permission levels for chat tool calls (Ask, Approve for me, Off, Full access) Replace the Bypass permissions on/off toggle with a four level permission selector, available in Settings > General (new Permissions section above Notifications), the chat settings panel, the composer plus menu, and a new always visible composer pill. Levels: - Ask for approval: every local tool call pauses for allow/deny. - Approve for me: only calls detected as potentially unsafe pause; the python/terminal sandbox stays on. - Off: never pauses; sandbox stays on (previous default behavior). - Full access: never pauses and the sandbox is disabled. Still requires the danger confirmation and is never restored across reloads. Backend adds permission_mode to the OpenAI compatible and Anthropic passthrough payloads and threads it through both tool loops. Auto mode uses a fail closed classifier in tools.py: terminal commands must be on a read only allowlist with no redirection or substitution, python code is AST scanned for writes, exec, process and network use, MCP tools auto run only with read only style names. Unknown tools always ask. Legacy bypass_permissions and confirm_tool_calls keep their exact behavior for existing API callers. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio permissions: Off is a plain toggle below Full access Off moves to the bottom of the level menu with a short description and acts as the feature-off state: the composer pill is hidden entirely while Off, and reselecting the active level toggles back to Off. * Studio permissions: higher contrast composer pill text The permission pill uses a foreground based grey instead of the shared muted pill color, so it reads darker in light mode and lighter in dark mode. Full access keeps the danger yellow. * Studio permissions: panel dropdown layout and shorter tooltip Chat settings panel: the Bypass permissions label sits on one line with a full width dropdown underneath, styled like the other panel selects. Tooltip shortened and wording uses Unsloth instead of Studio. * Studio permissions: harden auto-mode unsafe detection Extend the Approve for me classifier to catch write and exec paths that slipped through: - terminal: sort -o, tree -o, xxd -r, find -exec/-execdir/-ok/-delete and find -fprint/-fprintf/-fls now ask; plain read-only forms still auto-run. awk is no longer allowlisted since its program can write and call system(). - python: from-imports of mutating names (from os import remove [as rm]) and star imports now ask. Found by a fuzz and edge-case simulation matrix; pinned in test_permission_mode.py. * Studio permissions: split multi-line terminal commands in auto detection A shell runs each line as its own command, but shlex reads newlines as whitespace, so "ls\nrm -rf x" demoted rm to argument position and auto-ran. Normalize newlines and CR to separators, and treat any all separator token as a command boundary so runs of blank lines still split. Found by the simulation matrix; pinned in tests. * Studio permissions: address review feedback on auto-mode detection Auto-mode (Approve for me) safety classifier hardening: - Python: flag any reference to a mutating attribute, not only direct calls, so indirect refs (f = os.remove; f(x)) and aliases ask. Detect Path.open(mode) write modes and wrap the AST walk to fail closed. - Terminal: match attached short output flags (sort -o/tmp/out) and keep find context across grouping parens so find ( -delete ) asks. - Both: ask before reads that escape the sandbox workdir via parent traversal or hit credential paths (.ssh, .aws, id_rsa, .pem, etc.). permission_mode plumbing: - Fold permission_mode=full into bypass_permissions at the request model so route-level confirm-gate guards see it as bypass. - Reject ask/auto on the Anthropic Messages server-tools path, which has no confirmation channel (mirrors the confirm_tool_calls rejection). - Keep forced RAG autoinject in auto mode: the safe search_knowledge_base retrieval never gates, so derive the skip from the real confirm need. - Reset all local preferences now also clears the legacy confirm key so a reset restores the fresh default instead of the old level. Regression tests added for each case. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio permissions: close auto-mode classifier gaps from review round 2 Auto mode ("Approve for me") let a few mutating calls through as safe: - os.open(...) always creates/writes a descriptor, so treat it as unsafe even though builtin open in read mode stays safe. - fd -x/--exec/-X/--exec-batch runs a command per match; scan for these alongside find's -exec/-delete. - tempfile writes artefacts and hands back writable handles, so importing it now asks. - Calling the result of a call (getattr(os, "remove")("x"), partials) is a dynamic target the AST can't vet, so fail closed. - An MCP tool whose name pairs a read verb with a mutating one (get_or_create_issue, read_and_delete_file) no longer auto-runs on the read prefix alone. Also fold permission_mode="off" into confirm_tool_calls=False on both request models so the non-stream route guard sees the disabled gate, and drive the Confirm tool calls toggle off permission_mode="ask" so auto no longer shows it on. * Harden auto-mode classifier and normalize bypass to full for PR #7079 Approve for me now asks for a few cases it previously auto-ran: - os.open via an os alias (import os as o; o.open(path, O_CREAT)) - pathlib symlink_to / hardlink_to / link_to - importlib.import_module dynamic imports - os.mkfifo / os.mknod / os.utime Also fold bypass_permissions into full when a stale ask/auto permission_mode is sent alongside it, so the Anthropic route guard no longer 400s those legacy callers. Adds classifier and request-model regression tests. * Close more auto-mode classifier gaps for PR #7079 Approve for me now asks for cases the review surfaced: - builtin open aliased to a name (f = open; from builtins import open as w) or looked up dynamically (globals()['open']) - pickle / marshal / shelve / dill deserialization - io.FileIO write handles - sort --compress-program (runs an external program) - MCP names carrying save/archive/submit/commit/push/sync/register verbs Also refine the attribute open() write check so an explicit read mode (ZipFile.open(name, "r")) stays auto while os.open flags still ask. Adds test coverage for each case. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Close three more auto-mode gaps for PR #7079 - rg runs an arbitrary program per file via --pre / --hostname-bin, so "Approve for me" now asks for those flags (rg is on the read-only allowlist). - A path-qualified command token (./ls, /tmp/cat) is an arbitrary executable, not the trusted utility its basename matches, so it asks before running. - A direct /chat/completions caller that sets permission_mode ask/auto but omits the legacy confirm_tool_calls flag now self-enables the confirmation gate, so tools can no longer run ungated on that path. Adds classifier and request-model tests for each case. * Close auto-mode classifier gaps from review round 3 for PR #7079 Approve for me now asks for cases the latest pass surfaced: - short-option clusters bundling a write flag (sort -uo out => -u -o) - procfs reads that leak a process env/args/memory (cat /proc/self/environ, /proc/PID/cmdline, maps) - env-assignment prefixes that change command lookup/loading (LD_PRELOAD=x ls, PATH=. ls, IFS=x ls); benign FOO=1 cmd stays auto - os.open imported as a bare callable (from os import open as o) Also drops ps from the safe terminal allowlist: its BSD environment flags (ps auxe, ps eww) dump a parent process's unscrubbed env and cannot be flag-parsed reliably, so ps always asks now. Adds classifier tests for each case. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Close auto-mode classifier gaps from review round 4 for PR #7079 Terminal (Approve for me now asks for these): - cd dropped from the safe allowlist: cd /; cat etc/passwd moves the shell out of the session workdir so a later relative read escapes it - env -C/--chdir (workdir escape) and -S/--split-string (builds a fresh command line); wrapper flags are now checked - /etc//passwd and /etc/./passwd normalize to /etc/passwd before the sensitive-path scan - a sensitive path split across an assignment and an argument (p=/etc; cat $p/passwd) via best-effort NAME=value expansion Python: - builtins.exec / builtins.eval attribute calls (dynamic code execution) - destructured open aliases (f, _ = (open, print); f('out', 'w')) - a sensitive path composed from literals (os.path.join('/etc','passwd'), '/etc' + '/passwd') - ZipFile/TarFile write modes (ZipFile(name, 'w')); the reader stays auto Adds classifier tests for each case. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Close auto-mode classifier gaps from review round 5 for PR #7079 Terminal (Approve for me now asks for these): - procfs reads hidden by shell quotes (cat /proc/$PPID/enviro''n) or quoted/nested-variable assignments (p="/proc/$PPID"; cat $p/environ): quotes are stripped and NAME=value prefixes expanded before the scan - LESSOPEN/LESSCLOSE, which make less run an input preprocessor command Python: - os.chdir / os.fchdir, which move the cwd so a later relative read escapes the sandbox workdir - sensitive paths composed via a pathlib / chain (Path('/etc') / 'passwd') or an f-string of literals (f'/proc/{pid}/environ') - runpy (import) and runpy.run_path / run_module, which run arbitrary code Adds classifier tests for each case. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Close auto-mode classifier gaps from review round 6 for PR #7079 Approve for me now asks for these: - a mutating callable reached through a getattr alias (rm = getattr(os, "remove"); rm("f")): calls through a getattr-bound name fail closed - compound MCP tool names carrying clone/checkout/comment/fork/tag/ invite/share, which start with a read verb but still mutate - a sensitive path hidden behind a glob (cat /e??/passwd, cat /e[t]c/passwd): a ? / * / [..] token is matched against the sensitive-file set and bracket classes are de-obfuscated; benign globs (ls *.py) stay auto Also run first-pass RAG retrieval in off mode: like auto, off never prompts, so a direct caller passing a stale confirm flag should not lose document retrieval (both tool loops). Adds classifier tests for each case. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Close auto-mode classifier gaps from review round 7 for PR #7079 Approve for me now asks for these: - __builtins__.exec / __builtins__.eval (dynamic code via the dunder) - terminal reads that hide a credential path behind a backslash escape (cat /et\c/passwd) - read-named MCP filesystem calls pointed at a credential path (mcp__fs__read_file {"path": "/etc/passwd"}) - compound MCP names carrying append / prepend - open aliased through a subscript or builtins attribute (f = globals()["open"]; f = builtins.open) then called to write - open(..., **{"mode": "w"}) where a kwargs splat hides the write mode - a sensitive path with a dynamic segment (open(f"/etc/{name}"), os.path.join("/etc", name)); /tmp/{name} stays auto - urllib3 networking Also stop folding permission_mode ask/auto into confirm_tool_calls for external-provider requests: that branch rejects confirm_tool_calls with tools, and the mode only governs local tool calls. Local requests still self-gate. Adds tests for each case. * Close auto-mode classifier gaps from review round 8 for PR #7079 Approve for me now asks for these: - dbm on the unsafe-module list: dbm.open(file, "c"/"n") creates files, and importing the family signals a persistence writer - reads of ~/.azure and ~/.config/gh credential stores (Azure/GitHub tokens), in terminal, MCP arguments, and Python literals - compound MCP names carrying upsert / assign Adds classifier tests for each case. * Gate secret mounts and fix the composer pill count for PR #7079 - Add Docker/Kubernetes secret mount dirs (/run/secrets, /var/run/secrets) to the sensitive-path checks, so Approve for me asks before reading injected credentials (terminal, MCP args, Python). - Count the always-visible permission pill in the composer's compact threshold so labels collapse at the intended width instead of overflowing by one pill. Adds classifier tests for the secret mount paths. * Close auto-mode classifier gaps from review round 10 for PR #7079 Approve for me now asks for these: - qualified pathlib constructors (pathlib.Path('/etc') / name), folded the same as bare Path(...), so a dynamic sensitive path is detected - open aliased through an annotated assignment (f: object = open; f('out', 'w')), tracked like a plain assignment - recursive searches rooted at an absolute path (grep -R TOKEN /home, rg TOKEN /, fd pattern /etc), which read host files outside the sandbox tree; sandbox-relative searches stay auto Adds classifier tests for each case. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Close auto-mode classifier gaps from review round 11 for PR #7079 Approve for me now asks for these terminal reads, which bash would expand into a sensitive path only after the classifier had approved: - a glob that resolves into a secret mount or credential dir (cat /r?n/secrets/hf_token, cat /root/.s??/id_rsa) - a recursive search rooted at a tilde home (grep -R TOKEN ~root, grep -R TOKEN ~/logs) - a brace expansion that builds a credential path (cat /etc/pass{w,}d) - a default/alternate parameter expansion that builds one (cat /etc/pass${x:-wd}) - an input redirection that hides a glob (cat still fails closed. The path folder previously handled only tuple/scalar % right-hand sides and returned None for a dict, hiding the sensitive segment. - A read-named MCP database tool carrying PostgreSQL COPY. COPY ... FROM bulk-loads a table and COPY ... TO writes a server-side file, so both are matched as mutating queries like DELETE/UPDATE already were. A 'copy' substring in a column name stays safe (word boundary). - logging file handlers. logging.FileHandler('out.log', mode='w') (and the default append mode, RotatingFileHandler/TimedRotatingFileHandler/ WatchedFileHandler, and the bare from-import form) create or truncate a file like open(..., 'w'), so they are classified as writer calls. StreamHandler / NullHandler and logging reads stay safe. Adds regression rows for each gap and its safe counterpart. * Fix writer aliases, GraphQL mutations, and auto server tools (review round 29) - Auto-mode Python: an aliased writer or archive constructor is tracked like the existing open alias, so from numpy import save; s = save; s('out.npy', arr) (and z = ZipFile; z('a.zip', 'w'), incl. the destructured forms) ask instead of running the write unprompted. A benign builtin alias (x = len) stays safe. - Auto-mode MCP: a read-named tool carrying a GraphQL mutation now asks. query_graphql {"query": "mutation { deleteIssue(id: 1) }"} matches a leading mutation keyword (GraphQL uses # comments, so it scans the raw payload); GraphQL read queries stay safe. - Anthropic /v1/messages: permission_mode "auto" no longer 400s a safe-only server-tool selection. auto only needs a confirmation channel for an unsafe call, so like the omitted default it runs for web_search / RAG / render and rejects only when a gate-needing local terminal/python tool is selected. ask still always rejects (it asks per call, which this passthrough cannot honor). The rejection stays ahead of the model auto-switch. Adds regression rows/cases for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Gate asyncio spawn, net clients, default-captured open; allow safe-only auto (round 30) Auto-mode Python now asks for more process/network/write vectors: - asyncio process spawners (asyncio.create_subprocess_exec/shell and a loop's subprocess_exec/shell) run an arbitrary program without the terminal blocklist, so they gate like os.system/subprocess. - stdlib network clients imaplib / poplib / nntplib / xmlrpc(.client) / webbrowser open outbound connections the sandbox does not namespace off, so their import asks like the other network modules. - a callable captured as a function or lambda parameter default (def f(o=open): o('out', 'w')) now binds that parameter into the same alias set, so the later write through it is gated. A benign default (o=len) stays safe. Also, permission_mode "auto" no longer 400s a non-streaming local tool request whose selection is always-safe-only (web_search / RAG / render). auto only prompts for a classifier-flagged call, so a safe-only auto request needs no stream, while ask, an explicit confirm_tool_calls=true, MCP, and an unrestricted or unsafe selection still require it. Applied via a shared _confirm_gate_needs_stream helper at the pre-switch, GGUF, and safetensors confirm-stream guards; the loop's per-call confirm flag is unchanged. Adds regression rows/cases for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Catch brace-glob paths and attribute writer aliases; unfold auto (round 31) - Terminal auto mode now runs the glob-sensitive scan over every expansion candidate, so a brace-expanded glob (cat /e{t,}c/pass?d, which bash expands to /etc/pass?d and then globs to /etc/passwd) asks. Brace expansion alone spells no literal /etc/passwd and the glob only resolves once the brace group is expanded, so scanning both together is required. A benign brace + glob stays safe. - Python auto mode now tracks a mutating attribute captured as a plain name: s = np.save; s('out.npy', arr) binds a writer alias, a captured .open bound method (p = Path('out').open; p('w')) fails closed on any call since its mode position varies, and z = zipfile.ZipFile is gated like the bare import. A benign attribute alias (x = np.mean) stays safe. - permission_mode "auto" is no longer folded to confirm_tool_calls=true on the request model. Folding it defeated the safe-only-selection exception in _confirm_gate_needs_stream (an explicit confirm forces stream=true), so a non-streaming safe-only auto request was rejected. Leaving it unset lets the route apply the exception; the mode still drives the loop's per-call gate. "ask" still folds (it gates every call). Adds regression rows/cases for each. * Harden SQL/GraphQL/writer classification and passthrough guards (round 32) MCP argument mutation detection (read-named query tools): - CREATE DDL now matches modifiers and the broader object set, so CREATE OR REPLACE VIEW, CREATE UNIQUE INDEX, CREATE TEMP TABLE, CREATE MATERIALIZED VIEW and CREATE FUNCTION ask. - Stored-procedure invocation (CALL proc(...), EXEC/EXECUTE) and VACUUM ask; a natural-language "call me back" stays safe via the trailing "(" / ";" / end lookahead. - GraphQL # comments are stripped before the mutation match, so mutation # note\n { deleteIssue(id: 1) } no longer hides the mutation. Python auto-mode classification: - numpy.memmap / open_memmap and pandas ExcelWriter / HDFStore create or truncate a file on construction, so they gate like open(..., "w"). - asyncio networking (asyncio.open_connection, loop.create_connection / create_server and unix variants) opens outbound connections/listeners the sandbox does not isolate, so it gates like socket.connect. Terminal auto-mode: file -C / --compile writes a compiled magic database. Routing: - A JSON-schema response_format is guided-decoding passthrough, not a local tool loop, so a --enable-tools policy no longer 400s a non-streaming ask/auto structured-output request at the confirm guard. - An explicit confirm_tool_calls=False opts out of the Anthropic Messages server-tool gate entirely (it wins over the mode, mirroring _permission_mode_confirm and the GGUF path), so it runs even under ask. Adds regression rows/cases for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Track path-ctor aliases, exempt empty selection and safe safetensors card (round 33) - Python auto mode now propagates path constructor / join aliases, so assigning Path or os.path.join to another local name is still folded: P = Path; (P('/etc') / 'passwd').read_text() and j = os.path.join; open(j('/etc', 'passwd')) ask, while a benign /tmp alias stays safe. - _confirm_gate_needs_stream now distinguishes an omitted enabled_tools (None, all tools) from an explicit empty list ([], no tools). An empty selection runs no built-in tool and cannot prompt, so a non-streaming auto request with enable_tools=true, enabled_tools=[] is no longer 400ed under a --enable-tools policy. - The safetensors provisional render_html card now uses permission_mode: render_html is always safe and never prompts, so its early canvas card streams under auto (which ships confirm_tool_calls=true) instead of being suppressed, matching the GGUF path's is_always_safe_tool exemption. Adds regression rows/cases for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Extend auto-mode classifier: SQLite mutations, more net/xattr/compressed writers Additional fail-closed gaps found by a fresh adversarial pass, each with a reproduction and a benign control: - MCP read-named tools now ask on SQLite-flavored writes the base DML/DDL regex missed: ATTACH / DETACH DATABASE, a write-form PRAGMA (PRAGMA journal_mode=WAL / user_version=42 / foreign_keys(0), while the read-form PRAGMA journal_mode stays safe), and load_extension() which loads and runs an arbitrary shared library. - Python auto mode now gates the remaining asyncio network entry points (start_server, open_unix_connection, loop.create_datagram_endpoint, sock_connect), os.setxattr / os.removexattr metadata writes, the gzip / bz2 / lzma single-stream writers (GzipFile / BZ2File / LZMAFile, mode-gated like ZipFile so a read stays safe), pandas to_xml, and the websockets client. Benign controls (SELECT 1, read-form PRAGMA, asyncio.sleep, gzip read, numpy read, natural-language "attach"/"analyze") stay safe. Regression rows added to test_permission_mode.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Close follow-up auto-mode gaps: SQLite/GraphQL variants, more writers and net A fresh adversarial pass on the previous round found consistent extensions of the same fail-closed rules, each reproduced with a benign control: - MCP read-named tools: DROP / ALTER now cover the same broad object set as CREATE (DROP FUNCTION, ALTER INDEX, DROP MATERIALIZED VIEW); ATTACH is caught without the optional DATABASE keyword via its quoted-path form; a schema-qualified write PRAGMA (PRAGMA main.user_version=1) is matched; and a GraphQL mutation carrying directives (mutation M @audit { ... }) is treated as a mutation. - Python auto mode: os.startfile (Windows program launch), asyncio start_unix_server, and the socketserver framework now ask; a gzip/bz2/lzma open imported under an alias (from gzip import open as gopen) is gated like builtin open; and a dynamic path prefix that can form a sensitive absolute root (open(chr(47) + "etc/passwd"), open(os.sep + "etc/passwd")) is treated as sensitive, while a dynamic prefix with a benign suffix stays safe. Benign controls (read-form PRAGMA, natural-language "attach ... as", "drop the idea", SELECT dropped_at, query @cached, gzip read alias, dynamic prefix + data/file suffix) stay safe. Regression rows added to test_permission_mode.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Gate GNU time -o, basicConfig/methodcaller/fileinput, and more SQL mutations Another adversarial pass surfaced further consistent fail-closed gaps, each reproduced with a benign control: - Terminal: GNU time -o/--output/-a/--append truncate or append to a file with timing output; time is a wrapper, so the flag is checked before the wrapped command like env -C. - Python auto mode: logging.basicConfig(filename=...) opens a log file for write; operator.methodcaller("write_text"/...) hides a writer method behind a string and is now treated as dynamic dispatch (like getattr/partial); fileinput.input(..., inplace=True) rewrites a file in place (the default read form stays safe). - MCP read-named tools: UPDATE now matches quoted, bracketed, and schema-qualified targets (UPDATE "users" / public.users / ONLY public.users / [users] / `users` SET); SELECT ... INTO OUTFILE/DUMPFILE writes a server file; and state-changing SQL functions inside a SELECT (pg_terminate_backend, setval, pg_write_file, lo_export, ...) ask. Benign controls (time ls / time -p, basicConfig(level=), methodcaller("upper"), fileinput read, NL "update ... set", setval_col column, PL/pgSQL SELECT INTO var) stay safe. Regression rows added to test_permission_mode.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten auto-mode classifier comments Collapse the multi-line rationale blocks in the permission classifier to one or two lines each without dropping the exploit each branch closes. Comments and whitespace only (no code change); the classifier tests are unchanged and pass. * Retry transient SSE stalls in the tool-calling smoke probes The tool-calling job flaked with a bare "TimeoutError: timed out": the server-side python/bash probes stream over post_sse(), which (unlike post()) had no transport-level retry, so a single stalled stream on a shared CI runner hard-failed the whole step even though function calling had already passed. post_sse() now mirrors post(): a transport-level stall (stream open or a mid-stream read timing out) is retried once with a fresh request capped at 300s, while HTTP status errors still surface immediately. The Linux _run_tool_probe caps each attempt at 360s and treats a stall that outlives the retry as a failed attempt (rotate to the next seed) instead of raising, and the web_search probe uses the same 360s cap. A genuine server wedge still fails (the retry also times out), so real regressions are not masked. Applied to the Linux, macOS, and Windows inference-smoke workflows, which share the probe. * Close five more auto-mode classifier gaps from review Each reproduces with a benign control: - Path constructor aliased through an attribute (P = pathlib.Path) now folds like the bare-name alias, so (P('/etc') / 'passwd').read_text() asks while a /tmp alias stays safe. - Callable defaults that are not plain names now bind the parameter: an attribute writer (def f(s=np.save)), an archive constructor, a captured .open, and partial(open, mode='w') fold like the equivalent assignment; a benign default (np.mean) does not. - A dynamic piece inside a sensitive name (open('/et' + chr(99) + '/passwd'), which folds to '/et\x00/passwd') now asks: the literals around each dynamic segment are matched against a credential target with the segment as any run of non-separator chars, so an all-dynamic ('1 + 1') or segment-spanning (a + '/' + b) path stays safe. - MCP read-named tools now ask on REFRESH MATERIALIZED VIEW and REINDEX; a 'refresh' column or natural-language 'refresh' stays safe. - A writer/open alias handed to a higher-order invoker (map(open, names, modes), starmap(np.save, ...)) is gated even without a direct call site; a benign map(len, ...) is unaffected. Regression rows added to test_permission_mode.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Default tool pills off on model load so tool execution is opt-in resolveToolsEnabledOnLoad turned the web-search and code pills on for any tool-capable model when the user had expressed no preference. Default them off instead, so tool execution is enabled only when the person clicks the pill to turn it on; a saved preference (on or off) is still honoured, so a user who already enabled tools keeps them on. * Gate mark/subscribe MCP verbs and qualified higher-order writer invokers - A read-prefixed MCP tool name carrying mark / subscribe / unsubscribe (get_and_mark_read, get_and_subscribe) now asks; a 'mark' substring inside one token (list_bookmarks) stays safe. - The higher-order writer check now also fires for a qualified invoker (itertools.starmap(open, ...), functools.reduce(open, ...)), matching the bare-name map/filter form; the writer-check on the first arg keeps a benign itertools.starmap(len, ...) or itertools.chain(...) safe. Regression rows added to test_permission_mode.py. * Close more auto-mode gaps and align the ask confirm fold across paths Each classifier change reproduces with a benign control: - MCP read-named tools now ask on reply / notify verbs (get_and_reply_email, list_and_notify_users), on catalog writes COMMENT ON / SECURITY LABEL / LOCK TABLE and CREATE|DROP|ALTER POLICY, and on state-changing PostgreSQL functions inside a read-shaped SELECT (nextval, set_config, pg_notify, the advisory-lock family). A 'comment' column, a 'locks' table, and a 'nextval' column prefix stay safe; the natural-language NOTIFY/SET ROLE statement forms are left out because SET/NOTIFY overlap ordinary prose. - Python auto mode now gates loader.exec_module (runs a module's code), archive extractall (zip-slip file writes), the ensurepip / venv modules (install pip / build an environment), and pydoc.writedoc. The Hugging Face login token (~/.cache/huggingface/token and stored_tokens) is now a sensitive path, while the rest of that cache (model data) stays readable. - ChatCompletionRequest no longer overwrites an explicit confirm_tool_calls=false when permission_mode='ask': the fold only self-enables the gate when the flag is unset, so an explicit opt-out wins on the chat path exactly as it already does via _permission_mode_confirm and the Anthropic pre-switch guard. Regression rows added to test_permission_mode.py. * Gate sort -T, xxd outfile positional, and the legacy HF token path - sort -T / --temporary-directory writes spill files to a caller-chosen dir, so it joins -o / --output in sort's unsafe-flag set. - xxd [infile [outfile]] writes its second positional, like uniq; xxd now uses the same second-positional-write handling (xxd in.bin out.hex asks, xxd in.bin and xxd -c 16 in.bin stay read-only). - The sensitive-path regex now also covers the legacy ~/.huggingface/token location (optional leading dot), not just ~/.cache/huggingface/token; an unrelated dir like myhuggingface/token stays safe. Regression rows added to test_permission_mode.py. * Catch multi-char SQL mutation targets, globbed credential names, digit outfiles Three fail-open gaps in the auto-mode classifier, each with a benign control: - SQL: the trailing word boundary on the MCP mutation regex meant a bare \w stopped at the first character, so TRUNCATE users, GRANT SELECT ON t, and REVOKE ALL ON t (multi-character names) slipped through while single-letter targets matched. Match the whole identifier instead, and accept an explicit AS alias on UPDATE (UPDATE users AS u SET). The implicit-alias form is left out because it is indistinguishable from the prose "update set". A truncate_log column and a grants table stay safe. - A glob that resolves to a credential basename anywhere (cat ~/.huggingface/tok?n -> token, cat proj/.netr? -> .netrc, cat repo/.aws/cred*) now asks; the fixed target list only covered a handful of home paths. notes/dra?t.txt and token_counts.tx? stay safe. - uniq / xxd counted file positionals but skipped every numeric token to ignore a flag value, so a file literally named with digits (uniq 123 out) hid the output positional. Track each command's value-taking flags and consume only the value, so uniq -f 2 in stays safe while uniq 123 out asks. Regression rows added to test_permission_mode.py. * Isolate the permission-mode loop tests from process-global state The loop-driving tests (auto/off/full/bypass) drove run_safetensors_tool_loop against a process-global approval registry (state.tool_approvals._pending) keyed by a single shared session id, and read os.environ. Other backend test modules mutate both, some at import time, so in the full-suite ordering a stale pending approval or a leaked env var could make the loop deny or skip a call these tests expect to run. It passed when the file ran alone but failed only in the complete tests/ run on CI. Add an autouse fixture that snapshots and restores os.environ and the approval registry around each test, and give every _drive call a unique session id so a leaked approval can never collide. Attach a compact event-stream dump to the loop assertions so any residual full-suite-only failure reports what the loop actually did instead of a bare diff. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: harden auto-mode classifier for recursive listers, sort file lists, aliased invokers, single-member extract Close four fail-open gaps in is_potentially_unsafe_tool_call: - terminal: tree/du (always recursive) and ls -R rooted at an absolute or tilde path now ask, matching the existing grep/rg/find recursive-read gate; relative walks stay safe. - terminal: sort --files0-from=F reads the file list named in F, so it can read arbitrary host files indirectly; added to sort's unsafe flags. - python: track aliases of the higher-order invokers (m = map; from itertools import starmap as sm) so an aliased invoker handed open/a writer is still gated; a benign callable (map(len, ...)) stays safe. - python: single-member archive extract (ZipFile/TarFile.extract) writes to disk like extractall and is vulnerable to a crafted member path, so gate it. Also update the stale _FakeExecuteTool in test_permission_mode.py to accept the thread_id keyword that run_safetensors_tool_loop now forwards to execute_tool after the main merge, which had broken the five tool-loop tests. Adds regression rows covering each gap plus benign controls. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: normalize unknown permission_mode to 'ask' instead of a 422 The request models validated permission_mode with Literal[ask, auto, off, full], so an unrecognized value from a newer UI/client was rejected with a 422 before the tool loops could apply their unknown -> ask fallback (safetensors_agentic.py:464, llama_cpp.py:9001). That made the intended forward-compat degradation unreachable at the API boundary for both Chat Completions and the analogous Anthropic field. Accept a plain string on both ChatCompletionRequest and AnthropicMessagesRequest and normalize in a before-validator: None stays unset, the four known modes pass through, and any other value degrades to the safest gate ('ask'), matching the loops. Adds a regression test covering unknown/None/known across both models. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: close five more auto-mode classifier gaps - terminal: xargs is no longer a safe wrapper. It appends arguments read from stdin that the scan never sees, so `echo -o out /etc/passwd | xargs sort` forwards to `sort -o out /etc/passwd` (a write + sensitive read) while only the allow-listed literals are visible. Any xargs command now asks. - terminal: ionice -p/-P/-u change the I/O priority of an already running process / group / user instead of forwarding to a wrapped read-only command, so `ionice -c 3 -p ` now asks. ionice -c 3 stays safe. - MCP: gate ALTER SYSTEM, which persists PostgreSQL server configuration and was not one of the DDL objects the mutation detector matched. - MCP: a credential noun in a read-named tool (read_secret, list_tokens, get_credentials, fetch_api_key) is a sensitive disclosure, so it asks even without a mutating verb or a path/SQL argument. Scoped *_key nouns keep a primary_key / keyboard lookup safe. - render_html: no longer unconditionally safe. A static canvas still auto-runs, but one whose HTML/JS reaches the network (fetch/WebSocket/remote script) asks, since it can egress under the canvas CSP when artifact network access is on. Its early provisional card is suppressed under the auto confirm gate, and the confirm-without-stream guard now requires a stream when render_html is selectable. Adds regression rows and benign controls for each, and updates the render_html provisional-card and confirm-gate tests to the new behavior. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: extend auto-mode gates for indirect file lists, dynamic lookups, HTML network loads, and Anthropic render_html Follow-ups on the previous classifier round: - terminal: wc/du/find --files0-from (and find's -files0-from primary) read a NUL-separated list of input paths from a file, the same indirect mechanism as sort --files0-from, so a crafted list reads arbitrary host files past the literal path/root checks. Gate them like sort. - python: a namespace lookup through a dict-style call (f = __builtins__.__dict__.get('open'), globals().get('open'), vars(x).get(...)) can return open/eval/a mutator, so poison the bound name like getattr/subscript lookups already are. An ordinary dict .get or os.environ.get stays safe. - render_html: broaden the network detector so a canvas that loads a resource via CSS url()/@import, srcset, or a root-relative (/path) or protocol-relative (//host) src/href is treated as networked, not just fetch/WebSocket/remote script. Relative ./x and url(#id)/data: refs stay static/safe. - Anthropic /v1/messages: drop render_html from the unprompted-safe server-tool set. Since it can prompt (networked canvas) and this channel invokes the loop without confirm, selecting it under ask/auto/omitted now rejects like terminal/python; off/full (or an explicit confirm opt-out) run it. Adds regression rows and benign controls for each, plus an Anthropic route test. * Studio: close six more auto-mode classifier gaps - terminal: a glob that expands to a project .env (cat .e?v) now asks; .env joins the sensitive glob-basename set, matching the literal-path gate. - python: an open bound onto an attribute (box.f = open; box.f('out','w')) is tracked by attribute name, and open invoked via .__call__ (open.__call__('out','w'), unwrapped to the underlying callable) is gated, so neither slips past the name-based open-alias checks. Benign attribute callables and .__call__ on non-writers stay safe. - python: a namespace lookup via .get/.pop/.setdefault already covered the builtins case; unchanged here. - MCP: a mutating HTTP verb in a method/verb argument (get_url {"method": "DELETE"|"POST"|"PUT"|"PATCH"}) now asks, so a generic HTTP tool cannot mutate an external service unprompted; GET/HEAD stay safe. - MCP: a credential/secret environment-variable value (get_env {"name": "OPENAI_API_KEY"}) is treated as a sensitive read via the same credential-noun match used for tool names; PATH/HOME stay safe. - render_html: self-navigation sinks (location.assign/replace, window.open, assigning a URL to (window.)location(.href)) join the network detector, so a canvas that navigates itself to an external URL asks; location.reload() / history.back() stay static. Adds regression rows and benign controls for each. * Studio: gate obfuscated canvas egress, sensitive-dir iteration, and MCP metadata-host reads - render_html: strip block comments before the network scan so fetch/*x*/(...) cannot hide egress, and match bracket-access forms (window['fetch'](...), self['open'](...)). Line // comments are left alone so the // in an https URL is not eaten. A comment-only canvas stays static. - python: enumerating a directory outside the sandbox (Path('/etc').iterdir(), os.scandir('/etc'), os.listdir('/home'), os.walk('/')) reads host filenames the direct /etc/passwd checks would prompt for, so gate it when the target dir folds to an absolute/tilde/sensitive path; a relative dir stays safe and an unresolved dynamic dir is left to other checks. - MCP: a read-named HTTP tool pointed at a cloud-metadata / link-local host (fetch_url {"url": "http://169.254.169.254/..."}, metadata.google.internal) reads instance credentials, so classify those URL arguments as sensitive, mirroring the sandbox SSRF blocklist; ordinary and localhost URLs stay safe. Adds regression rows and benign controls for each. * Studio: gate meta-refresh navigation, pandas HTML/markdown exporters, absolute glob roots, and checksum verify mode * Studio: gate starred open writes, builtins.__import__, computed render_html sinks, and procfs fd reads in auto mode * Studio: gate remote worker canvases, huggingface_hub downloads, and write callables passed to user helpers in auto mode * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Unsloth Co-authored-by: Daniel Han --- studio/backend/core/inference/llama_cpp.py | 56 +- studio/backend/core/inference/orchestrator.py | 2 + .../core/inference/safetensors_agentic.py | 46 +- studio/backend/core/inference/tools.py | 2200 +++++++++++++++++ studio/backend/models/inference.py | 101 + studio/backend/routes/inference.py | 210 +- .../backend/tests/test_anthropic_messages.py | 124 +- .../backend/tests/test_llama_cpp_tool_loop.py | 33 + .../tests/test_openai_tool_passthrough.py | 154 ++ studio/backend/tests/test_permission_mode.py | 1594 ++++++++++++ .../tests/test_safetensors_tool_loop.py | 64 + .../src/components/assistant-ui/thread.tsx | 43 +- .../src/features/chat/api/chat-adapter.ts | 17 +- .../chat/bypass-permissions-menu-item.tsx | 78 +- .../src/features/chat/chat-settings-sheet.tsx | 92 +- studio/frontend/src/features/chat/index.ts | 1 + .../features/chat/permission-mode-select.tsx | 338 +++ .../src/features/chat/shared-composer.tsx | 39 +- .../chat/stores/chat-runtime-store.ts | 119 +- .../frontend/src/features/chat/types/api.ts | 9 + .../features/settings/tabs/general-tab.tsx | 15 +- studio/frontend/src/i18n/locales/en.ts | 6 + studio/frontend/src/index.css | 9 + 23 files changed, 5131 insertions(+), 219 deletions(-) create mode 100644 studio/backend/tests/test_permission_mode.py create mode 100644 studio/frontend/src/features/chat/permission-mode-select.tsx diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 093a92e38d..6b6c5373eb 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -8969,16 +8969,37 @@ class LlamaCppBackend: disable_parallel_tool_use: bool = False, confirm_tool_calls: bool = False, bypass_permissions: bool = False, + permission_mode: Optional[str] = None, ) -> Generator[dict, None, None]: """ Agentic loop: let the model call tools, execute them, and continue. + permission_mode: "ask" confirms every call (with confirm_tool_calls), + "auto" only pauses calls detected as potentially unsafe, "off" never + pauses (sandbox stays on), "full" is the same as bypass_permissions. + Unset/unknown behaves as "ask". + Yields dicts: {"type": "status", "text": "Searching: ..."/"Reading: ..."} -- tool status updates {"type": "content", "text": "token"} -- streamed content tokens (cumulative) {"type": "reasoning", "text": "token"} -- streamed reasoning tokens (cumulative) """ - from core.inference.tools import build_rag_autoinject, execute_tool + from core.inference.tools import ( + build_rag_autoinject, + execute_tool, + is_always_safe_tool, + is_potentially_unsafe_tool_call, + ) + + # Normalize the mode: "full" and bypass_permissions are the same + # switch, whichever arrives first wins toward the permissive side. + # "off" keeps the sandbox but never prompts. + if permission_mode == "full": + bypass_permissions = True + elif bypass_permissions: + permission_mode = "full" + elif permission_mode not in ("ask", "auto", "off"): + permission_mode = "ask" if not self.is_loaded: raise RuntimeError("llama-server is not loaded") @@ -8986,8 +9007,14 @@ class LlamaCppBackend: conversation = list(messages) # Forced first-pass RAG so a doc question doesn't lose to web_search. Emits - # the same tool card + citations a real call would. - _auto = None if confirm_tool_calls else build_rag_autoinject(conversation, rag_scope) + # the same tool card + citations a real call would. Skip it only when a + # retrieval call would actually prompt (ask mode); auto never gates the + # safe search_knowledge_base tool, so retrieval must still run there. + # off never prompts either, so it also keeps first-pass retrieval. + _skip_autoinject = ( + confirm_tool_calls and not bypass_permissions and permission_mode not in ("auto", "off") + ) + _auto = None if _skip_autoinject else build_rag_autoinject(conversation, rag_scope) if _auto: for _ev in _auto["events"]: yield _ev @@ -9357,8 +9384,16 @@ class LlamaCppBackend: in provisional_started_tool_calls.values() ) # Later parallel cards only reconcile when parallel use is enabled. + # In auto mode an always-safe tool (render_html) never + # prompts, so it must stream its early card too; mirror + # that here instead of gating on the raw confirm flag. _confirm_gated = ( - confirm_tool_calls and not bypass_permissions + confirm_tool_calls + and not bypass_permissions + and not ( + permission_mode == "auto" + and is_always_safe_tool(current_name) + ) ) # Keep small-argument tools on the normal path. _args_len = len( @@ -9925,7 +9960,18 @@ class LlamaCppBackend: # Bypass wins over the confirm gate at the loop level too, # so a direct internal caller with both flags never prompts. - needs_confirm = bool(confirm_tool_calls) and not bypass_permissions + # In "auto" mode only calls detected as potentially unsafe + # pause; read-only calls run straight through. "off" never + # prompts (sandbox stays on). + needs_confirm = ( + bool(confirm_tool_calls) + and not bypass_permissions + and permission_mode != "off" + ) + if needs_confirm and permission_mode == "auto": + needs_confirm = is_potentially_unsafe_tool_call( + decision.tool_name, decision.arguments + ) approval_id = new_approval_id() if needs_confirm else "" decision_slot = ( begin_tool_decision(session_id, approval_id) if needs_confirm else None diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 6d0b13ced9..c2082bc198 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -1372,6 +1372,7 @@ class InferenceOrchestrator: rag_scope: Optional[dict] = None, confirm_tool_calls: bool = False, bypass_permissions: bool = False, + permission_mode: Optional[str] = None, use_adapter: Optional[Union[bool, str]] = None, stats_holder: Optional[dict] = None, presence_penalty: float = 0.0, @@ -1439,6 +1440,7 @@ class InferenceOrchestrator: rag_scope = rag_scope, confirm_tool_calls = confirm_tool_calls, bypass_permissions = bypass_permissions, + permission_mode = permission_mode, ) def generate_with_adapter_control( diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index a18d2758ba..c1fffb71cb 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -428,6 +428,7 @@ def run_safetensors_tool_loop( rag_scope: Optional[dict] = None, confirm_tool_calls: bool = False, bypass_permissions: bool = False, + permission_mode: Optional[str] = None, ) -> Generator[dict, None, None]: """Drive an agentic tool loop on top of a cumulative-text generator. @@ -453,10 +454,27 @@ def run_safetensors_tool_loop( """ conversation = list(messages) - # Forced first-pass RAG (mirrors the GGUF loop) so doc Qs don't lose to web_search. + # Normalize the mode (mirrors the GGUF loop): "full" and + # bypass_permissions are the same switch; unset/unknown behaves as "ask". + # "off" keeps the sandbox but never prompts. + if permission_mode == "full": + bypass_permissions = True + elif bypass_permissions: + permission_mode = "full" + elif permission_mode not in ("ask", "auto", "off"): + permission_mode = "ask" + + # Forced first-pass RAG (mirrors the GGUF loop) so doc Qs don't lose to + # web_search. Skip only when a retrieval call would actually prompt (ask + # mode); auto never gates the safe search_knowledge_base tool. from core.inference.tools import build_rag_autoinject - _auto = None if confirm_tool_calls else build_rag_autoinject(conversation, rag_scope) + # off never prompts, so (like auto) it must not lose first-pass retrieval + # even if a direct caller passes a stale confirm_tool_calls flag. + _skip_autoinject = ( + confirm_tool_calls and not bypass_permissions and permission_mode not in ("auto", "off") + ) + _auto = None if _skip_autoinject else build_rag_autoinject(conversation, rag_scope) if _auto: for _ev in _auto["events"]: yield _ev @@ -539,7 +557,16 @@ def run_safetensors_tool_loop( # provisional card (keyed by tool_call_id, no approval) would show the # tool as "running" before the user has approved it. Suppress the early # card in that case and let the gated tool_start be the first signal. - _provisional_confirm_gated = bool(confirm_tool_calls) and not bypass_permissions + # In auto mode render_html is always safe and never prompts, so keep its + # early canvas card (the frontend sends confirm_tool_calls=true alongside + # auto); mirrors the GGUF path's _confirm_gated exemption. + from core.inference.tools import is_always_safe_tool + + _provisional_confirm_gated = ( + bool(confirm_tool_calls) + and not bypass_permissions + and not (permission_mode == "auto" and is_always_safe_tool("render_html")) + ) gen = _call_single_turn(single_turn, conversation, active_tools) prev_cumulative = "" @@ -1056,8 +1083,17 @@ def run_safetensors_tool_loop( assistant_msg.setdefault("tool_calls", []).append(decision.as_assistant_tool_call()) # Bypass wins over the confirm gate at the loop level too, so a - # direct internal caller passing both flags never prompts. - needs_confirm = bool(confirm_tool_calls) and not bypass_permissions + # direct internal caller passing both flags never prompts. In + # "auto" mode only calls detected as potentially unsafe pause. + # "off" never prompts (sandbox stays on). + needs_confirm = ( + bool(confirm_tool_calls) and not bypass_permissions and permission_mode != "off" + ) + if needs_confirm and permission_mode == "auto": + from core.inference.tools import is_potentially_unsafe_tool_call + needs_confirm = is_potentially_unsafe_tool_call( + decision.tool_name, decision.arguments + ) approval_id = new_approval_id() if needs_confirm else "" decision_slot = begin_tool_decision(session_id, approval_id) if needs_confirm else None start_event = decision.tool_start_event() diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 8d8e7ef3dd..84a67c6ad4 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -5,6 +5,7 @@ (DuckDuckGo), Python code execution, and terminal commands.""" import ast +import fnmatch import http.client import os import signal @@ -151,6 +152,42 @@ _COMMAND_PREFIXES = frozenset( } ) _ASSIGNMENT_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=") +# Env-assignment prefixes that change command lookup or code loading, so +# `LD_PRELOAD=x ls` / `PATH=. ls` run attacker code before the read-only +# utility. LD_*/DYLD_* and any *PATH are covered by the prefix/suffix check. +_AUTO_UNSAFE_ENV_ASSIGN = frozenset( + { + "IFS", + "BASH_ENV", + "ENV", + "SHELLOPTS", + "BASHOPTS", + "GLOBIGNORE", + "PROMPT_COMMAND", + "PS4", + "PYTHONSTARTUP", + "PYTHONHOME", + "NODE_OPTIONS", + "PERL5OPT", + "PERL5LIB", + "RUBYOPT", + "RUBYLIB", + # LESSOPEN/LESSCLOSE run an input preprocessor command for less. + "LESSOPEN", + "LESSCLOSE", + } +) + + +def _env_assignment_is_unsafe(name: str) -> bool: + """True if a NAME=value prefix affects command lookup/loading.""" + return ( + name in _AUTO_UNSAFE_ENV_ASSIGN + or name.startswith(("LD_", "DYLD_")) + or name.endswith("PATH") + ) + + _FIND_EXEC_FLAGS = frozenset({"-exec", "-execdir", "-ok", "-okdir"}) @@ -272,6 +309,2169 @@ def _find_blocked_commands(command: str) -> set[str]: return blocked +# ── "Approve for me" (permission_mode="auto") safety detection ────────────── +# Auto mode pauses only calls classified here as potentially unsafe. The sandbox +# and hard blocks (blocklist, rlimits) still apply at run time; this gate only +# decides prompting, and fails closed: anything not provably read-only asks. + +# Read-only commands allowed to run without confirmation in auto mode. +_AUTO_SAFE_TERMINAL_COMMANDS = frozenset( + { + "ls", + "dir", + "pwd", + # cd absent: `cd /; cat etc/passwd` escapes the workdir for a later + # relative read the path scan cannot see, so cd always asks. + "cat", + "head", + "tail", + # less/more absent: their pager escapes (+cmd, !shell, -o, LESSOPEN) can + # run a command or write a file, so they always ask. + "grep", + "egrep", + "fgrep", + "rg", + "find", + "fd", + "wc", + "sort", + "uniq", + "cut", + "tr", + "diff", + "cmp", + "file", + "stat", + "du", + "df", + # ps absent: BSD env flags (ps auxe, ps eww) dump a parent's unscrubbed + # env and can't be flag-parsed reliably, so ps always asks. + "date", + "cal", + "whoami", + "id", + "uname", + "hostname", + "uptime", + "which", + "whereis", + "type", + "basename", + "dirname", + "realpath", + "readlink", + "md5", + "md5sum", + "shasum", + "sha1sum", + "sha256sum", + "cksum", + "tree", + "printenv", + "echo", + "printf", + "true", + "false", + "test", + "[", + "seq", + "nl", + "od", + "xxd", + "hexdump", + "strings", + "column", + "paste", + "join", + "comm", + "expand", + "unexpand", + "fold", + "fmt", + "rev", + "tac", + "locale", + "arch", + "nproc", + "sw_vers", + "jq", + } +) +# Flags that turn an otherwise read-only command into a writer or executor +# (sort -o FILE, tree -o FILE, xxd -r IN OUT, find -exec/-delete/...). +_AUTO_UNSAFE_COMMAND_FLAGS = { + # --files0-from=F makes sort read the NUL-separated list of input files + # named in F, so a crafted list reads arbitrary host files indirectly. + "sort": frozenset( + {"-o", "--output", "--compress-program", "-T", "--temporary-directory", "--files0-from"} + ), + "tree": frozenset({"-o"}), + "xxd": frozenset({"-r"}), + # -c/--check makes a checksum tool read a manifest file and then read every + # path it names, so a manifest listing /etc/passwd turns `sha256sum -c list` + # into an indirect host-file read; the digest form (sha256sum file) only reads + # the named files. + "md5sum": frozenset({"-c", "--check"}), + "sha1sum": frozenset({"-c", "--check"}), + "sha256sum": frozenset({"-c", "--check"}), + "shasum": frozenset({"-c", "--check"}), + "cksum": frozenset({"-c", "--check"}), + # GNU time -o/--output/-a/--append FILE writes timing output; time is a + # wrapper, so the flag is checked before the wrapped command like env -C. + "time": frozenset({"-o", "--output", "-a", "--append"}), + # rg runs an arbitrary program per file with --pre/--hostname-bin. + "rg": frozenset({"--pre", "--hostname-bin"}), + # env -C/--chdir escapes the workdir; -S/--split-string builds a command. + "env": frozenset({"-C", "--chdir", "-S", "--split-string"}), + # ionice -p/-P/-u change the I/O priority of an already running process / + # group / user instead of forwarding to a wrapped read-only command, so a + # bare `ionice -c 3 -p ` mutates another process. ionice stays a safe + # wrapper for `ionice -c 3 `; only the process-target flags ask. + "ionice": frozenset({"-p", "-P", "-u"}), + # printf -v NAME assigns to a shell var, so `printf -v PATH %s .; ls` runs + # ./ls from the workdir. + "printf": frozenset({"-v"}), + # wc/du/find --files0-from=F read the NUL-separated list of input paths named + # in F, so a crafted list reads arbitrary host files past the literal path / + # root checks, like sort --files0-from. find spells it -files0-from (a primary). + "wc": frozenset({"--files0-from"}), + "du": frozenset({"--files0-from"}), + "find": frozenset( + { + "-exec", + "-execdir", + "-ok", + "-okdir", + "-delete", + "-fprint", + "-fprint0", + "-fprintf", + "-fls", + "-files0-from", + } + ), + # fd -x/--exec/-X/--exec-batch run a command per result; + # --base-directory/--search-path move the search root outside the workdir. + "fd": frozenset({"-x", "--exec", "-X", "--exec-batch", "--base-directory", "--search-path"}), + # date -s/--set writes the clock; display forms (+FORMAT, -d/-u/-R/-r) read. + "date": frozenset({"-s", "--set"}), + # file -C/--compile writes a compiled .mgc magic database; ident forms read. + "file": frozenset({"-C", "--compile"}), + # hostname -F/--file, -b/--boot set the hostname; display flags only read. + "hostname": frozenset({"-F", "--file", "-b", "--boot"}), +} +# Commands safe only without a mutating positional: `hostname NAME` sets the +# hostname, `date MMDDhhmm...` sets the clock (a +FORMAT token or a display +# flag's value stays read-only), so any other positional asks. +_AUTO_ARG_SENSITIVE_COMMANDS = frozenset({"hostname", "date"}) +# date display flags taking a value token (-d STRING, -r FILE, -f FILE); the +# value is not a clock-setting positional, so it is skipped. +_DATE_DISPLAY_VALUE_FLAGS = frozenset({"-d", "--date", "-r", "--reference", "-f", "--file"}) +# Commands that write their 2nd positional (uniq [INPUT [OUTPUT]], xxd [infile +# [outfile]]): the 1st file reads to stdout, but a second file positional +# overwrites it, like `sort -o`. +_AUTO_SECOND_POSITIONAL_WRITES = frozenset({"uniq", "xxd"}) +# Value-taking option flags for those commands whose argument is a separate token +# (uniq -f 2, xxd -c 16). The value must be consumed so a numeric option value is +# not miscounted as the output-file positional, and, conversely, a file that is +# literally named with digits (uniq 123 out) is still counted. +_SECOND_POSITIONAL_VALUE_FLAGS = { + "uniq": frozenset({"-f", "--skip-fields", "-s", "--skip-chars", "-w", "--check-chars"}), + "xxd": frozenset( + {"-c", "--cols", "-s", "--seek", "-l", "--len", "-g", "--groupsize", "-o", "--offset"} + ), +} +# find/fd group with (...) which resets command context, so scan every token for +# these once find/fd appears anywhere. +_AUTO_UNSAFE_FIND_LIKE_FLAGS = _AUTO_UNSAFE_COMMAND_FLAGS["find"] | _AUTO_UNSAFE_COMMAND_FLAGS["fd"] +# Recursive readers with an absolute-path target escape the workdir onto host +# files (grep -R TOKEN /home, rg TOKEN /), so they ask. +_AUTO_RECURSIVE_SEARCH = frozenset({"grep", "egrep", "fgrep", "rg", "ug", "find", "fd"}) +# Directory walkers that always recurse (tree /home, du /) read the whole host +# subtree under an absolute/tilde root, like a recursive search. ls only recurses +# with -R/--recursive, so it is gated separately when that flag is present. +_AUTO_RECURSIVE_LISTERS = frozenset({"tree", "du"}) +# Benign wrappers: safe AND forward command position to their target (checked in +# turn). sudo/su/chroot/etc. are absent, so they classify as unsafe. xargs is +# absent too: it appends arguments read from stdin that this scan never sees, so +# `echo -o out /etc/passwd | xargs sort` forwards to `sort -o out /etc/passwd` +# (a write + sensitive read) while only the allow-listed literals are visible. +_AUTO_SAFE_WRAPPERS = frozenset( + {"env", "command", "time", "timeout", "nice", "ionice", "stdbuf", "nohup"} +) + +# MCP tools whose names look read-only auto-run; anything else asks. +_AUTO_SAFE_MCP_TOOL_RE = re.compile( + r"^(get|list|search|read|fetch|query|find|describe|show|view|lookup|" + r"retrieve|count|status|info|help|check)(?:[_\-].*)?$", + re.IGNORECASE, +) +# A mutating verb anywhere in the name overrides a read-only prefix, so a +# compound name like get_or_create_issue or read_and_delete_file still asks. +_AUTO_UNSAFE_MCP_VERB_RE = re.compile( + r"(?:^|[_\-])(?:create|update|delete|remove|write|set|add|send|post|put|" + r"patch|insert|drop|kill|exec|execute|run|deploy|publish|move|rename|edit|" + r"modify|upload|replace|revoke|grant|approve|merge|close|cancel|pay|" + r"transfer|buy|sell|reset|clear|purge|destroy|terminate|revert|rollback|" + r"trigger|enable|disable|install|uninstall|restart|stop|start|" + r"save|archive|submit|commit|push|sync|register|" + r"clone|checkout|comment|fork|tag|invite|share|append|prepend|" + r"copy|duplicate|import|export|download|backup|restore|snapshot|mirror|" + r"upsert|assign|mark|subscribe|unsubscribe|reply|notify)(?:[_\-]|$)", + re.IGNORECASE, +) +# A read-named MCP tool that returns a secret is still a sensitive read, so a +# credential noun anywhere in the name (read_secret, list_tokens, +# get_credentials, fetch_api_key) asks even without a mutating verb or a path/SQL +# argument. Scoped nouns (api/access/private/... _key) avoid flagging benign +# keys like a primary_key or keyboard lookup. +_AUTO_SENSITIVE_MCP_NOUN_RE = re.compile( + r"(?:^|[_\-])(?:" + r"secret|token|credential|password|passwd|passphrase|apikey|" + r"(?:api|access|private|secret|signing|encryption|auth|session)[_\-]?keys?" + r")s?(?:[_\-]|$)", + re.IGNORECASE, +) + +# Python: modules whose import alone signals side effects auto mode should ask +# about (process spawning, network, bulk file ops, low-level memory). +_AUTO_UNSAFE_PY_MODULES = frozenset( + { + "subprocess", + "shutil", + "socket", + "ctypes", + "multiprocessing", + "pty", + "fcntl", + "requests", + "urllib", + "urllib3", + "http", + "httpx", + "aiohttp", + # huggingface_hub.hf_hub_download / snapshot_download fetch remote repo + # files over the network and write them to an on-disk cache. + "huggingface_hub", + # websockets opens a network connection; socketserver binds a listener. + "websockets", + "socketserver", + "ftplib", + "smtplib", + "telnetlib", + "paramiko", + # mail/news/rpc/browser stdlib clients open outbound connections + # (imaplib, poplib, xmlrpc.client, webbrowser.open). + "imaplib", + "poplib", + "nntplib", + "xmlrpc", + "webbrowser", + "tempfile", + # deserialization that can execute arbitrary code on load. + "pickle", + "marshal", + "shelve", + "dill", + # dbm.open(file, "c"/"n") creates files; treat the family as writers. + "dbm", + # sqlite3.connect(path) creates/mutates a database file (and runs DDL/DML + # without an open()/writer attribute), like dbm. + "sqlite3", + # runpy runs a script/module as code. + "runpy", + # ensurepip.bootstrap installs pip and venv.create builds an environment; + # both write to disk and can fetch/install packages. + "ensurepip", + "venv", + } +) +# Attribute calls that mutate the filesystem / spawn processes (os.remove, +# Path.write_text, sock.connect, ...) regardless of how the module was bound. +_AUTO_UNSAFE_PY_ATTRS = frozenset( + { + "remove", + "unlink", + "rmdir", + "removedirs", + "rename", + "renames", + "replace", + "rmtree", + "move", + "copy", + "copy2", + "copyfile", + "copytree", + "chmod", + "chown", + "system", + "popen", + "execv", + "execve", + "execl", + "execlp", + "execvp", + "spawnl", + "spawnv", + # os.startfile launches a program via its Windows association. + "startfile", + "fork", + "kill", + "killpg", + "symlink", + "link", + "mkdir", + "makedirs", + "truncate", + "touch", + "write_text", + "write_bytes", + "urlopen", + "urlretrieve", + "connect", + "bind", + "sendall", + # pathlib link creators, os node/metadata mutators, dynamic import. + "symlink_to", + "hardlink_to", + "link_to", + "mkfifo", + "mknod", + "utime", + # os.setxattr / os.removexattr mutate extended attributes, like chmod. + "setxattr", + "removexattr", + "import_module", + # loader.exec_module runs a module's code like import_module; archive + # extractall/extract write arbitrary files (zip-slip): extract takes a + # single member but an attacker-controlled member path still escapes. + "exec_module", + "extractall", + "extract", + "FileIO", + # asyncio subprocess spawners run a program past the terminal blocklist. + "create_subprocess_exec", + "create_subprocess_shell", + "subprocess_exec", + "subprocess_shell", + # asyncio outbound connections / listeners (open_connection, + # create_connection/server and unix variants), like socket.connect. + "open_connection", + "create_connection", + "create_server", + "create_unix_connection", + "create_unix_server", + # more asyncio listen/connect + UDP/raw socket helpers. + "start_server", + "start_unix_server", + "open_unix_connection", + "create_datagram_endpoint", + "sock_connect", + # os.chdir escapes the workdir; runpy helpers run arbitrary code. + "chdir", + "fchdir", + "run_path", + "run_module", + # types.FunctionType wraps a compiled code object into a callable, a + # dynamic-execution vector; pandas read_pickle deserializes (runs code). + "FunctionType", + "read_pickle", + } +) +# Pickle-backed loaders that can execute code embedded in the file; gated by +# receiver module (torch.load, joblib.load) since bare `load` is too common. +_AUTO_UNSAFE_PY_LOAD_MODULES = frozenset({"torch", "joblib", "cloudpickle"}) +# Writer methods that persist to disk without going through open() (numpy.save, +# Image.save, plt.savefig, DataFrame.to_csv, json.dump). Gated as method calls +# only, so a bare attribute reference is not mistaken for a write. +_AUTO_UNSAFE_PY_WRITE_METHODS = frozenset( + { + "save", + "savefig", + "savez", + "savez_compressed", + "savetxt", + "tofile", + "dump", + "to_csv", + "to_parquet", + "to_pickle", + "to_json", + "to_feather", + "to_hdf", + "to_excel", + "to_stata", + "to_sql", + "to_xml", + # pandas text exporters that write when given a path/buffer (to_html / + # to_markdown / to_latex mirror to_csv); to_clipboard / to_gbq persist + # off-process. to_string is omitted: it is overwhelmingly display-only. + "to_html", + "to_markdown", + "to_latex", + "to_clipboard", + "to_gbq", + "imwrite", + "imsave", + "write_image", + "write_html", + # ML persistence helpers (transformers/peft/safetensors/keras) that + # export adapters or weights to disk without an open()/writer attribute. + "save_pretrained", + "save_file", + "save_model", + "save_weights", + "save_lora", + "save_checkpoint", + # logging file handlers open a log file for write on construction (even + # default mode "a" creates); matched as attribute call and bare import. + "FileHandler", + "WatchedFileHandler", + "RotatingFileHandler", + "TimedRotatingFileHandler", + # numpy.memmap(..., mode="w+") and pandas writers create/truncate a file + # on construction, like open(..., "w"). + "memmap", + "open_memmap", + "ExcelWriter", + "HDFStore", + # pydoc.writedoc(name) writes name.html to the workdir. + "writedoc", + } +) +# Archive / compressed-file constructors taking the mode as their 2nd arg like +# open: ZipFile(name, "w") / gzip.GzipFile(name, "w") write, so gated only in +# write mode (reading a .gz is fine, so the modules are not blanket-unsafe). +_ARCHIVE_CTOR_NAMES = frozenset({"ZipFile", "TarFile", "GzipFile", "BZ2File", "LZMAFile"}) +# The stdlib module each archive constructor is imported from. +_ARCHIVE_CTOR_MODULES = { + "zipfile": "ZipFile", + "tarfile": "TarFile", + "gzip": "GzipFile", + "bz2": "BZ2File", + "lzma": "LZMAFile", +} +# Modules whose top-level open() takes the mode as its 2nd arg like builtin open, +# so `from gzip import open as gopen` binds an open alias gated on write mode. +_OPEN_ALIAS_MODULES = frozenset({"gzip", "bz2", "lzma"}) +# Builtins/itertools helpers that call their first argument once per item, so a +# writer/open alias handed to one runs without a direct call(...) site +# (list(map(open, names, modes)), starmap(np.save, ...)). filter's predicate is +# also invoked, so a writer smuggled there runs too. +_HIGHER_ORDER_INVOKERS = frozenset({"map", "filter", "starmap", "reduce"}) +_PY_WRITE_MODE_RE = re.compile(r"[wax+]") +# A file-mode literal ("w", "rb", "a+"): letters/flags only, no path chars. +# Used to tell a Path.open("w") mode from a ZipFile.open("name.txt") filename. +_PY_MODE_LITERAL_RE = re.compile(r"^[rwxa][btru+]*$") + +# Reading these off the host escapes the intent of "read-only is safe": they +# hold credentials. Path traversal (../) escapes the per-session workdir. +_SENSITIVE_PATH_RE = re.compile( + r"(?:^|[/\\])\.(?:ssh|aws|azure|gnupg|docker|kube|config/gcloud|config/gh)(?:[/\\]|$)" + r"|\.(?:netrc|npmrc|pypirc|git-credentials|env)(?:$|[/\\.\s'\"])" + r"|id_rsa|id_ed25519|id_ecdsa|id_dsa" + # Hugging Face stores the login token at ~/.cache/huggingface/token and the + # legacy ~/.huggingface/token (plus the multi-token store stored_tokens); the + # rest of that cache is model data, so only the credential files match. The + # optional leading dot covers the .huggingface dotdir form. + r"|(?:^|[/\\])\.?huggingface[/\\](?:token|stored_tokens)(?:$|[/\\.\s'\"])" + # /etc/ssh holds the host private keys (ssh_host_*_key); the whole dir is + # sensitive, not just passwd/shadow/sudoers. + r"|credentials|/etc/(?:passwd|shadow|sudoers|ssh(?:[/\\]|$))" + # Bash opens /dev/tcp/host/port and /dev/udp/host/port as network sockets, + # so a redirection to one reaches the network without the confirm prompt. + r"|/dev/(?:tcp|udp)/" + # Docker/Kubernetes secret mounts hold injected credentials. + r"|/(?:var/)?run/secrets(?:[/\\]|$)" + # procfs leaks a (possibly parent) process env/args/memory to a read, + # including the per-thread aliases under /proc//task//. The fd/ + # dir holds symlinks to a process's open files (a held credential/db file). + r"|/proc/[^/\s'\"]+/(?:task/[^/\s'\"]+/)?(?:environ|cmdline|mem|maps|fd)\b" + # A .pem/.key file (basename before the extension), not a bare ".key" + # (e.g. a jq '.key' filter). + r"|\w[\w.-]*\.(?:pem|key)(?:$|[\s'\"])", + re.IGNORECASE, +) +# A shell redirection with no following space (cat <../../notes) keeps `..` +# adjacent to `<`/`>`, so those count as leading delimiters here too. +_PARENT_TRAVERSAL_RE = re.compile(r"(?:^|[\s/\\'\"=:<>])\.\.(?:[/\\]|$|[\s'\"])") +# A sensitive directory: a dynamic segment under it (open(f"/etc/{name}")) is +# not provably safe, so fail closed when a folded path has a dynamic piece here. +_SENSITIVE_DIR_RE = re.compile( + r"/etc/|/(?:var/)?run/secrets[/\\]|(?:^|[/\\])\.(?:ssh|aws|azure|gnupg|docker|kube)[/\\]" + r"|(?:^|[/\\])\.config/(?:gcloud|gh)[/\\]", + re.IGNORECASE, +) +# Collapse /./ and repeated slashes so /etc/./passwd and /etc//passwd, which +# the OS resolves to /etc/passwd, still match the sensitive-path regex. +_REDUNDANT_SLASH_RE = re.compile(r"/\.?(?=/)") +# $name, ${name}, and operator/substring forms (${name:-x}, ${name:0:6}) all +# reference `name`; substituting the assigned value catches paths hidden behind +# a substring expansion (p=passwd; cat /etc/${p:0:6}). +_SHELL_VAR_RE = re.compile(r"\$\{(\w+)(?::[^{}]*)?\}|\$(\w+)") +# Pattern replacement (${p/X/w}, global ${p//X/w}) transforms the value before +# the path is used; apply it so p=passXd; cat /etc/${p/X/w} is scanned. +_SHELL_PARAM_REPL_RE = re.compile(r"\$\{(\w+)/(/)?([^/{}]*)/([^{}]*)\}") +# Case modification (${p^^} upper, ${p,,} lower, ${p^}/${p,} first char) also +# transforms the value, so p=PASSWD; cat /etc/${p,,} builds /etc/passwd. +_SHELL_PARAM_CASE_RE = re.compile(r"\$\{(\w+)(\^\^|,,|\^|,)\}") +# Indirect expansion ${!p} yields the value of the variable *named* by $p, so +# x=passwd; p=x; cat /etc/${!p} builds /etc/passwd. +_SHELL_PARAM_INDIRECT_RE = re.compile(r"\$\{!(\w+)\}") +_SHELL_ASSIGN_RE = re.compile(r"(?:^|[\s;&|(])([A-Za-z_]\w*)=([^\s;&|)]+)") +# Bash ANSI-C quoting ($'\x77' -> 'w') is expanded after this classifier, so +# decode $'...' bodies before the sensitive-path scan. +_ANSI_C_RE = re.compile(r"\$'((?:[^'\\]|\\.)*)'") +# Shell quotes only delimit; bash concatenates the pieces (cat /proc/x/enviro''n +# reads .../environ), so strip them before the sensitive-path scan. +_SHELL_QUOTE_RE = re.compile(r"['\"]") +# A glob bracket class [s] -> s, so .s[s]h de-obfuscates to .ssh for the scan. +_GLOB_BRACKET_RE = re.compile(r"\[([^!\]][^\]]*)\]") +# Bash POSIX character classes ([[:lower:]]) each match one char; Python fnmatch +# does not understand them, so normalize to `?` before the glob check. +_POSIX_CLASS_RE = re.compile(r"\[\[:\w+:\]\]") +# Canonical sensitive files a ? / * / [..] glob could expand to; fnmatch tests +# whether the pattern reaches one (cat /e??/passwd -> /etc/passwd). +_SENSITIVE_GLOB_TARGETS = ( + "/etc/passwd", + "/etc/shadow", + "/etc/sudoers", + "/root/.ssh/id_rsa", + "/root/.aws/credentials", + "/home/u/.ssh/id_rsa", + "/home/u/.ssh/id_ed25519", + "/home/u/.aws/credentials", + "/home/u/.netrc", + "/home/u/.git-credentials", +) +# Directories whose every file is a credential/secret; a glob resolving into one +# (cat /r?n/secrets/hf_token, cat /root/.s??/id_rsa) reads a secret even though +# the exact filename is never enumerated, so a globbed token here asks. +_SENSITIVE_GLOB_DIRS = ( + "/run/secrets", + "/var/run/secrets", + "/root/.ssh", + "/root/.aws", + "/root/.azure", + "/root/.gnupg", + "/root/.docker", + "/root/.kube", + "/root/.config/gcloud", + "/root/.config/gh", + "/home/u/.ssh", + "/home/u/.aws", + "/home/u/.azure", + "/home/u/.gnupg", + "/home/u/.docker", + "/home/u/.kube", + "/home/u/.config/gcloud", + "/home/u/.config/gh", +) +# Credential basenames a glob can reach even when the directory is not wholly +# sensitive (cat ~/.huggingface/tok?n -> token, cat ~/.netr? -> .netrc); the +# canonical-target list only covers a few fixed home paths, so match the globbed +# basename against these directly. +_SENSITIVE_GLOB_BASENAMES = frozenset( + { + "token", + "stored_tokens", + "credentials", + ".netrc", + "netrc", + ".pypirc", + ".npmrc", + ".git-credentials", + "id_rsa", + "id_ed25519", + "id_ecdsa", + "id_dsa", + "passwd", + "shadow", + # A project .env holds secrets; the literal path is gated elsewhere, so a + # glob that expands to it (cat .e?v) must be too. + ".env", + } +) +# A leading shell redirection (<, >, 2>, >>) hides the path from a plain glob +# scan (cat ]+") +# Bash brace expansion (cat /etc/pass{w,}d -> /etc/passwd /etc/passd, and the +# sequence form cat /etc/pass{w..w}d -> /etc/passwd) runs after this classifier; +# expand comma groups and .. sequences to scan each result. +_BRACE_COMMA_RE = re.compile(r"^\{([^{}]*,[^{}]*)\}$") +_BRACE_SEQ_RE = re.compile(r"^\{([^{}]+)\.\.([^{}]+)(?:\.\.(-?\d+))?\}$") +_BRACE_ANY_RE = re.compile(r"\{[^{}]*,[^{}]*\}|\{[^{}]+\.\.[^{}]+(?:\.\.-?\d+)?\}") +# Parameter expansion with a default/alternate operator (${x:-passwd}, +# ${x:+passwd}, ${x=passwd}) can synthesize a path after approval; the operand +# is substituted so the resulting path is scanned. +_SHELL_PARAM_OP_RE = re.compile(r"\$\{[A-Za-z_]\w*:?[-=+]([^{}]*)\}") + + +def _references_sensitive_path(text: str) -> bool: + """True if a command or string literal reads a credential path or escapes + the sandbox workdir via parent traversal.""" + norm = _REDUNDANT_SLASH_RE.sub("", text) + debracket = _GLOB_BRACKET_RE.sub(lambda m: m.group(1)[0], text) + return bool( + _PARENT_TRAVERSAL_RE.search(text) + or _SENSITIVE_PATH_RE.search(text) + or _SENSITIVE_PATH_RE.search(norm) + or _SENSITIVE_PATH_RE.search(debracket) + ) + + +def _pattern_matches_dir(pattern: str, target: str) -> bool: + """Segment-wise fnmatch so a glob segment does not cross a '/' boundary + (`/home/*` must not match `/home/u/.ssh`).""" + p = pattern.split("/") + t = target.split("/") + if len(p) != len(t): + return False + return all(fnmatch.fnmatch(tseg, pseg) for pseg, tseg in zip(p, t)) + + +def _glob_token_sensitive(token: str) -> bool: + """True if a single ? / * / [..] glob token could expand to a sensitive file + or a file under a secret/credential directory. Shared by the terminal scan + and the Python glob check (glob.glob('/e??/passwd')).""" + token = _REDIR_PREFIX_RE.sub("", _SHELL_QUOTE_RE.sub("", token)) + # A POSIX class ([[:lower:]]) matches one char, like `?`, but fnmatch treats + # it as a literal set; normalize so cat /etc/pass[[:lower:]]d resolves. + token = _POSIX_CLASS_RE.sub("?", token) + if not any(c in token for c in "?*["): + return False + if any(fnmatch.fnmatch(target, token) for target in _SENSITIVE_GLOB_TARGETS): + return True + # A glob that resolves to a credential basename is sensitive wherever it + # lives (cat ~/.huggingface/tok?n -> token, cat proj/.netr? -> .netrc); the + # fixed-target list only covers a handful of home paths. + base = token.rsplit("/", 1)[-1] + if any(c in base for c in "?*[") and any( + fnmatch.fnmatch(name, base) for name in _SENSITIVE_GLOB_BASENAMES + ): + return True + # A globbed directory that resolves into a secret/credential dir makes every + # file below it sensitive (cat /r?n/secrets/hf_token). + head = token.rsplit("/", 1)[0] if "/" in token else token + return any( + _pattern_matches_dir(token, d) or _pattern_matches_dir(head, d) + for d in _SENSITIVE_GLOB_DIRS + ) + + +def _glob_hits_sensitive(command: str) -> bool: + """True if any glob token in a command could expand to a sensitive file, so + `cat /e??/passwd` and `cat /r?n/secrets/hf_token` ask even without a literal + sensitive path.""" + return any( + _glob_token_sensitive(token) + for token in command.replace(";", " ").replace("|", " ").split() + ) + + +def _expand_shell_assignments(command: str) -> str: + """Best-effort substitution of `NAME=value ... $NAME`, so a sensitive path + split across an assignment and an argument (p=/etc; cat $p/passwd) is still + visible to the sensitive-path scan. Also applies pattern replacement + (p=passXd; cat /etc/${p/X/w}). Fail-open: only adds detections.""" + env = dict(_SHELL_ASSIGN_RE.findall(command)) + if not env: + return command + + def repl_pattern(m): + var, is_global, pat, rep = m.group(1), m.group(2), m.group(3), m.group(4) + if var not in env or not pat: + return m.group(0) + return env[var].replace(pat, rep) if is_global else env[var].replace(pat, rep, 1) + + def repl_case(m): + var, op = m.group(1), m.group(2) + if var not in env: + return m.group(0) + v = env[var] + if op == ",,": + return v.lower() + if op == "^^": + return v.upper() + if op == ",": + return v[:1].lower() + v[1:] + return v[:1].upper() + v[1:] + + def repl_indirect(m): + # ${!p} -> value of the variable named by $p (env[env[p]]). + pointed = env.get(m.group(1)) + return env.get(pointed, m.group(0)) if pointed is not None else m.group(0) + + command = _SHELL_PARAM_INDIRECT_RE.sub(repl_indirect, command) + command = _SHELL_PARAM_REPL_RE.sub(repl_pattern, command) + command = _SHELL_PARAM_CASE_RE.sub(repl_case, command) + return _SHELL_VAR_RE.sub(lambda m: env.get(m.group(1) or m.group(2), m.group(0)), command) + + +def _expand_param_defaults(command: str) -> str: + """Substitute the operand of a default/alternate parameter expansion + (cat /etc/pass${x:-wd} -> cat /etc/passwd), which bash applies after this + classifier. Fail-open: only adds detections.""" + return _SHELL_PARAM_OP_RE.sub(lambda m: m.group(1), command) + + +def _decode_ansi_c(command: str) -> str: + """Decode bash ANSI-C quoted words (cat $'/etc/pass\\x77d' -> cat /etc/passwd) + so an escape-obfuscated path is visible to the scan. Fail-open: only adds + detections.""" + + def dec(m): + try: + return bytes(m.group(1), "utf-8").decode("unicode_escape") + except (UnicodeDecodeError, ValueError): + return m.group(0) + + return _ANSI_C_RE.sub(dec, command) + + +def _brace_range(lo: str, hi: str, step: "str | None") -> "list[str]": + """Expand a bash sequence brace endpoint pair ({1..3}, {a..c}, {w..w}).""" + try: + istep = abs(int(step)) if step else 1 + istep = istep or 1 + if re.fullmatch(r"-?\d+", lo) and re.fullmatch(r"-?\d+", hi): + a, b = int(lo), int(hi) + rng = range(a, b + 1, istep) if a <= b else range(a, b - 1, -istep) + return [str(x) for x in rng][:64] + if len(lo) == 1 and len(hi) == 1 and lo.isalpha() and hi.isalpha(): + a, b = ord(lo), ord(hi) + rng = range(a, b + 1, istep) if a <= b else range(a, b - 1, -istep) + return [chr(x) for x in rng][:64] + except (ValueError, TypeError): + pass + return [] + + +def _brace_options(text: str) -> "list[str]": + """Options a single brace group expands to (comma list or .. sequence).""" + m = _BRACE_COMMA_RE.match(text) + if m: + return m.group(1).split(",") + m = _BRACE_SEQ_RE.match(text) + if m: + return _brace_range(m.group(1), m.group(2), m.group(3)) or [text] + return [text] + + +def _expand_braces(command: str) -> str: + """Best-effort bash brace expansion (cat /etc/pass{w,}d -> cat /etc/passwd + /etc/passd, cat /etc/pass{w..w}d -> cat /etc/passwd) so a sensitive path + split across a brace group is scanned. Bounded. Fail-open: only detects.""" + results = [command] + for _ in range(6): + if not any(_BRACE_ANY_RE.search(s) for s in results): + break + expanded = [] + for s in results: + m = _BRACE_ANY_RE.search(s) + if not m: + expanded.append(s) + continue + for opt in _brace_options(m.group(0)): + expanded.append(s[: m.start()] + opt + s[m.end() :]) + results = expanded[:64] + return " ".join(results) + + +def _mode_arg_writes(mode_node) -> bool: + """True if an AST node used as a file mode requests write/append.""" + if mode_node is None: + return False # default "r" + if isinstance(mode_node, ast.Constant) and isinstance(mode_node.value, str): + return bool(_PY_WRITE_MODE_RE.search(mode_node.value)) + return True # dynamic mode: cannot prove read-only + + +def _has_kwarg_splat(node) -> bool: + """True if the call has a ``**kwargs`` splat, which can hide a write mode.""" + return any(kw.arg is None for kw in node.keywords or []) + + +def _builtin_open_writes(node) -> bool: + """Write check for builtin ``open(file, mode)`` (mode is the 2nd arg).""" + if _has_kwarg_splat(node): + return True # **{"mode": "w"} could request a write + if any(isinstance(a, ast.Starred) for a in node.args): + return True # *("f", "w") could splat a write mode into the positionals + mode = node.args[1] if len(node.args) >= 2 else None + for kw in node.keywords or []: + if kw.arg == "mode": + mode = kw.value + return _mode_arg_writes(mode) + + +def _attr_open_writes(node) -> bool: + """Write check for ``x.open(...)`` (e.g. ``Path.open(mode)`` where mode is + the 1st arg). Only a mode-looking string is read as the mode, so a + ``ZipFile.open("name.txt")`` read is not mistaken for a write.""" + if _has_kwarg_splat(node): + return True # **{"mode": "w"} could request a write + for kw in node.keywords or []: + if kw.arg == "mode": + return _mode_arg_writes(kw.value) + if node.args: + first = node.args[0] + if isinstance(first, ast.Constant) and isinstance(first.value, str): + if _PY_MODE_LITERAL_RE.match(first.value): + return bool(_PY_WRITE_MODE_RE.search(first.value)) + # A 2nd positional arg is either a mode (x.open(name, "w")) or + # os.open(path, O_CREAT) flags via an alias: honor a string mode, + # otherwise cannot prove read-only, so ask. + if len(node.args) >= 2: + second = node.args[1] + if isinstance(second, ast.Constant) and isinstance(second.value, str): + return _mode_arg_writes(second) + return True + return False + return True # dynamic first arg: cannot prove read-only + return False # no args: read + + +_PATH_CTORS = ( + "Path", + "PurePath", + "PurePosixPath", + "PureWindowsPath", + "PosixPath", + "WindowsPath", +) +# Deterministic path pass-through/normalizer calls that return the same location +# (os.path.abspath('/etc') -> /etc, Path('/etc').resolve() -> /etc), so folding +# through them keeps a sensitive root visible to the scan. +_PATH_PASSTHROUGH_ATTRS = frozenset( + {"abspath", "normpath", "realpath", "expanduser", "expandvars", "resolve", "absolute"} +) +# pathlib methods that rewrite only the final path component, so the sensitive +# target is never spelled out as a literal (Path('/etc/x').with_name('passwd') +# -> /etc/passwd). Folded below so the rewritten path is still scanned. +_PATH_NAME_REWRITES = frozenset({"with_name", "with_stem", "with_suffix"}) +# Mapping-style %-format conversion specifier: %(name)s / %(n)5.2f. Used to fold +# '/etc/%(f)s' % {'f': 'passwd'} to /etc/passwd (a dynamic value becomes NUL). +_PERCENT_NAMED_RE = re.compile(r"%\((\w+)\)[-#0 +]*\d*(?:\.\d+)?[a-zA-Z]") + + +def _folded_path( + node, + literals = None, + ctors = None, + join_names = None, +) -> "str | None": + """Best-effort value of a path built from string literals, so a sensitive + path assembled from pieces (os.path.join('/etc', 'passwd'), '/etc'+'/passwd', + Path('/etc') / 'passwd', f'/proc/{pid}/environ', f'/etc/{name}') is still + visible to the scan. A dynamic piece becomes NUL, a non-slash placeholder, + so a dynamic segment under a sensitive dir (/etc/NUL) is still detectable. + ``literals`` maps names bound to string literals (base = '/etc'); ``ctors`` + is the set of pathlib constructor names (incl. import aliases); ``join_names`` + are bare names bound to os.path.join (from os.path import join).""" + literals = literals or {} + ctors = ctors or _PATH_CTORS + join_names = join_names or frozenset() + + def fold(node) -> "str | None": + if isinstance(node, ast.Constant) and isinstance(node.value, (str, bytes)): + # bytes paths are valid too (open(b'/etc/passwd')); decode for scan. + return ( + node.value.decode("latin-1", "ignore") + if isinstance(node.value, bytes) + else node.value + ) + if isinstance(node, ast.Name): + return literals.get(node.id) + if isinstance(node, ast.Attribute) and node.attr in ("parent", "parents"): + # A pathlib .parent/.parents walks above the current dir, escaping + # the per-session workdir without a literal '..'; mark it so a read + # folds to unsafe (\x02 is a non-slash escape sentinel). + return "\x02" + if ( + isinstance(node, ast.Subscript) + and isinstance(node.value, ast.Attribute) + and (node.value.attr == "parents") + ): + return "\x02" # Path(...).parents[1] + if isinstance(node, ast.JoinedStr): + return "".join( + v.value + if isinstance(v, ast.Constant) and isinstance(v.value, str) + else (fold(v.value) or "\x00") + if isinstance(v, ast.FormattedValue) + else "\x00" + for v in node.values + ) + if isinstance(node, ast.BinOp) and isinstance(node.op, (ast.Add, ast.Div)): + left = fold(node.left) + right = fold(node.right) + left = "\x00" if left is None else left + right = "\x00" if right is None else right + # Path('/etc') / 'passwd' joins with a separator; '+' concatenates. + return left + "/" + right if isinstance(node.op, ast.Div) else left + right + if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Mod): + # Old-style formatting: '%s/%s' % ('/etc', 'passwd') -> /etc/passwd. + template = fold(node.left) + if template is not None and "%" in template: + rhs = node.right + if "%(" in template: + # Mapping-style: '/etc/%(f)s' % {'f': 'passwd'} -> /etc/passwd. + # A literal dict resolves each name; an unresolved value or a + # non-literal mapping leaves the NUL marker so /etc/ + # still fails closed under a sensitive dir. + mapping: "dict[str, str]" = {} + if isinstance(rhs, ast.Dict): + for k, v in zip(rhs.keys, rhs.values): + if isinstance(k, ast.Constant) and isinstance(k.value, str): + fv = fold(v) + mapping[k.value] = fv if fv is not None else "\x00" + return _PERCENT_NAMED_RE.sub( + lambda m: mapping.get(m.group(1), "\x00"), template + ) + if isinstance(rhs, ast.Tuple): + args = tuple((fold(e) or "\x00") for e in rhs.elts) + else: + single = fold(rhs) + args = (single if single is not None else "\x00",) + try: + return template % args + except (TypeError, ValueError, KeyError): + return None + return None + if isinstance(node, ast.Call): + func = node.func + if isinstance(func, ast.Attribute) and func.attr == "joinpath": + # Path('/etc').joinpath('passwd') -> receiver and args are pieces. + base = fold(func.value) + parts = [base if base is not None else "\x00"] + parts += [(fold(a) or "\x00") for a in node.args] + return "/".join(parts) + if isinstance(func, ast.Attribute) and func.attr in ("glob", "rglob", "iglob"): + # Path('/etc').glob('passw?') -> the receiver dir joined with the + # glob pattern; _glob_token_sensitive then tests /etc/passw?. + base = fold(func.value) + pattern = fold(node.args[0]) if node.args else "\x00" + return (base if base is not None else "\x00") + "/" + (pattern or "\x00") + if isinstance(func, ast.Attribute) and func.attr in _PATH_NAME_REWRITES: + # Path('/etc/x').with_name('passwd') -> /etc/passwd; with_stem / + # with_suffix rewrite only the final component. Fold to the + # rewritten path so a sensitive target that no literal spells out + # is still caught. An unresolved receiver stays None (untracked, + # like a bare variable), and a dynamic arg becomes the NUL marker. + base = fold(func.value) + if base is None: + return None + arg = fold(node.args[0]) if node.args else None + arg = "\x00" if arg is None else arg + idx = base.rfind("/") + head = base[: idx + 1] if idx >= 0 else "" + name = base[idx + 1 :] if idx >= 0 else base + dot = name.rfind(".") + stem = name[:dot] if dot > 0 else name + suffix = name[dot:] if dot > 0 else "" + if func.attr == "with_name": + name = arg + elif func.attr == "with_stem": + name = arg + suffix + else: # with_suffix + name = stem + arg + return head + name + if isinstance(func, ast.Attribute) and func.attr in _PATH_PASSTHROUGH_ATTRS: + # Deterministic normalizers keep the same path: os.path.abspath( + # '/etc') -> /etc, Path('/etc').resolve() -> /etc. When called with + # a path arg fold it, else fold the receiver (Path method form). + return fold(node.args[0]) if node.args else fold(func.value) + if isinstance(func, ast.Attribute) and func.attr == "join": + # str.join has the separator as the receiver and the pieces in + # one iterable arg ("".join(['/etc', '/passwd']) -> /etc/passwd); + # tell it apart from os.path.join(*pieces). + sep = fold(func.value) + if ( + sep is not None + and len(node.args) == 1 + and isinstance(node.args[0], (ast.List, ast.Tuple)) + ): + pieces = [(fold(e) or "\x00") for e in node.args[0].elts] + return sep.join(pieces) + parts = [(fold(a) or "\x00") for a in node.args] + return "/".join(parts) + # A bare os.path.join alias (from os.path import join): join(*pieces). + if isinstance(func, ast.Name) and func.id in join_names: + parts = [(fold(a) or "\x00") for a in node.args] + return "/".join(parts) + # A bare/qualified/aliased pathlib constructor (Path(...), P(...)). + if (isinstance(func, ast.Attribute) and func.attr in ctors) or ( + isinstance(func, ast.Name) and func.id in ctors + ): + parts = [(fold(a) or "\x00") for a in node.args] + return "/".join(parts) + # '/etc/{}'.format('passwd') -> /etc/passwd (literal template + args). + if isinstance(func, ast.Attribute) and func.attr == "format": + template = fold(func.value) + if template is not None and "{" in template: + parts = [] + for a in node.args: + if isinstance(a, ast.Constant): + parts.append(str(a.value)) + else: + folded = fold(a) + parts.append("\x00" if folded is None else folded) + try: + return template.format(*parts) + except (IndexError, KeyError, ValueError): + return None + return None + + return fold(node) + + +def _dynamic_name_hits_sensitive(folded) -> bool: + """True if a folded path with a dynamic piece (NUL) inside a path segment + could spell a credential target, e.g. open('/et' + chr(99) + '/passwd') + folds to '/et\\x00/passwd'. NUL matches any run of non-separator chars so the + dynamic split of a sensitive name resolves, while an all-dynamic ('\\x00\\x00') + or segment-spanning ('\\x00/\\x00') path cannot form a single credential name + and stays safe.""" + if not folded or "\x00" not in folded: + return False + pattern = "".join(r"[^/\\]*" if ch == "\x00" else re.escape(ch) for ch in folded) + try: + rx = re.compile(pattern + r"\Z") + except re.error: + return True # pathological pattern: fail closed + return any(rx.match(t) for t in _SENSITIVE_GLOB_TARGETS) + + +def _folded_is_sensitive(folded) -> bool: + """A folded path is sensitive if it names a credential file, has a dynamic + segment (NUL) directly under a sensitive directory (/etc/NUL), walks out of + the sandbox via a pathlib .parent/.parents escape (\\x02), or is a glob that + could resolve to a credential path (glob.glob('/e??/passwd')).""" + if not folded: + return False + return ( + "\x02" in folded + or _references_sensitive_path(folded) + or ("\x00" in folded and bool(_SENSITIVE_DIR_RE.search(folded))) + # A dynamic segment (NUL) can be the "/" forming a sensitive root: + # open(os.sep + "etc/passwd") folds to "\x00etc/passwd", so re-scan with + # NUL as "/" (a benign "\x00data/file" -> "/data/file" stays safe). + or ("\x00" in folded and _references_sensitive_path(folded.replace("\x00", "/"))) + # A dynamic piece can also sit INSIDE a sensitive name: open('/et' + + # chr(99) + '/passwd') folds to "/et\x00/passwd", which none of the above + # catch. Match the literals around each NUL against a credential target, + # treating NUL as "any run of non-separator chars" so /et/passwd + # resolves while an all-dynamic ("\x00\x00" from 1 + 1) or segment-spanning + # ("\x00/\x00" from a + '/' + b) path stays safe. + or _dynamic_name_hits_sensitive(folded) + or _glob_token_sensitive(folded) + ) + + +def _terminal_is_potentially_unsafe(command: str) -> bool: + """Classify a terminal command for auto mode (fail closed).""" + if not command or not command.strip(): + return False + # Redirections and substitutions can hide writes or nested commands; a + # quoted ">" false-positives into a prompt, which is the safe direction. + if ">" in command or "`" in command or "$(" in command or "<(" in command: + return True + # Reads that escape the sandbox workdir (../) or hit credential paths are + # not "safe" reads; ask before running them. Strip shell quotes/backslash + # escapes and expand NAME=value prefixes first so `cat /proc/$PPID/enviro''n`, + # `cat /et\c/passwd`, and `p="/proc/$PPID"; cat $p/environ` are caught too. + stripped = _SHELL_QUOTE_RE.sub("", command).replace("\\", "") + # Bash applies brace/parameter/ANSI-C expansion after this classifier, so a + # path split across a brace group (/etc/pass{w,}d), a default/substring param + # (${x:-wd}, ${p:0:6}), or an escape ($'...') is invisible to the raw scan; + # expand first (ANSI-C decoded from the raw command, before backslash strip). + candidates = [] + for c in (command, stripped, _decode_ansi_c(command)): + c_param = _expand_param_defaults(c) + candidates.extend((c, c_param, _expand_braces(c_param), _expand_shell_assignments(c_param))) + # Run both the literal and glob-sensitive scans over every candidate, so a + # brace-expanded glob (cat /e{t,}c/pass?d -> /etc/pass?d) is caught. + if any(_glob_hits_sensitive(c) or _references_sensitive_path(c) for c in candidates): + return True + # Newlines (and CR) separate commands in a shell but read as plain + # whitespace to shlex, which would demote "ls\nrm x" to argument position. + command = command.replace("\r\n", ";").replace("\n", ";").replace("\r", ";") + try: + lexer = shlex.shlex(command, posix = True, punctuation_chars = ";&|()") + lexer.whitespace_split = True + tokens = list(lexer) + except ValueError: + return True + # A root can also hide behind an assignment (p=/; grep -R TOKEN $p) or a + # default parameter (grep -R TOKEN ${root:-/home}); re-lex the fully expanded + # command so the find/fd and recursive-search scans see the resolved token. + expanded_command = _expand_shell_assignments(_expand_param_defaults(command)) + if expanded_command != command: + try: + elexer = shlex.shlex(expanded_command, posix = True, punctuation_chars = ";&|()") + elexer.whitespace_split = True + scan_tokens = list(elexer) + except ValueError: + return True + else: + scan_tokens = tokens + # find/fd group with (...) which resets command context, so a trailing + # -delete/-exec could slip past; scan every token when find/fd appears. + if any(os.path.basename(t.strip(";&|()`{}")).lower() in ("find", "fd") for t in scan_tokens): + if any(t.split("=", 1)[0] in _AUTO_UNSAFE_FIND_LIKE_FLAGS for t in scan_tokens): + return True + # A recursive reader rooted outside the sandbox reads host files (grep -R + # TOKEN /home, rg TOKEN /, grep -R TOKEN ~root, p=/; grep -R TOKEN $p, and + # the always-recursive walkers tree /home / du /); ask. Bash expands + # ~/~user to a home dir after this decision, so a tilde root is a sandbox + # escape too. A path-qualified command token starts with "/" as well, but + # that already asks below. + if any(t.startswith("/") or t.startswith("~") for t in scan_tokens): + token_bases = [os.path.basename(t.strip(";&|()`{}")).lower() for t in tokens] + if any(b in _AUTO_RECURSIVE_SEARCH or b in _AUTO_RECURSIVE_LISTERS for b in token_bases): + return True + # ls only walks the whole subtree with -R/--recursive (ls -R /home, + # ls -laR /); a non-recursive ls /home lists one level and stays here. + if "ls" in token_bases and any( + t.split("=", 1)[0] in ("-R", "--recursive") + or (t[:1] == "-" and t[:2] != "--" and "=" not in t and "R" in t[1:]) + for t in tokens + ): + return True + expect_command = True + prefix_pending = False + current_command = "" + positional_args = 0 + pending_flag_value = False + for token in tokens: + # Runs of punctuation (";;", ";&") lex as one token; any token made + # purely of separator characters still separates commands. + if ( + token in _SHELL_SEPARATORS + or token in _SHELL_KEYWORDS_AS_SEP + or not set(token) - set(";&|()") + ): + expect_command = True + prefix_pending = False + current_command = "" + positional_args = 0 + pending_flag_value = False + continue + if token.startswith("-"): + # A write/exec flag on an otherwise read-only command asks + # (sort -o, tree -o, xxd -r, find -exec/-delete/...). Match + # "--output=x", an attached short option "-o/tmp/out", and a short + # option bundled in a cluster (sort -uo out => -u -o). + flag_head = token.split("=", 1)[0] + cluster = token[1:] if token[:2] != "--" and "=" not in token else "" + # GNU tools accept unambiguous abbreviations of a long option, so + # `sort --out=` reaches --output and `env --ch=/` reaches --chdir; + # a "--x" prefix of an unsafe long flag fails closed. + is_long_abbrev = flag_head.startswith("--") and len(flag_head) > 2 + for uf in _AUTO_UNSAFE_COMMAND_FLAGS.get(current_command, ()): + if flag_head == uf or (len(uf) == 2 and (token.startswith(uf) or uf[1] in cluster)): + return True + if is_long_abbrev and uf.startswith("--") and uf.startswith(flag_head): + return True + # A flag that takes a following value (date -d STRING / -r FILE; + # uniq -f N; xxd -c N) so the value token is not mistaken for a + # clock-setting positional or an output-file positional. + pending_flag_value = "=" not in token and ( + (current_command == "date" and flag_head in _DATE_DISPLAY_VALUE_FLAGS) + or flag_head in _SECOND_POSITIONAL_VALUE_FLAGS.get(current_command, ()) + ) + if not prefix_pending: + expect_command = False + continue + if not expect_command: + raw_pos = token.strip(";&|()`{}") + # uniq [INPUT [OUTPUT]] writes its second file positional; count file + # positionals and ask on the second one. A preceding option's value + # (uniq -f 2) is consumed via pending_flag_value, so a file literally + # named with digits (uniq 123 out) is still counted. + if current_command in _AUTO_SECOND_POSITIONAL_WRITES: + if pending_flag_value: + pending_flag_value = False + elif raw_pos: + positional_args += 1 + if positional_args >= 2: + return True + # hostname NAME sets the hostname; date sets the clock. A + # positional past a display flag's value therefore mutates state and + # asks (date's +FORMAT display token stays read-only). + elif current_command in _AUTO_ARG_SENSITIVE_COMMANDS: + if pending_flag_value: + pending_flag_value = False + elif raw_pos and not (current_command == "date" and raw_pos.startswith("+")): + return True + continue + if _ASSIGNMENT_RE.match(token): + # Benign NAME=value prefixes are skipped, but ones that change + # command lookup/loading (PATH, LD_PRELOAD, ...) fail closed. + if _env_assignment_is_unsafe(token.split("=", 1)[0]): + return True + continue + if prefix_pending and token.lstrip("-").isdigit(): + continue + raw = token.strip(";&|()`{}") + # A path-qualified command (./ls, /tmp/cat) is an arbitrary executable, + # not the trusted system utility its basename matches; ask first. + if "/" in raw or "\\" in raw: + return True + base = os.path.basename(raw).lower() + stem, ext = os.path.splitext(base) + if ext in {".exe", ".com", ".bat", ".cmd"}: + base = stem + if base in _AUTO_SAFE_WRAPPERS: + prefix_pending = True + # Track the wrapper so its own flags (env --chdir) are checked; + # the real command overwrites this when it is reached. + current_command = base + pending_flag_value = False + continue + if base not in _AUTO_SAFE_TERMINAL_COMMANDS: + return True + current_command = base + expect_command = False + prefix_pending = False + positional_args = 0 + pending_flag_value = False + return False + + +def _python_is_potentially_unsafe(code: str) -> bool: + """Classify python-tool code for auto mode (fail closed).""" + if not code or not code.strip(): + return False + # Anything the sandbox's static analysis already objects to would be + # refused at execution time; surface it as a confirmation first. + if _check_code_safety(code) is not None: + return True + try: + tree = ast.parse(code) + except SyntaxError: + return False # runs into a normal traceback; nothing to guard + # Names bound to the builtin open (f = open; from builtins import open as f; + # f, _ = (open, print)) so an aliased writer call is still checked below. + # builtins_aliases tracks `import builtins [as b]` for builtins.exec/eval. + open_aliases = {"open"} + # Attribute names bound to open (box.f = open), so a later box.f('out', 'w') + # write is still gated even though the callable is an attribute, not a name. + attr_open_aliases: "set[str]" = set() + builtins_aliases = {"builtins", "__builtins__"} + # Names bound to a dynamic lookup (rm = getattr(os, "remove"); + # f = globals()["open"]) whose calls cannot be proven read-only, so they + # fail closed. + dynamic_aliases = set() + # Names bound to a dynamic-code builtin, including aliased ones + # (from builtins import eval as e; e = builtins.exec), so a call or + # reference through the alias fails closed too. compile() builds a code + # object that FunctionType/exec can then run. + code_exec_aliases = {"exec", "eval", "__import__", "breakpoint", "compile"} + # Names bound to a string literal (base = '/etc'), so a sensitive path + # split through a variable (base + '/passwd') folds and is caught. + literal_str_vars: "dict[str, str]" = {} + # Pathlib constructor names incl. import aliases (from pathlib import Path as + # P), os.path.join names bound directly (from os.path import join as j), and + # writer functions imported as bare names (from numpy import save). + path_ctor_aliases = set(_PATH_CTORS) + pathjoin_aliases: "set[str]" = set() + writer_aliases: "set[str]" = set() + # Module names bound to os/posix (import os as o), so o.open(...) is still + # recognized as the low-level create/write that os.open is. + os_aliases = {"os", "posix"} + # Module names bound to a pickle-backed loader (import torch as t), so + # t.load(...) is still gated as a code-executing deserialize. + load_module_aliases = set(_AUTO_UNSAFE_PY_LOAD_MODULES) + # Names bound to the builtin getattr (g = getattr), so a dynamic lookup + # aliased through it (rm = g(os, "remove"); rm("f")) still fails closed. + getattr_aliases = {"getattr"} + # Names bound to functools.partial, so a partial that wraps open/a writer + # (w = partial(open, mode="w"); w("out.txt")) fails closed when w is called. + partial_aliases: "set[str]" = set() + # Archive constructors imported bare (from zipfile import ZipFile), so + # ZipFile(name, "w") is gated like the zipfile.ZipFile attribute call. + archive_ctor_aliases: "set[str]" = set() + # operator.methodcaller("write_text") is dynamic dispatch, like getattr. + operator_aliases = {"operator"} + methodcaller_aliases: "set[str]" = set() + # logging.basicConfig(filename=...) opens a log file for write. + basicconfig_aliases: "set[str]" = set() + # fileinput.input(..., inplace=True) rewrites a file in place. + fileinput_aliases = {"fileinput"} + # Higher-order invokers (map/filter/starmap/reduce) call their first arg, so + # one handed a writer (map(open, ...)) writes without a direct open() site. + # Track aliases (m = map; from itertools import starmap as sm) so an aliased + # invoker is still checked; the write-callable gate keeps map(len, ...) safe. + invoker_aliases = set(_HIGHER_ORDER_INVOKERS) + + def _is_dynamic_namespace(node) -> bool: + # A namespace mapping whose .get/.pop/.setdefault (or subscript) can return + # open/eval/a mutator: globals()/locals()/vars(...), any X.__dict__, + # __builtins__, sys.modules. Looking a name up through one is as dynamic as + # getattr, so a value fetched from it fails closed. + if isinstance(node, ast.Attribute): + if node.attr == "__dict__": + return True + return ( + node.attr == "modules" + and isinstance(node.value, ast.Name) + and node.value.id == "sys" + ) + if isinstance(node, ast.Name): + return node.id in builtins_aliases + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name): + return node.func.id in ("globals", "locals", "vars") + return False + + def _methodcaller_writes(call) -> bool: + # operator.methodcaller("write_text", ...) / methodcaller(name): unsafe + # when the method name is a known writer/mutator, or non-constant (cannot + # be proven read-only). + if not call.args: + return False + first = call.args[0] + if not (isinstance(first, ast.Constant) and isinstance(first.value, str)): + return True + return first.value in _AUTO_UNSAFE_PY_ATTRS or first.value in _AUTO_UNSAFE_PY_WRITE_METHODS + + def _fileinput_inplace(call) -> bool: + # fileinput.input(..., inplace=True) opens each file for in-place rewrite. + if _has_kwarg_splat(call): + return True + for kw in call.keywords or []: + if kw.arg == "inplace": + v = kw.value + if isinstance(v, ast.Constant): + return bool(v.value) + return True # dynamic inplace flag: cannot prove read-only + return False + + def _basicconfig_writes(call) -> bool: + # logging.basicConfig(filename=...) creates/opens a log file for writing. + if _has_kwarg_splat(call): + return True + return any(kw.arg == "filename" for kw in call.keywords or []) + + def _wraps_write_callable(arg) -> bool: + # The callable a partial wraps (partial(open, ...)); True when calling it + # could create/overwrite a file or resolve a dynamic/mutating function. + if isinstance(arg, ast.Name): + return ( + arg.id in open_aliases + or arg.id in dynamic_aliases + or arg.id in code_exec_aliases + or arg.id in getattr_aliases + or arg.id in writer_aliases + or arg.id in archive_ctor_aliases + ) + if isinstance(arg, ast.Attribute): + return ( + arg.attr == "open" + or arg.attr in _AUTO_UNSAFE_PY_ATTRS + or arg.attr in _AUTO_UNSAFE_PY_WRITE_METHODS + or arg.attr in _ARCHIVE_CTOR_NAMES + ) + return False + + def _passed_write_callable(arg) -> bool: + # A concrete write callable handed as an argument to another call: a + # name bound to open / a writer / an archive constructor, or an + # attribute reference to a writer method / mutating os attr / archive + # ctor / .open. Unlike _wraps_write_callable this omits the fail-closed + # dynamic / getattr / code-exec poison aliases, which are already gated + # where they are *called* and would over-trigger when a benign alias is + # merely passed or printed (print(getattr(o, 'name'))). + if isinstance(arg, ast.Name): + return ( + arg.id in open_aliases or arg.id in writer_aliases or arg.id in archive_ctor_aliases + ) + if isinstance(arg, ast.Attribute): + return ( + arg.attr == "open" + or arg.attr in _AUTO_UNSAFE_PY_ATTRS + or arg.attr in _AUTO_UNSAFE_PY_WRITE_METHODS + or arg.attr in _ARCHIVE_CTOR_NAMES + ) + return False + + # Names bound more than once cannot be folded to a single literal: this scan + # visits every assignment before any call is checked, so a later benign + # reassignment (base = '/etc'; open(base + '/passwd'); base = 'data') would + # otherwise mask the earlier sensitive value and auto-approve. Count every + # binding target up front and poison multiply-bound names to the escape + # sentinel so any path folded from them fails closed (asks) instead. + assign_counts: "dict[str, int]" = {} + for node in ast.walk(tree): + binding_targets = [] + if isinstance(node, ast.Assign): + binding_targets = node.targets + elif isinstance(node, (ast.AnnAssign, ast.AugAssign)): + binding_targets = [node.target] + for target in binding_targets: + for sub in ast.walk(target): + if isinstance(sub, ast.Name): + assign_counts[sub.id] = assign_counts.get(sub.id, 0) + 1 + multi_assigned_names = {name for name, count in assign_counts.items() if count > 1} + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.name == "builtins": + builtins_aliases.add(alias.asname or "builtins") + elif alias.name in ("os", "posix"): + os_aliases.add(alias.asname or alias.name) + elif alias.name in _AUTO_UNSAFE_PY_LOAD_MODULES: + load_module_aliases.add(alias.asname or alias.name) + elif alias.name == "operator": + operator_aliases.add(alias.asname or "operator") + elif alias.name == "fileinput": + fileinput_aliases.add(alias.asname or "fileinput") + elif isinstance(node, ast.ImportFrom): + if node.module == "operator": + for alias in node.names: + if alias.name == "methodcaller": + methodcaller_aliases.add(alias.asname or "methodcaller") + if node.module == "logging": + for alias in node.names: + if alias.name == "basicConfig": + basicconfig_aliases.add(alias.asname or "basicConfig") + if node.module == "builtins": + for alias in node.names: + if alias.name == "open": + open_aliases.add(alias.asname or "open") + elif alias.name in code_exec_aliases: + code_exec_aliases.add(alias.asname or alias.name) + if node.module in _OPEN_ALIAS_MODULES: + for alias in node.names: + if alias.name == "open": + # gzip/bz2/lzma open(file, mode) writes on "w"/"a"/"x", + # mode in the 2nd arg like builtin open. + open_aliases.add(alias.asname or "open") + if node.module == "pathlib": + for alias in node.names: + if alias.name in _PATH_CTORS: + path_ctor_aliases.add(alias.asname or alias.name) + if node.module in ("os.path", "posixpath", "ntpath"): + for alias in node.names: + if alias.name == "join": + pathjoin_aliases.add(alias.asname or "join") + if node.module == "functools": + for alias in node.names: + if alias.name == "partial": + partial_aliases.add(alias.asname or "partial") + if node.module in _ARCHIVE_CTOR_MODULES: + _ctor = _ARCHIVE_CTOR_MODULES[node.module] + for alias in node.names: + if alias.name == _ctor: + archive_ctor_aliases.add(alias.asname or _ctor) + for alias in node.names: + if alias.name in _AUTO_UNSAFE_PY_WRITE_METHODS: + writer_aliases.add(alias.asname or alias.name) + # from itertools import starmap as sm / from functools import + # reduce as r: an aliased higher-order invoker. + if alias.name in _HIGHER_ORDER_INVOKERS: + invoker_aliases.add(alias.asname or alias.name) + elif isinstance(node, (ast.Assign, ast.AnnAssign)) and node.value is not None: + value = node.value + # AnnAssign (f: object = open) has a single target, no destructuring. + if isinstance(node, ast.AnnAssign): + assign_targets = [node.target] + else: + assign_targets = node.targets + targets = [t.id for t in assign_targets if isinstance(t, ast.Name)] + attr_targets = [t.attr for t in assign_targets if isinstance(t, ast.Attribute)] + if isinstance(value, ast.Name) and value.id in open_aliases: + open_aliases.update(targets) + attr_open_aliases.update(attr_targets) # box.f = open + elif isinstance(value, ast.Name) and value.id in getattr_aliases: + getattr_aliases.update(targets) # g = getattr + elif isinstance(value, ast.Name) and value.id in partial_aliases: + partial_aliases.update(targets) # p = partial + elif isinstance(value, ast.Name) and value.id in writer_aliases: + writer_aliases.update(targets) # s = save (numpy save alias) + elif isinstance(value, ast.Name) and value.id in archive_ctor_aliases: + archive_ctor_aliases.update(targets) # z = ZipFile + elif isinstance(value, ast.Name) and value.id in invoker_aliases: + invoker_aliases.update(targets) # m = map + elif isinstance(value, ast.Name) and value.id in path_ctor_aliases: + path_ctor_aliases.update(targets) # P = Path + elif isinstance(value, ast.Name) and value.id in pathjoin_aliases: + pathjoin_aliases.update(targets) # j = join + elif isinstance(value, ast.Attribute) and value.attr == "join": + pathjoin_aliases.update(targets) # j = os.path.join + elif isinstance(value, ast.Attribute) and value.attr in _PATH_CTORS: + path_ctor_aliases.update(targets) # P = pathlib.Path + elif ( + isinstance(value, ast.Attribute) + and value.attr == "open" + and isinstance(value.value, ast.Name) + and value.value.id in builtins_aliases + ): + open_aliases.update(targets) # f = builtins.open + elif ( + isinstance(value, ast.Attribute) + and value.attr in code_exec_aliases + and isinstance(value.value, ast.Name) + and value.value.id in builtins_aliases + ): + code_exec_aliases.update(targets) # e = builtins.eval + elif isinstance(value, ast.Attribute) and value.attr in _AUTO_UNSAFE_PY_WRITE_METHODS: + writer_aliases.update(targets) # s = np.save + elif isinstance(value, ast.Attribute) and value.attr == "open": + # A captured .open bound method (p = Path('out').open) opens a file + # on any call; its mode position varies (Path.open mode is 1st arg, + # builtin open's is 2nd), so fail closed on the call rather than + # guess the write mode. + dynamic_aliases.update(targets) # p = Path('out').open; p('w') + elif isinstance(value, ast.Attribute) and value.attr in _ARCHIVE_CTOR_NAMES: + archive_ctor_aliases.update(targets) # z = zipfile.ZipFile + elif isinstance(value, ast.Subscript): + dynamic_aliases.update(targets) # f = globals()["open"] + elif ( + isinstance(value, ast.Call) + and isinstance(value.func, ast.Name) + and value.func.id in getattr_aliases + ): + dynamic_aliases.update(targets) # rm = getattr(os, "remove") / g(...) + elif ( + isinstance(value, ast.Call) + and isinstance(value.func, ast.Attribute) + and value.func.attr in ("get", "pop", "setdefault") + and _is_dynamic_namespace(value.func.value) + ): + # f = __builtins__.__dict__.get("open") / globals().get("open"): + # a namespace lookup can return open/eval, so poison like getattr. + dynamic_aliases.update(targets) + elif ( + isinstance(value, ast.Call) + and ( + (isinstance(value.func, ast.Name) and value.func.id in partial_aliases) + or (isinstance(value.func, ast.Attribute) and value.func.attr == "partial") + ) + and value.args + and _wraps_write_callable(value.args[0]) + ): + dynamic_aliases.update(targets) # w = partial(open, mode="w") + elif ( + isinstance(value, ast.Call) + and ( + (isinstance(value.func, ast.Name) and value.func.id in methodcaller_aliases) + or ( + isinstance(value.func, ast.Attribute) + and value.func.attr == "methodcaller" + and isinstance(value.func.value, ast.Name) + and value.func.value.id in operator_aliases + ) + ) + and _methodcaller_writes(value) + ): + dynamic_aliases.update(targets) # w = methodcaller("write_text", ...) + elif isinstance(value, ast.Constant) and isinstance(value.value, str): + # base = '/etc' -> resolve base in a later folded path. A name + # bound more than once is poisoned (\x02) so it fails closed. + for t in targets: + literal_str_vars[t] = "\x02" if t in multi_assigned_names else value.value + elif isinstance(value, (ast.Call, ast.BinOp, ast.Name, ast.JoinedStr)): + # p = Path('/etc'); q = p; r = os.path.join('/etc','x'): record a + # fully-literal folded path so a later reuse (p / 'passwd') folds. + folded = _folded_path(value, literal_str_vars, path_ctor_aliases, pathjoin_aliases) + if folded is not None and "\x00" not in folded and "\x02" not in folded: + for t in targets: + literal_str_vars[t] = "\x02" if t in multi_assigned_names else folded + elif isinstance(value, (ast.Tuple, ast.List)): + # Destructuring binds each element like a single assignment, so an + # aliased callable (f, _ = (open, print)) AND a string / path + # literal (base, leaf = ('/etc', 'passwd')) both propagate; without + # the latter a path folded from base/leaf would miss the sensitive + # target and auto-approve. + for target in assign_targets: + if isinstance(target, (ast.Tuple, ast.List)) and len(target.elts) == len( + value.elts + ): + for tgt_el, val_el in zip(target.elts, value.elts): + if not isinstance(tgt_el, ast.Name): + continue + tid = tgt_el.id + if isinstance(val_el, ast.Name) and val_el.id in open_aliases: + open_aliases.add(tid) + elif isinstance(val_el, ast.Name) and val_el.id in getattr_aliases: + getattr_aliases.add(tid) + elif isinstance(val_el, ast.Name) and val_el.id in partial_aliases: + partial_aliases.add(tid) + elif isinstance(val_el, ast.Name) and val_el.id in writer_aliases: + writer_aliases.add(tid) # s, _ = (save, 1) + elif isinstance(val_el, ast.Name) and val_el.id in archive_ctor_aliases: + archive_ctor_aliases.add(tid) # z, _ = (ZipFile, 1) + elif isinstance(val_el, ast.Constant) and isinstance(val_el.value, str): + literal_str_vars[tid] = ( + "\x02" if tid in multi_assigned_names else val_el.value + ) + elif isinstance(val_el, (ast.Call, ast.BinOp, ast.Name, ast.JoinedStr)): + folded = _folded_path( + val_el, literal_str_vars, path_ctor_aliases, pathjoin_aliases + ) + if ( + folded is not None + and "\x00" not in folded + and "\x02" not in folded + ): + literal_str_vars[tid] = ( + "\x02" if tid in multi_assigned_names else folded + ) + elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)): + # A callable captured as a parameter default (def f(o=open): o('x','w')) + # binds that parameter to the same alias set, so a later call through + # the parameter is still gated. defaults align to the tail of + # posonlyargs+args; kw_defaults align 1:1 with kwonlyargs (None = none). + _a = node.args + _defaulted = list( + zip( + (_a.posonlyargs + _a.args)[ + len(_a.posonlyargs) + len(_a.args) - len(_a.defaults) : + ], + _a.defaults, + ) + ) + [(p, d) for p, d in zip(_a.kwonlyargs, _a.kw_defaults) if d is not None] + for _param, _default in _defaulted: + if isinstance(_default, ast.Name): + _did = _default.id + if _did in open_aliases: + open_aliases.add(_param.arg) + elif _did in writer_aliases: + writer_aliases.add(_param.arg) + elif _did in archive_ctor_aliases: + archive_ctor_aliases.add(_param.arg) + elif _did in getattr_aliases: + getattr_aliases.add(_param.arg) + elif _did in partial_aliases: + partial_aliases.add(_param.arg) + elif _did in code_exec_aliases: + code_exec_aliases.add(_param.arg) + elif _did in dynamic_aliases: + dynamic_aliases.add(_param.arg) + elif isinstance(_default, ast.Attribute): + # An attribute writer / archive ctor / captured .open used as + # a default (def f(s=np.save), def f(z=zipfile.ZipFile), + # def f(o=Path('x').open)) binds the parameter like the + # equivalent assignment; a benign attribute (np.mean) does not. + if _default.attr in _AUTO_UNSAFE_PY_WRITE_METHODS: + writer_aliases.add(_param.arg) + elif _default.attr in _ARCHIVE_CTOR_NAMES: + archive_ctor_aliases.add(_param.arg) + elif _default.attr == "open": + dynamic_aliases.add(_param.arg) + elif ( + isinstance(_default, ast.Call) + and ( + ( + isinstance(_default.func, ast.Name) + and _default.func.id in partial_aliases + ) + or ( + isinstance(_default.func, ast.Attribute) + and _default.func.attr == "partial" + ) + ) + and _default.args + and _wraps_write_callable(_default.args[0]) + ): + dynamic_aliases.add(_param.arg) # def f(w=partial(open, mode="w")) + try: + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.name.split(".")[0] in _AUTO_UNSAFE_PY_MODULES: + return True + elif isinstance(node, ast.ImportFrom): + if node.module and node.module.split(".")[0] in _AUTO_UNSAFE_PY_MODULES: + return True + # from-imports can bind mutating callables to bare names + # (from os import remove [as rm]); star imports hide anything. + for alias in node.names: + if alias.name == "*" or alias.name in _AUTO_UNSAFE_PY_ATTRS: + return True + # os.open imported as a bare callable is a low-level + # create/write, like the os.open attribute call below. + if alias.name == "open" and node.module in ("os", "posix"): + return True + elif isinstance(node, ast.Attribute): + # Any reference to a mutating attribute fails closed, even + # without an immediate call (rm = os.remove; rm("x")). + if node.attr in _AUTO_UNSAFE_PY_ATTRS: + return True + # builtins.exec / builtins.eval / builtins.__import__ (and + # compile/breakpoint) are dynamic code execution, matching the + # bare-name code_exec_aliases path; __builtins__.__import__(...) + # is a dynamic import that dodges the static import check. + if ( + node.attr in ("exec", "eval", "__import__", "breakpoint", "compile") + and isinstance(node.value, ast.Name) + and node.value.id in builtins_aliases + ): + return True + elif isinstance(node, ast.Name): + if node.id in code_exec_aliases: + return True + elif isinstance(node, ast.Constant): + # Credential paths / parent traversal in a string or bytes + # literal (open('/etc/passwd') and open(b'/etc/passwd')), or a + # glob that resolves to one (glob.glob('/e??/passwd')). + val = node.value + if isinstance(val, bytes): + val = val.decode("latin-1", "ignore") + if isinstance(val, str) and ( + _references_sensitive_path(val) or _glob_token_sensitive(val) + ): + return True + elif isinstance(node, (ast.BinOp, ast.JoinedStr)): + # A sensitive path concatenated from literals ('/etc'+'/passwd'), + # a pathlib / chain, an f-string (f'/proc/{pid}/environ'), a + # dynamic segment under a sensitive dir (f'/etc/{name}'), or one + # split through a literal variable (base = '/etc'; base+'/passwd'). + if _folded_is_sensitive( + _folded_path(node, literal_str_vars, path_ctor_aliases, pathjoin_aliases) + ): + return True + elif isinstance(node, ast.Call): + # A sensitive path composed via os.path.join('/etc', name). + if _folded_is_sensitive( + _folded_path(node, literal_str_vars, path_ctor_aliases, pathjoin_aliases) + ): + return True + func = node.func + # x.__call__(args) is just x(args): unwrap so open.__call__('o', + # 'w') / save.__call__(...) reach the open/writer checks below + # instead of looking like a harmless ".__call__" attribute call. + if isinstance(func, ast.Attribute) and func.attr == "__call__": + func = func.value + if isinstance(func, (ast.Call, ast.Subscript)): + return True # calling a call/subscript result is dynamic + # A concrete write callable (open/writer/archive-ctor alias, or a + # writer/mutating attribute) handed as an argument to any call + # escapes into a helper that can invoke it without a direct + # open()/writer site -- the same bypass the map/starmap/reduce + # branches below gate, but through a user-defined helper + # (def run(fn): fn('o','w').write('x'); run(open)). A benign + # callable argument (run(len)) is unaffected. + if any(_passed_write_callable(a) for a in node.args) or any( + _passed_write_callable(kw.value) for kw in node.keywords + ): + return True + if isinstance(func, ast.Name): + if func.id in dynamic_aliases: + return True # call through a getattr alias is dynamic + if func.id in open_aliases and _builtin_open_writes(node): + return True + # A writer imported as a bare name (from numpy import save). + if func.id in writer_aliases: + return True + # A bare archive constructor (from zipfile import ZipFile) + # takes the mode as its 2nd arg like open, so ZipFile(x, "w") + # writes but ZipFile(x) reads. + if func.id in archive_ctor_aliases and _builtin_open_writes(node): + return True + # A bare-imported logging.basicConfig(filename=...) opens a + # log file for writing (from logging import basicConfig). + if func.id in basicconfig_aliases and _basicconfig_writes(node): + return True + # A writer/open alias handed to a higher-order invoker + # (map(open, names, modes), starmap(np.save, ...), or an + # aliased m = map / sm = starmap) is called without a direct + # open(...)/save(...) site; the callable is the first + # positional arg. A benign map(len, ...) is unaffected. + if ( + func.id in invoker_aliases + and node.args + and _wraps_write_callable(node.args[0]) + ): + return True + elif isinstance(func, ast.Attribute): + # Writer methods persist to disk without open() (np.save, + # img.save, plt.savefig, df.to_csv, json.dump); ask before + # they mutate the workdir in auto mode. + if func.attr in _AUTO_UNSAFE_PY_WRITE_METHODS: + return True + # logging.basicConfig(filename=...) opens a log file for write. + if func.attr == "basicConfig" and _basicconfig_writes(node): + return True + # A qualified higher-order invoker (itertools.starmap(open, ...), + # functools.reduce(open, ...)) calls its first arg like the bare + # map/filter form; the writer-check on that arg keeps a benign + # itertools.starmap(len, ...) / df.map(transform) safe. + if ( + func.attr in _HIGHER_ORDER_INVOKERS + and node.args + and _wraps_write_callable(node.args[0]) + ): + return True + # fileinput.input(..., inplace=True) rewrites a file in place; + # the default fileinput.input(...) only reads, so gate inplace. + if ( + func.attr == "input" + and isinstance(func.value, ast.Name) + and func.value.id in fileinput_aliases + and _fileinput_inplace(node) + ): + return True + # os.open() always creates/writes a file descriptor + # (tracked through import aliases: import os as o; o.open()). + if ( + func.attr == "open" + and isinstance(func.value, ast.Name) + and func.value.id in os_aliases + ): + return True + # A pickle-backed loader (torch.load, joblib.load) can execute + # code embedded in the file it deserializes. + if ( + func.attr == "load" + and isinstance(func.value, ast.Name) + and func.value.id in load_module_aliases + ): + return True + if func.attr == "open" and _attr_open_writes(node): + return True + # An open bound onto an attribute (box.f = open; box.f('o','w')) + # writes on 'w'/'a'/'x' like the builtin, so gate the attr name. + if func.attr in attr_open_aliases and _builtin_open_writes(node): + return True + # ZipFile/TarFile/GzipFile/BZ2File/LZMAFile take the mode as + # the 2nd arg (like builtin open), so ZipFile(name, "w") writes + # but ZipFile(name) reads. + if func.attr in _ARCHIVE_CTOR_NAMES and _builtin_open_writes(node): + return True + # Enumerating a directory outside the sandbox reads host + # filenames (and enables reading their contents) the direct + # /etc/passwd checks would prompt for: Path('/etc').iterdir(), + # os.scandir('/etc'), os.listdir('/home'), os.walk('/'), + # Path('/home').glob('*'), glob.glob('/home/*'). Gate when the + # target dir folds to an absolute/tilde/sensitive path; a + # relative dir (Path('.').iterdir(), glob.glob('src/*')) stays + # safe, and an unresolved dynamic dir is left to other checks. + _enum_dir = None + if func.attr == "iterdir": + _enum_dir = func.value + elif func.attr in ("glob", "rglob", "iglob"): + # Path('/home').glob('*') enumerates the receiver dir; + # glob.glob('/home/*') enumerates the pattern's root dir. + _recv = _folded_path( + func.value, literal_str_vars, path_ctor_aliases, pathjoin_aliases + ) + if isinstance(_recv, str) and _recv not in ("", "\x00"): + _enum_dir = func.value + elif node.args: + _enum_dir = node.args[0] + elif ( + func.attr in ("scandir", "listdir", "walk") + and isinstance(func.value, ast.Name) + and func.value.id in os_aliases + and node.args + ): + _enum_dir = node.args[0] + if _enum_dir is not None: + _folded_dir = _folded_path( + _enum_dir, literal_str_vars, path_ctor_aliases, pathjoin_aliases + ) + if isinstance(_folded_dir, str) and ( + _folded_dir.startswith("/") + or _folded_dir.startswith("~") + or _folded_is_sensitive(_folded_dir) + ): + return True + except Exception: + return True # unexpected AST shape: fail closed + return False + + +# Cloud-metadata / link-local hosts (mirrors the sandbox SSRF blocklist): a +# read-named HTTP MCP tool pointed at one (fetch_url +# {"url": "http://169.254.169.254/..."}) reads instance credentials, so it asks. +_MCP_METADATA_HOST_RE = re.compile( + r"169\.254\.\d{1,3}\.\d{1,3}|" + r"100\.100\.100\.\d{1,3}|" + r"fd00:ec2::254|" + r"metadata\.google\.internal|" + r"metadata\.tencentyun\.com|" + r"://metadata(?=[:/])", + re.IGNORECASE, +) + + +def _mcp_arguments_reference_sensitive(arguments) -> bool: + """True if any string in an MCP call's arguments names a credential path, a + credential/secret environment variable (get_env {"name": "OPENAI_API_KEY"}), + or a cloud-metadata host (fetch_url {"url": "http://169.254.169.254/..."}).""" + + def walk(value) -> bool: + if isinstance(value, str): + return ( + _references_sensitive_path(value) + or bool(_AUTO_SENSITIVE_MCP_NOUN_RE.search(value)) + or bool(_MCP_METADATA_HOST_RE.search(value)) + ) + if isinstance(value, dict): + return any(walk(v) for v in value.values()) + if isinstance(value, (list, tuple)): + return any(walk(v) for v in value) + return False + + return walk(arguments) + + +# DDL object types CREATE / DROP / ALTER share (DROP FUNCTION and ALTER INDEX +# mutate just like CREATE INDEX). +_SQL_DDL_OBJECTS = ( + r"table|database|schema|index|view|function|procedure|trigger|" + r"sequence|role|user|extension|type|domain|aggregate|policy" +) +# Modifiers between the DDL verb and object (CREATE OR REPLACE VIEW, DROP +# MATERIALIZED VIEW, CREATE UNIQUE INDEX). +_SQL_DDL_MODIFIERS = ( + r"(?:(?:or\s+replace|unique|temp|temporary|global|local|materialized|recursive)\s+)*" +) +# A SQL identifier (bare, "quoted", `quoted`, [bracketed]), optionally +# schema-qualified, so UPDATE "users"/public.users/ONLY .../[users] SET all hit. +_SQL_IDENT = r'(?:\w+|"(?:[^"]|"")*"|`(?:[^`]|``)*`|\[[^\]]+\])' +_SQL_UPDATE_TARGET = r"(?:only\s+)?" + _SQL_IDENT + r"(?:\s*\.\s*" + _SQL_IDENT + r")*" +# A read-named MCP tool (query_database, run_query) can still carry a mutating +# SQL statement; match DML/DDL as whole statements (DELETE FROM, DROP TABLE) so +# a natural-language query that merely contains the word "delete" stays safe. +_MCP_ARG_MUTATION_RE = re.compile( + r"\b(?:delete\s+from|" + r"drop\s+" + _SQL_DDL_MODIFIERS + r"(?:" + _SQL_DDL_OBJECTS + r")|" + # Match the whole identifier (the outer trailing \b needs the alternative to + # end on a word boundary, so a bare \w stops mid-name and TRUNCATE users slips + # through); the optional opening quote/bracket/backtick covers "users"/[users]. + r"truncate\s+(?:table\s+)?[\"\[`]?\w+|" + # UPDATE [AS alias] SET: allow an explicit AS alias before SET so + # UPDATE users AS u SET is caught, not just the bare form. The implicit-alias + # form (UPDATE users u SET) is left out because it is indistinguishable from + # the prose "update set" and would flag natural language. + r"update\s+" + _SQL_UPDATE_TARGET + r"(?:\s+as\s+" + _SQL_IDENT + r")?\s+set\b|" + r"insert\s+into|replace\s+into|" + # SELECT ... INTO OUTFILE/DUMPFILE writes a file (MySQL); bare SELECT INTO + # is left out (PL/pgSQL uses it to read into a variable). + r"select\s+[^;]*?\binto\s+(?:outfile|dumpfile)\b|" + # ALTER SYSTEM persists PostgreSQL server configuration; SYSTEM is not one of + # the DDL objects above, so match it explicitly. + r"alter\s+system\b|" + r"alter\s+" + _SQL_DDL_MODIFIERS + r"(?:" + _SQL_DDL_OBJECTS + r")|" + r"create\s+" + _SQL_DDL_MODIFIERS + r"(?:" + _SQL_DDL_OBJECTS + r")|" + r"grant\s+\w+|revoke\s+\w+|merge\s+into|" + # Catalog mutations: COMMENT ON , SECURITY LABEL, and LOCK TABLE change + # metadata or take a lock. Each needs a following keyword, so a "comment" + # column (SELECT comment FROM t) or "locks" table stays safe. + r"comment\s+on\b|security\s+label\b|lock\s+table\b|" + # PostgreSQL maintenance writes: REFRESH MATERIALIZED VIEW rewrites the view, + # REINDEX rebuilds an index. Both need a following object keyword/name, so a + # column or word "refresh"/"reindex" in prose stays safe. + r"refresh\s+materialized\s+view|reindex\s+\w+|" + # CALL proc(...) / EXEC[UTE] name / VACUUM mutate; CALL needs a following + # "(", ";", or end so natural-language "call me back" stays safe. + r"call\s+\w+(?=\s*[(;]|\s*$)|exec(?:ute)?\s+\w+|vacuum|" + # COPY ... FROM bulk-loads and COPY ... TO writes a file ([^;] stays in one + # statement). + r"copy\s+[^;]*?\b(?:from|to)\b)\b", + re.IGNORECASE, +) +# SQLite statements the base regex misses: ATTACH/DETACH a database (DATABASE +# optional via the quoted-path form), a write-form PRAGMA (name=value / name(...), +# unlike the read-form PRAGMA name), and load_extension() which runs a shared +# library. These tokens are not natural language, so benign text does not trip. +_MCP_ARG_SQLITE_MUTATION_RE = re.compile( + r"\b(?:attach|detach)\s+database\b" + r"|\battach\s+(?:database\s+)?['\"]" + r"|\bpragma\s+\w+(?:\.\w+)?\s*(?:=|\()" + r"|\bload_extension\s*\(", + re.IGNORECASE, +) +# State-changing SQL functions that mutate or write files inside a read-shaped +# SELECT (pg_terminate_backend, setval, pg_write_file, lo_export, ...). The +# trailing "(" is required, so a column named setval_count stays safe. +_MCP_ARG_SQL_FUNCTION_RE = re.compile( + r"\b(?:pg_terminate_backend|pg_cancel_backend|pg_write_file|lo_export|" + r"lo_import|setval|nextval|set_config|pg_notify|dblink_exec|pg_reload_conf|" + r"pg_rotate_logfile|" + # advisory locks change session/transaction lock state (read-shaped SELECT). + r"pg_advisory_(?:lock|lock_shared|unlock|unlock_shared|unlock_all|" + r"xact_lock|xact_lock_shared)|" + r"pg_try_advisory_(?:lock|lock_shared|xact_lock|xact_lock_shared))\s*\(", + re.IGNORECASE, +) +# SQL engines treat /* */ and -- comments as whitespace, so DELETE/**/FROM and +# UPDATE/**/users evade the \s+ in the mutation regex; collapse comments to a +# space before matching. +_SQL_COMMENT_RE = re.compile(r"/\*.*?\*/|--[^\n]*", re.DOTALL) +# A GraphQL mutation on a read-named tool. Directives are valid between the name +# and body (mutation M @audit { ... }), so allow @directive[(args)] before ( or {. +_GRAPHQL_MUTATION_RE = re.compile( + r"\bmutation\b\s*\w*\s*(?:@\w+(?:\s*\([^)]*\))?\s*)*[({]", re.IGNORECASE +) +# GraphQL # comments run to end-of-line and count as whitespace, so a comment +# between `mutation` and the body (mutation # note\n { ... }) would otherwise +# hide it; collapse them to a space before matching. +_GRAPHQL_COMMENT_RE = re.compile(r"#[^\n]*") + + +# HTTP verbs that mutate the target resource; a generic HTTP MCP tool +# (mcp__http__get_url {"method": "DELETE"}) mutates an external service even +# though its name looks read-only. GET/HEAD/OPTIONS/TRACE only read. +_MUTATING_HTTP_METHODS = frozenset({"POST", "PUT", "PATCH", "DELETE"}) +_HTTP_METHOD_KEYS = frozenset({"method", "http_method", "httpmethod", "verb", "http_verb"}) + + +def _mcp_arguments_mutate(arguments) -> bool: + """True if an MCP call's arguments carry a mutating command, so a read-named + but write-capable tool (query_database {"query": "DELETE FROM runs"}, + query_graphql {"query": "mutation { deleteIssue(id: 1) }"}, or an HTTP tool + {"method": "DELETE"}) asks.""" + + def walk(value) -> bool: + if isinstance(value, str): + _sql = _SQL_COMMENT_RE.sub(" ", value) + return ( + bool(_MCP_ARG_MUTATION_RE.search(_sql)) + or bool(_MCP_ARG_SQLITE_MUTATION_RE.search(_sql)) + or bool(_MCP_ARG_SQL_FUNCTION_RE.search(_sql)) + or bool(_GRAPHQL_MUTATION_RE.search(_GRAPHQL_COMMENT_RE.sub(" ", value))) + ) + if isinstance(value, dict): + for k, v in value.items(): + if ( + isinstance(k, str) + and k.lower() in _HTTP_METHOD_KEYS + and isinstance(v, str) + and v.strip().upper() in _MUTATING_HTTP_METHODS + ): + return True + return any(walk(v) for v in value.values()) + if isinstance(value, (list, tuple)): + return any(walk(v) for v in value) + return False + + return walk(arguments) + + +# Tools that are read-only / non state-mutating regardless of their arguments, +# so auto mode never has to pause them (their safety needs no argument scan). +# render_html is NOT unconditionally safe: it runs arbitrary HTML/JS in the +# canvas preview frame. A static canvas (charts, layout, inline SVG) never +# reaches the network, but code that calls out can exfiltrate or fetch under the +# preview's CSP when artifact network access is enabled, so those ask; a canvas +# with no network construct still auto-runs. Matches JS egress APIs, a remote or +# root-relative ") is False + ) + assert rh("") is False + assert rh("") is False + assert rh("") is True + assert rh("") is True + assert rh("") is True + assert rh("") is True + assert rh("") is True + # Worker / SharedWorker constructors run an off-thread script the scan cannot + # see (a module worker from a CORS CDN, or a blob/same-origin worker that + # fetches/importScripts) under worker-src http: https: blob:, so they ask. + assert rh("") is True + assert rh("") is True + assert rh("") is True + assert rh("") is False # not a ctor + assert rh("") is False # unrelated class, not a real Worker + # Resource-loading forms beyond a direct fetch also reach the network. + assert rh("") is True + assert rh("") is True + assert rh("") is True + assert rh("") is True # root-relative resolves to origin + assert rh("") is True # protocol-relative + # Self-navigation sinks exfiltrate by navigating the frame away. + assert rh("") is True + assert rh("") is True + assert rh("") is True + assert rh("") is True + assert rh("") is True + assert rh("") is False # reload is not navigation + assert rh("") is False + # Obfuscated egress: a block comment splitting fetch(, or bracket access. + assert rh("") is True + assert rh("") is True + # A computed bracket key spliced from string fragments on a global host object. + assert rh("") is True + assert rh("") is True + # A computed key on a plain object (not a global host) stays a static canvas. + assert rh("") is False + assert rh("") is False # comment only + # A meta-refresh with a url navigates the frame to an external origin. + assert rh('') is True + assert rh("") is True + assert rh('') is False # self-reload, no url + assert rh('

Hi

') is False # ordinary meta stays safe + + +def test_unknown_tools_fail_closed(): + assert is_potentially_unsafe_tool_call("mystery_tool", {}) is True + + +def test_is_always_safe_tool(): + from core.inference.tools import is_always_safe_tool + for name in ("web_search", "search_knowledge_base"): + assert is_always_safe_tool(name) is True + # render_html is no longer unconditionally safe: a networked canvas can prompt, + # which cannot be judged before its arguments stream. + for name in ("python", "terminal", "mystery_tool", "mcp__srv__read", "render_html"): + assert is_always_safe_tool(name) is False + + +@pytest.mark.parametrize( + ("tool", "unsafe"), + [ + ("get_weather", False), + ("list_files", False), + ("search", False), + ("send_email", True), + ("create_issue", True), + ("delete_row", True), + ("get_or_create_issue", True), # mutating verb overrides read prefix + ("read_and_delete_file", True), + ("find_and_update_row", True), + ("get_and_commit_changes", True), # commit/save/archive are mutating + ("read_and_save_file", True), + ("list_and_archive", True), + ("list_and_clone_repo", True), # clone/checkout/comment are mutating + ("fetch_and_comment_issue", True), + ("get_and_checkout_branch", True), + ("read_and_append_file", True), # append/prepend are mutating + ("prepend_line", True), + ("get_and_upsert_row", True), # upsert/assign are mutating + ("list_and_assign_issue", True), + ("read_and_copy_file", True), # copy-style verbs create/overwrite state + ("get_and_copy_resource", True), + ("read_and_duplicate_entry", True), + ("fetch_and_download_asset", True), # download writes local state + ("list_and_export_data", True), # import/export/backup/restore/snapshot + ("get_and_snapshot_volume", True), + ("get_and_mark_read", True), # mark/subscribe change external state + ("get_and_subscribe", True), + ("list_and_unsubscribe", True), + ("get_and_reply_email", True), # reply/notify send/change external state + ("list_and_notify_users", True), + ("read_secret", True), # credential noun: a read that discloses a secret + ("list_tokens", True), + ("get_credentials", True), + ("fetch_api_key", True), # scoped *_key noun + ("read_access_key", True), + ("get_password", True), + ("read_passphrase", True), + ("read_report", False), # plain read stays safe + ("get_primary_key", False), # a schema key is not a credential + ("search_keyboard_shortcuts", False), # 'key' inside another word stays safe + ("list_bookmarks", False), # 'mark' substring in a token stays safe + ("list_notifications", False), # 'notify' is a different token than 'notifications' + ], +) +def test_mcp_classifier(tool, unsafe): + name = f"{MCP_TOOL_PREFIX}srv1__{tool}" + assert is_potentially_unsafe_tool_call(name, {}) is unsafe + + +@pytest.mark.parametrize( + ("args", "unsafe"), + [ + ({"path": "/etc/passwd"}, True), # read-named tool at a credential path + ({"path": "../../.ssh/id_rsa"}, True), + ({"nested": {"file": "~/.aws/credentials"}}, True), + ({"name": "OPENAI_API_KEY"}, True), # explicit credential env-var read + ({"name": "AWS_SECRET_ACCESS_KEY"}, True), + ({"key": "DATABASE_PASSWORD"}, True), + ( + {"url": "http://169.254.169.254/latest/meta-data/iam/security-credentials/"}, + True, + ), # AWS instance-metadata host + ( + {"url": "http://metadata.google.internal/computeMetadata/v1/"}, + True, + ), # GCP metadata host + ({"path": "notes.txt"}, False), # ordinary path stays safe + ({"path": "data/report.csv"}, False), + ({"name": "PATH"}, False), # a non-secret env var stays safe + ({"name": "HOME"}, False), + ({"url": "https://example.com/api"}, False), # ordinary URL stays safe + ({"url": "http://localhost:8080/health"}, False), # localhost app stays safe + ], +) +def test_mcp_sensitive_arguments(args, unsafe): + name = f"{MCP_TOOL_PREFIX}fs__read_file" + assert is_potentially_unsafe_tool_call(name, args) is unsafe + + +@pytest.mark.parametrize( + ("args", "unsafe"), + [ + ({"query": "DELETE FROM runs"}, True), # read-named tool, mutating query + ({"sql": "DROP TABLE users"}, True), + ({"query": "UPDATE t SET x=1"}, True), + ({"query": "INSERT INTO t VALUES (1)"}, True), + ({"query": "SELECT * FROM runs"}, False), # read query stays safe + ({"query": "how to delete old files"}, False), # NL text with 'delete' stays safe + ({"query": "find the created_at column"}, False), # 'created' substring stays safe + ({"query": "DELETE/**/FROM runs"}, True), # inline SQL comment as whitespace + ({"query": "UPDATE/**/t SET x=1"}, True), + ({"query": "DROP/**/TABLE users"}, True), + ({"query": "SELECT * FROM runs -- delete later"}, False), # trailing comment stays safe + ({"query": "COPY users FROM '/tmp/u.csv'"}, True), # bulk load writes the table + ({"query": "COPY users (id, name)\nFROM STDIN"}, True), # multiline COPY FROM + ({"query": "COPY (SELECT 1) TO '/tmp/o.csv'"}, True), # COPY TO writes a server file + ({"query": "SELECT copy_count FROM t"}, False), # 'copy' substring column stays safe + ({"query": "mutation { deleteIssue(id: 1) }"}, True), # GraphQL mutation + ({"query": "mutation DelIssue { deleteIssue(id: 1) }"}, True), # named GraphQL mutation + ({"query": "mutation # note\n { deleteIssue(id: 1) }"}, True), # comment before body + ({"query": "mutation # c\n Del { deleteIssue(id: 1) }"}, True), # comment before name + ({"query": "query { issue(id: 1) { title } }"}, False), # GraphQL read query stays safe + ({"query": "{ issue(id: 1) { title } }"}, False), # shorthand GraphQL query stays safe + ({"query": "query # note\n { issue(id: 1) }"}, False), # commented read query stays safe + ({"query": "CREATE OR REPLACE VIEW v AS SELECT 1"}, True), # DDL with a modifier + ({"query": "CREATE UNIQUE INDEX idx ON t(x)"}, True), # DDL with UNIQUE + ({"query": "CREATE TEMP TABLE t (id int)"}, True), # DDL with TEMP + ({"query": "CREATE MATERIALIZED VIEW mv AS SELECT 1"}, True), # materialized view DDL + ({"query": "CREATE FUNCTION f() RETURNS int AS $$ $$"}, True), # function DDL + ({"query": "ALTER SYSTEM SET work_mem = '1GB'"}, True), # persists server config + ({"query": "alter system reset all"}, True), # ALTER SYSTEM RESET + ({"query": "SELECT * FROM system_logs"}, False), # 'system' as a table name stays safe + ({"query": "SELECT * FROM created_view"}, False), # 'create' substring stays safe + ({"query": "CALL delete_all_users()"}, True), # stored procedure invocation + ({"query": "EXEC purge_queue"}, True), # EXEC procedure + ({"query": "EXECUTE sp_drop"}, True), # EXECUTE procedure + ({"query": "VACUUM INTO 'backup.db'"}, True), # VACUUM rewrites the database + ({"query": "please call me back later"}, False), # NL 'call' stays safe + ({"query": "ATTACH DATABASE '/tmp/x.db' AS x"}, True), # attaches a database file + ({"query": "DETACH DATABASE x"}, True), # detaches a database + ({"query": "PRAGMA user_version = 42"}, True), # write-form PRAGMA + ({"query": "PRAGMA journal_mode=WAL"}, True), # write-form PRAGMA (no spaces) + ({"query": "PRAGMA foreign_keys(0)"}, True), # call-form PRAGMA write + ({"query": "SELECT load_extension('/tmp/evil.so')"}, True), # loads native code + ({"query": "PRAGMA journal_mode"}, False), # read-form PRAGMA stays safe + ({"query": "can you attach the report to the email"}, False), # NL 'attach' stays safe + ({"query": "ATTACH '/tmp/x.db' AS x"}, True), # ATTACH without DATABASE keyword + ({"query": "PRAGMA main.user_version = 1"}, True), # schema-qualified write PRAGMA + ({"query": "attach it as draft"}, False), # NL 'attach ... as' stays safe + ({"query": "DROP FUNCTION f()"}, True), # DROP of a non-table object + ({"query": "ALTER INDEX idx RENAME TO idx2"}, True), # ALTER of a non-table object + ({"query": "DROP MATERIALIZED VIEW mv"}, True), # DROP with a modifier + ({"query": "ALTER USER bob WITH PASSWORD 'x'"}, True), # ALTER USER mutates + ({"query": "SELECT dropped_at FROM t"}, False), # 'drop' substring column stays safe + ({"query": "mutation M @audit { deleteIssue(id: 1) }"}, True), # directive GraphQL mutation + ( + {"query": "query Q @cached { issue(id: 1) { title } }"}, + False, + ), # directive GraphQL read stays safe + ({"query": 'UPDATE "users" SET admin=1'}, True), # double-quoted UPDATE target + ({"query": "UPDATE public.users SET admin=1"}, True), # schema-qualified UPDATE + ({"query": "UPDATE ONLY public.users SET admin=1"}, True), # ONLY-qualified UPDATE + ({"query": "UPDATE `users` SET admin=1"}, True), # backtick-quoted UPDATE + ({"query": "UPDATE [users] SET admin=1"}, True), # bracket-quoted UPDATE + ({"query": "please update the documentation set"}, False), # NL 'update ... set' stays safe + ({"query": "SELECT pg_terminate_backend(123)"}, True), # state-changing SQL function + ({"query": "SELECT setval('s', 1)"}, True), # sequence mutation function + ({"query": "SELECT pg_write_file('/tmp/p', 'x')"}, True), # server-side file write + ({"query": "SELECT lo_export(123, '/tmp/p')"}, True), # large-object export to a file + ({"query": "SELECT setval_col FROM t"}, False), # 'setval' column prefix stays safe + ( + {"query": "SELECT secret INTO OUTFILE '/tmp/leak' FROM users"}, + True, + ), # INTO OUTFILE write + ({"query": "SELECT x INTO DUMPFILE '/tmp/d' FROM t"}, True), # INTO DUMPFILE write + ( + {"query": "SELECT count(*) INTO cnt FROM t"}, + False, + ), # PL/pgSQL SELECT INTO var stays safe + ({"query": "REFRESH MATERIALIZED VIEW mv"}, True), # materialized view rewrite + ({"query": "REINDEX INDEX idx"}, True), # index rebuild + ({"query": "REINDEX TABLE t"}, True), # table reindex + ({"query": "SELECT refresh_count FROM t"}, False), # 'refresh' column stays safe + ({"query": "please refresh the page"}, False), # NL 'refresh' stays safe + ({"query": "COMMENT ON TABLE users IS 'owned'"}, True), # catalog metadata write + ({"query": "LOCK TABLE users IN ACCESS EXCLUSIVE MODE"}, True), # explicit lock + ({"query": "SECURITY LABEL FOR x ON TABLE t IS 'z'"}, True), # security label write + ({"query": "CREATE POLICY p ON accounts USING (true)"}, True), # row-security policy DDL + ({"query": "SELECT comment FROM t"}, False), # 'comment' column stays safe + ({"query": "SELECT * FROM locks"}, False), # 'locks' table stays safe + ({"query": "SELECT nextval('billing_seq')"}, True), # sequence advance mutates + ({"query": "SELECT pg_advisory_lock(42)"}, True), # advisory lock changes state + ({"query": "SELECT pg_notify('jobs', 'wake')"}, True), # server-side notification + ({"query": "SELECT set_config('x', 'y', false)"}, True), # session config write + ({"query": "SELECT nextval_col FROM t"}, False), # 'nextval' column prefix stays safe + ({"query": "TRUNCATE users"}, True), # multi-char table name (bare TRUNCATE) + ({"query": "TRUNCATE TABLE accounts"}, True), # multi-char TRUNCATE TABLE + ({"query": 'TRUNCATE TABLE "users"'}, True), # quoted TRUNCATE target + ({"query": "TRUNCATE accounts RESTART IDENTITY"}, True), # TRUNCATE with options + ({"query": "SELECT truncate_log FROM t"}, False), # 'truncate' column stays safe + ({"query": "UPDATE users AS u SET admin=1"}, True), # aliased UPDATE target (AS) + ({"query": 'UPDATE "users" AS u SET x=1'}, True), # quoted+aliased UPDATE + ({"query": "UPDATE public.users AS u SET x=1"}, True), # schema-qualified aliased UPDATE + ({"query": "SELECT * FROM users AS u"}, False), # aliased SELECT stays safe + ({"query": "please update the documentation set"}, False), # NL, no AS, stays safe + ({"query": "GRANT SELECT ON t TO u"}, True), # privilege grant (multi-word) + ({"query": "REVOKE ALL ON t FROM u"}, True), # privilege revoke (multi-word) + ({"query": "SELECT * FROM grants"}, False), # 'grants' table stays safe + ({"url": "http://x", "method": "DELETE"}, True), # mutating HTTP verb arg + ({"method": "POST"}, True), + ({"verb": "PUT"}, True), # alternate method-key name + ({"method": "GET"}, False), # read HTTP verb stays safe + ({"method": "HEAD"}, False), + ], +) +def test_mcp_mutating_arguments(args, unsafe): + name = f"{MCP_TOOL_PREFIX}db__query_database" + assert is_potentially_unsafe_tool_call(name, args) is unsafe + + +# ── loop behavior ─────────────────────────────────────────────────── + +_DEFAULT_TOOLS = [ + {"type": "function", "function": {"name": "python"}}, + {"type": "function", "function": {"name": "web_search"}}, +] + + +class _FakeExecuteTool: + def __init__(self): + self.calls = [] + self.disable_sandbox_seen = [] + + def __call__( + self, + name, + arguments, + *, + cancel_event = None, + timeout = None, + session_id = None, + thread_id = None, + rag_scope = None, + disable_sandbox = False, + ): + self.calls.append((name, arguments)) + self.disable_sandbox_seen.append(disable_sandbox) + return f"RESULT[{name}]" + + +def _tool_call(name, args_json): + return f'{{"name": "{name}", "arguments": {args_json}}}' + + +def _multi_turn(turns): + turn_iter = iter(turns) + + def _gen(_messages): + try: + yield next(turn_iter) + except StopIteration: + return + + return _gen + + +def _drive(turns, decisions, **loop_kwargs): + """Run the loop, resolving each gated tool_start with the next decision.""" + decision_iter = iter(decisions) + exec_fn = _FakeExecuteTool() + # A per-call session id so a leaked pending approval from another test can + # never collide with this run's approval registry entries. + session = f"{_SESSION}-{uuid.uuid4().hex}" + gen = run_safetensors_tool_loop( + single_turn = _multi_turn(turns), + messages = [{"role": "user", "content": "hi"}], + tools = _DEFAULT_TOOLS, + execute_tool = exec_fn, + session_id = session, + **loop_kwargs, + ) + events = [] + for ev in gen: + events.append(ev) + if ev["type"] == "tool_start" and ev.get("awaiting_confirmation"): + resolve_tool_decision(ev["approval_id"], next(decision_iter), session_id = session) + return events, exec_fn + + +def _tool_starts(events): + return [e for e in events if e["type"] == "tool_start"] + + +def _diag(events, exec_fn): + """A compact dump of what the loop actually did, attached to the loop-driving + assertions so a full-suite-only failure on CI (which does not reproduce when + the file runs alone) reports the real event stream instead of a bare diff.""" + return ( + f"calls={exec_fn.calls} sandbox_seen={exec_fn.disable_sandbox_seen} " + f"events={[(e.get('type'), e.get('awaiting_confirmation'), e.get('tool_name')) for e in events]}" + ) + + +def test_auto_mode_does_not_gate_safe_calls(): + events, exec_fn = _drive( + [_tool_call("python", '{"code": "print(1)"}'), "final"], + [], + confirm_tool_calls = True, + permission_mode = "auto", + ) + starts = _tool_starts(events) + assert starts and starts[0]["awaiting_confirmation"] is False, _diag(events, exec_fn) + assert starts[0]["approval_id"] == "" + assert exec_fn.calls == [("python", {"code": "print(1)"})], _diag(events, exec_fn) + assert exec_fn.disable_sandbox_seen == [False], _diag( + events, exec_fn + ) # sandbox stays on in auto + + +def test_auto_mode_gates_unsafe_calls(): + events, exec_fn = _drive( + [_tool_call("python", '{"code": "import os; os.remove(\\"x\\")"}'), "final"], + ["allow"], + confirm_tool_calls = True, + permission_mode = "auto", + ) + starts = _tool_starts(events) + assert starts and starts[0]["awaiting_confirmation"] is True, _diag(events, exec_fn) + assert starts[0]["approval_id"] + assert len(exec_fn.calls) == 1, _diag(events, exec_fn) + assert exec_fn.disable_sandbox_seen == [False], _diag(events, exec_fn) + + +def test_ask_mode_gates_even_safe_calls(): + events, _ = _drive( + [_tool_call("python", '{"code": "print(1)"}'), "final"], + ["allow"], + confirm_tool_calls = True, + permission_mode = "ask", + ) + starts = _tool_starts(events) + assert starts and starts[0]["awaiting_confirmation"] is True + + +def test_unset_mode_behaves_as_ask(): + events, _ = _drive( + [_tool_call("python", '{"code": "print(1)"}'), "final"], + ["allow"], + confirm_tool_calls = True, + ) + starts = _tool_starts(events) + assert starts and starts[0]["awaiting_confirmation"] is True + + +def test_off_mode_never_gates_and_keeps_sandbox(): + # "Off": no prompts even for unsafe calls, but the sandbox stays on. + events, exec_fn = _drive( + [_tool_call("python", '{"code": "import os; os.remove(\\"x\\")"}'), "final"], + [], + confirm_tool_calls = True, # off must win over a stray confirm flag + permission_mode = "off", + ) + starts = _tool_starts(events) + assert starts and starts[0]["awaiting_confirmation"] is False, _diag(events, exec_fn) + assert starts[0]["approval_id"] == "" + assert exec_fn.disable_sandbox_seen == [False], _diag(events, exec_fn) + + +def test_full_mode_never_gates_and_drops_sandbox(): + events, exec_fn = _drive( + [_tool_call("python", '{"code": "import os; os.remove(\\"x\\")"}'), "final"], + [], + confirm_tool_calls = True, # full must win over the confirm gate + permission_mode = "full", + ) + starts = _tool_starts(events) + assert starts and starts[0]["awaiting_confirmation"] is False, _diag(events, exec_fn) + assert exec_fn.disable_sandbox_seen == [True], _diag(events, exec_fn) + + +def test_bypass_flag_implies_full_mode(): + # Legacy callers that only set bypass_permissions keep the same behavior. + events, exec_fn = _drive( + [_tool_call("python", '{"code": "print(1)"}'), "final"], + [], + confirm_tool_calls = True, + bypass_permissions = True, + ) + starts = _tool_starts(events) + assert starts and starts[0]["awaiting_confirmation"] is False, _diag(events, exec_fn) + assert exec_fn.disable_sandbox_seen == [True], _diag(events, exec_fn) + + +def test_bypass_permissions_folds_to_full_on_request_models(): + # A legacy bypass caller that also sends a stale ask/auto mode normalizes to + # full, so the route guards (which reject ask/auto) don't 400 the request. + for cls in (ChatCompletionRequest, AnthropicMessagesRequest): + req = cls( + messages = [{"role": "user", "content": "hi"}], + bypass_permissions = True, + permission_mode = "auto", + ) + assert req.permission_mode == "full" + assert req.bypass_permissions is True + + +def test_unknown_permission_mode_normalizes_to_ask_on_request_models(): + # An unrecognized mode from a newer UI/client must degrade to the safest gate + # ("ask") at the API boundary instead of a 422, so the forward-compat fallback + # the tool loops already apply (unknown -> ask) is reachable. None stays unset; + # the four known modes pass through untouched. + for cls in (ChatCompletionRequest, AnthropicMessagesRequest): + for unknown in ("paranoid", "readonly", "bogus", ""): + req = cls( + messages = [{"role": "user", "content": "hi"}], + permission_mode = unknown, + ) + assert req.permission_mode == "ask", (cls.__name__, unknown) + assert ( + cls(messages = [{"role": "user", "content": "hi"}], permission_mode = None).permission_mode + is None + ) + for known in ("ask", "auto", "off", "full"): + req = cls( + messages = [{"role": "user", "content": "hi"}], + permission_mode = known, + ) + # 'full' folds to bypass but the mode string is preserved. + assert req.permission_mode == known, (cls.__name__, known) + + +def test_ask_auto_self_enable_confirm_on_chat_request(): + # "Ask" gates every call, so a direct /chat/completions caller that requests + # ask but omits the legacy confirm flag self-enables it when Studio's own tool + # loop is requested. Only the router's loop-entry signals count (enable_tools / + # mcp_enabled); enabled_tools alone never starts the loop. + for loop in ({"enable_tools": True}, {"mcp_enabled": True}): + req = ChatCompletionRequest( + messages = [{"role": "user", "content": "hi"}], + permission_mode = "ask", + **loop, + ) + assert req.confirm_tool_calls is True + # "auto" is NOT folded: it only prompts for a classifier-flagged call, so + # leaving confirm unset lets the route apply the safe-only-selection exception + # (a safe-only auto request needs no stream) instead of an explicit confirm + # forcing stream=true. The mode still drives the loop's per-call gate. + for loop in ({"enable_tools": True}, {"mcp_enabled": True}): + req = ChatCompletionRequest( + messages = [{"role": "user", "content": "hi"}], + permission_mode = "auto", + **loop, + ) + assert req.confirm_tool_calls is None + # enabled_tools by itself is a passthrough filter, not a loop-entry signal: + # a client-tool passthrough that also lists enabled_tools must route verbatim + # (confirm stays unset), else the confirm-without-stream guard 400s it. + for mode in ("ask", "auto"): + req = ChatCompletionRequest( + messages = [{"role": "user", "content": "hi"}], + permission_mode = mode, + enabled_tools = ["terminal"], + tools = [{"type": "function", "function": {"name": "f"}}], + ) + assert req.confirm_tool_calls is None + # An explicit confirm_tool_calls=False wins over the ask mode (opts out of the + # gate), matching _permission_mode_confirm and the Anthropic pre-switch guard; + # the fold only self-enables when the flag is unset, so a caller cannot get a + # different answer on the chat path than the Anthropic path for the same body. + req = ChatCompletionRequest( + messages = [{"role": "user", "content": "hi"}], + permission_mode = "ask", + enable_tools = True, + confirm_tool_calls = False, + ) + assert req.confirm_tool_calls is False + # A plain client-tool passthrough (client-supplied tools that Studio does not + # execute) must NOT self-enable confirm, or the route rejects the passthrough. + req = ChatCompletionRequest( + messages = [{"role": "user", "content": "hi"}], + permission_mode = "ask", + tools = [{"type": "function", "function": {"name": "f"}}], + ) + assert req.confirm_tool_calls is None + # ask/auto without any tool request has nothing to gate; confirm stays unset. + req = ChatCompletionRequest( + messages = [{"role": "user", "content": "hi"}], + permission_mode = "ask", + ) + assert req.confirm_tool_calls is None + # Legacy callers with no permission_mode keep their confirm flag untouched. + req = ChatCompletionRequest( + messages = [{"role": "user", "content": "hi"}], + confirm_tool_calls = False, + ) + assert req.confirm_tool_calls is False + # External-provider requests are not folded (the provider branch rejects + # confirm_tool_calls with tools, and permission_mode is a local concept). + for extra in ({"provider_id": "p1"}, {"provider_type": "openai"}): + req = ChatCompletionRequest( + messages = [{"role": "user", "content": "hi"}], + permission_mode = "ask", + enable_tools = True, + **extra, + ) + assert req.confirm_tool_calls is None + + +def test_permission_mode_confirm_derivation(): + # The route derives the effective confirm gate from permission_mode so that a + # tool loop forced on by CLI policy (no request-level tool flag) still honors + # the documented "unset behaves as ask" default. + from routes.inference import _permission_mode_confirm + + def req(**kw): + return ChatCompletionRequest(messages = [{"role": "user", "content": "hi"}], **kw) + + # An explicit confirm flag always wins (True gates, False opts out). + assert _permission_mode_confirm(req(confirm_tool_calls = True, stream = False)) is True + assert _permission_mode_confirm(req(confirm_tool_calls = False, permission_mode = "ask")) is False + # Explicit ask/auto always engage the gate (a non-streaming one is rejected + # by the guard that reads this). + assert _permission_mode_confirm(req(permission_mode = "ask", stream = False)) is True + assert _permission_mode_confirm(req(permission_mode = "auto", stream = False)) is True + # off/full never prompt. + assert _permission_mode_confirm(req(permission_mode = "off")) is False + assert _permission_mode_confirm(req(permission_mode = "full")) is False + # An unset mode defaults to ask, but only realizably on a streaming request; + # a non-streaming unset request keeps the legacy run-without-gate behavior. + assert _permission_mode_confirm(req(stream = True)) is True + assert _permission_mode_confirm(req(stream = False)) is False + + +def test_confirm_gate_needs_stream(): + # auto only prompts for a classifier-flagged call, so an auto request that can + # only select always-safe tools (web_search / RAG) needs no stream and must not + # be rejected by the confirm-without-stream guard. + from routes.inference import _confirm_gate_needs_stream + + def req(**kw): + return ChatCompletionRequest(messages = [{"role": "user", "content": "hi"}], **kw) + + safe = ["web_search", "search_knowledge_base"] + # auto + a safe-only selection never prompts -> no stream needed. + assert _confirm_gate_needs_stream(req(permission_mode = "auto", enabled_tools = safe)) is False + assert ( + _confirm_gate_needs_stream(req(permission_mode = "auto", enabled_tools = ["web_search"])) + is False + ) + # render_html can prompt when its canvas reaches the network, so a selection + # that includes it needs a stream to deliver that prompt. + assert ( + _confirm_gate_needs_stream( + req(permission_mode = "auto", enabled_tools = ["web_search", "render_html"]) + ) + is True + ) + # But a selectable unsafe tool, an unrestricted (omitted) selection, MCP, or an + # explicit confirm flag all still require streaming under auto. + assert ( + _confirm_gate_needs_stream(req(permission_mode = "auto", enabled_tools = ["terminal"])) is True + ) + assert _confirm_gate_needs_stream(req(permission_mode = "auto", enable_tools = True)) is True + assert ( + _confirm_gate_needs_stream( + req(permission_mode = "auto", enabled_tools = ["web_search"], mcp_enabled = True) + ) + is True + ) + assert ( + _confirm_gate_needs_stream( + req(permission_mode = "auto", enabled_tools = ["web_search"], confirm_tool_calls = True) + ) + is True + ) + # An explicit empty selection runs no built-in tool, so nothing can prompt and + # no stream is needed (distinct from an omitted list, which means all tools). + assert ( + _confirm_gate_needs_stream(req(permission_mode = "auto", enable_tools = True, enabled_tools = [])) + is False + ) + # ask prompts for every call, so even a safe-only selection needs streaming. + assert _confirm_gate_needs_stream(req(permission_mode = "ask", enabled_tools = safe)) is True + # off/full never prompt; unset non-streaming keeps the legacy run-without-gate. + assert _confirm_gate_needs_stream(req(permission_mode = "off", enabled_tools = safe)) is False + assert _confirm_gate_needs_stream(req(permission_mode = "full", enabled_tools = safe)) is False + assert _confirm_gate_needs_stream(req(enabled_tools = safe, stream = False)) is False diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 63fdbbd8e9..eae1a75161 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -2592,6 +2592,50 @@ class TestLoopBasic: assert tool_starts[0]["arguments"] == {} assert "" in tool_starts[1]["arguments"]["code"] + def test_render_html_auto_mode_static_runs_without_prompt(self): + """permission_mode="auto" ships confirm_tool_calls=true. render_html is no + longer unconditionally safe (a networked canvas must ask), so its early + provisional card is suppressed under the confirm gate; a static canvas is + still classified safe and runs without an approval prompt.""" + exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) + turn_iter = iter( + [ + [ + "", + "", + "Hi", + ], + ["Done."], + ] + ) + + def _gen(_messages): + chunks = next(turn_iter) + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "make html"}], + tools = [{"type": "function", "function": {"name": "render_html"}}], + execute_tool = exec_fn, + confirm_tool_calls = True, + permission_mode = "auto", + session_id = "sess", + max_tool_iterations = 3, + ) + events = _collect_events(loop) + tool_starts = [e for e in events if e["type"] == "tool_start"] + + # No early provisional card under the auto confirm gate; just the real call. + assert len(tool_starts) == 1 + assert tool_starts[0]["tool_name"] == "render_html" + assert "" in tool_starts[0]["arguments"]["code"] + # A static canvas is classified safe, so it runs without an approval gate. + assert tool_starts[0].get("awaiting_confirmation") in (False, None) + def test_render_html_provisional_card_closed_on_generator_exception(self): """If the model generator raises mid-stream after a provisional render_html card was surfaced, the loop must close that card as errored before the @@ -3674,6 +3718,26 @@ class TestGuardrails: assert any(e.get("type") == "content" and e.get("text") == "plain answer" for e in events) assert exec_fn.calls == [] + def test_auto_mode_still_runs_rag_autoinject(self, monkeypatch): + # "auto" sends confirm_tool_calls=true so unsafe calls gate, but the + # safe search_knowledge_base retrieval never gates, so autoinject must + # still run (unlike ask mode above). + ran = {"called": False} + + def fake_autoinject(*_args, **_kwargs): + ran["called"] = True + return None + + monkeypatch.setattr("core.inference.tools.build_rag_autoinject", fake_autoinject) + loop, _exec_fn = _make_loop( + turns = [["plain answer"]], + confirm_tool_calls = True, + permission_mode = "auto", + rag_scope = {"thread_id": "t1"}, + ) + _collect_events(loop) + assert ran["called"] is True + def test_auto_heal_disabled_preserves_xml_on_final_no_tools_pass(self): turns = iter( [ diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 5b6264c6d7..9b502a5000 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -79,6 +79,7 @@ import { McpComposerButton } from "@/features/chat/mcp-composer-button"; import { getExternalReasoningCapabilities } from "@/features/chat/provider-capabilities"; import { useRagToolDisabled } from "@/features/chat/hooks/use-rag-tool-disabled"; import { BypassPermissionsMenuItem } from "@/features/chat/bypass-permissions-menu-item"; +import { PermissionModeComposerPill } from "@/features/chat/permission-mode-select"; import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; import { useExternalProvidersStore } from "@/features/chat/stores/external-providers-store"; import { PROMPT_QUEUE_STOP_EVENT } from "@/features/chat/utils/prompt-queue-boundary"; @@ -131,7 +132,6 @@ import { Image03Icon, McpServerIcon, PencilRulerIcon, - ShieldBanIcon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { useNavigate } from "@tanstack/react-router"; @@ -1428,11 +1428,14 @@ const Composer: FC<{ const artifactsEnabled = useChatRuntimeStore((s) => s.artifactsEnabled); const mcpEnabledForChat = useChatRuntimeStore((s) => s.mcpEnabledForChat); const ragEnabled = useChatRuntimeStore((s) => s.ragEnabled); + const permissionMode = useChatRuntimeStore((s) => s.permissionMode); const bypassPermissions = useChatRuntimeStore((s) => s.bypassPermissions); - // More than 4 pills: collapse to icons only. Search and Code always show; + // More than 4 pills: collapse to icons only. Search and Code always show; the + // permission pill shows in every mode except "off" (it renders null there); // Images, RAG, Canvas and MCP are conditional. const pillsCompact = 2 + + (permissionMode !== "off" ? 1 : 0) + (ragEnabled ? 1 : 0) + (supportsBuiltinImageGeneration ? 1 : 0) + (artifactsEnabled ? 1 : 0) + @@ -1856,9 +1859,9 @@ const Composer: FC<{ data-pill-compact={pillsCompact ? "true" : undefined} > - {/* Active-mode badge: always visible when bypass is on, even while - the pill row is collapsed (returns null when off). */} - + {/* Permission-level pill: always visible, even while the pill row + is collapsed; opens the permission level dropdown. */} + {composerExpanded ? ( <> @@ -2620,36 +2623,6 @@ const ArtifactsToggle: FC = () => { ); }; -// Claude gold pill shown while Bypass permissions is on; click to turn it off. -// Mirror of shared-composer's badge so both composers surface the state. -const BypassPermissionsToggle: FC = () => { - const bypassPermissions = useChatRuntimeStore((s) => s.bypassPermissions); - const setBypassPermissions = useChatRuntimeStore( - (s) => s.setBypassPermissions, - ); - if (!bypassPermissions) return null; - return ( - - ); -}; - const ToolStatusDisplay: FC = () => { const toolStatus = useChatRuntimeStore((s) => s.toolStatus); const isThreadRunning = useAuiState(({ thread }) => thread.isRunning); diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 12c1c3b385..c7dd6372aa 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -173,6 +173,7 @@ interface ResponseDetailsMetadata { artifacts: boolean; confirmToolCalls: boolean; bypassPermissions: boolean; + permissionMode?: string; }; } @@ -1951,6 +1952,7 @@ export function createOpenAIStreamAdapter( mcpEnabledForChat, confirmToolCalls, bypassPermissions, + permissionMode, webFetchToolsEnabled, ragEnabled, ragSource, @@ -2642,6 +2644,7 @@ export function createOpenAIStreamAdapter( artifacts: renderHtmlToolEnabledForThisTurn, confirmToolCalls, bypassPermissions, + permissionMode, }, }); const externalCapabilities = getProviderCapabilities( @@ -2953,6 +2956,16 @@ export function createOpenAIStreamAdapter( ...(supportsPreserveThinking ? { preserve_thinking: preserveThinking } : {}), + // Permission level for local tool calls is sent for every local + // chat, not only when a tool pill is on: a process policy + // (unsloth run --enable-tools) can open the tool loop with no pill, + // and the backend must still see the selected gate. ask/auto request + // the confirm gate ("auto" only pauses calls flagged unsafe); off + // and full never prompt, full also drops the sandbox. + permission_mode: permissionMode, + confirm_tool_calls: + permissionMode === "ask" || permissionMode === "auto", + bypass_permissions: bypassPermissions, ...(supportsTools && (toolsEnabled || codeToolsEnabled || @@ -2974,10 +2987,6 @@ export function createOpenAIStreamAdapter( : []), ], mcp_enabled: mcpEnabledForChat, - // Bypass Permissions wins: never request the confirm gate - // while bypassing, and tell the backend to drop the sandbox. - confirm_tool_calls: confirmToolCalls && !bypassPermissions, - bypass_permissions: bypassPermissions, // Scope: thread_id = this thread's docs, kb_id = a KB, // project_id = the thread's project sources (auto-on whenever // the project has indexed sources, no Docs pill needed). diff --git a/studio/frontend/src/features/chat/bypass-permissions-menu-item.tsx b/studio/frontend/src/features/chat/bypass-permissions-menu-item.tsx index 14cb6747e9..b35317b2fa 100644 --- a/studio/frontend/src/features/chat/bypass-permissions-menu-item.tsx +++ b/studio/frontend/src/features/chat/bypass-permissions-menu-item.tsx @@ -14,45 +14,49 @@ import { AlertDialogHeader, AlertDialogTitle, } from "@/components/ui/alert-dialog"; -import { DropdownMenuItem } from "@/components/ui/dropdown-menu"; +import { + DropdownMenuSub, + DropdownMenuSubContent, + DropdownMenuSubTrigger, +} from "@/components/ui/dropdown-menu"; import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; -import { Tick02Icon } from "@/lib/tick-icon"; +import { PermissionModeMenuItems } from "./permission-mode-select"; -// "Bypass permissions" entry for the composer "+" -> More menu. Mirrors the -// settings toggle: enabling demands the danger warning, disabling is immediate. -// The menu closes normally on select (no preventDefault) -- the warning dialog -// lives outside the menu (BypassPermissionsConfirmDialog, mounted once at the -// chat-page root and driven by the store), so it survives the menu unmounting -// and the "+"/More popovers don't stay frozen. +// "Bypass permissions" entry for the composer "+" -> More menu. Like the MCP +// pill, it opens a submenu where the user picks the permission level (Ask for +// approval / Approve for me / Full access). Picking Full access demands the +// danger warning; the other levels apply immediately. The menu closes normally +// on select (no preventDefault) -- the warning dialog lives outside the menu +// (BypassPermissionsConfirmDialog, mounted once at the chat-page root and +// driven by the store), so it survives the menu unmounting and the "+"/More +// popovers don't stay frozen. export function BypassPermissionsMenuItem() { - const bypassPermissions = useChatRuntimeStore((s) => s.bypassPermissions); - const setBypassPermissions = useChatRuntimeStore( - (s) => s.setBypassPermissions, - ); + const permissionMode = useChatRuntimeStore((s) => s.permissionMode); const setBypassConfirmOpen = useChatRuntimeStore( (s) => s.setBypassConfirmOpen, ); return ( - { - if (bypassPermissions) { - setBypassPermissions(false); - } else { - // Defer past Radix's menu-close focus restoration: opening the dialog - // synchronously here lets the dropdown grab focus back and breaks the - // dialog's focus trap. - setTimeout(() => setBypassConfirmOpen(true), 0); + + - - Bypass permissions - {bypassPermissions ? ( - - ) : null} - + > + + Bypass permissions + + + + setTimeout(() => setBypassConfirmOpen(true), 0) + } + /> + + ); } @@ -63,19 +67,17 @@ export function BypassPermissionsMenuItem() { export function BypassPermissionsConfirmDialog() { const open = useChatRuntimeStore((s) => s.bypassConfirmOpen); const setOpen = useChatRuntimeStore((s) => s.setBypassConfirmOpen); - const setBypassPermissions = useChatRuntimeStore( - (s) => s.setBypassPermissions, - ); + const setPermissionMode = useChatRuntimeStore((s) => s.setPermissionMode); return ( - Enable Bypass permissions? + Enable Full access? - Bypass permissions is dangerous since the AI model might delete, - corrupt your machine, and or cause real world damage to you or the - world - only accept if you are certain + Full access (Bypass permissions) is dangerous since the AI model + might delete, corrupt your machine, and or cause real world damage + to you or the world - only accept if you are certain @@ -84,7 +86,7 @@ export function BypassPermissionsConfirmDialog() { variant="destructive" className="!bg-destructive !text-destructive-foreground hover:!bg-destructive/90" onClick={() => { - setBypassPermissions(true); + setPermissionMode("full"); setOpen(false); }} > diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 07ddffdd59..cedd298ecf 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -6,16 +6,6 @@ import { AlertDescription, AlertTitle, } from "@/components/ui/alert"; -import { - AlertDialog, - AlertDialogAction, - AlertDialogCancel, - AlertDialogContent, - AlertDialogDescription, - AlertDialogFooter, - AlertDialogHeader, - AlertDialogTitle, -} from "@/components/ui/alert-dialog"; import { Button } from "@/components/ui/button"; import { Checkbox } from "@/components/ui/checkbox"; import { @@ -81,6 +71,7 @@ import { Fragment, type ReactNode } from "react"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { toast } from "@/lib/toast"; import { OpenAICodeExecSection } from "./components/openai-code-exec-section"; +import { PermissionModeDropdown } from "./permission-mode-select"; import { resyncInferenceStatusAfterServerModelChange } from "./hooks/use-chat-model-runtime"; import { type ExternalProviderConfig, @@ -2037,9 +2028,8 @@ function NudgeToolCallsToggle() { } function ConfirmToolCallsToggle() { - const confirmToolCalls = useChatRuntimeStore((s) => s.confirmToolCalls); const setConfirmToolCalls = useChatRuntimeStore((s) => s.setConfirmToolCalls); - const bypassPermissions = useChatRuntimeStore((s) => s.bypassPermissions); + const permissionMode = useChatRuntimeStore((s) => s.permissionMode); return (
@@ -2049,85 +2039,49 @@ function ConfirmToolCallsToggle() { Confirm tool calls - When on, local Studio tool calls pause for your approval before they - run. Provider-hosted tools are not gated here. + When on, every local Unsloth tool call pauses for your approval + before it runs (the "Ask for approval" level). When off, tool calls + run without prompts inside the sandbox (the "Off" level). + Provider-hosted tools are not gated here.
- {bypassPermissions ? ( + {permissionMode === "full" ? ( - Overridden by Bypass permissions + Overridden by Full access (Bypass permissions) ) : null} ); } function BypassPermissionsToggle() { - const bypassPermissions = useChatRuntimeStore((s) => s.bypassPermissions); - const setBypassPermissions = useChatRuntimeStore( - (s) => s.setBypassPermissions, - ); - const [dialogOpen, setDialogOpen] = useState(false); + const permissionMode = useChatRuntimeStore((s) => s.permissionMode); return ( -
-
-
- - Bypass permissions - - - Dangerous. Runs every tool call with no confirmation and disables - the python/terminal sandbox. Environment secrets are stripped, but - code can still read files and credentials on your machine. - -
- { - if (next) setDialogOpen(true); - else setBypassPermissions(false); - }} - /> +
+
+ + Bypass permissions + + + How Unsloth approves tool calls before they run. Full access is + dangerous: it disables confirmations and the code sandbox. +
- {bypassPermissions ? ( + {/* Full width, styled like the panel selects/preset input. */} + + {permissionMode === "full" ? ( Tool calls run with no confirmation and no sandbox. ) : null} - - - - Enable Bypass permissions? - - Bypass permissions is dangerous since the AI model might delete, - corrupt your machine, and or cause real world damage to you or the - world - only accept if you are certain - - - - Cancel - { - setBypassPermissions(true); - setDialogOpen(false); - }} - > - I understand - - - -
); } diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts index 3099884645..7e894bb92e 100644 --- a/studio/frontend/src/features/chat/index.ts +++ b/studio/frontend/src/features/chat/index.ts @@ -17,6 +17,7 @@ export { type Preset, } from "./chat-settings-sheet"; export { useChatRuntimeStore } from "./stores/chat-runtime-store"; +export { PermissionModeDropdown } from "./permission-mode-select"; export { useChatSearchStore } from "./stores/chat-search-store"; export { usePinnedChatsStore } from "./stores/pinned-chats-store"; export { useChatPreferencesStore } from "./stores/chat-preferences-store"; diff --git a/studio/frontend/src/features/chat/permission-mode-select.tsx b/studio/frontend/src/features/chat/permission-mode-select.tsx new file mode 100644 index 0000000000..4277c1bfcf --- /dev/null +++ b/studio/frontend/src/features/chat/permission-mode-select.tsx @@ -0,0 +1,338 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { + ChevronDown, + CircleAlert, + CircleOff, + Hand, + ShieldCheck, + XIcon, +} from "lucide-react"; +import { useState } from "react"; + +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; +import { Button } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuLabel, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { ChevronDownStandardIcon } from "@/lib/chevron-icons"; +import { Tick02Icon } from "@/lib/tick-icon"; +import { cn } from "@/lib/utils"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { + type PermissionMode, + useChatRuntimeStore, +} from "./stores/chat-runtime-store"; + +/** + * Permission levels for the Bypass permissions dropdowns (General settings, + * chat settings sheet, composer "+" menu). Off sits last as the toggle that + * turns the feature off entirely. + */ +export const PERMISSION_MODE_OPTIONS: readonly { + value: PermissionMode; + label: string; + description: string; + icon: typeof Hand; +}[] = [ + { + value: "ask", + label: "Ask for approval", + description: "Always ask before tool calls edit files or use the internet", + icon: Hand, + }, + { + value: "auto", + label: "Approve for me", + description: "Only ask for actions detected as potentially unsafe", + icon: ShieldCheck, + }, + { + value: "full", + label: "Full access", + description: + "Unrestricted: no approval prompts and the code sandbox is disabled", + icon: CircleAlert, + }, + { + value: "off", + label: "Off", + description: "Turn off bypass permissions", + icon: CircleOff, + }, +] as const; + +export function permissionModeOption(mode: PermissionMode) { + return ( + PERMISSION_MODE_OPTIONS.find((option) => option.value === mode) ?? + PERMISSION_MODE_OPTIONS[0] + ); +} + +/** The option rows shared by every permission dropdown/submenu. Non-full + * levels apply directly; picking Full access must go through the caller's + * danger confirmation, so it's a separate callback. */ +export function PermissionModeMenuItems({ + onRequestFullAccess, +}: { + onRequestFullAccess: () => void; +}) { + const permissionMode = useChatRuntimeStore((s) => s.permissionMode); + const setPermissionMode = useChatRuntimeStore((s) => s.setPermissionMode); + + return ( + <> + {PERMISSION_MODE_OPTIONS.map((option) => ( + { + // Reselecting the active level toggles the feature off. + if (option.value === permissionMode) { + setPermissionMode("off"); + } else if (option.value === "full") { + onRequestFullAccess(); + } else { + setPermissionMode(option.value); + } + }} + className={cn( + "items-start gap-2 py-2", + permissionMode === option.value && "font-medium", + option.value === "full" && + permissionMode === "full" && + "text-bypass", + )} + > + + + {option.label} + + {option.description} + + + {permissionMode === option.value ? ( + + ) : null} + + ))} + + ); +} + +/** Danger confirmation shown before Full access turns on. Self-contained so + * the dropdown works outside the chat page (e.g. the Settings dialog). */ +export function FullAccessConfirmDialog({ + open, + onOpenChange, +}: { + open: boolean; + onOpenChange: (open: boolean) => void; +}) { + const setPermissionMode = useChatRuntimeStore((s) => s.setPermissionMode); + + return ( + + + + Enable Full access? + + Full access (Bypass permissions) is dangerous since the AI model + might delete, corrupt your machine, and or cause real world damage + to you or the world - only accept if you are certain + + + + Cancel + { + setPermissionMode("full"); + onOpenChange(false); + }} + > + I understand + + + + + ); +} + +/** + * Select-style dropdown (like the MCP composer menu) for picking the + * permission level. Used in General settings and the chat settings sheet. + */ +export function PermissionModeDropdown({ + side = "bottom", + align = "end", + triggerClassName, +}: { + side?: "top" | "bottom"; + align?: "start" | "end"; + triggerClassName?: string; +} = {}) { + const permissionMode = useChatRuntimeStore((s) => s.permissionMode); + const [confirmOpen, setConfirmOpen] = useState(false); + const active = permissionModeOption(permissionMode); + const ActiveIcon = active.icon; + + return ( + <> + + + + + + + How should tool calls be approved? + + + setTimeout(() => setConfirmOpen(true), 0) + } + /> + + + + + ); +} + +/** + * Composer pill (mirrors the MCP pill) showing the current permission level + * in the chat box; clicking opens the level dropdown. Danger-styled while + * Full access is on. The Full access pick routes through the store-driven + * BypassPermissionsConfirmDialog mounted at the chat-page root, so the + * warning survives this menu unmounting. + */ +export function PermissionModeComposerPill({ + side = "bottom", +}: { + side?: "top" | "bottom"; +} = {}) { + const permissionMode = useChatRuntimeStore((s) => s.permissionMode); + const setBypassConfirmOpen = useChatRuntimeStore( + (s) => s.setBypassConfirmOpen, + ); + const setPermissionMode = useChatRuntimeStore((s) => s.setPermissionMode); + const active = permissionModeOption(permissionMode); + const ActiveIcon = active.icon; + const fullAccess = permissionMode === "full"; + + // Off means the feature is off: no pill (re-enable via the "+" menu or + // settings, like the pre-levels bypass badge). + if (permissionMode === "off") return null; + + return ( + + + + + + + How should tool calls be approved? + + + setTimeout(() => setBypassConfirmOpen(true), 0) + } + /> + + + ); +} diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index 47a0720dac..2ed9589461 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -48,7 +48,6 @@ import { Image03Icon, McpServerIcon, PencilRulerIcon, - ShieldBanIcon, } from "@hugeicons/core-free-icons"; import { useNavigate } from "@tanstack/react-router"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -62,6 +61,7 @@ import { import { listPromptEntries, type PromptEntry } from "./api/prompts-api"; import { McpComposerButton } from "./mcp-composer-button"; import { BypassPermissionsMenuItem } from "./bypass-permissions-menu-item"; +import { PermissionModeComposerPill } from "./permission-mode-select"; import { reasoningCapsFromLoad } from "./lib/apply-inference-status-to-store"; import { KnowledgeBaseComposerButton } from "@/features/rag/components/knowledge-base-composer-button"; import { NewProjectDialog } from "./components/new-project-dialog"; @@ -510,6 +510,7 @@ export function SharedComposer({ ); const artifactsEnabled = useChatRuntimeStore((s) => s.artifactsEnabled); const setArtifactsEnabled = useChatRuntimeStore((s) => s.setArtifactsEnabled); + const permissionMode = useChatRuntimeStore((s) => s.permissionMode); const mcpEnabledForChat = useChatRuntimeStore((s) => s.mcpEnabledForChat); const setMcpEnabledForChat = useChatRuntimeStore( (s) => s.setMcpEnabledForChat, @@ -529,10 +530,6 @@ export function SharedComposer({ const setWebFetchToolsEnabled = useChatRuntimeStore( (s) => s.setWebFetchToolsEnabled, ); - const bypassPermissions = useChatRuntimeStore((s) => s.bypassPermissions); - const setBypassPermissions = useChatRuntimeStore( - (s) => s.setBypassPermissions, - ); const ragEnabled = useChatRuntimeStore((s) => s.ragEnabled); const setRagEnabled = useChatRuntimeStore((s) => s.setRagEnabled); const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); @@ -685,9 +682,12 @@ export function SharedComposer({ const ragDisabled = modelLoaded && (isExternalModel || !supportsTools); const showRagPill = !isExternalModel; // Above 4 pills, collapse to icons only to cut clutter. Compare, Search and - // Code always show; the rest are conditional. + // Code always show; the permission pill shows in every mode except "off" + // (it renders null there); the rest are conditional. + const permissionPillVisible = permissionMode !== "off"; const pillsCompact = 3 + + (permissionPillVisible ? 1 : 0) + (showImagePill ? 1 : 0) + (showRagPill && ragEnabled && !ragDisabled ? 1 : 0) + (showWebFetchPill ? 1 : 0) + @@ -1656,29 +1656,10 @@ export function SharedComposer({ Compare - {/* Bypass sits immediately after Compare and ahead of every other - tool pill (Search, Code, ...) so the active danger state reads - first; only Compare outranks it. */} - {bypassPermissions && ( - - )} + {/* Permission-level pill sits immediately after Compare and ahead + of every other tool pill (Search, Code, ...) so the Full access + danger state reads first; only Compare outranks it. */} +