From af4464a0e0492e6bb3a37bb255b2b14bf1c165dc Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Fri, 24 Jul 2026 14:41:34 -0700 Subject: [PATCH] Studio: harden MTP companion pairing --- studio/backend/routes/inference.py | 17 +++++- .../tests/test_mtp_drafter_companion.py | 14 +++++ .../tests/test_native_gguf_companion.py | 58 +++++++++++++++++++ studio/backend/utils/models/model_config.py | 6 +- studio/backend/utils/native_path_leases.py | 16 +++-- 5 files changed, 103 insertions(+), 8 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index bc48f0b0bf..f4b8a94af5 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1020,6 +1020,7 @@ try: from utils.models import ModelConfig from utils.inference import load_inference_config from utils.models.model_config import ( + _local_gguf_companion_search_root, detect_mtp_file, load_model_defaults, ) @@ -1059,6 +1060,7 @@ except ImportError: from utils.models import ModelConfig from utils.inference import load_inference_config from utils.models.model_config import ( + _local_gguf_companion_search_root, detect_mtp_file, load_model_defaults, ) @@ -3077,6 +3079,7 @@ def _validate_native_gguf_companion( label: str, *, allow_mtp_subdir: bool = False, + mtp_search_root: str | Path | None = None, ) -> None: """Reject a companion GGUF (mmproj / MTP drafter) that a native-lease load would otherwise hand to llama-server: must be a regular file (no symlink @@ -3103,7 +3106,10 @@ def _validate_native_gguf_companion( ) try: if not native_gguf_companion_parent_allowed( - companion, gguf, allow_mtp_subdir = allow_mtp_subdir + companion, + gguf, + allow_mtp_subdir = allow_mtp_subdir, + mtp_search_root = mtp_search_root, ): location = ( "beside the selected GGUF or in its MTP directory" @@ -3337,7 +3343,10 @@ def _request_matches_loaded_settings( else llama_backend.extra_args ) if not _extra_args_set_spec_type(effective_extras): - detected = detect_mtp_file(llama_backend.gguf_path) + companion_root = _local_gguf_companion_search_root( + llama_backend.gguf_path, llama_backend.gguf_path + ) + detected = detect_mtp_file(llama_backend.gguf_path, search_root = companion_root) stored = llama_backend.mtp_draft_path try: detected_resolved = Path(detected).resolve() if detected else None @@ -4656,11 +4665,15 @@ async def _load_model_impl( # The drafter is optional (unlike mmproj for a vision # model): drop it rather than fail the load. try: + mtp_search_root = _local_gguf_companion_search_root( + config.gguf_file, config.gguf_file + ) _validate_native_gguf_companion( config.gguf_mtp_file, config.gguf_file, "MTP drafter", allow_mtp_subdir = True, + mtp_search_root = mtp_search_root, ) except HTTPException as exc: logger.warning("Dropping MTP drafter for native load: %s", exc.detail) diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index ed8a6a5d0e..5c1e7fb035 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -259,6 +259,20 @@ def test_detect_mtp_file_subdir_skips_foreign_drafter(tmp_path): assert detect_mtp_file(str(weight)) is None +@pytest.mark.parametrize( + "companion_path", + ["mtp-gemma-4-E4B-it-Q4_0.gguf", "MTP/mtp-gemma-4-E4B-it-Q4_0.gguf"], +) +def test_detect_mtp_file_requires_model_name_boundary(tmp_path, companion_path): + weight = tmp_path / "gemma-4-E4B-item-qat-Q4_0.gguf" + weight.write_bytes(b"x") + companion = tmp_path / companion_path + companion.parent.mkdir(parents = True, exist_ok = True) + companion.write_bytes(b"x") + + assert detect_mtp_file(str(weight)) is None + + def test_detect_mtp_file_accepts_case_variant_subdir(tmp_path): weight = tmp_path / "gemma-4-E4B-it-qat-Q4_0.gguf" weight.write_bytes(b"x") diff --git a/studio/backend/tests/test_native_gguf_companion.py b/studio/backend/tests/test_native_gguf_companion.py index d4fe8ba831..9ed260aa90 100644 --- a/studio/backend/tests/test_native_gguf_companion.py +++ b/studio/backend/tests/test_native_gguf_companion.py @@ -16,9 +16,13 @@ if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) from routes.inference import _validate_native_gguf_companion +from routes.inference import _request_matches_loaded_settings +from core.inference.llama_cpp import LlamaCppBackend +from models.inference import LoadRequest def _write_pair(tmp_path: Path, folder: str | None = None) -> tuple[Path, Path]: + tmp_path.mkdir(parents = True, exist_ok = True) weight = tmp_path / "model.gguf" weight.write_bytes(b"model") parent = tmp_path if folder is None else tmp_path / folder @@ -41,6 +45,60 @@ def test_native_mtp_companion_allows_mtp_directory(tmp_path, folder): ) +def test_native_mtp_companion_allows_repo_root_mtp_directory(tmp_path): + quant_dir = tmp_path / "Q4_0" + weight, _ = _write_pair(quant_dir) + companion_dir = tmp_path / "MTP" + companion_dir.mkdir() + companion = companion_dir / "mtp-model.gguf" + companion.write_bytes(b"draft") + + _validate_native_gguf_companion( + str(companion), + str(weight), + "MTP drafter", + allow_mtp_subdir = True, + mtp_search_root = str(tmp_path), + ) + + +def test_native_mtp_companion_rejects_unrelated_search_root(tmp_path): + quant_dir = tmp_path / "repo" / "Q4_0" + weight, _ = _write_pair(quant_dir) + companion_dir = tmp_path / "MTP" + companion_dir.mkdir() + companion = companion_dir / "mtp-model.gguf" + companion.write_bytes(b"draft") + + with pytest.raises(HTTPException, match = "must live beside"): + _validate_native_gguf_companion( + str(companion), + str(weight), + "MTP drafter", + allow_mtp_subdir = True, + mtp_search_root = str(tmp_path), + ) + + +def test_reload_dedup_finds_repo_root_mtp_companion(tmp_path, monkeypatch): + quant_dir = tmp_path / "Q4_0" + quant_dir.mkdir() + weight = quant_dir / "model.gguf" + weight.write_bytes(b"model") + companion_dir = tmp_path / "MTP" + companion_dir.mkdir() + companion = companion_dir / "mtp-model.gguf" + companion.write_bytes(b"draft") + + monkeypatch.setattr(LlamaCppBackend, "_kill_orphaned_servers", staticmethod(lambda: 0)) + backend = LlamaCppBackend() + backend._gguf_path = str(weight) + backend._mtp_draft_path = str(companion) + + request = LoadRequest(model_path = str(weight)) + assert _request_matches_loaded_settings(request, backend) + + def test_native_vision_companion_rejects_mtp_directory(tmp_path): weight, companion = _write_pair(tmp_path, "MTP") with pytest.raises(HTTPException, match = "must live next to"): diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 4c71414559..0d66515182 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1471,7 +1471,11 @@ def detect_mtp_file(path: str, search_root: Optional[str] = None) -> Optional[st if weight_name is None: return True stem = _pairing_stem(candidate.name) - return bool(stem) and weight_name.startswith(stem) + return ( + bool(stem) + and weight_name.startswith(stem) + and (len(weight_name) == len(stem) or not weight_name[len(stem)].isalnum()) + ) def _precision_rank(candidate: Path) -> tuple[int, str]: name = candidate.name.lower() diff --git a/studio/backend/utils/native_path_leases.py b/studio/backend/utils/native_path_leases.py index 090cedfeb6..d17a96eb00 100644 --- a/studio/backend/utils/native_path_leases.py +++ b/studio/backend/utils/native_path_leases.py @@ -52,15 +52,21 @@ def native_gguf_companion_parent_allowed( gguf_path: str | Path, *, allow_mtp_subdir: bool = False, + mtp_search_root: str | Path | None = None, ) -> bool: """Check whether a GGUF companion is in an allowed directory.""" companion_parent = Path(companion_path).resolve(strict = True).parent gguf_parent = Path(gguf_path).resolve(strict = True).parent - return companion_parent == gguf_parent or bool( - allow_mtp_subdir - and companion_parent.parent == gguf_parent - and companion_parent.name.casefold() == "mtp" - ) + if companion_parent == gguf_parent: + return True + if not allow_mtp_subdir or companion_parent.name.casefold() != "mtp": + return False + allowed_roots = {gguf_parent} + if mtp_search_root is not None: + search_root = Path(mtp_search_root).resolve(strict = True) + if search_root in {gguf_parent, gguf_parent.parent}: + allowed_roots.add(search_root) + return companion_parent.parent in allowed_roots @dataclass(frozen = True)