Fix cache_key collapse of non-pydantic private/__main__ objects

ccflow/tests/utils/test_tokenize.py

@@ -1254,3 +1254,126 @@ def test_enum_includes_module(self):
         E1 = _enum.Enum("Color", {"RED": 1}, module="pkg.one")
         E2 = _enum.Enum("Color", {"RED": 1}, module="pkg.two")
         assert tokenize(E1.RED) != tokenize(E2.RED)
+
+
+class TestPrivateModuleNoStateCollapse:
+    """Objects in ``__main__``/private modules must fold state, not name-only-collapse.
+
+    Regression for a fallback that name-only-tokenized any object whose type module
+    started with ``_`` or contained ``._``, silently dropping instance state and
+    collapsing genuinely-different values onto one cache key.
+    """
+
+    @staticmethod
+    def _make(module):
+        class _Add:
+            def __init__(self, n):
+                self.n = n
+
+            def __call__(self, x):
+                return x + self.n
+
+        _Add.__module__ = module
+        return _Add
+
+    @pytest.mark.parametrize("module", ["_secret", "pkg._internal", "__main__"])
+    def test_picklable_private_module_instances_distinct(self, module):
+        cls = self._make(module)
+        assert compute_data_token(cls(5)) != compute_data_token(cls(7))
+
+    @pytest.mark.parametrize("module", ["_secret", "pkg._internal", "__main__"])
+    def test_picklable_private_module_instances_deterministic(self, module):
+        cls = self._make(module)
+        assert compute_data_token(cls(5)) == compute_data_token(cls(5))
+
+    def test_stateful_frozen_dataclass_in_private_module_distinct(self):
+        import dataclasses
+
+        @dataclasses.dataclass(frozen=True)
+        class _Config:
+            values: tuple
+
+        _Config.__module__ = "pkg._private"
+        assert compute_data_token(_Config((1, 2))) != compute_data_token(_Config((1, 3)))
+
+    def test_unpicklable_user_object_in_main_fails_loud(self):
+        import threading
+
+        class _Holder:
+            def __init__(self):
+                self.lock = threading.Lock()
+
+        _Holder.__module__ = "__main__"
+        with pytest.raises(TypeError):
+            normalize_token(_Holder())
+
+    def test_unpicklable_user_object_in_private_module_fails_loud(self):
+        # An unpicklable object in a private *user* package must fail loud rather than
+        # name-only-collapse: two distinct instances would otherwise share a key. Only
+        # curated interpreter-internal modules are exempt, never arbitrary ``pkg._x``.
+        import threading
+
+        class _Resource:
+            def __init__(self):
+                self.lock = threading.Lock()
+
+        _Resource.__module__ = "company._models"
+        with pytest.raises(TypeError):
+            normalize_token(_Resource())
+
+    def test_unpicklable_interpreter_internal_is_name_only_stable(self):
+        # _thread.lock is an allowlisted interpreter internal: unpicklable, but a behavior-
+        # irrelevant primitive with no semantic identity, so it degrades to a stable name-only
+        # token rather than crashing behavior hashing (same treatment as 3.14 _abc._abc_data).
+        import threading
+
+        lock = threading.Lock()
+        token = normalize_token(lock)
+        assert token == ("__internal__", "_thread", type(lock).__qualname__)
+        assert normalize_token(threading.Lock()) == token
+
+    def test_picklable_framework_internal_is_name_only(self):
+        # Picklable-but-behavior-irrelevant framework internals (pydantic compiled
+        # validators) fold only module + qualname, never volatile runtime state.
+        class _Fake:
+            pass
+
+        _Fake.__module__ = "pydantic_core._pydantic_core"
+        assert normalize_token(_Fake()) == ("__internal__", "pydantic_core._pydantic_core", _Fake.__qualname__)
+
+    @pytest.mark.parametrize(
+        "module",
+        [
+            "pydantic._internal._model_construction",
+            "pydantic_core._pydantic_core",
+            "pydantic.plugin._schema_validator",
+        ],
+    )
+    def test_private_pydantic_submodules_are_name_only(self, module):
+        # Private submodules across the whole pydantic namespace stay name-only, matching the
+        # original behavior; public modules like pydantic.fields must not (guarded below).
+        class _Fake:
+            pass
+
+        _Fake.__module__ = module
+        assert normalize_token(_Fake()) == ("__internal__", module, _Fake.__qualname__)
+
+    def test_public_pydantic_module_is_not_name_only(self):
+        class _Fake:
+            pass
+
+        _Fake.__module__ = "pydantic.fields"
+        assert normalize_token(_Fake())[0] == "__cloudpickle__"
+
+    def test_real_pydantic_internal_objects_are_name_only(self):
+        # Guards the concrete regression: a model's compiled validator/serializer live in a
+        # private pydantic submodule and must be keyed by name, whatever module path the
+        # installed pydantic version puts them in.
+        import pydantic
+
+        class M(pydantic.BaseModel):
+            x: int = 1
+
+        for obj in (M.__pydantic_validator__, M.__pydantic_serializer__):
+            token = normalize_token(obj)
+            assert token[0] == "__internal__", (type(obj).__module__, token)

