From a01c1224521d60b5cffc635d4e8378ca4fd49ce2 Mon Sep 17 00:00:00 2001 From: "open-swe[bot]" Date: Wed, 15 Apr 2026 14:19:20 +0000 Subject: [PATCH] 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> --- .../langgraph/checkpoint/serde/__init__.py | 3 + .../langgraph/checkpoint/serde/_msgpack.py | 12 ++ .../langgraph/checkpoint/serde/jsonplus.py | 92 ++++++++------- libs/checkpoint/tests/test_jsonplus.py | 106 ++++++++++++++++++ 4 files changed, 172 insertions(+), 41 deletions(-) diff --git a/libs/checkpoint/langgraph/checkpoint/serde/__init__.py b/libs/checkpoint/langgraph/checkpoint/serde/__init__.py index e69de29bb..dd268d5ec 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/__init__.py @@ -0,0 +1,3 @@ +from langgraph.checkpoint.serde._msgpack import register_safe_types + +__all__ = ["register_safe_types"] diff --git a/libs/checkpoint/langgraph/checkpoint/serde/_msgpack.py b/libs/checkpoint/langgraph/checkpoint/serde/_msgpack.py index 0b09a9a4e..1aab277f2 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/_msgpack.py +++ b/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)) diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 4ef4a40b2..ba78fb2ba 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/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: diff --git a/libs/checkpoint/tests/test_jsonplus.py b/libs/checkpoint/tests/test_jsonplus.py index 9242c59bf..e4126ecd6 100644 --- a/libs/checkpoint/tests/test_jsonplus.py +++ b/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__) + )