fix: schema funkiness, again

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

@@ -37,6 +37,46 @@ LC_REVIVER = Reviver()
 EMPTY_BYTES = b""
 logger = logging.getLogger(__name__)
 
+SAFE_MSGPACK_TYPES: frozenset[tuple[str, ...]] = frozenset(
+    {
+        # datetime types
+        ("datetime", "datetime"),
+        ("datetime", "date"),
+        ("datetime", "time"),
+        ("datetime", "timedelta"),
+        ("datetime", "timezone"),
+        # uuid
+        ("uuid", "UUID"),
+        # numeric
+        ("decimal", "Decimal"),
+        # collections
+        ("builtins", "set"),
+        ("builtins", "frozenset"),
+        ("collections", "deque"),
+        # ip addresses
+        ("ipaddress", "IPv4Address"),
+        ("ipaddress", "IPv4Interface"),
+        ("ipaddress", "IPv4Network"),
+        ("ipaddress", "IPv6Address"),
+        ("ipaddress", "IPv6Interface"),
+        ("ipaddress", "IPv6Network"),
+        # pathlib
+        ("pathlib", "Path"),
+        ("pathlib", "PosixPath"),
+        ("pathlib", "WindowsPath"),
+        # pathlib in Python 3.13+
+        ("pathlib._local", "Path"),
+        ("pathlib._local", "PosixPath"),
+        ("pathlib._local", "WindowsPath"),
+        # langgraph
+        ("langgraph.types", "Send"),
+        ("langgraph.types", "Interrupt"),
+        ("langgraph.types", "Command"),
+        ("langgraph.types", "StateSnapshot"),
+        ("langgraph.types", "PregelTask"),
+    }
+)
+
 
 class JsonPlusSerializer(SerializerProtocol):
     """Serializer that uses ormsgpack, with optional fallbacks.
@@ -54,18 +94,29 @@ class JsonPlusSerializer(SerializerProtocol):
         *,
         pickle_fallback: bool = False,
         allowed_json_modules: Sequence[tuple[str, ...]] | Literal[True] | None = None,
+        # TODO: change default to None once users have had time to configure allowlists
+        allowed_msgpack_modules: Sequence[tuple[str, ...]]
+        | Literal[True]
+        | None = True,
         __unpack_ext_hook__: Callable[[int, bytes], Any] | None = None,
     ) -> None:
         self.pickle_fallback = pickle_fallback
-        self._allowed_modules = (
+        # JSON allowlist
+        self._allowed_json_modules: set[tuple[str, ...]] | Literal[True] | None = (
             {mod_and_name for mod_and_name in allowed_json_modules}
             if allowed_json_modules and allowed_json_modules is not True
             else (allowed_json_modules if allowed_json_modules is True else None)
         )
+        # Msgpack allowlist
+        self._allowed_msgpack_modules: set[tuple[str, ...]] | Literal[True] | None = (
+            {mod_and_name for mod_and_name in allowed_msgpack_modules}
+            if allowed_msgpack_modules and allowed_msgpack_modules is not True
+            else (allowed_msgpack_modules if allowed_msgpack_modules is True else None)
+        )
         self._unpack_ext_hook = (
             __unpack_ext_hook__
             if __unpack_ext_hook__ is not None
-            else _msgpack_ext_hook
+            else _create_msgpack_ext_hook(self._allowed_msgpack_modules)
         )
 
     def _encode_constructor_args(
@@ -90,7 +141,7 @@ class JsonPlusSerializer(SerializerProtocol):
         return out
 
     def _reviver(self, value: dict[str, Any]) -> Any:
-        if self._allowed_modules and (
+        if self._allowed_json_modules and (
             value.get("lc", None) == 2
             and value.get("type", None) == "constructor"
             and value.get("id", None) is not None
@@ -107,7 +158,7 @@ class JsonPlusSerializer(SerializerProtocol):
         return LC_REVIVER(value)
 
     def _revive_lc2(self, value: dict[str, Any]) -> Any:
-        self._check_allowed_modules(value)
+        self._check_allowed_json_modules(value)
 
         [*module, name] = value["id"]
         try:
@@ -139,7 +190,7 @@ class JsonPlusSerializer(SerializerProtocol):
         except Exception:
             return None
 
-    def _check_allowed_modules(self, value: dict[str, Any]) -> None:
+    def _check_allowed_json_modules(self, value: dict[str, Any]) -> None:
         needed = tuple(value["id"])
         method = value.get("method")
         if isinstance(method, list):
@@ -150,7 +201,7 @@ class JsonPlusSerializer(SerializerProtocol):
             method_display = "<init>"
 
         dotted = ".".join(needed)
-        if not self._allowed_modules:
+        if not self._allowed_json_modules:
             raise InvalidModuleError(
                 f"Refused to deserialize JSON constructor: {dotted} (method: {method_display}). "
                 "No allowed_json_modules configured.\n\n"
@@ -161,9 +212,9 @@ class JsonPlusSerializer(SerializerProtocol):
                 "or plain-JSON representations revived without import-time side effects."
             )
 
-        if self._allowed_modules is True:
+        if self._allowed_json_modules is True:
             return
-        if needed in self._allowed_modules:
+        if needed in self._allowed_json_modules:
             return
 
         raise InvalidModuleError(
@@ -448,92 +499,158 @@ def _msgpack_default(obj: Any) -> str | ormsgpack.Ext:
         raise TypeError(f"Object of type {obj.__class__.__name__} is not serializable")
 
 
-def _msgpack_ext_hook(code: int, data: bytes) -> Any:
-    if code == EXT_CONSTRUCTOR_SINGLE_ARG:
-        try:
-            tup = ormsgpack.unpackb(
-                data, ext_hook=_msgpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
-            )
-            # module, name, arg
-            return getattr(importlib.import_module(tup[0]), tup[1])(tup[2])
-        except Exception:
-            return
-    elif code == EXT_CONSTRUCTOR_POS_ARGS:
-        try:
-            tup = ormsgpack.unpackb(
-                data, ext_hook=_msgpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
-            )
-            # module, name, args
-            return getattr(importlib.import_module(tup[0]), tup[1])(*tup[2])
-        except Exception:
-            return
-    elif code == EXT_CONSTRUCTOR_KW_ARGS:
-        try:
-            tup = ormsgpack.unpackb(
-                data, ext_hook=_msgpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
-            )
-            # module, name, args
-            return getattr(importlib.import_module(tup[0]), tup[1])(**tup[2])
-        except Exception:
-            return
-    elif code == EXT_METHOD_SINGLE_ARG:
-        try:
-            tup = ormsgpack.unpackb(
-                data, ext_hook=_msgpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
-            )
-            # module, name, arg, method
-            return getattr(getattr(importlib.import_module(tup[0]), tup[1]), tup[3])(
-                tup[2]
-            )
-        except Exception:
-            return
-    elif code == EXT_PYDANTIC_V1:
-        try:
-            tup = ormsgpack.unpackb(
-                data, ext_hook=_msgpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
-            )
-            # module, name, kwargs
-            cls = getattr(importlib.import_module(tup[0]), tup[1])
-            try:
-                return cls(**tup[2])
-            except Exception:
-                return cls.construct(**tup[2])
-        except Exception:
-            # for pydantic objects we can't find/reconstruct
-            # let's return the kwargs dict instead
-            try:
-                return tup[2]
-            except NameError:
-                return
-    elif code == EXT_PYDANTIC_V2:
-        try:
-            tup = ormsgpack.unpackb(
-                data, ext_hook=_msgpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
-            )
-            # module, name, kwargs, method
-            cls = getattr(importlib.import_module(tup[0]), tup[1])
-            try:
-                return cls(**tup[2])
-            except Exception:
-                return cls.model_construct(**tup[2])
-        except Exception:
-            # for pydantic objects we can't find/reconstruct
-            # let's return the kwargs dict instead
-            try:
-                return tup[2]
-            except NameError:
-                return
-    elif code == EXT_NUMPY_ARRAY:
-        try:
-            import numpy as _np
+def _create_msgpack_ext_hook(
+    allowed_modules: set[tuple[str, ...]] | Literal[True] | None,
+) -> Callable[[int, bytes], Any]:
+    """Create msgpack ext hook with allowlist.
 
