fix: dedup deserialization warnings and add register_safe_types utility

Warnings in _check_allowed and _check_allowed_method now track already-warned
types/methods per ext_hook closure, so the same warning is only emitted once.

Added register_safe_types() to _msgpack.py as a public API for users to add
custom types to the global safe-types set, exported from
langgraph.checkpoint.serde.

Co-authored-by: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com>

libs/checkpoint/langgraph/checkpoint/serde/__init__.py

@@ -0,0 +1,3 @@
+from langgraph.checkpoint.serde._msgpack import register_safe_types
+
+__all__ = ["register_safe_types"]

libs/checkpoint/langgraph/checkpoint/serde/_msgpack.py

@@ -1,5 +1,6 @@
 import os
 from collections.abc import Iterable
+from inspect import isclass
 from typing import cast
 
 STRICT_MSGPACK_ENABLED = os.getenv("LANGGRAPH_STRICT_MSGPACK", "false").lower() in (
@@ -87,3 +88,14 @@ SAFE_MSGPACK_METHODS: frozenset[tuple[str, str, str]] = frozenset(
 
 
 AllowedMsgpackModules = Iterable[tuple[str, ...] | type]
+
+_CUSTOM_SAFE_MSGPACK_TYPES: set[tuple[str, ...]] = set()
+
+
+def register_safe_types(types: Iterable[type | tuple[str, ...]]) -> None:
+    """Add types to the global safe-types allowlist so they deserialize without warnings."""
+    for t in types:
+        if isclass(t):
+            _CUSTOM_SAFE_MSGPACK_TYPES.add((t.__module__, t.__name__))
+        else:
+            _CUSTOM_SAFE_MSGPACK_TYPES.add(cast(tuple[str, ...], t))

libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py

@@ -511,51 +511,59 @@ def _create_msgpack_ext_hook(
         An ext_hook function for use with ormsgpack.unpackb.
     """
 
+    _warned_types: set[tuple[str, ...]] = set()
+    _warned_methods: set[tuple[str, str, str]] = set()
+
     def _check_allowed(module: str, name: str) -> bool:
         """Check if type is allowed. Returns True if allowed, False if blocked."""
         key = (module, name)
 
-        if key in _lg_msgpack.SAFE_MSGPACK_TYPES:
+        if (
+            key in _lg_msgpack.SAFE_MSGPACK_TYPES
+            or key in _lg_msgpack._CUSTOM_SAFE_MSGPACK_TYPES
+        ):
             return True
 
         if allowed_modules is True:
-            # default is to warn but allow unregistered types
+            if key not in _warned_types:
+                _warned_types.add(key)
+                emit_serde_event(
+                    {
+                        "kind": "msgpack_unregistered_allowed",
+                        "module": module,
+                        "name": name,
+                    }
+                )
+                logger.warning(
+                    "Deserializing unregistered type %s.%s from checkpoint. "
+                    "This will be blocked in a future version. "
+                    "Add to allowed_msgpack_modules to silence: [(%r, %r)]",
+                    module,
+                    name,
+                    module,
+                    name,
+                )
+            return True
+        if allowed_modules is not None:
+            if key in allowed_modules:
+                return True
+        if key not in _warned_types:
+            _warned_types.add(key)
             emit_serde_event(
                 {
-                    "kind": "msgpack_unregistered_allowed",
+                    "kind": "msgpack_blocked",
                     "module": module,
                     "name": name,
                 }
             )
             logger.warning(
-                "Deserializing unregistered type %s.%s from checkpoint. "
-                "This will be blocked in a future version. "
-                "Add to allowed_msgpack_modules to silence: [(%r, %r)]",
+                "Blocked deserialization of %s.%s - not in allowed_msgpack_modules. "
+                "Add to allowed_msgpack_modules to allow: [(%r, %r)]",
                 module,
                 name,
                 module,
                 name,
             )
-            return True
-        if allowed_modules is not None:
-            if key in allowed_modules:
-                return True
-        # strict mode blocks unregistered types
-        emit_serde_event(
-            {
-                "kind": "msgpack_blocked",
-                "module": module,
-                "name": name,
-            }
-        )
-        logger.warning(
-            "Blocked deserialization of %s.%s - not in allowed_msgpack_modules. "
-            "Add to allowed_msgpack_modules to allow: [(%r, %r)]",
-            module,
-            name,
-            module,
-            name,
-        )
         return False
 
     def _check_allowed_method(module: str, name: str, method: str) -> bool:
@@ -563,21 +571,23 @@ def _create_msgpack_ext_hook(
         key = (module, name, method)
         if key in _lg_msgpack.SAFE_MSGPACK_METHODS:
             return True
-        emit_serde_event(
-            {
-                "kind": "msgpack_method_blocked",
-                "module": module,
-                "name": name,
-                "method": method,
-            }
-        )
-        logger.warning(
-            "Blocked deserialization of method call %s.%s.%s - "
-            "not in allowed methods set.",
-            module,
-            name,
-            method,
-        )
+        if key not in _warned_methods:
+            _warned_methods.add(key)
+            emit_serde_event(
+                {
+                    "kind": "msgpack_method_blocked",
+                    "module": module,
+                    "name": name,
+                    "method": method,
+                }
+            )
+            logger.warning(
+                "Blocked deserialization of method call %s.%s.%s - "
+                "not in allowed methods set.",
+                module,
+                name,
+                method,
+            )
         return False
 
     def ext_hook(code: int, data: bytes) -> Any:

libs/checkpoint/tests/test_jsonplus.py

@@ -983,3 +983,109 @@ def test_msgpack_nested_pydantic_serializes_as_dict(
     # No blocking should occur - inner is serialized as dict, not ext
     assert "blocked" not in caplog.text.lower()
     assert result == obj
+
+
+def test_msgpack_warning_dedup_default_mode(caplog: pytest.LogCaptureFixture) -> None:
+    """Repeated deserialization of the same unregistered type should only warn once."""
+    current = _lg_msgpack.STRICT_MSGPACK_ENABLED
+    _lg_msgpack.STRICT_MSGPACK_ENABLED = False
+    serde = JsonPlusSerializer()
+
+    obj = AnotherPydantic(foo="test")
+    dumped = serde.dumps_typed(obj)
+
+    caplog.clear()
+    serde.loads_typed(dumped)
+    first_count = caplog.text.lower().count("unregistered type")
+    assert first_count >= 1
+
+    caplog.clear()
+    serde.loads_typed(dumped)
+    assert caplog.text.lower().count("unregistered type") == 0
+
+    _lg_msgpack.STRICT_MSGPACK_ENABLED = current
+
+
+def test_msgpack_warning_dedup_strict_mode(caplog: pytest.LogCaptureFixture) -> None:
+    """Repeated deserialization of the same blocked type should only warn once."""
+    serde = JsonPlusSerializer(allowed_msgpack_modules=None)
+
+    obj = AnotherPydantic(foo="test")
+    dumped = serde.dumps_typed(obj)
+
+    caplog.clear()
+    serde.loads_typed(dumped)
+    first_count = caplog.text.lower().count("blocked")
+    assert first_count >= 1
+
+    caplog.clear()
+    serde.loads_typed(dumped)
+    assert caplog.text.lower().count("blocked") == 0
+
+
+def test_register_safe_types_silences_warning(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Types added via register_safe_types should deserialize without warnings."""
+    _lg_msgpack.register_safe_types([AnotherPydantic])
+    try:
+        serde = JsonPlusSerializer()
+
+        obj = AnotherPydantic(foo="test")
+
+        caplog.clear()
+        dumped = serde.dumps_typed(obj)
+        result = serde.loads_typed(dumped)
+
+        assert "unregistered type" not in caplog.text.lower()
+        assert "blocked" not in caplog.text.lower()
+        assert result == obj
+    finally:
+        _lg_msgpack._CUSTOM_SAFE_MSGPACK_TYPES.discard(
+            (AnotherPydantic.__module__, AnotherPydantic.__name__)
+        )
+
+
+def test_register_safe_types_with_tuples(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """register_safe_types should accept (module, name) tuples."""
+    _lg_msgpack.register_safe_types([("tests.test_jsonplus", "AnotherPydantic")])
+    try:
+        serde = JsonPlusSerializer()
+
+        obj = AnotherPydantic(foo="test")
+
+        caplog.clear()
+        dumped = serde.dumps_typed(obj)
+        result = serde.loads_typed(dumped)
+
+        assert "unregistered type" not in caplog.text.lower()
+        assert "blocked" not in caplog.text.lower()
+        assert result == obj
+    finally:
+        _lg_msgpack._CUSTOM_SAFE_MSGPACK_TYPES.discard(
+            ("tests.test_jsonplus", "AnotherPydantic")
+        )
+
+
+def test_register_safe_types_strict_mode(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Types added via register_safe_types should also work in strict mode."""
+    _lg_msgpack.register_safe_types([AnotherPydantic])
+    try:
+        serde = JsonPlusSerializer(allowed_msgpack_modules=None)
+
+        obj = AnotherPydantic(foo="test")
+
+        caplog.clear()
+        dumped = serde.dumps_typed(obj)
+        result = serde.loads_typed(dumped)
+
+        assert "blocked" not in caplog.text.lower()
+        assert result == obj
+    finally:
+        _lg_msgpack._CUSTOM_SAFE_MSGPACK_TYPES.discard(
+            (AnotherPydantic.__module__, AnotherPydantic.__name__)
+        )