Merge pull request #551 from willmiao/codex/refactor-model-metadata-saving-logic

fix: hydrate cached metadata before persisting updates

py/services/metadata_sync_service.py

@@ -153,7 +153,12 @@ class MetadataSyncService:
         model_data: Dict[str, Any],
         update_cache_func: Callable[[str, str, Dict[str, Any]], Awaitable[bool]],
     ) -> tuple[bool, Optional[str]]:
-        """Fetch metadata for a model and update both disk and cache state."""
+        """Fetch metadata for a model and update both disk and cache state.
+
+        Callers should hydrate ``model_data`` via ``MetadataManager.hydrate_model_data``
+        before invoking this method so that the persisted payload retains all known
+        metadata fields.
+        """
 
         if not isinstance(model_data, dict):
             error = f"Invalid model_data type: {type(model_data)}"
@@ -176,6 +181,7 @@ class MetadataSyncService:
                 metadata_provider = await self._get_default_provider()
 
             civitai_metadata, error = await metadata_provider.get_model_by_hash(sha256)
+
             if not civitai_metadata:
                 if error == "Model not found":
                     model_data["from_civitai"] = False

tests/services/test_metadata_sync_service.py

@@ -1,4 +1,5 @@
 from types import SimpleNamespace
+from typing import Any, Dict
 from unittest.mock import AsyncMock
 
 import pytest
@@ -20,7 +21,10 @@ def build_service(
     default_provider: SimpleNamespace | None = None,
     provider_selector: AsyncMock | None = None,
 ):
-    metadata_manager = SimpleNamespace(save_metadata=AsyncMock())
+    metadata_manager = SimpleNamespace(
+        save_metadata=AsyncMock(),
+        hydrate_model_data=AsyncMock(side_effect=lambda payload: payload),
+    )
     preview_service = SimpleNamespace(ensure_preview_for_metadata=AsyncMock())
     settings = DummySettings(settings_values)
 
@@ -105,12 +109,27 @@ async def test_fetch_and_update_model_success_updates_cache(tmp_path):
     }
     helpers.default_provider.get_model_by_hash.return_value = (civitai_payload, None)
 
-    model_data = {"model_name": "Local", "folder": "root"}
+    model_path = tmp_path / "model.safetensors"
+
+    async def hydrate(payload: Dict[str, Any]) -> Dict[str, Any]:
+        payload["hydrated"] = True
+        return payload
+
+    helpers.metadata_manager.hydrate_model_data.side_effect = hydrate
+
+    model_data = {
+        "model_name": "Local",
+        "folder": "root",
+        "file_path": str(model_path),
+    }
     update_cache = AsyncMock(return_value=True)
 
+    await hydrate(model_data)
+    helpers.metadata_manager.hydrate_model_data.reset_mock()
+
     ok, error = await helpers.service.fetch_and_update_model(
         sha256="abc",
-        file_path=str(tmp_path / "model.safetensors"),
+        file_path=str(model_path),
         model_data=model_data,
         update_cache_func=update_cache,
     )
@@ -120,10 +139,15 @@ async def test_fetch_and_update_model_success_updates_cache(tmp_path):
     assert model_data["civitai_deleted"] is False
     assert "civitai" in model_data
 
-    metadata_path = str(tmp_path / "model.metadata.json")
+    helpers.metadata_manager.hydrate_model_data.assert_not_awaited()
+    assert model_data["hydrated"] is True
+
+    metadata_path = str(model_path.with_suffix(".metadata.json"))
     await_args = helpers.metadata_manager.save_metadata.await_args_list
     assert await_args, "expected metadata to be persisted"
-    assert await_args[-1][0][0] == metadata_path
+    last_call = await_args[-1]
+    assert last_call.args[0] == metadata_path
+    assert last_call.args[1]["hydrated"] is True
     update_cache.assert_awaited_once()
 
 
@@ -132,14 +156,26 @@ async def test_fetch_and_update_model_handles_missing_remote_metadata(tmp_path):
     helpers = build_service()
     helpers.default_provider.get_model_by_hash.return_value = (None, "Model not found")
 
+    model_path = tmp_path / "model.safetensors"
+
+    async def hydrate(payload: Dict[str, Any]) -> Dict[str, Any]:
+        payload["hydrated"] = True
+        return payload
+
+    helpers.metadata_manager.hydrate_model_data.side_effect = hydrate
+
     model_data = {
         "model_name": "Local",
         "folder": "sub",
+        "file_path": str(model_path),
     }
 
+    await hydrate(model_data)
+    helpers.metadata_manager.hydrate_model_data.reset_mock()
+
     ok, error = await helpers.service.fetch_and_update_model(
         sha256="missing",
-        file_path=str(tmp_path / "model.safetensors"),
+        file_path=str(model_path),
         model_data=model_data,
         update_cache_func=AsyncMock(),
     )
@@ -149,17 +185,24 @@ async def test_fetch_and_update_model_handles_missing_remote_metadata(tmp_path):
     assert model_data["from_civitai"] is False
     assert model_data["civitai_deleted"] is True
 
+    helpers.metadata_manager.hydrate_model_data.assert_not_awaited()
+    assert model_data["hydrated"] is True
+
     helpers.metadata_manager.save_metadata.assert_awaited_once()
-    args, _ = helpers.metadata_manager.save_metadata.await_args
-    assert args[0].endswith("model.safetensors")
-    assert "folder" not in args[1]
+    call_args = helpers.metadata_manager.save_metadata.await_args
+    assert call_args.args[0].endswith("model.safetensors")
+    assert "folder" not in call_args.args[1]
+    assert call_args.args[1]["hydrated"] is True
 
 
 @pytest.mark.asyncio
 async def test_fetch_and_update_model_respects_deleted_without_archive():
     helpers = build_service(settings_values={"enable_metadata_archive_db": False})
 
-    model_data = {"civitai_deleted": True}
+    model_data = {
+        "civitai_deleted": True,
+        "file_path": "/tmp/model.safetensors",
+    }
     update_cache = AsyncMock()
 
     ok, error = await helpers.service.fetch_and_update_model(
@@ -172,6 +215,7 @@ async def test_fetch_and_update_model_respects_deleted_without_archive():
     assert not ok
     assert "metadata archive DB is not enabled" in error
     helpers.default_provider_factory.assert_not_awaited()
+    helpers.metadata_manager.hydrate_model_data.assert_not_awaited()
     update_cache.assert_not_awaited()
 
 

tests/services/test_route_support_services.py

@@ -30,6 +30,9 @@ class RecordingMetadataManager:
         Path(metadata_path).write_text(json.dumps(metadata))
         return True
 
+    async def hydrate_model_data(self, model_data: Dict[str, Any]) -> Dict[str, Any]:
+        return model_data
+
 
 class RecordingPreviewService:
     def __init__(self) -> None:
@@ -114,6 +117,7 @@ def test_metadata_sync_fetch_and_update_updates_cache(tmp_path: Path) -> None:
     )
 
     model_data = {"sha256": "abc", "file_path": str(tmp_path / "model.safetensors")}
+    asyncio.run(manager.hydrate_model_data(model_data))
     success, error = asyncio.run(
         service.fetch_and_update_model(
             sha256="abc",