-            dtype_str, shape, order, buf = ormsgpack.unpackb(
-                data, ext_hook=_msgpack_ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
+    Args:
+        allowed_modules: Set of (module, name) tuples that are allowed to be
+        deserialized, or True to allow all with warnings for unregistered types, or None to only allow safe types.
+
+    Returns:
+        An ext_hook function for use with ormsgpack.unpackb.
+    """
+
+    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 SAFE_MSGPACK_TYPES:
+            return True
+
+        if allowed_modules is not None and allowed_modules is not True:
+            if key in allowed_modules:
+                return True
+
+        if allowed_modules is True:
+            # default is to warn but allow unregistered types
+            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,
             )
-            arr = _np.frombuffer(buf, dtype=_np.dtype(dtype_str))
-            return arr.reshape(shape, order=order)
-        except Exception:
-            return
+            return True
+        else:
+            # strict mode blocks unregistered types
+            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 ext_hook(code: int, data: bytes) -> Any:
+        if code == EXT_CONSTRUCTOR_SINGLE_ARG:
+            try:
+                tup = ormsgpack.unpackb(
+                    data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
+                )
+                if not _check_allowed(tup[0], tup[1]):
+                    return None
+                # module, name, arg
+                return getattr(importlib.import_module(tup[0]), tup[1])(tup[2])
+            except Exception:
+                return None
+        elif code == EXT_CONSTRUCTOR_POS_ARGS:
+            try:
+                tup = ormsgpack.unpackb(
+                    data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
+                )
+                if not _check_allowed(tup[0], tup[1]):
+                    return None
+                # module, name, args
+                return getattr(importlib.import_module(tup[0]), tup[1])(*tup[2])
+            except Exception:
+                return None
+        elif code == EXT_CONSTRUCTOR_KW_ARGS:
+            try:
+                tup = ormsgpack.unpackb(
+                    data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
+                )
+                if not _check_allowed(tup[0], tup[1]):
+                    return None
+                # module, name, kwargs
+                return getattr(importlib.import_module(tup[0]), tup[1])(**tup[2])
+            except Exception:
+                return None
+        elif code == EXT_METHOD_SINGLE_ARG:
+            try:
+                tup = ormsgpack.unpackb(
+                    data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
+                )
+                if not _check_allowed(tup[0], tup[1]):
+                    return None
+                # module, name, arg, method
+                return getattr(
+                    getattr(importlib.import_module(tup[0]), tup[1]), tup[3]
+                )(tup[2])
+            except Exception:
+                return None
+        elif code == EXT_PYDANTIC_V1:
+            try:
+                tup = ormsgpack.unpackb(
+                    data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
+                )
+                if not _check_allowed(tup[0], tup[1]):
+                    return None
+                # module, name, kwargs
+                cls = getattr(importlib.import_module(tup[0]), tup[1])
+                try:
+                    return cls(**tup[2])
+                except Exception:
+                    return cls.construct(**tup[2])
+            except Exception:
+                # for pydantic objects we can't find/reconstruct
+                # let's return the kwargs dict instead
+                try:
+                    return tup[2]
+                except NameError:
+                    return None
+        elif code == EXT_PYDANTIC_V2:
+            try:
+                tup = ormsgpack.unpackb(
+                    data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
+                )
+                if not _check_allowed(tup[0], tup[1]):
+                    return None
+                # module, name, kwargs, method
+                cls = getattr(importlib.import_module(tup[0]), tup[1])
+                try:
+                    return cls(**tup[2])
+                except Exception:
+                    return cls.model_construct(**tup[2])
+            except Exception:
+                # for pydantic objects we can't find/reconstruct
+                # let's return the kwargs dict instead
+                try:
+                    return tup[2]
+                except NameError:
+                    return None
+        elif code == EXT_NUMPY_ARRAY:
+            try:
+                import numpy as _np
+
+                dtype_str, shape, order, buf = ormsgpack.unpackb(
+                    data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
+                )
+                arr = _np.frombuffer(buf, dtype=_np.dtype(dtype_str))
+                return arr.reshape(shape, order=order)
+            except Exception:
+                return None
+        return None
+
+    return ext_hook
+
+
+_msgpack_ext_hook = _create_msgpack_ext_hook(allowed_modules=None)
 
 
 def _msgpack_ext_hook_to_json(code: int, data: bytes) -> Any:

libs/checkpoint/tests/test_jsonplus.py

@@ -512,5 +512,88 @@ def test_serde_jsonplus_pandas_series(series: pd.Series) -> None:
 
     assert dumped[0] == "pickle"
     result = serde.loads_typed(dumped)
-
     assert result.equals(series)
+
+
+def test_msgpack_safe_types_no_warning(caplog: pytest.LogCaptureFixture) -> None:
+    """Test safe types deserialize without warnings."""
+
+    serde = JsonPlusSerializer()
+
+    safe_objects = [
+        datetime.now(),
+        date.today(),
+        time(12, 30),
+        timezone.utc,
+        uuid.uuid4(),
+        Decimal("123.45"),
+        {1, 2, 3},
+        frozenset([1, 2, 3]),
+        deque([1, 2, 3]),
+        IPv4Address("192.168.1.1"),
+        pathlib.Path("/tmp/test"),
+    ]
+
+    for obj in safe_objects:
+        caplog.clear()
+        dumped = serde.dumps_typed(obj)
+        result = serde.loads_typed(dumped)
+        assert "unregistered type" not in caplog.text.lower(), (
+            f"Unexpected warning for {type(obj)}"
+        )
+        assert result is not None
+
+
+def test_msgpack_pydantic_warns_by_default(caplog: pytest.LogCaptureFixture) -> None:
+    """Pydantic models not in allowlist should log warning but still deserialize.
+
+    TODO: We'll want to change this to block unregistered types in the future."""
+
+    serde = JsonPlusSerializer()
+
+    obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world"))
+
+    caplog.clear()
+    dumped = serde.dumps_typed(obj)
+    result = serde.loads_typed(dumped)
+
+    assert "unregistered type" in caplog.text.lower()
+    assert "allowed_msgpack_modules" in caplog.text
+    assert result == obj
+
+
+def test_msgpack_allowlist_silences_warning(caplog: pytest.LogCaptureFixture) -> None:
+    """Types in allowed_msgpack_modules should deserialize without warnings."""
+
+    serde = JsonPlusSerializer(
+        allowed_msgpack_modules=[
+            ("tests.test_jsonplus", "MyPydantic"),
+            ("tests.test_jsonplus", "InnerPydantic"),
+        ]
+    )
+
+    obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world"))
+
+    caplog.clear()
+    dumped = serde.dumps_typed(obj)
+    result = serde.loads_typed(dumped)
+
+    assert "unregistered type" not in caplog.text.lower()
+    assert result == obj
+
+
+def test_msgpack_none_blocks_unregistered(caplog: pytest.LogCaptureFixture) -> None:
+    """allowed_msgpack_modules=None should block unregistered types.
+
+    TODO: This will be the default behavior in the future."""
+
+    serde = JsonPlusSerializer(allowed_msgpack_modules=None)
+
+    obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world"))
+
+    caplog.clear()
+    dumped = serde.dumps_typed(obj)
+    result = serde.loads_typed(dumped)
+
+    assert "blocked" in caplog.text.lower()
+    assert result is None

libs/cli/langgraph_cli/schemas.py

@@ -126,7 +126,7 @@ class SerdeConfig(TypedDict, total=False):
     If omitted, no serde is set up (the object store will still be present, however)."""
 
     allowed_json_modules: list[list[str]] | bool | None
-    """Optional. List of allowed python modules to de-serialize custom objects from.
+    """Optional. List of allowed python modules to de-serialize custom objects from JSON.
     
     If provided, only the specified modules will be allowed to be deserialized.
     If omitted, no modules are allowed, and the object returned will simply be a json object OR
@@ -146,7 +146,34 @@ class SerdeConfig(TypedDict, total=False):
     Example:
     {...
         "serde": {
-            "allowed_json_modules": true
+            "allowed_json_modules": True
+        }
+    }
+
+    """
+    allowed_msgpack_modules: list[list[str]] | bool | None
+    """Optional. List of allowed python modules to de-serialize custom objects from msgpack.
+
+    Known safe types (langgraph.checkpoint.serde.jsonplus.SAFE_MSGPACK_TYPES) are always
+    allowed regardless of this setting. Use this to allowlist your custom Pydantic models,
+    dataclasses, and other user-defined types.
+
+    If True (default), unregistered types will log a warning but still be deserialized.
+    If None, only known safe types will be deserialized; unregistered types will be blocked.
+
+    Example - allowlist specific types (no warnings for these):
+    {...
+        "serde": {
+            "allowed_msgpack_modules": [
+                ["my_agent", "models", "MyState"],
+            ]
+        }
+    }
+
+    Example - strict mode (only safe types allowed):
+    {...
+        "serde": {
+            "allowed_msgpack_modules": null
         }
     }
     
@@ -306,8 +333,7 @@ class EncryptionConfig(TypedDict, total=False):
     """Configuration for custom at-rest encryption logic.
 
     Allows you to implement custom encryption for sensitive data stored in the database,
-    including metadata fields and checkpoint blobs.
-    """
+    including metadata fields and checkpoint blobs."""
 
     path: str
     """Required. Path to an instance of the Encryption() class that implements custom encryption handlers.

libs/cli/schemas/schema.json

@@ -591,7 +591,27 @@
               "type": "null"
             }
           ],
-          "description": "Optional. List of allowed python modules to de-serialize custom objects from.\n\nIf provided, only the specified modules will be allowed to be deserialized.\nIf omitted, no modules are allowed, and the object returned will simply be a json object OR\na deserialized langchain object.\n"
+          "description": "Optional. List of allowed python modules to de-serialize custom objects from JSON.\n\nIf provided, only the specified modules will be allowed to be deserialized.\nIf omitted, no modules are allowed, and the object returned will simply be a json object OR\na deserialized langchain object.\n"
+        },
+        "allowed_msgpack_modules": {
+          "anyOf": [
+            {
+              "type": "array",
+              "items": {
+                "type": "array",
+                "items": {
+                  "type": "string"
+                }
+              }
+            },
+            {
+              "type": "boolean"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "description": "Optional. List of allowed python modules to de-serialize custom objects from msgpack.\n\nKnown safe types (langgraph.checkpoint.serde.jsonplus.SAFE_MSGPACK_TYPES) are always\nallowed regardless of this setting. Use this to allowlist your custom Pydantic models,\ndataclasses, and other user-defined types.\n\nIf True (default), unregistered types will log a warning but still be deserialized.\nIf None, only known safe types will be deserialized; unregistered types will be blocked.\n\n{...\n[\"my_agent\", \"models\", \"MyState\"],\n]\n}\n}\n\n{...\n}\n}\n\n"
         },
         "pickle_fallback": {
           "type": "boolean",

libs/cli/schemas/schema.v0.json

@@ -591,7 +591,27 @@
               "type": "null"
             }
           ],
-          "description": "Optional. List of allowed python modules to de-serialize custom objects from.\n\nIf provided, only the specified modules will be allowed to be deserialized.\nIf omitted, no modules are allowed, and the object returned will simply be a json object OR\na deserialized langchain object.\n"
+          "description": "Optional. List of allowed python modules to de-serialize custom objects from JSON.\n\nIf provided, only the specified modules will be allowed to be deserialized.\nIf omitted, no modules are allowed, and the object returned will simply be a json object OR\na deserialized langchain object.\n"
+        },
+        "allowed_msgpack_modules": {
+          "anyOf": [
+            {
+              "type": "array",
+              "items": {
+                "type": "array",
+                "items": {
+                  "type": "string"
+                }
+              }
+            },
+            {
+              "type": "boolean"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "description": "Optional. List of allowed python modules to de-serialize custom objects from msgpack.\n\nKnown safe types (langgraph.checkpoint.serde.jsonplus.SAFE_MSGPACK_TYPES) are always\nallowed regardless of this setting. Use this to allowlist your custom Pydantic models,\ndataclasses, and other user-defined types.\n\nIf True (default), unregistered types will log a warning but still be deserialized.\nIf None, only known safe types will be deserialized; unregistered types will be blocked.\n\n{...\n[\"my_agent\", \"models\", \"MyState\"],\n]\n}\n}\n\n{...\n}\n}\n\n"
         },
         "pickle_fallback": {
           "type": "boolean",