ccflow/utils/tokenize.py

@@ -74,6 +74,32 @@ def _with_cycle_check(obj: Any, build: Callable[[], Any]) -> Any:
             _visited.reset(token)
 
 
+# Frameworks whose *private* submodules expose compiled/derived objects that are picklable but carry
+# volatile runtime state (e.g. ``pydantic._internal._model_construction``, ``pydantic_core._pydantic_core``,
+# ``pydantic.plugin._schema_validator``). Key these by module + qualname only, never fold their bytes.
+# Matched by top-level package plus any underscore-prefixed path component, so public modules such as
+# ``pydantic.fields`` are unaffected and unrelated user packages are never captured.
+_FRAMEWORK_INTERNAL_TOP_LEVEL = ("pydantic", "pydantic_core")
+
+# Unpicklable interpreter internals that are behavior-irrelevant and safe to key by name alone: either
+# derived state already captured elsewhere in the hash (``_abc._abc_data``, surfaced in ABC-derived
+# classes' closures on Python 3.14) or primitives with no semantic identity (``_thread`` lock/RLock).
+# Anything else that fails to serialize raises loudly rather than silently collapsing distinct objects
+# (e.g. DB connections, sockets, or user classes in a private ``pkg._internal`` module) onto one key.
+_NAME_ONLY_UNPICKLABLE_MODULES = ("_abc", "_thread")
+
+
+def _is_framework_internal(module: str) -> bool:
+    """Return whether ``module`` is a private submodule of a known framework (see the list above)."""
+    parts = module.split(".")
+    return parts[0] in _FRAMEWORK_INTERNAL_TOP_LEVEL and any(part.startswith("_") for part in parts)
+
+
+def _module_in(module: str, prefixes: tuple[str, ...]) -> bool:
+    """Return whether ``module`` equals or is a submodule of any prefix in ``prefixes``."""
+    return any(module == prefix or module.startswith(prefix + ".") for prefix in prefixes)
+
+
 @singledispatch
 def normalize_token(obj: Any) -> Any:
     """Produce a canonical, deterministically hashable representation of ``obj``.
@@ -85,13 +111,17 @@ def _(obj):
             return ("mytype", ...)
 
     Unknown types fall back to a ``cloudpickle``-based digest, raising ``TypeError`` on pickling failure.
+
+    Serializability, not the module name, is the primary signal: any object cloudpickle can serialize
+    has its state folded in, so two genuinely different values (e.g. instances authored in ``__main__``
+    or a private module) never collapse to one key. The module name is consulted only via two small
+    curated allowlists -- one for private framework internals whose bytes are volatile, and one for
+    unpicklable interpreter internals that are behavior-irrelevant -- so that any other object that
+    cannot be serialized fails loud instead of silently sharing a key.
     """
-    # Python 3.14 exposes internal objects (e.g. _abc._abc_data, pydantic._internal.*,
-    # pydantic_core._pydantic_core.*) in function closures of ABC-derived classes.
-    # These are not behavior-relevant; produce a stable token keyed by module + qualname
-    # so they don't affect hashing.
     obj_module = getattr(type(obj), "__module__", "") or ""
-    if obj_module.startswith("_") or "._" in obj_module:
+
+    if _is_framework_internal(obj_module):
         return ("__internal__", obj_module, type(obj).__qualname__)
 
     try:
@@ -102,6 +132,8 @@ def _(obj):
     try:
         pickled = cloudpickle.dumps(obj)
     except Exception as exc:
+        if _module_in(obj_module, _NAME_ONLY_UNPICKLABLE_MODULES):
+            return ("__internal__", obj_module, type(obj).__qualname__)
         raise TypeError(f"Cannot tokenize object of type {type(obj).__qualname__}. Register a normalize_token handler for this type.") from exc
     return ("__cloudpickle__", hashlib.sha256(pickled).hexdigest())