From e6538c83bb7c466407444ee6e7a673e43786f466 Mon Sep 17 00:00:00 2001
From: Will Miao <willmiao.dev@gmail.com>
Date: Fri, 24 Jul 2026 12:07:18 +0800
Subject: [PATCH] fix(metadata): restore sha256 after hydrate_model_data to
 prevent KeyError in CivitAI fetch

hydrate_model_data replaces model_data with .metadata.json content which
may lack sha256 (corrupted file, concurrent write, etc.). Restore the
cached sha256 after hydration and persist the fix back to disk so
subsequent lookups don't hit the same error.

Also improve error log to include file_path for debugging.
---
 py/routes/handlers/model_handlers.py          | 34 ++++++++++++++++++-
 .../bulk_metadata_refresh_use_case.py         | 11 ++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/py/routes/handlers/model_handlers.py b/py/routes/handlers/model_handlers.py
index 408df67f..180a5068 100644
--- a/py/routes/handlers/model_handlers.py
+++ b/py/routes/handlers/model_handlers.py
@@ -537,6 +537,7 @@ class ModelManagementHandler:
                     # Update model_data with new hash
                     model_data["sha256"] = sha256
                     model_data["hash_status"] = "completed"
+                    hash_status = "completed"
                 else:
                     return web.json_response(
                         {"success": False, "error": "No SHA256 hash found"}, status=400
@@ -544,6 +545,32 @@ class ModelManagementHandler:
 
             await MetadataManager.hydrate_model_data(model_data)
 
+            # hydrate_model_data replaces model_data with .metadata.json content,
+            # which may lack sha256. Restore from cache and persist the fix.
+            if not model_data.get("sha256"):
+                if sha256:
+                    model_data["sha256"] = sha256
+                    model_data["hash_status"] = model_data.get("hash_status", hash_status)
+                    data_to_save = model_data.copy()
+                    data_to_save.pop("folder", None)
+                    await MetadataManager.save_metadata(file_path, data_to_save)
+                else:
+                    sha256 = await calculate_sha256(file_path)
+                    if sha256:
+                        model_data["sha256"] = sha256.lower()
+                        model_data["hash_status"] = "completed"
+                        data_to_save = model_data.copy()
+                        data_to_save.pop("folder", None)
+                        await MetadataManager.save_metadata(file_path, data_to_save)
+                    else:
+                        return web.json_response(
+                            {
+                                "success": False,
+                                "error": "Failed to compute SHA256 hash for model",
+                            },
+                            status=500,
+                        )
+
             success, error = await self._metadata_sync.fetch_and_update_model(
                 sha256=model_data["sha256"],
                 file_path=file_path,
@@ -566,7 +593,12 @@ class ModelManagementHandler:
                     {"success": False, "error": OFFLINE_FRIENDLY_MESSAGE},
                     status=503,
                 )
-            self._logger.error("Error fetching from CivitAI: %s", exc, exc_info=True)
+            self._logger.error(
+                "Error fetching from CivitAI for %s: %s",
+                locals().get("file_path", "unknown"),
+                exc,
+                exc_info=True,
+            )
             return web.json_response({"success": False, "error": str(exc)}, status=500)
 
     async def relink_civitai(self, request: web.Request) -> web.Response:
diff --git a/py/services/use_cases/bulk_metadata_refresh_use_case.py b/py/services/use_cases/bulk_metadata_refresh_use_case.py
index b037fde0..98136afd 100644
--- a/py/services/use_cases/bulk_metadata_refresh_use_case.py
+++ b/py/services/use_cases/bulk_metadata_refresh_use_case.py
@@ -126,6 +126,7 @@ class BulkMetadataRefreshUseCase:
                         if sha256:
                             model["sha256"] = sha256
                             model["hash_status"] = "completed"
+                            hash_status = "completed"
                         else:
                             self._logger.error(f"Failed to calculate hash for {file_path}")
                             failures.append({"name": model.get("model_name", file_path or "Unknown"), "error": "Failed to calculate hash"})
@@ -148,6 +149,16 @@ class BulkMetadataRefreshUseCase:
                     continue
 
                 await MetadataManager.hydrate_model_data(model)
+
+                # hydrate_model_data replaces model with .metadata.json content,
+                # which may lack sha256. Restore from cache and persist the fix.
+                if not model.get("sha256"):
+                    model["sha256"] = sha256
+                    model["hash_status"] = model.get("hash_status", hash_status)
+                    data_to_save = model.copy()
+                    data_to_save.pop("folder", None)
+                    await MetadataManager.save_metadata(file_path, data_to_save)
+
                 result, error_msg = await self._metadata_sync.fetch_and_update_model(
                     sha256=model["sha256"],
                     file_path=model["file_path"],