diff --git a/services/model_manager.py b/services/model_manager.py index 9b57c46..64417e1 100644 --- a/services/model_manager.py +++ b/services/model_manager.py @@ -114,6 +114,7 @@ MODEL_TYPE_EXCLUSION_REASONS = { "diffusers": "diffusers is folder-valued and needs a directory-tree install design", "classifiers": "classifiers are extensionless and need a dedicated content policy", "custom_nodes": "custom_nodes are executable plugin code, not managed model files", + "datasets": "datasets contain user-managed training data, not managed model weights", } MODEL_TYPE_ALIASES = { "config": "configs", diff --git a/services/preflight.py b/services/preflight.py index 0ace3ed..e1fd727 100644 --- a/services/preflight.py +++ b/services/preflight.py @@ -48,7 +48,9 @@ _LEGACY_INVENTORY_CACHE_KEY = "inventory" _INVENTORY_LOCK = threading.RLock() _INVENTORY_SCAN_THREAD: threading.Thread | None = None _INVENTORY_ERROR_RETRY_SEC = 5 -_INVENTORY_EXCLUDED_MODEL_TYPES = {"custom_nodes"} +# CRITICAL: datasets are user training data and executable custom_nodes are code; neither +# may cross the model inventory filename boundary. +_INVENTORY_EXCLUDED_MODEL_TYPES = {"custom_nodes", "datasets"} # Heuristic mapping: input_key -> folder_paths type _INPUT_KEY_MAP = { @@ -147,7 +149,11 @@ def _scan_model_inventory(checkpoint: List[str] | None = None) -> Dict[str, List def _copy_inventory_snapshot(models: Dict[str, List[str]]) -> Dict[str, List[str]]: - return {key: list(value) for key, value in (models or {}).items()} + return { + key: list(value) + for key, value in (models or {}).items() + if key not in _INVENTORY_EXCLUDED_MODEL_TYPES + } def _inventory_snapshot_stale_locked(now: float | None = None) -> bool: diff --git a/tests/test_dataset_inventory_boundary.py b/tests/test_dataset_inventory_boundary.py new file mode 100644 index 0000000..274fb03 --- /dev/null +++ b/tests/test_dataset_inventory_boundary.py @@ -0,0 +1,255 @@ +from __future__ import annotations + +import json +import tempfile +import time +import unittest +from pathlib import Path +from unittest.mock import MagicMock, patch + +import services.preflight +from services.model_manager import ( + MODEL_TYPE_EXCLUSION_REASONS, + ModelManager, + ModelManagerError, + _model_type_exclusion_reason, + _norm_model_type, +) + +try: + from aiohttp import web + from aiohttp.test_utils import AioHTTPTestCase, unittest_run_loop +except Exception: # pragma: no cover + web = None # type: ignore + AioHTTPTestCase = unittest.TestCase # type: ignore + + def unittest_run_loop(fn): # type: ignore + return fn + + +DATASET_SENTINEL = "private-subject-caption-0001.txt" +DATASET_REASON = ( + "datasets contain user-managed training data, not managed model weights" +) + + +class DatasetInventoryServiceBoundaryTests(unittest.TestCase): + def setUp(self): + services.preflight._reset_inventory_state_for_tests() + + def tearDown(self): + services.preflight._reset_inventory_state_for_tests() + + @staticmethod + def _folder_paths() -> MagicMock: + folder_paths = MagicMock() + folder_paths.folder_names_and_paths = { + "checkpoints": [], + "clip": [], + "unet": [], + "datasets": [], + "custom_nodes": [], + } + folder_paths.get_filename_list.side_effect = lambda model_type: ( + [DATASET_SENTINEL] if model_type == "datasets" else [] + ) + return folder_paths + + def test_dynamic_resolution_and_scan_exclude_dataset_user_data(self): + folder_paths = self._folder_paths() + + with patch.object(services.preflight, "folder_paths", folder_paths): + model_types = services.preflight._resolve_inventory_model_types() + snapshot = services.preflight._scan_model_inventory() + + self.assertNotIn("datasets", model_types) + self.assertNotIn("custom_nodes", model_types) + self.assertIn("text_encoders", model_types) + self.assertIn("diffusion_models", model_types) + self.assertNotIn("datasets", snapshot) + self.assertNotIn(DATASET_SENTINEL, json.dumps(snapshot)) + scanned_types = { + call.args[0] for call in folder_paths.get_filename_list.call_args_list + } + self.assertNotIn("datasets", scanned_types) + self.assertNotIn("custom_nodes", scanned_types) + + def test_preexisting_cached_dataset_snapshot_is_filtered_on_copy(self): + folder_paths = self._folder_paths() + services.preflight._CACHE[services.preflight._INVENTORY_SNAPSHOT_KEY] = { + "checkpoints": ["safe.safetensors"], + "datasets": [DATASET_SENTINEL], + } + services.preflight._CACHE[services.preflight._INVENTORY_SNAPSHOT_TS_KEY] = 123.0 + + with patch.object(services.preflight, "folder_paths", folder_paths): + snapshot = services.preflight.get_model_inventory_snapshot( + trigger_refresh=False + ) + + self.assertEqual(snapshot["models"], {"checkpoints": ["safe.safetensors"]}) + self.assertNotIn(DATASET_SENTINEL, json.dumps(snapshot)) + self.assertEqual(snapshot["snapshot_ts"], 123.0) + self.assertIn(snapshot["scan_state"], {"idle", "error", "refreshing"}) + + def test_background_publication_filters_defensive_dataset_result(self): + folder_paths = self._folder_paths() + + with ( + patch.object(services.preflight, "folder_paths", folder_paths), + patch.object( + services.preflight, + "_scan_model_inventory", + return_value={ + "checkpoints": ["safe.safetensors"], + "datasets": [DATASET_SENTINEL], + }, + ), + ): + services.preflight._inventory_refresh_worker() + snapshot = services.preflight.get_model_inventory_snapshot( + trigger_refresh=False + ) + + self.assertEqual(snapshot["models"], {"checkpoints": ["safe.safetensors"]}) + self.assertEqual(snapshot["scan_state"], "idle") + self.assertFalse(snapshot["stale"]) + self.assertIsNone(snapshot["last_error"]) + self.assertNotIn(DATASET_SENTINEL, json.dumps(snapshot)) + + +class DatasetModelManagerBoundaryTests(unittest.TestCase): + def setUp(self): + self.temp_dir = tempfile.TemporaryDirectory() + root = Path(self.temp_dir.name) + self.manager = ModelManager( + state_root=root / "state", + install_root=root / "install", + ) + + def tearDown(self): + self.temp_dir.cleanup() + + def test_dataset_type_has_stable_user_data_exclusion(self): + self.assertEqual(MODEL_TYPE_EXCLUSION_REASONS["datasets"], DATASET_REASON) + for value in ("datasets", " DATASETS ", "Datasets"): + with self.subTest(value=value): + self.assertEqual(_norm_model_type(value), "other") + self.assertEqual(_model_type_exclusion_reason(value), DATASET_REASON) + self.assertNotIn("datasets", self.manager._model_type_to_subdir) + + def test_dataset_download_is_rejected_before_url_path_or_task_work(self): + payload = { + "model_id": "dataset-attempt", + "name": "Dataset Attempt", + "model_type": "datasets", + "source": "catalog", + "source_label": "Catalog", + "download_url": f"https://example.invalid/{DATASET_SENTINEL}", + "expected_sha256": "a" * 64, + "provenance": { + "publisher": "private-publisher-sentinel", + "license": "private-license-sentinel", + "source_url": "https://example.invalid/private-source", + }, + "destination_subdir": "../private-dataset", + "filename": DATASET_SENTINEL, + } + + with ( + patch.object(self.manager, "_assert_budget") as assert_budget, + patch.object( + self.manager, + "_validate_url_policy", + side_effect=AssertionError( + "dataset exclusion must run before URL validation" + ), + ) as validate_url, + patch.object(self.manager, "_validate_provenance") as validate_provenance, + patch.object(self.manager, "_sanitize_subdir") as sanitize_subdir, + patch.object(self.manager, "_sanitize_filename") as sanitize_filename, + self.assertRaises(ModelManagerError) as caught, + ): + self.manager.create_download_task(**payload) + + self.assertEqual(caught.exception.code, "unsupported_model_type") + self.assertEqual( + caught.exception.detail, + "model_type 'datasets' is not supported for managed install/import: " + + DATASET_REASON, + ) + assert_budget.assert_called_once_with() + validate_url.assert_not_called() + validate_provenance.assert_not_called() + sanitize_subdir.assert_not_called() + sanitize_filename.assert_not_called() + self.assertEqual(self.manager._tasks, {}) + self.assertEqual(list(self.manager.install_root.rglob("*")), []) + + +@unittest.skipIf(web is None, "aiohttp not installed") +class DatasetInventoryAliasApiTests(AioHTTPTestCase): + def setUp(self): + super().setUp() + services.preflight._reset_inventory_state_for_tests() + + def tearDown(self): + services.preflight._reset_inventory_state_for_tests() + super().tearDown() + + async def get_application(self): + from api.preflight_handler import inventory_handler + + app = web.Application() + for prefix in ( + "/openclaw", + "/moltbot", + "/api/openclaw", + "/api/moltbot", + ): + app.router.add_get(f"{prefix}/preflight/inventory", inventory_handler) + return app + + @patch("api.preflight_handler.check_rate_limit", return_value=True) + @patch("api.preflight_handler.require_admin_token", return_value=(True, None)) + @patch("api.preflight_handler._get_node_class_mappings", return_value={}) + @unittest_run_loop + async def test_every_inventory_alias_omits_cached_dataset_names( + self, + _nodes_mock, + _admin_mock, + _rate_limit_mock, + ): + services.preflight._CACHE[services.preflight._INVENTORY_SNAPSHOT_KEY] = { + "checkpoints": ["safe.safetensors"], + "datasets": [DATASET_SENTINEL], + } + services.preflight._CACHE[services.preflight._INVENTORY_SNAPSHOT_TS_KEY] = ( + time.time() + ) + folder_paths = DatasetInventoryServiceBoundaryTests._folder_paths() + + with patch.object(services.preflight, "folder_paths", folder_paths): + for prefix in ( + "/openclaw", + "/moltbot", + "/api/openclaw", + "/api/moltbot", + ): + with self.subTest(prefix=prefix): + response = await self.client.get(f"{prefix}/preflight/inventory") + payload = await response.json() + rendered = json.dumps(payload) + self.assertEqual(response.status, 200) + self.assertEqual( + payload["models"], {"checkpoints": ["safe.safetensors"]} + ) + self.assertNotIn("datasets", rendered) + self.assertNotIn(DATASET_SENTINEL, rendered) + self.assertIsInstance(payload["snapshot_ts"], float) + self.assertEqual(payload["scan_state"], "idle") + self.assertFalse(payload["stale"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_model_manager_api.py b/tests/test_model_manager_api.py index 37843b1..8907f7d 100644 --- a/tests/test_model_manager_api.py +++ b/tests/test_model_manager_api.py @@ -1,4 +1,5 @@ import hashlib +import json import tempfile import time import unittest @@ -146,6 +147,50 @@ class TestModelManagerAPI(AioHTTPTestCase): resp = await self.client.get("/openclaw/models/search") self.assertEqual(resp.status, 403) + @patch("api.model_manager.require_admin_token", return_value=(True, None)) + @patch( + "services.model_manager.validate_outbound_url", + side_effect=AssertionError( + "dataset exclusion must run before outbound URL validation" + ), + ) + @unittest_run_loop + async def test_dataset_download_returns_stable_unsupported_contract( + self, mock_validate_url, _mock_admin + ): + response = await self.client.post( + "/openclaw/models/downloads", + json={ + "model_id": "dataset-attempt", + "name": "Dataset Attempt", + "model_type": "datasets", + "source": "catalog", + "source_label": "Catalog", + "download_url": "https://example.invalid/private-caption.txt", + "expected_sha256": "a" * 64, + "provenance": { + "publisher": "private-publisher-sentinel", + "license": "private-license-sentinel", + "source_url": "https://example.invalid/private-source", + }, + "destination_subdir": "../private-dataset", + "filename": "private-caption.txt", + }, + ) + payload = await response.json() + + self.assertEqual(response.status, 400) + self.assertFalse(payload["ok"]) + self.assertEqual(payload["error"], "unsupported_model_type") + self.assertEqual( + payload["detail"], + "model_type 'datasets' is not supported for managed install/import: " + "datasets contain user-managed training data, not managed model weights", + ) + self.assertNotIn("private-caption", json.dumps(payload)) + self.assertFalse(mock_validate_url.called) + self.assertEqual(self.manager._tasks, {}) + @patch("api.model_manager.require_admin_token", return_value=(True, None)) @patch( "services.model_manager.validate_outbound_url", diff --git a/web/tests/unit/model_manager_tab.test.js b/web/tests/unit/model_manager_tab.test.js index f16415b..5fe6f7a 100644 --- a/web/tests/unit/model_manager_tab.test.js +++ b/web/tests/unit/model_manager_tab.test.js @@ -129,5 +129,6 @@ describe("model_manager_tab", () => { expect(options).not.toContain("diffusers"); expect(options).not.toContain("classifiers"); expect(options).not.toContain("custom_nodes"); + expect(options).not.toContain("datasets"); }); });