diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 0615485304..f8cc2d9a1b 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -3423,18 +3423,18 @@ def _request_matches_loaded_settings( # copy and, failing that, to no drafter at all. An ordinary # load reaches the root drafter, so it keeps root-first # detection and reloads when one appears. - if detected and not _native_gguf_companion_usable( - detected, llama_backend.gguf_path, mtp_search_root = companion_root - ): + def _usable(candidate: str) -> bool: + return _native_gguf_companion_usable( + candidate, llama_backend.gguf_path, mtp_search_root = companion_root + ) + + if detected and not _usable(detected): detected = detect_mtp_file( llama_backend.gguf_path, search_root = companion_root, skip_root = True, + accept = _usable, ) - if detected and not _native_gguf_companion_usable( - detected, llama_backend.gguf_path, mtp_search_root = companion_root - ): - detected = None stored = llama_backend.mtp_draft_path try: detected_resolved = Path(detected).resolve() if detected else None @@ -4777,21 +4777,21 @@ async def _load_model_impl( if not _mtp_allowed(config.gguf_mtp_file): # The preferred root drafter is out of bounds for a - # grant on a quant subdir, but its MTP/ copy may - # not be. Use that before dropping MTP entirely. + # grant on a quant subdir, but an MTP/ copy may not + # be. Scan them in preference order rather than + # dropping MTP on the first rejection. fallback = detect_mtp_file( config.gguf_file, search_root = mtp_search_root, skip_root = True, + accept = _mtp_allowed, ) - if fallback and _mtp_allowed(fallback): + if fallback: logger.info( "Using MTP subdirectory drafter for native load: %s", fallback, ) - config.gguf_mtp_file = fallback - else: - config.gguf_mtp_file = None + config.gguf_mtp_file = fallback _source_load_kwargs = dict( gguf_path = config.gguf_file, mmproj_path = config.gguf_mmproj_file, diff --git a/studio/backend/tests/test_native_gguf_companion.py b/studio/backend/tests/test_native_gguf_companion.py index 3b0df808dd..3a7b4b0caf 100644 --- a/studio/backend/tests/test_native_gguf_companion.py +++ b/studio/backend/tests/test_native_gguf_companion.py @@ -19,6 +19,8 @@ from routes.inference import _validate_native_gguf_companion from routes.inference import _request_matches_loaded_settings from routes.inference import _validate_native_mtp_drafter from routes.inference import _loaded_is_local_model +from routes.inference import _native_gguf_companion_usable +from utils.models.model_config import detect_mtp_file from core.inference.llama_cpp import LlamaCppBackend from models.inference import LoadRequest @@ -334,3 +336,51 @@ def test_status_provenance_survives_deleted_model_directory(tmp_path, monkeypatc stale = LlamaCppBackend() assert not _loaded_is_local_model(stale, False, "unsloth/gemma-4-12b") assert _loaded_is_local_model(stale, True, None) + + +def test_native_load_skips_rejected_mtp_candidate_for_next_one(tmp_path): + """MTP/ can hold several compatible copies. If the size-preferred one is + out of the grant, the next must be tried instead of disabling MTP.""" + quant_dir = tmp_path / "Q4_0" + quant_dir.mkdir() + weight = quant_dir / "model.gguf" + weight.write_bytes(b"model") + outside = tmp_path.parent / "outside-blob.gguf" + outside.write_bytes(b"d") + companion_dir = tmp_path / "MTP" + companion_dir.mkdir() + try: + (companion_dir / "mtp-model-Q4_0.gguf").symlink_to(outside) + except OSError as exc: + pytest.skip(f"symlinks unavailable: {exc}") + larger = companion_dir / "mtp-model-Q8_0.gguf" + larger.write_bytes(b"d" * 5000) + + def _usable(candidate: str) -> bool: + return _native_gguf_companion_usable(candidate, str(weight), mtp_search_root = str(tmp_path)) + + # Preferred by size, but it resolves out of the permitted directory. + assert not _usable(detect_mtp_file(str(weight), str(tmp_path), skip_root = True)) + assert detect_mtp_file(str(weight), str(tmp_path), skip_root = True, accept = _usable) == str( + larger.resolve() + ) + + +def test_native_load_returns_none_when_no_candidate_passes(tmp_path): + quant_dir = tmp_path / "Q4_0" + quant_dir.mkdir() + weight = quant_dir / "model.gguf" + weight.write_bytes(b"model") + outside = tmp_path.parent / "outside-only.gguf" + outside.write_bytes(b"d") + companion_dir = tmp_path / "MTP" + companion_dir.mkdir() + try: + (companion_dir / "mtp-model-Q4_0.gguf").symlink_to(outside) + except OSError as exc: + pytest.skip(f"symlinks unavailable: {exc}") + + def _usable(candidate: str) -> bool: + return _native_gguf_companion_usable(candidate, str(weight), mtp_search_root = str(tmp_path)) + + assert detect_mtp_file(str(weight), str(tmp_path), skip_root = True, accept = _usable) is None diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index e6c1d0b4d7..d379e7ad0a 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -29,7 +29,7 @@ import re import subprocess import sys from pathlib import Path -from typing import List, Tuple +from typing import Callable, List, Tuple import hashlib import json import threading @@ -1486,6 +1486,7 @@ def detect_mtp_file( path: str, search_root: Optional[str] = None, skip_root: bool = False, + accept: Optional[Callable[[str], bool]] = None, ) -> Optional[str]: """Find the separate MTP drafter (``mtp-*.gguf``) for a local GGUF model. @@ -1505,6 +1506,9 @@ def detect_mtp_file( ``skip_root`` scans only ``MTP/``, for callers that must discard an out-of-bounds root drafter and still want the subdir copy (native loads). + ``accept`` filters candidates in preference order, so a caller with extra + rules (a native lease) keeps scanning instead of treating the first + rejection as no drafter at all. """ def _pairing_stem(name: str) -> str: @@ -1594,10 +1598,14 @@ def detect_mtp_file( if not _matches_weight(f): continue try: - if f.is_file() and _launchable(f): - return _drafter_launch_path(f) + if not (f.is_file() and _launchable(f)): + continue + launch = _drafter_launch_path(f) except OSError: continue + if accept is not None and not accept(launch): + continue + return launch subdir_candidates: list[Path] = [] for d in dirs: @@ -1647,6 +1655,8 @@ def detect_mtp_file( resolved = _drafter_launch_path(candidate) except OSError: continue + if accept is not None and not accept(resolved): + continue logger.info(f"Detected MTP subdirectory drafter: {resolved}") return resolved return None diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 9409945072..7a7475406f 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -1538,7 +1538,17 @@ export const useChatRuntimeStore = create((set, get) => ({ maxTokens: nextMaxTokens, }, activeGgufVariant: ggufVariant ?? null, - ...(checkpointChanged ? { contextUsage: null } : {}), + // Provenance and the spec-fallback reason both describe the model + // being replaced, so they go together on a real change. Dropping only + // one leaves the settings sheet pairing a stale reason with the wrong + // recovery text. The load or status response reseeds both. + ...(checkpointChanged + ? { + contextUsage: null, + activeModelIsLocal: false, + specFallbackReason: null, + } + : {}), }; }), setActiveThreadId: (activeThreadId) => diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py index 35125f02dd..3450bfdc74 100644 --- a/tests/studio/test_model_picker_contracts.py +++ b/tests/studio/test_model_picker_contracts.py @@ -357,6 +357,14 @@ def test_local_mtp_warning_covers_path_and_native_gguf_sources(): # repository id. activeModelIsLocal is the backend's own answer for both. assert "activeNativePathToken" not in local.group(0) assert ".gguf" not in local.group(0) + + # Switching models must drop both together: a kept flag would classify the + # newly selected model by the old one's provenance. + store = _read("features/chat/stores/chat-runtime-store.ts") + reset = re.search(r"setCheckpoint: \(modelId, ggufVariant\) =>.*?\}\),", store, re.S) + assert reset + assert "activeModelIsLocal: false" in reset.group(0) + assert "specFallbackReason: null" in reset.group(0) assert "isLocalGguf" in src.split('specFallbackReason === "drafter_not_found"', 1)[1]