From 04097e6902f881fce66ebbc2f67edcb1b39dea1d Mon Sep 17 00:00:00 2001
From: Deep Mehta <deepme987@gmail.com>
Date: Tue, 3 Mar 2026 13:04:10 -0800
Subject: [PATCH] fix: address coderabbit review feedback

- Add try/except to _build_context, return None when hash fails
- Return None from _serialize_cache_key on total failure (no id()-based fallback)
- Replace hex-like test literal with non-secret placeholder

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 comfy_execution/cache_provider.py             | 10 ++++----
 comfy_execution/caching.py                    | 23 ++++++++++++-------
 .../execution_test/test_cache_provider.py     |  4 ++--
 3 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/comfy_execution/cache_provider.py b/comfy_execution/cache_provider.py
index 79a5b71e8..e689952a3 100644
--- a/comfy_execution/cache_provider.py
+++ b/comfy_execution/cache_provider.py
@@ -242,11 +242,12 @@ def _canonicalize(obj: Any) -> Any:
         return ("__repr__", repr(obj))
 
 
-def _serialize_cache_key(cache_key: Any) -> str:
+def _serialize_cache_key(cache_key: Any) -> Optional[str]:
     """
     Serialize cache key to a hex digest string for external storage.
 
-    Returns SHA256 hex string suitable for Redis/database keys.
+    Returns SHA256 hex string suitable for Redis/database keys,
+    or None if serialization fails entirely (fail-closed).
 
     Note: Uses canonicalize + JSON serialization instead of pickle because
     pickle is NOT deterministic across Python sessions due to hash randomization
@@ -263,8 +264,9 @@ def _serialize_cache_key(cache_key: Any) -> str:
         try:
             serialized = pickle.dumps(cache_key, protocol=4)
             return hashlib.sha256(serialized).hexdigest()
-        except Exception:
-            return hashlib.sha256(str(id(cache_key)).encode()).hexdigest()
+        except Exception as fallback_error:
+            _logger.warning(f"Failed pickle fallback for cache key: {fallback_error}")
+            return None
 
 
 def _contains_nan(obj: Any) -> bool:
diff --git a/comfy_execution/caching.py b/comfy_execution/caching.py
index 91669e583..47681f332 100644
--- a/comfy_execution/caching.py
+++ b/comfy_execution/caching.py
@@ -316,14 +316,21 @@ class BasicCache:
             return ''
 
     def _build_context(self, node_id, cache_key):
-        """Build CacheContext with hash. Returns None if hashing fails on NaN."""
-        from comfy_execution.cache_provider import CacheContext, _serialize_cache_key
-        return CacheContext(
-            prompt_id=self._current_prompt_id,
-            node_id=node_id,
-            class_type=self._get_class_type(node_id),
-            cache_key_hash=_serialize_cache_key(cache_key)
-        )
+        """Build CacheContext with hash. Returns None if hashing fails."""
+        from comfy_execution.cache_provider import CacheContext, _serialize_cache_key, _logger
+        try:
+            cache_key_hash = _serialize_cache_key(cache_key)
+            if cache_key_hash is None:
+                return None
+            return CacheContext(
+                prompt_id=self._current_prompt_id,
+                node_id=node_id,
+                class_type=self._get_class_type(node_id),
+                cache_key_hash=cache_key_hash,
+            )
+        except Exception as e:
+            _logger.warning(f"Failed to build cache context for node {node_id}: {e}")
+            return None
 
     async def _ensure_subcache(self, node_id, children_ids):
         subcache_key = self.cache_key_set.get_subcache_key(node_id)
diff --git a/tests-unit/execution_test/test_cache_provider.py b/tests-unit/execution_test/test_cache_provider.py
index 1d9ca095a..a11673610 100644
--- a/tests-unit/execution_test/test_cache_provider.py
+++ b/tests-unit/execution_test/test_cache_provider.py
@@ -333,13 +333,13 @@ class TestCacheContext:
             prompt_id="prompt-123",
             node_id="node-456",
             class_type="KSampler",
-            cache_key_hash="abcdef1234567890" * 4,
+            cache_key_hash="a" * 64,
         )
 
         assert context.prompt_id == "prompt-123"
         assert context.node_id == "node-456"
         assert context.class_type == "KSampler"
-        assert context.cache_key_hash == "abcdef1234567890" * 4
+        assert context.cache_key_hash == "a" * 64
 
 
 class TestCacheValue: