# SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """Offline RAG embedding-model handling (issue #6817). Offline the studio must never call the Hub (a DNS-dead session hangs on retries). Using a fake HF cache under a temp HF_HUB_CACHE, assert that offline: is_embedding_model classifies from the cached modules.json without the Hub; the file-security gate fails CLOSED on an unscanned pickle weight with no safetensors alternative and allows an inert cache; the embedder threads local_files_only into the load. Online behavior is unchanged (bounded timeout + cache fallback). """ import sys import types from pathlib import Path from types import SimpleNamespace from unittest.mock import patch import pytest from utils.security import evaluate_file_security from utils.utils import ( hf_cache_snapshot_dir, hf_cache_snapshot_is_loadable, hf_env_offline, st_repo_id_candidates, ) # Minimal sentence-transformers modules.json (the marker the gate keys on). MODULES_JSON = ( '[{"idx": 0, "name": "0", "path": "", "type": "sentence_transformers.models.Transformer"}]' ) def _modules_json(*paths): """modules.json listing one Transformer module per path (a load root).""" import json return json.dumps( [ { "idx": i, "name": str(i), "path": p, "type": "sentence_transformers.models.Transformer", } for i, p in enumerate(paths) ] ) _COMMIT = "0123456789abcdef0123456789abcdef01234567" def _fs_case_sensitive(root): """Whether root's filesystem is case-sensitive (Linux yes; macOS/Windows usually no). The gate mirrors the loader, whose file lookups follow the same rule, so some cases only exist on one.""" probe = Path(root) / "_case_probe" probe.write_text("x") try: return not (Path(root) / "_CASE_PROBE").exists() finally: probe.unlink() def _requires_case_sensitive_fs(root): if not _fs_case_sensitive(root): pytest.skip("requires a case-sensitive filesystem") def _requires_case_insensitive_fs(root): if _fs_case_sensitive(root): pytest.skip("requires a case-insensitive filesystem") def _make_cache( root, repo_id, files, commit = _COMMIT, ): """Build a canonical HF-cache snapshot (refs/main + snapshots//) for repo_id under root from {relpath: contents}; returns the snapshot dir.""" from huggingface_hub.file_download import repo_folder_name repo_dir = Path(root) / repo_folder_name(repo_id = repo_id, repo_type = "model") (repo_dir / "refs").mkdir(parents = True, exist_ok = True) (repo_dir / "refs" / "main").write_text(commit) snapshot = repo_dir / "snapshots" / commit snapshot.mkdir(parents = True, exist_ok = True) for rel, contents in files.items(): path = snapshot / rel path.parent.mkdir(parents = True, exist_ok = True) path.write_text(contents) return snapshot def _no_network(): """Patch model_info to fail loudly if any offline path reaches the network.""" return patch("huggingface_hub.model_info", side_effect = AssertionError("hit the network")) def _is_embedding_model(*args, **kwargs): from utils.models.model_config import is_embedding_model return is_embedding_model(*args, **kwargs) @pytest.fixture def hf_cache(tmp_path, monkeypatch): """Point the HF cache at a fresh temp dir. get_hf_cache_paths() reads an import-time env snapshot, not live os.environ, so point it (and thus active_hf_hub_cache + the snapshot lookup's selected root) at this temp cache too.""" root = tmp_path / "hub" root.mkdir() monkeypatch.setenv("HF_HOME", str(tmp_path)) monkeypatch.setenv("HF_HUB_CACHE", str(root)) monkeypatch.setattr( "utils.hf_cache_settings.get_hf_cache_paths", lambda: SimpleNamespace(hub_cache = root), ) return root @pytest.fixture(autouse = True) def _clean_env(monkeypatch): """Start each test online with an empty detection cache; offline tests opt in.""" monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) from utils.models import model_config as mc mc._embedding_detection_cache.clear() yield mc._embedding_detection_cache.clear() # ── hf_env_offline ─────────────────────────────────────────────── @pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "on", " On "]) def test_hf_env_offline_true(monkeypatch, value): monkeypatch.setenv("HF_HUB_OFFLINE", value) assert hf_env_offline() is True @pytest.mark.parametrize("value", ["0", "false", "no", "off", ""]) def test_hf_env_offline_false(monkeypatch, value): monkeypatch.setenv("HF_HUB_OFFLINE", value) assert hf_env_offline() is False def test_hf_env_offline_honors_transformers_flag(monkeypatch): monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") assert hf_env_offline() is True def test_hf_env_offline_default_false(): assert hf_env_offline() is False # ── st_repo_id_candidates ──────────────────────────────────────── def test_candidates_slashless_adds_st_alias(): assert st_repo_id_candidates("all-MiniLM-L6-v2") == [ "all-MiniLM-L6-v2", "sentence-transformers/all-MiniLM-L6-v2", ] def test_candidates_with_org_is_verbatim(): assert st_repo_id_candidates("org/model") == ["org/model"] def test_candidates_empty_name(): assert st_repo_id_candidates(" ") == [] # ── hf_cache_snapshot_dir ──────────────────────────────────────── def test_snapshot_dir_resolves_active_commit(hf_cache): snapshot = _make_cache(hf_cache, "org/emb", {"modules.json": MODULES_JSON}) assert hf_cache_snapshot_dir("org/emb") == snapshot def test_snapshot_dir_none_when_uncached(hf_cache): assert hf_cache_snapshot_dir("org/missing") is None def test_snapshot_dir_uses_st_alias_for_slashless(hf_cache): snapshot = _make_cache( hf_cache, "sentence-transformers/all-MiniLM-L6-v2", {"modules.json": MODULES_JSON} ) assert hf_cache_snapshot_dir("all-MiniLM-L6-v2") == snapshot def test_snapshot_dir_none_when_snapshot_missing(hf_cache): from huggingface_hub.file_download import repo_folder_name repo_dir = hf_cache / repo_folder_name(repo_id = "org/broken", repo_type = "model") (repo_dir / "refs").mkdir(parents = True) (repo_dir / "refs" / "main").write_text("deadbeef") # no snapshots/deadbeef dir assert hf_cache_snapshot_dir("org/broken") is None def test_snapshot_dir_expands_env_vars_in_cache_path(tmp_path, monkeypatch): # An unexpanded $VAR in HF_HUB_CACHE must resolve where the loader looks. real = tmp_path / "hub" real.mkdir() monkeypatch.setenv("MY_HF_CACHE", str(real)) monkeypatch.setenv("HF_HUB_CACHE", "$MY_HF_CACHE") monkeypatch.delenv("HF_HOME", raising = False) monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) snapshot = _make_cache(real, "org/emb", {"modules.json": MODULES_JSON}) assert hf_cache_snapshot_dir("org/emb") == snapshot def test_snapshot_dir_uses_sentence_transformers_home(tmp_path, monkeypatch): # ST uses SENTENCE_TRANSFORMERS_HOME as its cache_folder, so the gate must inspect it too. st_home = tmp_path / "st_home" st_home.mkdir() monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_home)) monkeypatch.delenv("HF_HUB_CACHE", raising = False) monkeypatch.delenv("HF_HOME", raising = False) snapshot = _make_cache(st_home, "org/emb", {"modules.json": MODULES_JSON}) assert hf_cache_snapshot_dir("org/emb") == snapshot def test_snapshot_dir_prefers_selected_cache_over_st_home(tmp_path, monkeypatch): # The RAG loader passes cache_folder=active_hf_hub_cache(), which overrides # SENTENCE_TRANSFORMERS_HOME, so the snapshot + offline security lookup must # search the selected cache even when ST_HOME points elsewhere. Otherwise the # gate scans a cache the model never loads from and a pickle weight in the # selected cache slips through. st_home = tmp_path / "st_home" st_home.mkdir() selected = tmp_path / "hub" selected.mkdir() monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_home)) monkeypatch.delenv("HF_HUB_CACHE", raising = False) monkeypatch.delenv("HF_HOME", raising = False) monkeypatch.setattr( "utils.hf_cache_settings.get_hf_cache_paths", lambda: SimpleNamespace(hub_cache = selected), ) snapshot = _make_cache(selected, "org/emb", {"modules.json": MODULES_JSON}) # only in selected assert hf_cache_snapshot_dir("org/emb") == snapshot def test_snapshot_is_loadable_with_config_and_weights(hf_cache): _make_cache(hf_cache, "org/emb", {"config.json": "{}", "model.safetensors": "x"}) assert hf_cache_snapshot_is_loadable("org/emb") is True def test_snapshot_is_not_loadable_when_metadata_only(hf_cache): # A partial cache (refs/main resolves but no weights) is not loadable. _make_cache(hf_cache, "org/partial", {"config.json": "{}", "modules.json": MODULES_JSON}) assert hf_cache_snapshot_is_loadable("org/partial") is False def test_snapshot_is_not_loadable_when_uncached(hf_cache): assert hf_cache_snapshot_is_loadable("org/missing") is False def test_gate_blocks_pickle_in_sentence_transformers_home(tmp_path, monkeypatch): # A pickle under SENTENCE_TRANSFORMERS_HOME must still fail closed offline. st_home = tmp_path / "st_home" st_home.mkdir() monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_home)) monkeypatch.delenv("HF_HUB_CACHE", raising = False) monkeypatch.delenv("HF_HOME", raising = False) _make_cache(st_home, "org/pk", {"config.json": "{}", "pytorch_model.bin": "x"}) with _no_network(): assert evaluate_file_security("org/pk", local_only_load = True).blocked is True # ── is_embedding_model: offline (no network) ───────────────────── def test_offline_true_for_cached_st_model(hf_cache, monkeypatch): monkeypatch.setenv("HF_HUB_OFFLINE", "1") _make_cache(hf_cache, "org/emb", {"modules.json": MODULES_JSON, "config.json": "{}"}) with _no_network(): assert _is_embedding_model("org/emb") is True def test_offline_false_for_cached_non_st_model(hf_cache, monkeypatch): monkeypatch.setenv("HF_HUB_OFFLINE", "1") _make_cache(hf_cache, "org/plain", {"config.json": "{}", "model.safetensors": "x"}) with _no_network(): assert _is_embedding_model("org/plain") is False def test_offline_false_when_uncached(hf_cache, monkeypatch): monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") with _no_network(): assert _is_embedding_model("org/missing") is False def test_offline_slashless_resolves_via_alias(hf_cache, monkeypatch): monkeypatch.setenv("HF_HUB_OFFLINE", "1") _make_cache(hf_cache, "sentence-transformers/all-MiniLM-L6-v2", {"modules.json": MODULES_JSON}) with _no_network(): assert _is_embedding_model("all-MiniLM-L6-v2") is True def test_offline_ignores_stale_online_memo(hf_cache, monkeypatch): # An online lookup memoizes True for an UNCACHED repo (tags say embedding, no weights). Once # offline, is_embedding_model must reclassify from the empty cache and return False, not the # stale online True that would make settings accept a repo _get() cannot load. with patch( "huggingface_hub.model_info", side_effect = lambda *a, **k: SimpleNamespace( tags = ["sentence-transformers"], pipeline_tag = None ), ): assert _is_embedding_model("org/uncached-emb") is True # memoized True online monkeypatch.setenv("HF_HUB_OFFLINE", "1") with _no_network(): assert _is_embedding_model("org/uncached-emb") is False # recomputed from empty cache def test_offline_recomputes_after_cache_materializes(hf_cache, monkeypatch): # Because the offline branch never records a memo, once an uncached repo's snapshot # materializes (another process populates the cache) the next call re-reports True. monkeypatch.setenv("HF_HUB_OFFLINE", "1") with _no_network(): assert _is_embedding_model("org/later") is False # uncached _make_cache(hf_cache, "org/later", {"modules.json": MODULES_JSON}) assert _is_embedding_model("org/later") is True # cache now present, no stale negative # ── is_embedding_model: online (bounded + fallback) ────────────── def test_online_passes_bounded_timeout(hf_cache): seen = {} def _mi( name, token = None, timeout = None, **kw, ): seen["timeout"] = timeout return SimpleNamespace(tags = ["sentence-transformers"], pipeline_tag = None) with patch("huggingface_hub.model_info", side_effect = _mi): assert _is_embedding_model("org/emb") is True assert seen["timeout"] == 15.0 def test_online_error_falls_back_to_cache_marker(hf_cache): _make_cache(hf_cache, "org/emb", {"modules.json": MODULES_JSON}) with patch("huggingface_hub.model_info", side_effect = RuntimeError("dns dead")): assert _is_embedding_model("org/emb") is True def test_online_error_without_cache_returns_false(hf_cache): with patch("huggingface_hub.model_info", side_effect = RuntimeError("dns dead")): assert _is_embedding_model("org/missing") is False # ── evaluate_file_security: offline fail-closed gate ───────────── def _offline_decision(name): return evaluate_file_security(name, local_only_load = True) def test_gate_allows_safetensors_only(hf_cache): _make_cache(hf_cache, "org/st", {"modules.json": MODULES_JSON, "model.safetensors": "x"}) with _no_network(): assert _offline_decision("org/st").blocked is False def test_gate_blocks_pickle_without_safetensors(hf_cache): _make_cache(hf_cache, "org/pk", {"config.json": "{}", "pytorch_model.bin": "x"}) with _no_network(): decision = _offline_decision("org/pk") assert decision.blocked is True assert any(u["path"] == "pytorch_model.bin" for u in decision.unsafe_files) def test_gate_allows_pickle_with_safetensors_sibling(hf_cache): _make_cache(hf_cache, "org/both", {"pytorch_model.bin": "x", "model.safetensors": "y"}) with _no_network(): assert _offline_decision("org/both").blocked is False def test_gate_blocks_sharded_pickle(hf_cache): _make_cache( hf_cache, "org/shard", { "pytorch_model-00001-of-00002.bin": "a", "pytorch_model-00002-of-00002.bin": "b", }, ) with _no_network(): assert _offline_decision("org/shard").blocked is True def test_gate_blocks_indexed_pickle_shard_in_subdirectory(hf_cache): # from_pretrained follows weight_map paths relative to the root index, so these nested shards # are deserialized even though they are not direct children of the load root (iterdir misses # them). The online gate blocks index-referenced subdir pickles; the offline gate must too. _make_cache( hf_cache, "org/indexed-shard", { "pytorch_model.bin.index.json": ( '{"weight_map": {"layer.weight": "shards/pytorch_model-00001-of-00001.bin"}}' ), "shards/pytorch_model-00001-of-00001.bin": "pickle", }, ) with _no_network(): decision = _offline_decision("org/indexed-shard") assert decision.blocked is True assert any( u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files ) def test_gate_blocks_indexed_pickle_shard_with_nonstandard_stem(hf_cache): # The index tells the loader to deserialize this file, so a pickle EXTENSION is enough -- the # shard's stem need not match the on-disk weight-name heuristic (which only guesses bare files). _make_cache( hf_cache, "org/indexed-odd", { "pytorch_model.bin.index.json": '{"weight_map": {"w": "shards/evil-00001-of-00001.bin"}}', "shards/evil-00001-of-00001.bin": "pickle", }, ) with _no_network(): decision = _offline_decision("org/indexed-odd") assert decision.blocked is True assert any(u["path"] == "shards/evil-00001-of-00001.bin" for u in decision.unsafe_files) def test_gate_blocks_safetensors_index_pointing_to_pickle_shard(hf_cache): # load_state_dict picks safetensors vs torch.load by each shard's own suffix, so a # model.safetensors.index.json that maps a weight to a .bin shard still deserializes it. The # index's own existence must not suppress the shard it names. _make_cache( hf_cache, "org/st-index-pickle", { "model.safetensors.index.json": ( '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' ), "shards/pytorch_model-00001-of-00001.bin": "pickle", }, ) with _no_network(): decision = _offline_decision("org/st-index-pickle") assert decision.blocked is True assert any( u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files ) def test_gate_blocks_indexed_shard_with_no_pickle_extension(hf_cache): # Transformers torch.loads any indexed shard not ending in .safetensors, so an unconventional # extensionless name is still a deserialization target. _make_cache( hf_cache, "org/indexed-noext", { "pytorch_model.bin.index.json": '{"weight_map": {"w": "shards/payload"}}', "shards/payload": "pickle", }, ) with _no_network(): decision = _offline_decision("org/indexed-noext") assert decision.blocked is True assert any(u["path"] == "shards/payload" for u in decision.unsafe_files) _UPPER_INDEX_FILES = { "PYTORCH_MODEL.BIN.INDEX.JSON": ( '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' ), "shards/pytorch_model-00001-of-00001.bin": "pickle", } def test_gate_blocks_uppercase_index_on_case_insensitive_fs(hf_cache): # On a case-insensitive volume (Windows/macOS) from_pretrained opens an oddly-cased index when it # requests the canonical lowercase name, so the loader-mirror lookup resolves it and blocks. _requires_case_insensitive_fs(hf_cache) _make_cache(hf_cache, "org/upper-index", _UPPER_INDEX_FILES) with _no_network(): decision = _offline_decision("org/upper-index") assert decision.blocked is True assert any( u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files ) def test_gate_allows_uppercase_index_on_case_sensitive_fs(hf_cache): # On a case-sensitive FS from_pretrained's os.path.isfile of the canonical lowercase name misses # the uppercase artifact and never loads its shard, so the gate must not over-block it. _requires_case_sensitive_fs(hf_cache) _make_cache(hf_cache, "org/upper-index", _UPPER_INDEX_FILES) with _no_network(): assert _offline_decision("org/upper-index").blocked is False def test_gate_blocks_indexed_shard_named_with_backslash(hf_cache): # On POSIX a backslash is a literal filename char, so from_pretrained joins the raw weight_map # value and deserializes a file actually named "dir\payload.bin"; the gate must probe it verbatim. import os if os.sep != "/": pytest.skip("backslash is a path separator off POSIX") _make_cache( hf_cache, "org/backslash", { "pytorch_model.bin.index.json": '{"weight_map": {"w": "dir\\\\payload.bin"}}', "dir\\payload.bin": "pickle", }, ) with _no_network(): decision = _offline_decision("org/backslash") assert decision.blocked is True assert any(u["path"] == "dir\\payload.bin" for u in decision.unsafe_files) def test_gate_blocks_indexed_shard_with_uppercase_safetensors_suffix(hf_cache): # load_state_dict's endswith(".safetensors") is case-sensitive, so a shard named payload.SAFETENSORS # falls to torch.load. The gate must classify shard suffixes case-sensitively to match it. _make_cache( hf_cache, "org/upper-suffix", { "pytorch_model.bin.index.json": '{"weight_map": {"w": "shards/payload.SAFETENSORS"}}', "shards/payload.SAFETENSORS": "pickle", }, ) with _no_network(): decision = _offline_decision("org/upper-suffix") assert decision.blocked is True assert any(u["path"] == "shards/payload.SAFETENSORS" for u in decision.unsafe_files) def test_gate_allows_stale_safetensors_index_beside_direct_safetensors(hf_cache): # A complete direct model.safetensors is selected before either index, so a stale # model.safetensors.index.json referencing a .bin shard never deserializes -> must not block. _make_cache( hf_cache, "org/direct-plus-stale-index", { "model.safetensors": "tensors", "model.safetensors.index.json": ( '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' ), "shards/pytorch_model-00001-of-00001.bin": "pickle", }, ) with _no_network(): assert _offline_decision("org/direct-plus-stale-index").blocked is False def test_gate_blocks_pytorch_index_with_uppercase_safetensors_decoy(hf_cache): # On a case-sensitive FS, from_pretrained asks for the canonical lowercase model.safetensors, does # not find an uppercase decoy, and selects the pytorch index instead. The decoy must not suppress. _requires_case_sensitive_fs(hf_cache) _make_cache( hf_cache, "org/upper-decoy", { "MODEL.SAFETENSORS": "decoy", "pytorch_model.bin.index.json": ( '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' ), "shards/pytorch_model-00001-of-00001.bin": "pickle", }, ) with _no_network(): decision = _offline_decision("org/upper-decoy") assert decision.blocked is True assert any( u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files ) def test_gate_blocks_direct_pickle_with_uppercase_safetensors_decoy(hf_cache): # Same decoy against a direct pytorch_model.bin: the loader selects the pickle, so the uppercase # safetensors must not suppress it on a case-sensitive FS. _requires_case_sensitive_fs(hf_cache) _make_cache( hf_cache, "org/upper-decoy-direct", {"MODEL.SAFETENSORS": "decoy", "pytorch_model.bin": "pickle"}, ) with _no_network(): decision = _offline_decision("org/upper-decoy-direct") assert decision.blocked is True assert any(u["path"] == "pytorch_model.bin" for u in decision.unsafe_files) def test_gate_blocks_indexed_pickle_shard_in_module_subdir(hf_cache): # A weight index inside a sentence-transformers module load root points at a nested pickle shard. _make_cache( hf_cache, "org/mod-indexed", { "modules.json": _modules_json("0_Transformer"), "0_Transformer/pytorch_model.bin.index.json": ( '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' ), "0_Transformer/shards/pytorch_model-00001-of-00001.bin": "pickle", }, ) with _no_network(): decision = _offline_decision("org/mod-indexed") assert decision.blocked is True assert any( u["path"] == "0_Transformer/shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files ) def test_gate_allows_indexed_pickle_shard_with_safetensors_sibling(hf_cache): # A base model.safetensors makes the loader ignore the pickle index entirely, so it must not # block (mirrors the direct-file safetensors-sibling suppression). _make_cache( hf_cache, "org/indexed-both", { "pytorch_model.bin.index.json": ( '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' ), "shards/pytorch_model-00001-of-00001.bin": "pickle", "model.safetensors": "y", }, ) with _no_network(): assert _offline_decision("org/indexed-both").blocked is False def test_gate_allows_indexed_safetensors_shard_in_subdirectory(hf_cache): # A safetensors index lists inert shards -- following it must never block (guards against a # scanner that flags every indexed shard regardless of format). _make_cache( hf_cache, "org/st-indexed", { "model.safetensors.index.json": ( '{"weight_map": {"w": "shards/model-00001-of-00001.safetensors"}}' ), "shards/model-00001-of-00001.safetensors": "tensors", }, ) with _no_network(): assert _offline_decision("org/st-indexed").blocked is False def test_gate_blocks_on_index_path_traversal(hf_cache): # A weight_map entry escaping the snapshot via ".." is abnormal/hostile -> fail closed. _make_cache( hf_cache, "org/escape", {"pytorch_model.bin.index.json": '{"weight_map": {"w": "../../../../etc/evil.bin"}}'}, ) with _no_network(): assert _offline_decision("org/escape").blocked is True def test_gate_allows_symlinked_sharded_safetensors(tmp_path, monkeypatch): # Real HF caches store snapshot files as symlinks into blobs/. A resolve()-based containment # check would escape the snapshot and false-block every sharded model; the lexical gate must not. import hashlib import os from huggingface_hub.file_download import repo_folder_name root = tmp_path / "hub" root.mkdir() monkeypatch.setenv("HF_HOME", str(tmp_path)) monkeypatch.setenv("HF_HUB_CACHE", str(root)) monkeypatch.setattr( "utils.hf_cache_settings.get_hf_cache_paths", lambda: SimpleNamespace(hub_cache = root), ) repo_dir = root / repo_folder_name(repo_id = "org/sym", repo_type = "model") (repo_dir / "refs").mkdir(parents = True) (repo_dir / "refs" / "main").write_text(_COMMIT) blobs = repo_dir / "blobs" blobs.mkdir() snapshot = repo_dir / "snapshots" / _COMMIT (snapshot / "shards").mkdir(parents = True) def _blobbed(rel, content): digest = hashlib.sha256(content.encode()).hexdigest() (blobs / digest).write_text(content) target = snapshot / rel target.parent.mkdir(parents = True, exist_ok = True) target.symlink_to(os.path.relpath(blobs / digest, target.parent)) _blobbed("config.json", "{}") _blobbed( "model.safetensors.index.json", '{"weight_map": {"w": "shards/model-00001-of-00001.safetensors"}}', ) _blobbed("shards/model-00001-of-00001.safetensors", "tensors") with _no_network(): assert _offline_decision("org/sym").blocked is False def test_gate_allows_index_without_weight_map(hf_cache): # An index whose top-level JSON has no dict weight_map lets the loader resolve no shards, so it # must not crash or block on its own (only inert safetensors are cached here). _make_cache( hf_cache, "org/no-wm", {"model.safetensors.index.json": "[]", "model.safetensors": "x"}, ) with _no_network(): assert _offline_decision("org/no-wm").blocked is False def test_gate_allows_nothing_cached(hf_cache): with _no_network(): assert _offline_decision("org/missing").blocked is False def test_gate_allows_gguf_only(hf_cache): _make_cache(hf_cache, "org/gg", {"model.gguf": "x"}) with _no_network(): assert _offline_decision("org/gg").blocked is False def test_gate_blocks_pickle_in_module_subdir(hf_cache): # 0_Transformer is a module load root (listed in modules.json), so its pickle blocks. _make_cache( hf_cache, "org/mod", {"modules.json": _modules_json("0_Transformer"), "0_Transformer/pytorch_model.bin": "x"}, ) with _no_network(): assert _offline_decision("org/mod").blocked is True def test_gate_allows_pickle_in_subdir_with_safetensors(hf_cache): _make_cache( hf_cache, "org/mod2", { "modules.json": _modules_json("0_Transformer"), "0_Transformer/pytorch_model.bin": "x", "0_Transformer/model.safetensors": "y", }, ) with _no_network(): assert _offline_decision("org/mod2").blocked is False def test_gate_allows_unreferenced_nested_pickle(hf_cache): # A pickle in a dir NOT referenced by modules.json (e.g. nemo/) is never deserialized, so it # must not block the offline load (matches the online gate). _make_cache( hf_cache, "org/aux", { "modules.json": MODULES_JSON, # Transformer at the root only "model.safetensors": "w", "nemo/pytorch_model.bin": "x", }, ) with _no_network(): assert _offline_decision("org/aux").blocked is False def test_gate_blocks_adapter_pickle_without_safetensors(hf_cache): _make_cache(hf_cache, "org/ad", {"config.json": "{}", "adapter_model.bin": "x"}) with _no_network(): decision = _offline_decision("org/ad") assert decision.blocked is True assert any(u["path"] == "adapter_model.bin" for u in decision.unsafe_files) def test_gate_allows_adapter_pickle_with_adapter_safetensors(hf_cache): _make_cache(hf_cache, "org/ad2", {"adapter_model.bin": "x", "adapter_model.safetensors": "y"}) with _no_network(): assert _offline_decision("org/ad2").blocked is False def test_gate_blocks_base_pickle_with_only_adapter_safetensors_decoy(hf_cache): # A decoy adapter_model.safetensors must NOT suppress a base pytorch_model.bin (the base # loader would still deserialize the unscanned pickle). _make_cache(hf_cache, "org/decoy", {"pytorch_model.bin": "x", "adapter_model.safetensors": "y"}) with _no_network(): assert _offline_decision("org/decoy").blocked is True def test_gate_blocks_adapter_pickle_with_only_base_safetensors_decoy(hf_cache): # Symmetric: a base model.safetensors must NOT suppress an adapter_model.bin. _make_cache(hf_cache, "org/decoy2", {"adapter_model.bin": "x", "model.safetensors": "y"}) with _no_network(): assert _offline_decision("org/decoy2").blocked is True def test_gate_reports_snapshot_relative_path(hf_cache): _make_cache( hf_cache, "org/mod3", {"modules.json": _modules_json("0_Transformer"), "0_Transformer/pytorch_model.bin": "x"}, ) with _no_network(): decision = _offline_decision("org/mod3") assert decision.blocked is True assert any(u["path"] == "0_Transformer/pytorch_model.bin" for u in decision.unsafe_files) # ── evaluate_file_security: online path unchanged ──────────────── def test_online_default_blocks_unsafe(): status = { "scansDone": True, "filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}], } with patch( "huggingface_hub.model_info", side_effect = lambda *a, **k: SimpleNamespace(security_repo_status = status), ): assert evaluate_file_security("org/x").blocked is True def test_online_default_allows_clean(): status = {"scansDone": True, "filesWithIssues": []} with patch( "huggingface_hub.model_info", side_effect = lambda *a, **k: SimpleNamespace(security_repo_status = status), ): assert evaluate_file_security("org/x").blocked is False # ── embeddings guard + loader ──────────────────────────────────── def test_guard_offline_blocks_pickle_only(hf_cache): from core.rag.embeddings import UnsafeEmbeddingModelError, _guard_model_security _make_cache(hf_cache, "org/pk", {"config.json": "{}", "pytorch_model.bin": "x"}) with _no_network(): with pytest.raises(UnsafeEmbeddingModelError): _guard_model_security("org/pk", local_only = True) def test_guard_offline_allows_safetensors(hf_cache): from core.rag.embeddings import _guard_model_security _make_cache(hf_cache, "org/st", {"modules.json": MODULES_JSON, "model.safetensors": "x"}) with _no_network(): _guard_model_security("org/st", local_only = True) # must not raise def _install_fake_sentence_transformers(monkeypatch, captured): class FakeSentenceTransformer: def __init__( self, name, *, device = None, model_kwargs = None, local_files_only = False, **kw, ): captured["name"] = name captured["device"] = device captured["local_files_only"] = local_files_only module = types.ModuleType("sentence_transformers") module.SentenceTransformer = FakeSentenceTransformer monkeypatch.setitem(sys.modules, "sentence_transformers", module) def test_get_offline_loads_from_local_snapshot(hf_cache, monkeypatch): from core.rag import embeddings snapshot = _make_cache( hf_cache, "org/st", {"modules.json": MODULES_JSON, "model.safetensors": "x"} ) # TRANSFORMERS_OFFLINE only: a cached model loads from its local snapshot dir (a local path, # never the Hub), offline-safe on ANY sentence-transformers version. monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) monkeypatch.setattr(embeddings, "_model", None, raising = False) monkeypatch.setattr(embeddings, "_name", None, raising = False) monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None) monkeypatch.setattr(embeddings, "_device", lambda: "cpu") captured = {} _install_fake_sentence_transformers(monkeypatch, captured) with _no_network(): embeddings._get("org/st") assert captured["name"] == str(snapshot) def test_get_offline_uncached_uses_local_files_only(tmp_path, monkeypatch): from core.rag import embeddings empty = tmp_path / "hub" empty.mkdir() monkeypatch.setenv("HF_HUB_CACHE", str(empty)) monkeypatch.delenv("HF_HOME", raising = False) monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) monkeypatch.setattr(embeddings, "_model", None, raising = False) monkeypatch.setattr(embeddings, "_name", None, raising = False) monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None) monkeypatch.setattr(embeddings, "_device", lambda: "cpu") # No cache -> repo-id load forced cache-only (fails fast offline, not a hang). monkeypatch.setattr(embeddings, "_guard_model_security", lambda name, local_only = False: None) captured = {} _install_fake_sentence_transformers(monkeypatch, captured) embeddings._get("org/uncached-xyz") assert captured["name"] == "org/uncached-xyz" assert captured["local_files_only"] is True def test_get_online_omits_local_files_only(monkeypatch): from core.rag import embeddings monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) monkeypatch.setattr(embeddings, "_model", None, raising = False) monkeypatch.setattr(embeddings, "_name", None, raising = False) monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None) monkeypatch.setattr(embeddings, "_device", lambda: "cpu") # Isolate the loader wiring from the online guard's network calls. monkeypatch.setattr(embeddings, "_guard_model_security", lambda name, local_only = False: None) captured = {} _install_fake_sentence_transformers(monkeypatch, captured) embeddings._get("org/online") assert captured["local_files_only"] is False