Adopt single-pass signature hashing

comfy_execution/caching.py

@@ -64,12 +64,13 @@ _PRIMITIVE_SIGNATURE_TYPES = (int, float, str, bool, bytes, type(None))
 _CONTAINER_SIGNATURE_TYPES = (dict, list, tuple, set, frozenset)
 _MAX_SIGNATURE_DEPTH = 32
 _MAX_SIGNATURE_CONTAINER_VISITS = 10_000
+_FAILED_SIGNATURE = object()
 
 
-def _mark_signature_tainted(taint_state):
+def _primitive_signature_sort_key(obj):
-    """Record that signature sanitization hit a fail-closed condition."""
+    """Return a deterministic ordering key for primitive signature values."""
-    if taint_state is not None:
+    obj_type = type(obj)
-        taint_state["tainted"] = True
+    return ("primitive", obj_type.__module__, obj_type.__qualname__, repr(obj))
 
 
 def _sanitized_sort_key(obj, depth=0, max_depth=_MAX_SIGNATURE_DEPTH, active=None, memo=None):
@@ -123,21 +124,10 @@ def _sanitized_sort_key(obj, depth=0, max_depth=_MAX_SIGNATURE_DEPTH, active=Non
     return result
 
 
-def _sanitize_signature_input(obj, depth=0, max_depth=_MAX_SIGNATURE_DEPTH, active=None, memo=None, budget=None, taint_state=None):
+def _signature_to_hashable_impl(obj, depth=0, max_depth=_MAX_SIGNATURE_DEPTH, active=None, memo=None, budget=None):
-    """Normalize signature inputs to safe built-in containers.
+    """Canonicalize signature inputs directly into their final hashable form."""
-
-    Preserves built-in container type, replaces opaque runtime values with
-    Unhashable(), stops safely on cycles or excessive depth, memoizes repeated
-    built-in substructures so shared DAG-like inputs do not explode into
-    repeated recursive work, and optionally records when sanitization had to
-    fail closed anywhere in the traversed structure.
-    """
-    if taint_state is not None and taint_state.get("tainted"):
-        return Unhashable()
-
     if depth >= max_depth:
-        _mark_signature_tainted(taint_state)
+        return _FAILED_SIGNATURE
-        return Unhashable()
 
     if active is None:
         active = set()
@@ -148,93 +138,102 @@ def _sanitize_signature_input(obj, depth=0, max_depth=_MAX_SIGNATURE_DEPTH, acti
 
     obj_type = type(obj)
     if obj_type in _PRIMITIVE_SIGNATURE_TYPES:
-        return obj
+        return obj, _primitive_signature_sort_key(obj)
-    if obj_type not in _CONTAINER_SIGNATURE_TYPES:
+    if obj_type is Unhashable or obj_type not in _CONTAINER_SIGNATURE_TYPES:
-        _mark_signature_tainted(taint_state)
+        return _FAILED_SIGNATURE
-        return Unhashable()
 
     obj_id = id(obj)
     if obj_id in memo:
         return memo[obj_id]
     if obj_id in active:
-        _mark_signature_tainted(taint_state)
+        return _FAILED_SIGNATURE
-        return Unhashable()
 
     budget["remaining"] -= 1
     if budget["remaining"] < 0:
-        _mark_signature_tainted(taint_state)
+        return _FAILED_SIGNATURE
-        return Unhashable()
 
     active.add(obj_id)
     try:
         if obj_type is dict:
             try:
                 items = list(obj.items())
-                sort_memo = {}
+            except RuntimeError:
-                sanitized_items = [
+                return _FAILED_SIGNATURE
-                    (
-                        _sanitize_signature_input(key, depth + 1, max_depth, active, memo, budget, taint_state),
-                        _sanitize_signature_input(value, depth + 1, max_depth, active, memo, budget, taint_state),
-                    )
-                    for key, value in items
-                ]
-                ordered_items = [
-                    (
-                        (
-                            _sanitized_sort_key(key, depth + 1, max_depth, memo=sort_memo),
-                            _sanitized_sort_key(value, depth + 1, max_depth, memo=sort_memo),
-                        ),
-                        (key, value),
-                    )
-                    for key, value in sanitized_items
-                ]
-                ordered_items.sort(key=lambda item: item[0])
 
-                result = Unhashable()
+            ordered_items = []
-                for index in range(1, len(ordered_items)):
+            for key, value in items:
-                    previous_sort_key, previous_item = ordered_items[index - 1]
+                key_result = _signature_to_hashable_impl(key, depth + 1, max_depth, active, memo, budget)
-                    current_sort_key, current_item = ordered_items[index]
+                if key_result is _FAILED_SIGNATURE:
-                    if previous_sort_key == current_sort_key and previous_item != current_item:
+                    return _FAILED_SIGNATURE
-                        _mark_signature_tainted(taint_state)
+                value_result = _signature_to_hashable_impl(value, depth + 1, max_depth, active, memo, budget)
-                        break
+                if value_result is _FAILED_SIGNATURE:
-                else:
+                    return _FAILED_SIGNATURE
-                    result = {key: value for _, (key, value) in ordered_items}
+                key_value, key_sort = key_result
-            except RuntimeError:
+                value_value, value_sort = value_result
-                _mark_signature_tainted(taint_state)
+                ordered_items.append((((key_sort, value_sort)), (key_value, value_value)))
-                result = Unhashable()
+
-        elif obj_type is list:
+            ordered_items.sort(key=lambda item: item[0])
+            for index in range(1, len(ordered_items)):
+                previous_sort_key, previous_item = ordered_items[index - 1]
+                current_sort_key, current_item = ordered_items[index]
+                if previous_sort_key == current_sort_key and previous_item != current_item:
+                    return _FAILED_SIGNATURE
+
+            value = ("dict", tuple(item for _, item in ordered_items))
+            sort_key = ("dict", tuple(sort_key for sort_key, _ in ordered_items))
+        elif obj_type is list or obj_type is tuple:
             try:
                 items = list(obj)
-                result = [_sanitize_signature_input(item, depth + 1, max_depth, active, memo, budget, taint_state) for item in items]
             except RuntimeError:
-                _mark_signature_tainted(taint_state)
+                return _FAILED_SIGNATURE
-                result = Unhashable()
+
-        elif obj_type is tuple:
+            child_results = []
-            try:
+            for item in items:
-                items = list(obj)
+                child_result = _signature_to_hashable_impl(item, depth + 1, max_depth, active, memo, budget)
-                result = tuple(_sanitize_signature_input(item, depth + 1, max_depth, active, memo, budget, taint_state) for item in items)
+                if child_result is _FAILED_SIGNATURE:
-            except RuntimeError:
+                    return _FAILED_SIGNATURE
-                _mark_signature_tainted(taint_state)
+                child_results.append(child_result)
-                result = Unhashable()
+
-        elif obj_type is set:
+            container_tag = "list" if obj_type is list else "tuple"
-            try:
+            value = (container_tag, tuple(child for child, _ in child_results))
-                items = list(obj)
+            sort_key = (container_tag, tuple(child_sort for _, child_sort in child_results))
-                result = {_sanitize_signature_input(item, depth + 1, max_depth, active, memo, budget, taint_state) for item in items}
-            except RuntimeError:
-                _mark_signature_tainted(taint_state)
-                result = Unhashable()
         else:
             try:
                 items = list(obj)
-                result = frozenset(_sanitize_signature_input(item, depth + 1, max_depth, active, memo, budget, taint_state) for item in items)
             except RuntimeError:
-                _mark_signature_tainted(taint_state)
+                return _FAILED_SIGNATURE
-                result = Unhashable()
+
+            ordered_items = []
+            for item in items:
+                child_result = _signature_to_hashable_impl(item, depth + 1, max_depth, active, memo, budget)
+                if child_result is _FAILED_SIGNATURE:
+                    return _FAILED_SIGNATURE
+                child_value, child_sort = child_result
+                ordered_items.append((child_sort, child_value))
+
+            ordered_items.sort(key=lambda item: item[0])
+            for index in range(1, len(ordered_items)):
+                previous_sort_key, previous_value = ordered_items[index - 1]
+                current_sort_key, current_value = ordered_items[index]
+                if previous_sort_key == current_sort_key and previous_value != current_value:
+                    return _FAILED_SIGNATURE
+
+            container_tag = "set" if obj_type is set else "frozenset"
+            value = (container_tag, tuple(child_value for _, child_value in ordered_items))
+            sort_key = (container_tag, tuple(child_sort for child_sort, _ in ordered_items))
     finally:
         active.discard(obj_id)
 
-    memo[obj_id] = result
+    memo[obj_id] = (value, sort_key)
-    return result
+    return memo[obj_id]
+
+
+def _signature_to_hashable(obj, max_nodes=_MAX_SIGNATURE_CONTAINER_VISITS):
+    """Build the final cache-signature representation in one fail-closed pass."""
+    result = _signature_to_hashable_impl(obj, budget={"remaining": max_nodes})
+    if result is _FAILED_SIGNATURE:
+        return Unhashable()
+    return result[0]
 
 
 def to_hashable(obj, max_nodes=_MAX_SIGNATURE_CONTAINER_VISITS):
@@ -397,11 +396,7 @@ class CacheKeySetInputSignature(CacheKeySet):
         signature.append(await self.get_immediate_node_signature(dynprompt, node_id, order_mapping))
         for ancestor_id in ancestors:
             signature.append(await self.get_immediate_node_signature(dynprompt, ancestor_id, order_mapping))
-        taint_state = {"tainted": False}
+        return _signature_to_hashable(signature)
-        signature = _sanitize_signature_input(signature, taint_state=taint_state)
-        if taint_state["tainted"]:
-            return Unhashable()
-        return to_hashable(signature)
 
     async def get_immediate_node_signature(self, dynprompt, node_id, ancestor_order_mapping):
         """Build the cache-signature fragment for a node's immediate inputs.
@@ -424,7 +419,7 @@ class CacheKeySetInputSignature(CacheKeySet):
                 ancestor_index = ancestor_order_mapping[ancestor_id]
                 signature.append((key,("ANCESTOR", ancestor_index, ancestor_socket)))
             else:
-                signature.append((key, _sanitize_signature_input(inputs[key])))
+                signature.append((key, inputs[key]))
         return signature
 
     # This function returns a list of all ancestors of the given node. The order of the list is

tests-unit/execution_test/caching_test.py

@@ -1,4 +1,4 @@
-"""Unit tests for cache-signature sanitization and hash conversion hardening."""
+"""Unit tests for cache-signature canonicalization hardening."""
 
 import asyncio
 import importlib
@@ -76,96 +76,91 @@ def caching_module(monkeypatch):
     return module, nodes_module
 
 
-def test_sanitize_signature_input_handles_shared_builtin_substructures(caching_module):
+def test_signature_to_hashable_handles_shared_builtin_substructures(caching_module):
-    """Shared built-in substructures should sanitize without collapsing to Unhashable."""
+    """Shared built-in substructures should canonicalize without collapsing to Unhashable."""
     caching, _ = caching_module
     shared = [{"value": 1}, {"value": 2}]
 
-    sanitized = caching._sanitize_signature_input([shared, shared])
+    signature = caching._signature_to_hashable([shared, shared])
 
-    assert isinstance(sanitized, list)
+    assert signature[0] == "list"
-    assert sanitized[0] == sanitized[1]
+    assert signature[1][0] == signature[1][1]
-    assert sanitized[0][0]["value"] == 1
+    assert signature[1][0][0] == "list"
-    assert sanitized[0][1]["value"] == 2
+    assert signature[1][0][1][0] == ("dict", (("value", 1),))
+    assert signature[1][0][1][1] == ("dict", (("value", 2),))
 
 
-def test_sanitize_signature_input_marks_tainted_on_opaque_values(caching_module):
+def test_signature_to_hashable_fails_closed_on_opaque_values(caching_module):
-    """Opaque values should mark the containing signature as tainted."""
+    """Opaque values should collapse the full signature to Unhashable immediately."""
     caching, _ = caching_module
-    taint_state = {"tainted": False}
 
-    sanitized = caching._sanitize_signature_input(["safe", object()], taint_state=taint_state)
+    signature = caching._signature_to_hashable(["safe", object()])
 
-    assert isinstance(sanitized, list)
+    assert isinstance(signature, caching.Unhashable)
-    assert taint_state["tainted"] is True
-    assert isinstance(sanitized[1], caching.Unhashable)
 
 
-def test_sanitize_signature_input_stops_descending_after_taint(caching_module, monkeypatch):
+def test_signature_to_hashable_stops_descending_after_failure(caching_module, monkeypatch):
-    """Once tainted, later recursive calls should return immediately without deeper descent."""
+    """Once canonicalization fails, later recursive descent should stop immediately."""
     caching, _ = caching_module
-    original = caching._sanitize_signature_input
+    original = caching._signature_to_hashable_impl
     marker = object()
     marker_seen = False
 
-    def tracking_sanitize(obj, *args, **kwargs):
+    def tracking_canonicalize(obj, *args, **kwargs):
-        """Track whether recursion reaches the nested marker after tainting."""
+        """Track whether recursion reaches the nested marker after failure."""
         nonlocal marker_seen
         if obj is marker:
             marker_seen = True
         return original(obj, *args, **kwargs)
 
-    monkeypatch.setattr(caching, "_sanitize_signature_input", tracking_sanitize)
+    monkeypatch.setattr(caching, "_signature_to_hashable_impl", tracking_canonicalize)
 
-    taint_state = {"tainted": False}
+    signature = caching._signature_to_hashable([object(), [marker]])
-    sanitized = original([object(), [marker]], taint_state=taint_state)
 
-    assert isinstance(sanitized, list)
+    assert isinstance(signature, caching.Unhashable)
-    assert taint_state["tainted"] is True
     assert marker_seen is False
-    assert isinstance(sanitized[1], caching.Unhashable)
 
 
-def test_sanitize_signature_input_snapshots_list_before_recursing(caching_module, monkeypatch):
+def test_signature_to_hashable_snapshots_list_before_recursing(caching_module, monkeypatch):
-    """List sanitization should read a point-in-time snapshot before recursive descent."""
+    """List canonicalization should read a point-in-time snapshot before recursive descent."""
     caching, _ = caching_module
-    original = caching._sanitize_signature_input
+    original = caching._signature_to_hashable_impl
-    marker = object()
+    marker = ("marker",)
     values = [marker, 2]
 
-    def mutating_sanitize(obj, *args, **kwargs):
+    def mutating_canonicalize(obj, *args, **kwargs):
         """Mutate the live list during recursion to verify snapshot-based traversal."""
         if obj is marker:
             values[1] = 3
         return original(obj, *args, **kwargs)
 
-    monkeypatch.setattr(caching, "_sanitize_signature_input", mutating_sanitize)
+    monkeypatch.setattr(caching, "_signature_to_hashable_impl", mutating_canonicalize)
 
-    sanitized = original(values)
+    signature = caching._signature_to_hashable(values)
 
-    assert isinstance(sanitized, list)
+    assert signature == ("list", (("tuple", ("marker",)), 2))
-    assert sanitized[1] == 2
+    assert values[1] == 3
 
 
-def test_sanitize_signature_input_snapshots_dict_before_recursing(caching_module, monkeypatch):
+def test_signature_to_hashable_snapshots_dict_before_recursing(caching_module, monkeypatch):
-    """Dict sanitization should read a point-in-time snapshot before recursive descent."""
+    """Dict canonicalization should read a point-in-time snapshot before recursive descent."""
     caching, _ = caching_module
-    original = caching._sanitize_signature_input
+    original = caching._signature_to_hashable_impl
-    marker = object()
+    marker = ("marker",)
     values = {"first": marker, "second": 2}
 
-    def mutating_sanitize(obj, *args, **kwargs):
+    def mutating_canonicalize(obj, *args, **kwargs):
         """Mutate the live dict during recursion to verify snapshot-based traversal."""
         if obj is marker:
             values["second"] = 3
         return original(obj, *args, **kwargs)
 
-    monkeypatch.setattr(caching, "_sanitize_signature_input", mutating_sanitize)
+    monkeypatch.setattr(caching, "_signature_to_hashable_impl", mutating_canonicalize)
 
-    sanitized = original(values)
+    signature = caching._signature_to_hashable(values)
 
-    assert isinstance(sanitized, dict)
+    assert signature == ("dict", (("first", ("tuple", ("marker",))), ("second", 2)))
-    assert sanitized["second"] == 2
+    assert values["second"] == 3
 
 
 @pytest.mark.parametrize(
@@ -178,31 +173,31 @@ def test_sanitize_signature_input_snapshots_dict_before_recursing(caching_module
         lambda marker: {marker: "value"},
     ],
 )
-def test_sanitize_signature_input_fails_closed_on_runtimeerror(caching_module, monkeypatch, container_factory):
+def test_signature_to_hashable_fails_closed_on_runtimeerror(caching_module, monkeypatch, container_factory):
-    """Traversal RuntimeError should degrade sanitization to Unhashable."""
+    """Traversal RuntimeError should degrade canonicalization to Unhashable."""
     caching, _ = caching_module
-    original = caching._sanitize_signature_input
+    original = caching._signature_to_hashable_impl
     marker = object()
 
-    def raising_sanitize(obj, *args, **kwargs):
+    def raising_canonicalize(obj, *args, **kwargs):
         """Raise a traversal RuntimeError for the marker value and delegate otherwise."""
         if obj is marker:
             raise RuntimeError("container changed during iteration")
         return original(obj, *args, **kwargs)
 
-    monkeypatch.setattr(caching, "_sanitize_signature_input", raising_sanitize)
+    monkeypatch.setattr(caching, "_signature_to_hashable_impl", raising_canonicalize)
 
-    sanitized = original(container_factory(marker))
+    signature = caching._signature_to_hashable(container_factory(marker))
 
-    assert isinstance(sanitized, caching.Unhashable)
+    assert isinstance(signature, caching.Unhashable)
 
 
 def test_to_hashable_handles_shared_builtin_substructures(caching_module):
-    """Repeated sanitized content should hash stably for shared substructures."""
+    """The legacy helper should still hash sanitized built-ins stably when used directly."""
     caching, _ = caching_module
     shared = [{"value": 1}, {"value": 2}]
 
-    sanitized = caching._sanitize_signature_input([shared, shared])
+    sanitized = [shared, shared]
     hashable = caching.to_hashable(sanitized)
 
     assert hashable[0] == "list"
@@ -232,7 +227,7 @@ def test_to_hashable_fails_closed_on_runtimeerror(caching_module, monkeypatch, c
     assert isinstance(hashable, caching.Unhashable)
 
 
-def test_sanitize_signature_input_fails_closed_for_ambiguous_dict_ordering(caching_module):
+def test_signature_to_hashable_fails_closed_for_ambiguous_dict_ordering(caching_module):
     """Ambiguous dict sort ties should fail closed instead of depending on input order."""
     caching, _ = caching_module
     ambiguous = {
@@ -240,7 +235,7 @@ def test_sanitize_signature_input_fails_closed_for_ambiguous_dict_ordering(cachi
         _OpaqueValue(): _OpaqueValue(),
     }
 
-    sanitized = caching._sanitize_signature_input(ambiguous)
+    sanitized = caching._signature_to_hashable(ambiguous)
 
     assert isinstance(sanitized, caching.Unhashable)