diff --git a/libs/checkpoint-postgres/uv.lock b/libs/checkpoint-postgres/uv.lock index d75989652..1cca56676 100644 --- a/libs/checkpoint-postgres/uv.lock +++ b/libs/checkpoint-postgres/uv.lock @@ -280,6 +280,7 @@ dev = [ { name = "numpy" }, { name = "pandas" }, { name = "pandas-stubs", specifier = ">=2.2.2.240807" }, + { name = "pycryptodome", specifier = ">=3.23.0" }, { name = "pytest" }, { name = "pytest-asyncio" }, { name = "pytest-mock" }, diff --git a/libs/checkpoint-sqlite/uv.lock b/libs/checkpoint-sqlite/uv.lock index cc8cb0241..9bde94ddd 100644 --- a/libs/checkpoint-sqlite/uv.lock +++ b/libs/checkpoint-sqlite/uv.lock @@ -289,6 +289,7 @@ dev = [ { name = "numpy" }, { name = "pandas" }, { name = "pandas-stubs", specifier = ">=2.2.2.240807" }, + { name = "pycryptodome", specifier = ">=3.23.0" }, { name = "pytest" }, { name = "pytest-asyncio" }, { name = "pytest-mock" }, diff --git a/libs/checkpoint/Makefile b/libs/checkpoint/Makefile index 66f7cf600..5c376936e 100644 --- a/libs/checkpoint/Makefile +++ b/libs/checkpoint/Makefile @@ -37,4 +37,4 @@ type: format format_diff: uv run ruff format $(PYTHON_FILES) - uv run ruff check --select I --fix $(PYTHON_FILES) + uv run ruff check --fix $(PYTHON_FILES) diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 5b59942e6..33495d5e3 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -1,6 +1,8 @@ from __future__ import annotations -from collections.abc import AsyncIterator, Iterator, Mapping, Sequence +import copy +import logging +from collections.abc import AsyncIterator, Collection, Iterator, Mapping, Sequence from typing import ( # noqa: UP035 Any, Generic, @@ -14,6 +16,7 @@ from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base.id import uuid6 from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_methods +from langgraph.checkpoint.serde.encrypted import EncryptedSerializer from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer from langgraph.checkpoint.serde.types import ( ERROR, @@ -25,6 +28,7 @@ from langgraph.checkpoint.serde.types import ( V = TypeVar("V", int, float, str) PendingWrite = tuple[str, str, Any] +logger = logging.getLogger(__name__) # Marked as total=False to allow for future expansion. @@ -474,6 +478,37 @@ class BaseCheckpointSaver(Generic[V]): else: return current + 1 + def with_allowlist( + self, extra_allowlist: Collection[tuple[str, ...]] + ) -> BaseCheckpointSaver[V]: + """Return a shallow clone with a derived msgpack allowlist.""" + serde = _with_msgpack_allowlist(self.serde, extra_allowlist) + if serde is self.serde: + return self + clone = copy.copy(self) + clone.serde = maybe_add_typed_methods(serde) + return clone + + +def _with_msgpack_allowlist( + serde: SerializerProtocol, extra_allowlist: Collection[tuple[str, ...]] +) -> SerializerProtocol: + if isinstance(serde, JsonPlusSerializer): + return serde.with_msgpack_allowlist(extra_allowlist) + if isinstance(serde, EncryptedSerializer): + inner = serde.serde + if isinstance(inner, JsonPlusSerializer): + updated_inner = inner.with_msgpack_allowlist(extra_allowlist) + if updated_inner is inner: + return serde + return EncryptedSerializer(serde.cipher, updated_inner) + logger.warning( + "Serializer %s does not support msgpack allowlist. " + "Strict msgpack deserialization will not be enforced.", + type(serde).__name__, + ) + return serde + class EmptyChannelError(Exception): """Raised when attempting to get the value of a channel that hasn't been updated diff --git a/libs/checkpoint/langgraph/checkpoint/serde/_msgpack.py b/libs/checkpoint/langgraph/checkpoint/serde/_msgpack.py new file mode 100644 index 000000000..9866806d4 --- /dev/null +++ b/libs/checkpoint/langgraph/checkpoint/serde/_msgpack.py @@ -0,0 +1,71 @@ +import os +from collections.abc import Iterable +from typing import cast + +STRICT_MSGPACK_ENABLED = os.getenv("LANGGRAPH_STRICT_MSGPACK", "false").lower() in ( + "1", + "true", + "yes", +) + + +_SENTINEL = cast(None, object()) + +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"), + # zoneinfo + ("zoneinfo", "ZoneInfo"), + # regex + ("re", "compile"), + # langgraph + ("langgraph.types", "Send"), + ("langgraph.types", "Interrupt"), + ("langgraph.types", "Command"), + ("langgraph.types", "StateSnapshot"), + ("langgraph.types", "PregelTask"), + ("langgraph.types", "Overwrite"), + ("langgraph.store.base", "Item"), + ("langgraph.store.base", "GetOp"), + } +) + +# Allowed (module, name, method) triples for EXT_METHOD_SINGLE_ARG. +# Only these specific method invocations are permitted during deserialization. +# This is separate from SAFE_MSGPACK_TYPES which only governs construction. +SAFE_MSGPACK_METHODS: frozenset[tuple[str, str, str]] = frozenset( + { + ("datetime", "datetime", "fromisoformat"), + } +) + + +AllowedMsgpackModules = Iterable[tuple[str, ...] | type] diff --git a/libs/checkpoint/langgraph/checkpoint/serde/encrypted.py b/libs/checkpoint/langgraph/checkpoint/serde/encrypted.py index 829c7dd2a..9f517e405 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/encrypted.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/encrypted.py @@ -41,7 +41,7 @@ class EncryptedSerializer(SerializerProtocol): ) -> "EncryptedSerializer": """Create an `EncryptedSerializer` using AES encryption.""" try: - from Crypto.Cipher import AES # type: ignore + from Crypto.Cipher import AES except ImportError: raise ImportError( "Pycryptodome is not installed. Please install it with `pip install pycryptodome`." diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index c4b550308..c267ac8be 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -10,7 +10,7 @@ import pickle import re import sys from collections import deque -from collections.abc import Callable, Sequence +from collections.abc import Callable, Iterable, Sequence from datetime import date, datetime, time, timedelta, timezone from enum import Enum from inspect import isclass @@ -22,17 +22,24 @@ from ipaddress import ( IPv6Interface, IPv6Network, ) -from typing import Any, Literal +from typing import TYPE_CHECKING, Any, Literal, cast from uuid import UUID from zoneinfo import ZoneInfo import ormsgpack from langchain_core.load.load import Reviver +from langgraph.checkpoint.serde import _msgpack as _lg_msgpack from langgraph.checkpoint.serde.base import SerializerProtocol from langgraph.checkpoint.serde.types import SendProtocol from langgraph.store.base import Item +if TYPE_CHECKING: + from langgraph.checkpoint.serde._msgpack import ( + AllowedMsgpackModules, + ) + from langgraph.checkpoint.serde.types import SendProtocol + LC_REVIVER = Reviver() EMPTY_BYTES = b"" logger = logging.getLogger(__name__) @@ -53,19 +60,59 @@ class JsonPlusSerializer(SerializerProtocol): self, *, pickle_fallback: bool = False, - allowed_json_modules: Sequence[tuple[str, ...]] | Literal[True] | None = None, + allowed_json_modules: Iterable[tuple[str, ...]] | Literal[True] | None = None, + allowed_msgpack_modules: ( + AllowedMsgpackModules | Literal[True] | None + ) = _lg_msgpack._SENTINEL, __unpack_ext_hook__: Callable[[int, bytes], Any] | None = None, ) -> None: + if allowed_msgpack_modules is _lg_msgpack._SENTINEL: + if _lg_msgpack.STRICT_MSGPACK_ENABLED: + allowed_msgpack_modules = None + else: + allowed_msgpack_modules = True self.pickle_fallback = pickle_fallback - self._allowed_modules = ( - {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) + self._allowed_json_modules: set[tuple[str, ...]] | Literal[True] | None = ( + _normalize_allowlist(allowed_json_modules) ) + self._allowed_msgpack_modules = _normalize_allowlist(allowed_msgpack_modules) + + self._custom_unpack_ext_hook = __unpack_ext_hook__ is not 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 with_msgpack_allowlist( + self, extra_allowlist: Iterable[tuple[str, ...] | type] + ) -> JsonPlusSerializer: + """Return a new serializer with a merged msgpack allowlist.""" + base_allowlist = self._allowed_msgpack_modules + if base_allowlist is True or base_allowlist is False: + return self + elif base_allowlist: + base_allowlist = set(base_allowlist) + else: + base_allowlist = set() + extra = _normalize_module_keys(tuple(extra_allowlist)) + merged = base_allowlist | extra + if merged == base_allowlist: + return self + allowed_msgpack_modules: AllowedMsgpackModules | Literal[True] | None + if merged: + allowed_msgpack_modules = tuple(merged) + elif isinstance(self._allowed_msgpack_modules, set): + allowed_msgpack_modules = tuple(self._allowed_msgpack_modules) + else: + allowed_msgpack_modules = self._allowed_msgpack_modules + return self.__class__( + pickle_fallback=self.pickle_fallback, + allowed_json_modules=self._allowed_json_modules, + allowed_msgpack_modules=allowed_msgpack_modules, + __unpack_ext_hook__=( + self._unpack_ext_hook if self._custom_unpack_ext_hook else None + ), ) def _encode_constructor_args( @@ -90,7 +137,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 +154,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 +186,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 +197,7 @@ class JsonPlusSerializer(SerializerProtocol): method_display = "" 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 +208,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 +495,174 @@ 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 _lg_msgpack.SAFE_MSGPACK_TYPES: + 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 + if allowed_modules is not None: + if key in allowed_modules: + return True + # 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 _check_allowed_method(module: str, name: str, method: str) -> bool: + """Check if a method invocation is allowed.""" + key = (module, name, method) + if key in _lg_msgpack.SAFE_MSGPACK_METHODS: + return True + 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: + 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]): + # We default to returning the raw data. If the user + # is using this in the context of a pydantic state, etc., then + # it would be validated upon construction. + return tup[2] + # 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 tup[2] + # 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 tup[2] + # 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_method(tup[0], tup[1], tup[3]): + return tup[2] + # 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 tup[2] + # 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 tup[2] + # 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 + + +# Aliasing in case anyone imported it directly +_msgpack_ext_hook = _create_msgpack_ext_hook(allowed_modules=None) def _msgpack_ext_hook_to_json(code: int, data: bytes) -> Any: @@ -648,3 +777,26 @@ _option = ( def _msgpack_enc(data: Any) -> bytes: return ormsgpack.packb(data, default=_msgpack_default, option=_option) + + +def _normalize_allowlist( + allowlist: AllowedMsgpackModules | Literal[True] | None, +) -> set[tuple[str, ...]] | Literal[True] | None: + if allowlist is True: + return allowlist + elif allowlist: + return _normalize_module_keys(allowlist) + else: + return None + + +def _normalize_module_keys( + modules: AllowedMsgpackModules, +) -> set[tuple[str, ...]]: + normalized: set[tuple[str, ...]] = set() + for module in modules: + if isclass(module): + normalized.add((module.__module__, module.__name__)) + else: + normalized.add(cast(tuple[str, ...], module)) + return normalized diff --git a/libs/checkpoint/pyproject.toml b/libs/checkpoint/pyproject.toml index 71d104de3..df8c1619f 100644 --- a/libs/checkpoint/pyproject.toml +++ b/libs/checkpoint/pyproject.toml @@ -42,6 +42,7 @@ lint = [ dev = [ {include-group = "test"}, {include-group = "lint"}, + "pycryptodome>=3.23.0", ] [tool.hatch.build.targets.wheel] diff --git a/libs/checkpoint/tests/test_encrypted.py b/libs/checkpoint/tests/test_encrypted.py new file mode 100644 index 000000000..696b93d32 --- /dev/null +++ b/libs/checkpoint/tests/test_encrypted.py @@ -0,0 +1,437 @@ +"""Tests for EncryptedSerializer with msgpack allowlist functionality. + +These tests mirror the msgpack allowlist tests in test_jsonplus.py but run them +through the EncryptedSerializer to ensure the allowlist behavior is preserved +when encryption is enabled. +""" + +from __future__ import annotations + +import logging +import pathlib +import re +import uuid +from collections import deque +from datetime import date, datetime, time, timezone +from decimal import Decimal +from ipaddress import IPv4Address +from typing import Literal, cast + +import ormsgpack +import pytest +from pydantic import BaseModel + +from langgraph.checkpoint.base import BaseCheckpointSaver, _with_msgpack_allowlist +from langgraph.checkpoint.serde import _msgpack as _lg_msgpack +from langgraph.checkpoint.serde.base import CipherProtocol +from langgraph.checkpoint.serde.encrypted import EncryptedSerializer +from langgraph.checkpoint.serde.jsonplus import ( + EXT_METHOD_SINGLE_ARG, + JsonPlusSerializer, + _msgpack_enc, +) + + +class InnerPydantic(BaseModel): + hello: str + + +class MyPydantic(BaseModel): + foo: str + bar: int + inner: InnerPydantic + + +class AnotherPydantic(BaseModel): + foo: str + + +class _PassthroughCipher(CipherProtocol): + def encrypt(self, plaintext: bytes) -> tuple[str, bytes]: + return "passthrough", plaintext + + def decrypt(self, ciphername: str, ciphertext: bytes) -> bytes: + assert ciphername == "passthrough" + return ciphertext + + +def _make_encrypted_serde( + allowed_msgpack_modules: ( + _lg_msgpack.AllowedMsgpackModules | Literal[True] | None | object + ) = _lg_msgpack._SENTINEL, +) -> EncryptedSerializer: + """Create an EncryptedSerializer with AES encryption for testing.""" + inner = JsonPlusSerializer( + allowed_msgpack_modules=cast( + _lg_msgpack.AllowedMsgpackModules | Literal[True] | None, + allowed_msgpack_modules, + ) + ) + return EncryptedSerializer.from_pycryptodome_aes( + serde=inner, key=b"1234567890123456" + ) + + +def test_msgpack_method_pathlib_blocked_encrypted_strict( + tmp_path: pathlib.Path, caplog: pytest.LogCaptureFixture +) -> None: + target = tmp_path / "secret.txt" + target.write_text("secret") + payload = ormsgpack.packb( + ormsgpack.Ext( + EXT_METHOD_SINGLE_ARG, + _msgpack_enc(("pathlib", "Path", target, "read_text")), + ), + option=ormsgpack.OPT_NON_STR_KEYS, + ) + serde = EncryptedSerializer( + _PassthroughCipher(), + JsonPlusSerializer(allowed_msgpack_modules=None), + ) + + caplog.set_level(logging.WARNING, logger="langgraph.checkpoint.serde.jsonplus") + caplog.clear() + result = serde.loads_typed(("msgpack+passthrough", payload)) + + assert result == target + assert "blocked deserialization of method call pathlib.path.read_text" in ( + caplog.text.lower() + ) + + +class TestEncryptedSerializerMsgpackAllowlist: + """Test msgpack allowlist behavior through EncryptedSerializer.""" + + def test_safe_types_no_warning(self, caplog: pytest.LogCaptureFixture) -> None: + """Test safe types deserialize without warnings through encryption.""" + serde = _make_encrypted_serde() + + 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) + # Verify encryption is happening + assert "+aes" in dumped[0], f"Expected encryption for {type(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_pydantic_warns_by_default(self, caplog: pytest.LogCaptureFixture) -> None: + """Pydantic models not in allowlist should log warning but still deserialize.""" + current = _lg_msgpack.STRICT_MSGPACK_ENABLED + _lg_msgpack.STRICT_MSGPACK_ENABLED = False + serde = _make_encrypted_serde() + + obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world")) + + caplog.clear() + dumped = serde.dumps_typed(obj) + assert "+aes" in dumped[0] + result = serde.loads_typed(dumped) + + assert "unregistered type" in caplog.text.lower() + assert "allowed_msgpack_modules" in caplog.text + assert result == obj + _lg_msgpack.STRICT_MSGPACK_ENABLED = current + + def test_strict_mode_blocks_unregistered( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """Strict mode should block unregistered types through encryption.""" + serde = _make_encrypted_serde(allowed_msgpack_modules=None) + + obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world")) + + caplog.clear() + dumped = serde.dumps_typed(obj) + assert "+aes" in dumped[0] + result = serde.loads_typed(dumped) + + assert "blocked" in caplog.text.lower() + expected = obj.model_dump() + assert result == expected + + def test_allowlist_silences_warning(self, caplog: pytest.LogCaptureFixture) -> None: + """Types in allowed_msgpack_modules should deserialize without warnings.""" + serde = _make_encrypted_serde( + allowed_msgpack_modules=[ + ("tests.test_encrypted", "MyPydantic"), + ("tests.test_encrypted", "InnerPydantic"), + ] + ) + + obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world")) + + caplog.clear() + dumped = serde.dumps_typed(obj) + assert "+aes" in dumped[0] + result = serde.loads_typed(dumped) + + assert "unregistered type" not in caplog.text.lower() + assert "blocked" not in caplog.text.lower() + assert result == obj + + def test_allowlist_blocks_non_listed( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """Allowlists should block unregistered types even through encryption.""" + serde = _make_encrypted_serde( + allowed_msgpack_modules=[("tests.test_encrypted", "MyPydantic")] + ) + + obj = AnotherPydantic(foo="nope") + + caplog.clear() + dumped = serde.dumps_typed(obj) + assert "+aes" in dumped[0] + result = serde.loads_typed(dumped) + + assert "blocked" in caplog.text.lower() + expected = obj.model_dump() + assert result == expected + + def test_safe_types_value_equality(self, caplog: pytest.LogCaptureFixture) -> None: + """Verify safe types are correctly restored with proper values through encryption.""" + serde = _make_encrypted_serde(allowed_msgpack_modules=None) + + test_cases = [ + datetime(2024, 1, 15, 12, 30, 45, 123456), + date(2024, 6, 15), + time(14, 30, 0), + uuid.UUID("12345678-1234-5678-1234-567812345678"), + Decimal("123.456789"), + {1, 2, 3, 4, 5}, + frozenset(["a", "b", "c"]), + deque([1, 2, 3]), + IPv4Address("10.0.0.1"), + pathlib.Path("/some/test/path"), + re.compile(r"\d+", re.MULTILINE), + ] + + for obj in test_cases: + caplog.clear() + dumped = serde.dumps_typed(obj) + assert "+aes" in dumped[0], f"Expected encryption for {type(obj)}" + result = serde.loads_typed(dumped) + + assert "blocked" not in caplog.text.lower(), f"Blocked for {type(obj)}" + if isinstance(obj, re.Pattern): + assert result.pattern == obj.pattern + assert result.flags == obj.flags + else: + assert result == obj, ( + f"Value mismatch for {type(obj)}: {result} != {obj}" + ) + + def test_regex_safe_type(self, caplog: pytest.LogCaptureFixture) -> None: + """re.compile patterns should deserialize without warnings as a safe type.""" + serde = _make_encrypted_serde(allowed_msgpack_modules=None) + pattern = re.compile(r"foo.*bar", re.IGNORECASE | re.DOTALL) + + caplog.clear() + dumped = serde.dumps_typed(pattern) + assert "+aes" in dumped[0] + result = serde.loads_typed(dumped) + + assert "blocked" not in caplog.text.lower() + assert "unregistered" not in caplog.text.lower() + assert result.pattern == pattern.pattern + assert result.flags == pattern.flags + + +class TestWithMsgpackAllowlistEncrypted: + """Test _with_msgpack_allowlist function with EncryptedSerializer.""" + + def test_propagates_allowlist_to_inner_serde(self) -> None: + """_with_msgpack_allowlist should propagate allowlist to inner JsonPlusSerializer.""" + inner = JsonPlusSerializer(allowed_msgpack_modules=None) + encrypted = EncryptedSerializer.from_pycryptodome_aes( + serde=inner, key=b"1234567890123456" + ) + + extra = [("my.module", "MyClass")] + result = _with_msgpack_allowlist(encrypted, extra) + + # Should return a new EncryptedSerializer + assert isinstance(result, EncryptedSerializer) + assert result is not encrypted + # Inner serde should have the allowlist + assert isinstance(result.serde, JsonPlusSerializer) + assert isinstance(result.serde._allowed_msgpack_modules, set) + assert ("my.module", "MyClass") in result.serde._allowed_msgpack_modules + + def test_preserves_cipher(self) -> None: + """_with_msgpack_allowlist should preserve the cipher from the original.""" + inner = JsonPlusSerializer(allowed_msgpack_modules=None) + encrypted = EncryptedSerializer.from_pycryptodome_aes( + serde=inner, key=b"1234567890123456" + ) + + result = _with_msgpack_allowlist(encrypted, [("my.module", "MyClass")]) + + assert isinstance(result, EncryptedSerializer) + # Should use the same cipher + assert result.cipher is encrypted.cipher + + def test_returns_same_if_not_jsonplus_inner(self) -> None: + """_with_msgpack_allowlist should return same serde if inner is not JsonPlusSerializer.""" + + class DummyInnerSerde: + def dumps_typed(self, obj: object) -> tuple[str, bytes]: + return ("dummy", b"") + + def loads_typed(self, data: tuple[str, bytes]) -> None: + return None + + from langgraph.checkpoint.serde.base import CipherProtocol + + class DummyCipher(CipherProtocol): + def encrypt(self, plaintext: bytes) -> tuple[str, bytes]: + return "dummy", plaintext + + def decrypt(self, ciphername: str, ciphertext: bytes) -> bytes: + return ciphertext + + encrypted = EncryptedSerializer(DummyCipher(), DummyInnerSerde()) + result = _with_msgpack_allowlist(encrypted, [("my.module", "MyClass")]) + + assert result is encrypted + + def test_warns_if_allowlist_unsupported( + self, caplog: pytest.LogCaptureFixture + ) -> None: + class DummySerde: + def dumps_typed(self, obj: object) -> tuple[str, bytes]: + return ("dummy", b"") + + def loads_typed(self, data: tuple[str, bytes]) -> object: + return data + + serde = DummySerde() + caplog.set_level(logging.WARNING, logger="langgraph.checkpoint.base") + caplog.clear() + + result = _with_msgpack_allowlist(serde, [("my.module", "MyClass")]) + + assert result is serde + assert "does not support msgpack allowlist" in caplog.text.lower() + + def test_noop_allowlist_returns_same_encrypted_instance(self) -> None: + inner = JsonPlusSerializer(allowed_msgpack_modules=None) + encrypted = EncryptedSerializer.from_pycryptodome_aes( + serde=inner, key=b"1234567890123456" + ) + + result = _with_msgpack_allowlist(encrypted, ()) + + assert result is encrypted + + def test_functional_roundtrip_with_allowlist( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """End-to-end test: allowlist applied via _with_msgpack_allowlist works.""" + inner = JsonPlusSerializer(allowed_msgpack_modules=None) + encrypted = EncryptedSerializer.from_pycryptodome_aes( + serde=inner, key=b"1234567890123456" + ) + + # Apply allowlist for MyPydantic + updated = _with_msgpack_allowlist( + encrypted, + [ + ("tests.test_encrypted", "MyPydantic"), + ("tests.test_encrypted", "InnerPydantic"), + ], + ) + + obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world")) + + caplog.clear() + dumped = updated.dumps_typed(obj) + assert "+aes" in dumped[0] + result = updated.loads_typed(dumped) + + # Should deserialize without blocking + assert "blocked" not in caplog.text.lower() + assert result == obj + + def test_original_still_blocks_after_with_allowlist( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """Original serde should still block after _with_msgpack_allowlist creates a new one.""" + inner = JsonPlusSerializer(allowed_msgpack_modules=None) + encrypted = EncryptedSerializer.from_pycryptodome_aes( + serde=inner, key=b"1234567890123456" + ) + + # Apply allowlist - this should create a NEW serde + _with_msgpack_allowlist( + encrypted, + [("tests.test_encrypted", "MyPydantic")], + ) + + # Original should still block + obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world")) + + caplog.clear() + dumped = encrypted.dumps_typed(obj) + result = encrypted.loads_typed(dumped) + + assert "blocked" in caplog.text.lower() + assert result == obj.model_dump() + + +class TestEncryptedSerializerUnencryptedFallback: + """Test that EncryptedSerializer handles unencrypted data correctly.""" + + def test_loads_unencrypted_data(self) -> None: + """EncryptedSerializer should handle unencrypted data for backwards compat.""" + plain = JsonPlusSerializer(allowed_msgpack_modules=None) + encrypted = _make_encrypted_serde(allowed_msgpack_modules=None) + + obj = {"key": "value", "number": 42} + + # Serialize with plain serde + dumped = plain.dumps_typed(obj) + assert "+aes" not in dumped[0] + + # Should still deserialize with encrypted serde + result = encrypted.loads_typed(dumped) + assert result == obj + + +def test_with_allowlist_uses_copy_protocol() -> None: + class CopyAwareSaver(BaseCheckpointSaver[str]): + def __init__(self) -> None: + super().__init__(serde=JsonPlusSerializer(allowed_msgpack_modules=None)) + self.copy_was_used = False + + def __copy__(self) -> object: + clone = object.__new__(self.__class__) + clone.__dict__ = self.__dict__.copy() + clone.copy_was_used = True + return clone + + saver = CopyAwareSaver() + + updated = saver.with_allowlist([("tests.test_encrypted", "MyPydantic")]) + + assert isinstance(updated, CopyAwareSaver) + assert updated is not saver + assert updated.copy_was_used is True + assert saver.copy_was_used is False diff --git a/libs/checkpoint/tests/test_jsonplus.py b/libs/checkpoint/tests/test_jsonplus.py index c2ff5f2d0..e2d9690d3 100644 --- a/libs/checkpoint/tests/test_jsonplus.py +++ b/libs/checkpoint/tests/test_jsonplus.py @@ -1,5 +1,6 @@ import dataclasses import json +import logging import pathlib import re import sys @@ -13,15 +14,20 @@ from zoneinfo import ZoneInfo import dataclasses_json import numpy as np +import ormsgpack import pandas as pd import pytest from pydantic import BaseModel, SecretStr from pydantic.v1 import BaseModel as BaseModelV1 from pydantic.v1 import SecretStr as SecretStrV1 +from langgraph.checkpoint.serde import _msgpack as _lg_msgpack +from langgraph.checkpoint.serde._msgpack import AllowedMsgpackModules from langgraph.checkpoint.serde.jsonplus import ( + EXT_METHOD_SINGLE_ARG, InvalidModuleError, JsonPlusSerializer, + _msgpack_enc, _msgpack_ext_hook_to_json, ) from langgraph.store.base import Item @@ -37,6 +43,10 @@ class MyPydantic(BaseModel): inner: InnerPydantic +class AnotherPydantic(BaseModel): + foo: str + + class InnerPydanticV1(BaseModelV1): hello: str @@ -138,7 +148,27 @@ def test_serde_jsonplus() -> None: ) to_serialize["my_secret_str_v1"] = SecretStrV1("meow") - serde = JsonPlusSerializer() + allowed_msgpack_modules: AllowedMsgpackModules = [ + InnerDataclass, + MyDataclass, + MyDataclassWSlots, + MyEnum, + InnerPydantic, + MyPydantic, + # Testing that it supports both. + (Person.__module__, Person.__name__), + (SecretStr.__module__, SecretStr.__name__), + ] + if sys.version_info < (3, 14): + allowed_msgpack_modules.extend( # type: ignore + [ + (InnerPydanticV1.__module__, InnerPydanticV1.__name__), + (MyPydanticV1.__module__, MyPydanticV1.__name__), + (SecretStrV1.__module__, SecretStrV1.__name__), + ] + ) + + serde = JsonPlusSerializer(allowed_msgpack_modules=allowed_msgpack_modules) dumped = serde.dumps_typed(to_serialize) @@ -512,5 +542,337 @@ 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.""" + current = _lg_msgpack.STRICT_MSGPACK_ENABLED + _lg_msgpack.STRICT_MSGPACK_ENABLED = False + 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 + _lg_msgpack.STRICT_MSGPACK_ENABLED = current + + +def test_msgpack_env_strict_default( + caplog: pytest.LogCaptureFixture, +) -> None: + """Strict msgpack env should default to blocking unregistered types.""" + current = _lg_msgpack.STRICT_MSGPACK_ENABLED + _lg_msgpack.STRICT_MSGPACK_ENABLED = True + 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 "blocked" in caplog.text.lower() + assert result == obj.model_dump() + _lg_msgpack.STRICT_MSGPACK_ENABLED = current + + +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.""" + + 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() + expected = obj.model_dump() + assert result == expected + + +def test_msgpack_allowlist_blocks_non_listed( + caplog: pytest.LogCaptureFixture, +) -> None: + """Allowlists should block unregistered types even if msgpack is enabled.""" + + serde = JsonPlusSerializer( + allowed_msgpack_modules=[("tests.test_jsonplus", "MyPydantic")] + ) + + obj = AnotherPydantic(foo="nope") + + caplog.clear() + dumped = serde.dumps_typed(obj) + result = serde.loads_typed(dumped) + + assert "blocked" in caplog.text.lower() + expected = obj.model_dump() + # It's not allowed, so we just leave it as a dict + assert result == expected + + +def test_msgpack_strict_allows_safe_types( + caplog: pytest.LogCaptureFixture, +) -> None: + """Safe types should still deserialize in strict mode without warnings.""" + + serde = JsonPlusSerializer(allowed_msgpack_modules=None) + safe = uuid.uuid4() + + caplog.clear() + dumped = serde.dumps_typed(safe) + result = serde.loads_typed(dumped) + + assert "blocked" not in caplog.text.lower() + assert result == safe + + +def test_msgpack_regex_safe_type(caplog: pytest.LogCaptureFixture) -> None: + """re.compile patterns should deserialize without warnings as a safe type.""" + + serde = JsonPlusSerializer(allowed_msgpack_modules=None) + pattern = re.compile(r"foo.*bar", re.IGNORECASE | re.DOTALL) + + caplog.clear() + dumped = serde.dumps_typed(pattern) + result = serde.loads_typed(dumped) + + assert "blocked" not in caplog.text.lower() + assert "unregistered" not in caplog.text.lower() + assert result.pattern == pattern.pattern + assert result.flags == pattern.flags + + +def test_msgpack_method_pathlib_blocked_in_strict( + tmp_path: pathlib.Path, caplog: pytest.LogCaptureFixture +) -> None: + target = tmp_path / "secret.txt" + target.write_text("secret") + serde = JsonPlusSerializer(allowed_msgpack_modules=None) + payload = ormsgpack.packb( + ormsgpack.Ext( + EXT_METHOD_SINGLE_ARG, + _msgpack_enc(("pathlib", "Path", target, "read_text")), + ), + option=ormsgpack.OPT_NON_STR_KEYS, + ) + + caplog.set_level(logging.WARNING, logger="langgraph.checkpoint.serde.jsonplus") + caplog.clear() + result = serde.loads_typed(("msgpack", payload)) + + assert result == target + assert "blocked deserialization of method call pathlib.path.read_text" in ( + caplog.text.lower() + ) + + +def test_msgpack_method_pathlib_blocked_default_mode( + tmp_path: pathlib.Path, caplog: pytest.LogCaptureFixture +) -> None: + target = tmp_path / "secret.txt" + target.write_text("secret") + serde = JsonPlusSerializer(allowed_msgpack_modules=True) + payload = ormsgpack.packb( + ormsgpack.Ext( + EXT_METHOD_SINGLE_ARG, + _msgpack_enc(("pathlib", "Path", target, "read_text")), + ), + option=ormsgpack.OPT_NON_STR_KEYS, + ) + + caplog.set_level(logging.WARNING, logger="langgraph.checkpoint.serde.jsonplus") + caplog.clear() + result = serde.loads_typed(("msgpack", payload)) + + assert result == target + assert "blocked deserialization of method call pathlib.path.read_text" in ( + caplog.text.lower() + ) + + +def test_msgpack_regex_still_works_strict(caplog: pytest.LogCaptureFixture) -> None: + serde = JsonPlusSerializer(allowed_msgpack_modules=None) + pattern = re.compile(r"pattern", re.IGNORECASE | re.MULTILINE) + + caplog.clear() + result = serde.loads_typed(serde.dumps_typed(pattern)) + + assert "blocked" not in caplog.text.lower() + assert result.pattern == pattern.pattern + assert result.flags == pattern.flags + + +def test_msgpack_path_constructor_still_works() -> None: + serde = JsonPlusSerializer(allowed_msgpack_modules=None) + path_obj = pathlib.Path("/tmp/foo") + + result = serde.loads_typed(serde.dumps_typed(path_obj)) + + assert result == path_obj + + +def test_with_msgpack_allowlist_noop_returns_same_instance() -> None: + serde = JsonPlusSerializer(allowed_msgpack_modules=None) + + result = serde.with_msgpack_allowlist(()) + + assert result is serde + + +@pytest.mark.skipif(sys.version_info >= (3, 14), reason="pydantic v1 not on 3.14+") +def test_msgpack_pydantic_v1_allowlist(caplog: pytest.LogCaptureFixture) -> None: + """Pydantic v1 models in allowlist should deserialize without warnings.""" + + serde = JsonPlusSerializer( + allowed_msgpack_modules=[ + ("tests.test_jsonplus", "MyPydanticV1"), + ("tests.test_jsonplus", "InnerPydanticV1"), + ] + ) + + obj = MyPydanticV1(foo="test", bar=42, inner=InnerPydanticV1(hello="world")) + + 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 + + +def test_msgpack_dataclass_allowlist(caplog: pytest.LogCaptureFixture) -> None: + """Dataclasses in allowlist should deserialize without warnings.""" + + serde = JsonPlusSerializer( + allowed_msgpack_modules=[ + ("tests.test_jsonplus", "MyDataclass"), + ("tests.test_jsonplus", "InnerDataclass"), + ] + ) + + obj = MyDataclass(foo="test", bar=42, inner=InnerDataclass(hello="world")) + + 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 + + +def test_msgpack_safe_types_value_equality(caplog: pytest.LogCaptureFixture) -> None: + """Verify safe types are correctly restored with proper values.""" + + serde = JsonPlusSerializer(allowed_msgpack_modules=None) + + test_cases = [ + datetime(2024, 1, 15, 12, 30, 45, 123456), + date(2024, 6, 15), + time(14, 30, 0), + uuid.UUID("12345678-1234-5678-1234-567812345678"), + Decimal("123.456789"), + {1, 2, 3, 4, 5}, + frozenset(["a", "b", "c"]), + deque([1, 2, 3]), + IPv4Address("10.0.0.1"), + pathlib.Path("/some/test/path"), + re.compile(r"\d+", re.MULTILINE), + ] + + for obj in test_cases: + caplog.clear() + dumped = serde.dumps_typed(obj) + result = serde.loads_typed(dumped) + + assert "blocked" not in caplog.text.lower(), f"Blocked for {type(obj)}" + # For regex patterns, compare pattern and flags + if isinstance(obj, re.Pattern): + assert result.pattern == obj.pattern + assert result.flags == obj.flags + else: + assert result == obj, f"Value mismatch for {type(obj)}: {result} != {obj}" + + +def test_msgpack_nested_pydantic_serializes_as_dict( + caplog: pytest.LogCaptureFixture, +) -> None: + """Nested Pydantic models are serialized via model_dump() as dicts. + + This means nested models don't go through the ext hook and don't need + to be in the allowlist - only the outer type does. + """ + + # Only allow outer type - inner is serialized as dict via model_dump() + serde = JsonPlusSerializer( + allowed_msgpack_modules=[("tests.test_jsonplus", "MyPydantic")] + ) + + obj = MyPydantic(foo="test", bar=42, inner=InnerPydantic(hello="world")) + + caplog.clear() + dumped = serde.dumps_typed(obj) + result = serde.loads_typed(dumped) + + # No blocking should occur - inner is serialized as dict, not ext + assert "blocked" not in caplog.text.lower() + assert result == obj diff --git a/libs/checkpoint/tests/test_memory.py b/libs/checkpoint/tests/test_memory.py index 7d85f4ed5..a68a23d90 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -1,7 +1,9 @@ +import logging from typing import Any import pytest from langchain_core.runnables import RunnableConfig +from pydantic import BaseModel from langgraph.checkpoint.base import ( Checkpoint, @@ -10,6 +12,11 @@ from langgraph.checkpoint.base import ( empty_checkpoint, ) from langgraph.checkpoint.memory import InMemorySaver +from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer + + +class MemoryPydantic(BaseModel): + foo: str class TestMemorySaver: @@ -199,3 +206,105 @@ async def test_memory_saver() -> None: with memory_saver as sync_memory_saver: assert sync_memory_saver is memory_saver + + +def test_memory_saver_warns_on_unregistered_msgpack( + caplog: pytest.LogCaptureFixture, +) -> None: + serde = JsonPlusSerializer() + memory_saver = InMemorySaver(serde=serde) + obj = MemoryPydantic(foo="bar") + + checkpoint = empty_checkpoint() + checkpoint["channel_values"] = {"foo": obj} + checkpoint["channel_versions"] = {"foo": 1} + + config: RunnableConfig = { + "configurable": {"thread_id": "thread-1", "checkpoint_ns": ""} + } + + caplog.set_level(logging.WARNING, logger="langgraph.checkpoint.serde.jsonplus") + new_config = memory_saver.put(config, checkpoint, {}, {"foo": 1}) + result = memory_saver.get_tuple(new_config) + + assert result is not None + assert "unregistered type" in caplog.text.lower() + assert result.checkpoint["channel_values"]["foo"] == obj + + +def test_memory_saver_allowlist_silences_warning( + caplog: pytest.LogCaptureFixture, +) -> None: + serde = JsonPlusSerializer( + allowed_msgpack_modules=[("tests.test_memory", "MemoryPydantic")] + ) + memory_saver = InMemorySaver(serde=serde) + obj = MemoryPydantic(foo="bar") + + checkpoint = empty_checkpoint() + checkpoint["channel_values"] = {"foo": obj} + checkpoint["channel_versions"] = {"foo": 1} + + config: RunnableConfig = { + "configurable": {"thread_id": "thread-1", "checkpoint_ns": ""} + } + + caplog.set_level(logging.WARNING, logger="langgraph.checkpoint.serde.jsonplus") + new_config = memory_saver.put(config, checkpoint, {}, {"foo": 1}) + result = memory_saver.get_tuple(new_config) + + assert result is not None + assert "unregistered type" not in caplog.text.lower() + assert result.checkpoint["channel_values"]["foo"] == obj + + +def test_memory_saver_strict_blocks_unregistered( + caplog: pytest.LogCaptureFixture, +) -> None: + serde = JsonPlusSerializer(allowed_msgpack_modules=None) + memory_saver = InMemorySaver(serde=serde) + obj = MemoryPydantic(foo="bar") + + checkpoint = empty_checkpoint() + checkpoint["channel_values"] = {"foo": obj} + checkpoint["channel_versions"] = {"foo": 1} + + config: RunnableConfig = { + "configurable": {"thread_id": "thread-1", "checkpoint_ns": ""} + } + + caplog.set_level(logging.WARNING, logger="langgraph.checkpoint.serde.jsonplus") + new_config = memory_saver.put(config, checkpoint, {}, {"foo": 1}) + result = memory_saver.get_tuple(new_config) + + assert result is not None + assert "blocked" in caplog.text.lower() + expected = obj.model_dump() if hasattr(obj, "model_dump") else obj.dict() + assert result.checkpoint["channel_values"]["foo"] == expected + + +def test_memory_saver_with_allowlist_proxy_isolated() -> None: + serde = JsonPlusSerializer(allowed_msgpack_modules=None) + memory_saver = InMemorySaver(serde=serde) + proxy = memory_saver.with_allowlist([("tests.test_memory", "MemoryPydantic")]) + + obj = MemoryPydantic(foo="bar") + + checkpoint = empty_checkpoint() + checkpoint["channel_values"] = {"foo": obj} + checkpoint["channel_versions"] = {"foo": 1} + + config: RunnableConfig = { + "configurable": {"thread_id": "thread-1", "checkpoint_ns": ""} + } + + new_config = proxy.put(config, checkpoint, {}, {"foo": 1}) + + proxied = proxy.get_tuple(new_config) + assert proxied is not None + assert proxied.checkpoint["channel_values"]["foo"] == obj + + direct = memory_saver.get_tuple(new_config) + assert direct is not None + expected = obj.model_dump() if hasattr(obj, "model_dump") else obj.dict() + assert direct.checkpoint["channel_values"]["foo"] == expected diff --git a/libs/checkpoint/uv.lock b/libs/checkpoint/uv.lock index 8490472a2..cc4384e3d 100644 --- a/libs/checkpoint/uv.lock +++ b/libs/checkpoint/uv.lock @@ -302,6 +302,7 @@ dev = [ { name = "numpy", version = "2.4.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "pandas" }, { name = "pandas-stubs" }, + { name = "pycryptodome" }, { name = "pytest" }, { name = "pytest-asyncio" }, { name = "pytest-mock" }, @@ -341,6 +342,7 @@ dev = [ { name = "numpy" }, { name = "pandas" }, { name = "pandas-stubs", specifier = ">=2.2.2.240807" }, + { name = "pycryptodome", specifier = ">=3.23.0" }, { name = "pytest" }, { name = "pytest-asyncio" }, { name = "pytest-mock" }, @@ -912,6 +914,41 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, ] +[[package]] +name = "pycryptodome" +version = "3.23.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/8e/a6/8452177684d5e906854776276ddd34eca30d1b1e15aa1ee9cefc289a33f5/pycryptodome-3.23.0.tar.gz", hash = "sha256:447700a657182d60338bab09fdb27518f8856aecd80ae4c6bdddb67ff5da44ef", size = 4921276, upload-time = "2025-05-17T17:21:45.242Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/5d/bdb09489b63cd34a976cc9e2a8d938114f7a53a74d3dd4f125ffa49dce82/pycryptodome-3.23.0-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:0011f7f00cdb74879142011f95133274741778abba114ceca229adbf8e62c3e4", size = 2495152, upload-time = "2025-05-17T17:20:20.833Z" }, + { url = "https://files.pythonhosted.org/packages/a7/ce/7840250ed4cc0039c433cd41715536f926d6e86ce84e904068eb3244b6a6/pycryptodome-3.23.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:90460fc9e088ce095f9ee8356722d4f10f86e5be06e2354230a9880b9c549aae", size = 1639348, upload-time = "2025-05-17T17:20:23.171Z" }, + { url = "https://files.pythonhosted.org/packages/ee/f0/991da24c55c1f688d6a3b5a11940567353f74590734ee4a64294834ae472/pycryptodome-3.23.0-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4764e64b269fc83b00f682c47443c2e6e85b18273712b98aa43bcb77f8570477", size = 2184033, upload-time = "2025-05-17T17:20:25.424Z" }, + { url = "https://files.pythonhosted.org/packages/54/16/0e11882deddf00f68b68dd4e8e442ddc30641f31afeb2bc25588124ac8de/pycryptodome-3.23.0-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:eb8f24adb74984aa0e5d07a2368ad95276cf38051fe2dc6605cbcf482e04f2a7", size = 2270142, upload-time = "2025-05-17T17:20:27.808Z" }, + { url = "https://files.pythonhosted.org/packages/d5/fc/4347fea23a3f95ffb931f383ff28b3f7b1fe868739182cb76718c0da86a1/pycryptodome-3.23.0-cp313-cp313t-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d97618c9c6684a97ef7637ba43bdf6663a2e2e77efe0f863cce97a76af396446", size = 2309384, upload-time = "2025-05-17T17:20:30.765Z" }, + { url = "https://files.pythonhosted.org/packages/6e/d9/c5261780b69ce66d8cfab25d2797bd6e82ba0241804694cd48be41add5eb/pycryptodome-3.23.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:9a53a4fe5cb075075d515797d6ce2f56772ea7e6a1e5e4b96cf78a14bac3d265", size = 2183237, upload-time = "2025-05-17T17:20:33.736Z" }, + { url = "https://files.pythonhosted.org/packages/5a/6f/3af2ffedd5cfa08c631f89452c6648c4d779e7772dfc388c77c920ca6bbf/pycryptodome-3.23.0-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:763d1d74f56f031788e5d307029caef067febf890cd1f8bf61183ae142f1a77b", size = 2343898, upload-time = "2025-05-17T17:20:36.086Z" }, + { url = "https://files.pythonhosted.org/packages/9a/dc/9060d807039ee5de6e2f260f72f3d70ac213993a804f5e67e0a73a56dd2f/pycryptodome-3.23.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:954af0e2bd7cea83ce72243b14e4fb518b18f0c1649b576d114973e2073b273d", size = 2269197, upload-time = "2025-05-17T17:20:38.414Z" }, + { url = "https://files.pythonhosted.org/packages/f9/34/e6c8ca177cb29dcc4967fef73f5de445912f93bd0343c9c33c8e5bf8cde8/pycryptodome-3.23.0-cp313-cp313t-win32.whl", hash = "sha256:257bb3572c63ad8ba40b89f6fc9d63a2a628e9f9708d31ee26560925ebe0210a", size = 1768600, upload-time = "2025-05-17T17:20:40.688Z" }, + { url = "https://files.pythonhosted.org/packages/e4/1d/89756b8d7ff623ad0160f4539da571d1f594d21ee6d68be130a6eccb39a4/pycryptodome-3.23.0-cp313-cp313t-win_amd64.whl", hash = "sha256:6501790c5b62a29fcb227bd6b62012181d886a767ce9ed03b303d1f22eb5c625", size = 1799740, upload-time = "2025-05-17T17:20:42.413Z" }, + { url = "https://files.pythonhosted.org/packages/5d/61/35a64f0feaea9fd07f0d91209e7be91726eb48c0f1bfc6720647194071e4/pycryptodome-3.23.0-cp313-cp313t-win_arm64.whl", hash = "sha256:9a77627a330ab23ca43b48b130e202582e91cc69619947840ea4d2d1be21eb39", size = 1703685, upload-time = "2025-05-17T17:20:44.388Z" }, + { url = "https://files.pythonhosted.org/packages/db/6c/a1f71542c969912bb0e106f64f60a56cc1f0fabecf9396f45accbe63fa68/pycryptodome-3.23.0-cp37-abi3-macosx_10_9_universal2.whl", hash = "sha256:187058ab80b3281b1de11c2e6842a357a1f71b42cb1e15bce373f3d238135c27", size = 2495627, upload-time = "2025-05-17T17:20:47.139Z" }, + { url = "https://files.pythonhosted.org/packages/6e/4e/a066527e079fc5002390c8acdd3aca431e6ea0a50ffd7201551175b47323/pycryptodome-3.23.0-cp37-abi3-macosx_10_9_x86_64.whl", hash = "sha256:cfb5cd445280c5b0a4e6187a7ce8de5a07b5f3f897f235caa11f1f435f182843", size = 1640362, upload-time = "2025-05-17T17:20:50.392Z" }, + { url = "https://files.pythonhosted.org/packages/50/52/adaf4c8c100a8c49d2bd058e5b551f73dfd8cb89eb4911e25a0c469b6b4e/pycryptodome-3.23.0-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:67bd81fcbe34f43ad9422ee8fd4843c8e7198dd88dd3d40e6de42ee65fbe1490", size = 2182625, upload-time = "2025-05-17T17:20:52.866Z" }, + { url = "https://files.pythonhosted.org/packages/5f/e9/a09476d436d0ff1402ac3867d933c61805ec2326c6ea557aeeac3825604e/pycryptodome-3.23.0-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c8987bd3307a39bc03df5c8e0e3d8be0c4c3518b7f044b0f4c15d1aa78f52575", size = 2268954, upload-time = "2025-05-17T17:20:55.027Z" }, + { url = "https://files.pythonhosted.org/packages/f9/c5/ffe6474e0c551d54cab931918127c46d70cab8f114e0c2b5a3c071c2f484/pycryptodome-3.23.0-cp37-abi3-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:aa0698f65e5b570426fc31b8162ed4603b0c2841cbb9088e2b01641e3065915b", size = 2308534, upload-time = "2025-05-17T17:20:57.279Z" }, + { url = "https://files.pythonhosted.org/packages/18/28/e199677fc15ecf43010f2463fde4c1a53015d1fe95fb03bca2890836603a/pycryptodome-3.23.0-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:53ecbafc2b55353edcebd64bf5da94a2a2cdf5090a6915bcca6eca6cc452585a", size = 2181853, upload-time = "2025-05-17T17:20:59.322Z" }, + { url = "https://files.pythonhosted.org/packages/ce/ea/4fdb09f2165ce1365c9eaefef36625583371ee514db58dc9b65d3a255c4c/pycryptodome-3.23.0-cp37-abi3-musllinux_1_2_i686.whl", hash = "sha256:156df9667ad9f2ad26255926524e1c136d6664b741547deb0a86a9acf5ea631f", size = 2342465, upload-time = "2025-05-17T17:21:03.83Z" }, + { url = "https://files.pythonhosted.org/packages/22/82/6edc3fc42fe9284aead511394bac167693fb2b0e0395b28b8bedaa07ef04/pycryptodome-3.23.0-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:dea827b4d55ee390dc89b2afe5927d4308a8b538ae91d9c6f7a5090f397af1aa", size = 2267414, upload-time = "2025-05-17T17:21:06.72Z" }, + { url = "https://files.pythonhosted.org/packages/59/fe/aae679b64363eb78326c7fdc9d06ec3de18bac68be4b612fc1fe8902693c/pycryptodome-3.23.0-cp37-abi3-win32.whl", hash = "sha256:507dbead45474b62b2bbe318eb1c4c8ee641077532067fec9c1aa82c31f84886", size = 1768484, upload-time = "2025-05-17T17:21:08.535Z" }, + { url = "https://files.pythonhosted.org/packages/54/2f/e97a1b8294db0daaa87012c24a7bb714147c7ade7656973fd6c736b484ff/pycryptodome-3.23.0-cp37-abi3-win_amd64.whl", hash = "sha256:c75b52aacc6c0c260f204cbdd834f76edc9fb0d8e0da9fbf8352ef58202564e2", size = 1799636, upload-time = "2025-05-17T17:21:10.393Z" }, + { url = "https://files.pythonhosted.org/packages/18/3d/f9441a0d798bf2b1e645adc3265e55706aead1255ccdad3856dbdcffec14/pycryptodome-3.23.0-cp37-abi3-win_arm64.whl", hash = "sha256:11eeeb6917903876f134b56ba11abe95c0b0fd5e3330def218083c7d98bbcb3c", size = 1703675, upload-time = "2025-05-17T17:21:13.146Z" }, + { url = "https://files.pythonhosted.org/packages/d9/12/e33935a0709c07de084d7d58d330ec3f4daf7910a18e77937affdb728452/pycryptodome-3.23.0-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:ddb95b49df036ddd264a0ad246d1be5b672000f12d6961ea2c267083a5e19379", size = 1623886, upload-time = "2025-05-17T17:21:20.614Z" }, + { url = "https://files.pythonhosted.org/packages/22/0b/aa8f9419f25870889bebf0b26b223c6986652bdf071f000623df11212c90/pycryptodome-3.23.0-pp310-pypy310_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d8e95564beb8782abfd9e431c974e14563a794a4944c29d6d3b7b5ea042110b4", size = 1672151, upload-time = "2025-05-17T17:21:22.666Z" }, + { url = "https://files.pythonhosted.org/packages/d4/5e/63f5cbde2342b7f70a39e591dbe75d9809d6338ce0b07c10406f1a140cdc/pycryptodome-3.23.0-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:14e15c081e912c4b0d75632acd8382dfce45b258667aa3c67caf7a4d4c13f630", size = 1664461, upload-time = "2025-05-17T17:21:25.225Z" }, + { url = "https://files.pythonhosted.org/packages/d6/92/608fbdad566ebe499297a86aae5f2a5263818ceeecd16733006f1600403c/pycryptodome-3.23.0-pp310-pypy310_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:a7fc76bf273353dc7e5207d172b83f569540fc9a28d63171061c42e361d22353", size = 1702440, upload-time = "2025-05-17T17:21:27.991Z" }, + { url = "https://files.pythonhosted.org/packages/d1/92/2eadd1341abd2989cce2e2740b4423608ee2014acb8110438244ee97d7ff/pycryptodome-3.23.0-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:45c69ad715ca1a94f778215a11e66b7ff989d792a4d63b68dc586a1da1392ff5", size = 1803005, upload-time = "2025-05-17T17:21:31.37Z" }, +] + [[package]] name = "pydantic" version = "2.12.5" diff --git a/libs/cli/langgraph_cli/schemas.py b/libs/cli/langgraph_cli/schemas.py index 011ce4752..1a1669e6f 100644 --- a/libs/cli/langgraph_cli/schemas.py +++ b/libs/cli/langgraph_cli/schemas.py @@ -128,7 +128,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 @@ -148,7 +148,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 } } @@ -328,8 +355,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. diff --git a/libs/cli/schemas/schema.json b/libs/cli/schemas/schema.json index 4ba2700e8..a29d82027 100644 --- a/libs/cli/schemas/schema.json +++ b/libs/cli/schemas/schema.json @@ -608,7 +608,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", diff --git a/libs/cli/schemas/schema.v0.json b/libs/cli/schemas/schema.v0.json index 4ba2700e8..a29d82027 100644 --- a/libs/cli/schemas/schema.v0.json +++ b/libs/cli/schemas/schema.v0.json @@ -608,7 +608,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", diff --git a/libs/langgraph/Makefile b/libs/langgraph/Makefile index 3b406baf4..e2c57e19d 100644 --- a/libs/langgraph/Makefile +++ b/libs/langgraph/Makefile @@ -87,15 +87,15 @@ integration_tests: WORKERS ?= auto XDIST_ARGS := $(if $(WORKERS),-n $(WORKERS) --dist worksteal,) -MAXFAIL ?= -MAXFAIL_ARGS := $(if $(MAXFAIL),--maxfail $(MAXFAIL),) +MAXFAIL ?= 1 +MAXFAIL_ARGS = $(if $(MAXFAIL),--maxfail $(MAXFAIL),) # Add an '-x' if xdist is enabled XDIST_ARGS := $(if $(WORKERS),-x $(XDIST_ARGS),) test_watch: make start-services &&\ make start-dev-server &&\ - uv run ptw . -- --ff -vv $(XDIST_ARGS) $(MAXFAIL_ARGS) $(TEST); \ + uv run ptw -- --ff -vv $(XDIST_ARGS) $(MAXFAIL_ARGS) $(TEST); \ EXIT_CODE=$$?; \ make stop-services; \ make stop-dev-server; \ @@ -130,7 +130,7 @@ type: format format_diff: uv run ruff format $(PYTHON_FILES) - uv run ruff check --select I --fix $(PYTHON_FILES) + uv run ruff check --fix $(PYTHON_FILES) spell_check: uv run codespell --toml pyproject.toml diff --git a/libs/langgraph/bench/__main__.py b/libs/langgraph/bench/__main__.py index 0bfa16ad1..d824abe5f 100644 --- a/libs/langgraph/bench/__main__.py +++ b/libs/langgraph/bench/__main__.py @@ -10,6 +10,7 @@ from bench.fanout_to_subgraph import fanout_to_subgraph, fanout_to_subgraph_sync from bench.pydantic_state import pydantic_state from bench.react_agent import react_agent from bench.sequential import create_sequential +from bench.serde_allowlist import collect_allowlist_large, collect_allowlist_small from bench.wide_dict import wide_dict from bench.wide_state import wide_state from langgraph.graph import StateGraph @@ -513,3 +514,7 @@ compilation_benchmarks = ( for name, graph in compilation_benchmarks: r.bench_func(name + "_compilation", compile_graph, graph) + +# Serde allowlist collection +r.bench_func("serde_allowlist_small", collect_allowlist_small) +r.bench_func("serde_allowlist_large", collect_allowlist_large) diff --git a/libs/langgraph/bench/serde_allowlist.py b/libs/langgraph/bench/serde_allowlist.py new file mode 100644 index 000000000..4278b9af5 --- /dev/null +++ b/libs/langgraph/bench/serde_allowlist.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass +from enum import Enum +from typing import Annotated + +from pydantic import BaseModel +from typing_extensions import NotRequired, TypedDict + +from langgraph._internal._serde import collect_allowlist_from_schemas + + +class Color(Enum): + RED = "red" + BLUE = "blue" + + +@dataclass +class InnerDataclass: + value: int + + +class InnerModel(BaseModel): + name: str + + +class InnerTyped(TypedDict): + payload: InnerDataclass + optional: NotRequired[InnerModel] + + +@dataclass +class Node: + value: int + child: Node | None = None + + +@dataclass +class NestedDataclass: + inner: InnerDataclass + items: list[InnerModel] + mapping: dict[str, InnerDataclass] + optional: InnerModel | None + union: InnerDataclass | InnerModel + queue: deque[InnerDataclass] + frozen: frozenset[InnerModel] + + +AnnotatedList = Annotated[list[InnerDataclass], "meta"] + + +class DummyChannel: + @property + def ValueType(self) -> type[InnerDataclass]: + return InnerDataclass + + @property + def UpdateType(self) -> type[InnerModel]: + return InnerModel + + +SCHEMAS_SMALL = [InnerDataclass, InnerModel, Color] +SCHEMAS_LARGE = [ + InnerDataclass, + InnerModel, + Color, + InnerTyped, + Node, + NestedDataclass, + AnnotatedList, +] +CHANNELS = {"a": DummyChannel(), "b": DummyChannel()} + + +def collect_allowlist_small() -> None: + collect_allowlist_from_schemas(schemas=SCHEMAS_SMALL, channels=CHANNELS) + + +def collect_allowlist_large() -> None: + collect_allowlist_from_schemas(schemas=SCHEMAS_LARGE, channels=CHANNELS) diff --git a/libs/langgraph/langgraph/_internal/_serde.py b/libs/langgraph/langgraph/_internal/_serde.py new file mode 100644 index 000000000..775242a87 --- /dev/null +++ b/libs/langgraph/langgraph/_internal/_serde.py @@ -0,0 +1,253 @@ +from __future__ import annotations + +import dataclasses +import logging +import sys +import types +from collections import deque +from enum import Enum +from typing import ( + Annotated, + Any, + Literal, + Union, + get_args, + get_origin, + get_type_hints, +) + +from langchain_core import messages as lc_messages +from langgraph.checkpoint.base import BaseCheckpointSaver +from pydantic import BaseModel +from typing_extensions import NotRequired, Required, is_typeddict + +try: + from langgraph.checkpoint.serde._msgpack import ( # noqa: F401 + STRICT_MSGPACK_ENABLED, + ) +except ImportError: + STRICT_MSGPACK_ENABLED = False + +_warned_allowlist_unsupported = False + +logger = logging.getLogger(__name__) + + +def _supports_checkpointer_allowlist() -> bool: + return hasattr(BaseCheckpointSaver, "with_allowlist") + + +_SUPPORTS_ALLOWLIST = _supports_checkpointer_allowlist() + + +def apply_checkpointer_allowlist( + checkpointer: Any, allowlist: set[tuple[str, ...]] | None +) -> Any: + if not checkpointer or allowlist is None or checkpointer in (True, False): + return checkpointer + if not _SUPPORTS_ALLOWLIST: + global _warned_allowlist_unsupported + if not _warned_allowlist_unsupported: + logger.warning( + "Checkpointer does not support with_allowlist; strict msgpack " + "allowlist will be skipped." + ) + _warned_allowlist_unsupported = True + return checkpointer + return checkpointer.with_allowlist(allowlist) + + +def curated_core_allowlist() -> set[tuple[str, ...]]: + allowlist: set[tuple[str, ...]] = set() + for name in ( + "BaseMessage", + "BaseMessageChunk", + "HumanMessage", + "HumanMessageChunk", + "AIMessage", + "AIMessageChunk", + "SystemMessage", + "SystemMessageChunk", + "ChatMessage", + "ChatMessageChunk", + "ToolMessage", + "ToolMessageChunk", + "FunctionMessage", + "FunctionMessageChunk", + "RemoveMessage", + ): + cls = getattr(lc_messages, name, None) + if cls is None: + continue + allowlist.add((cls.__module__, cls.__name__)) + + return allowlist + + +def build_serde_allowlist( + *, + schemas: list[type[Any]] | None = None, + channels: dict[str, Any] | None = None, +) -> set[tuple[str, ...]]: + allowlist = curated_core_allowlist() + if schemas: + schemas = [schema for schema in schemas if schema is not None] + return allowlist | collect_allowlist_from_schemas( + schemas=schemas, + channels=channels, + ) + + +def collect_allowlist_from_schemas( + *, + schemas: list[type[Any]] | None = None, + channels: dict[str, Any] | None = None, +) -> set[tuple[str, ...]]: + allowlist: set[tuple[str, ...]] = set() + seen: set[Any] = set() + seen_ids: set[int] = set() + + if schemas: + for schema in schemas: + _collect_from_type(schema, allowlist, seen, seen_ids) + + if channels: + for channel in channels.values(): + value_type = getattr(channel, "ValueType", None) + if value_type is not None: + _collect_from_type(value_type, allowlist, seen, seen_ids) + update_type = getattr(channel, "UpdateType", None) + if update_type is not None: + _collect_from_type(update_type, allowlist, seen, seen_ids) + + return allowlist + + +def _collect_from_type( + typ: Any, + allowlist: set[tuple[str, ...]], + seen: set[Any], + seen_ids: set[int], +) -> None: + if _already_seen(typ, seen, seen_ids): + return + + if typ is Any or typ is None: + return + + if typ is Literal: + return + + if isinstance(typ, types.UnionType): + for arg in typ.__args__: + _collect_from_type(arg, allowlist, seen, seen_ids) + return + + origin = get_origin(typ) + if origin is Union: + for arg in get_args(typ): + _collect_from_type(arg, allowlist, seen, seen_ids) + return + if origin is Annotated or origin in (Required, NotRequired): + args = get_args(typ) + if args: + _collect_from_type(args[0], allowlist, seen, seen_ids) + return + + if origin is Literal: + return + + if origin in (list, set, tuple, dict, deque, frozenset): + for arg in get_args(typ): + _collect_from_type(arg, allowlist, seen, seen_ids) + return + + if hasattr(typ, "__supertype__"): + _collect_from_type(typ.__supertype__, allowlist, seen, seen_ids) + return + + if is_typeddict(typ): + for field_type in _safe_get_type_hints(typ).values(): + _collect_from_type(field_type, allowlist, seen, seen_ids) + return + + if _is_pydantic_model(typ): + allowlist.add((typ.__module__, typ.__name__)) + field_types = _safe_get_type_hints(typ) + if field_types: + for field_type in field_types.values(): + _collect_from_type(field_type, allowlist, seen, seen_ids) + else: + for field_type in _pydantic_field_types(typ): + _collect_from_type(field_type, allowlist, seen, seen_ids) + return + + if dataclasses.is_dataclass(typ): + if typ_name := getattr(typ, "__name__", None): + allowlist.add((typ.__module__, typ_name)) + field_types = _safe_get_type_hints(typ) + if field_types: + for field_type in field_types.values(): + _collect_from_type(field_type, allowlist, seen, seen_ids) + else: + for field in dataclasses.fields(typ): + _collect_from_type(field.type, allowlist, seen, seen_ids) + return + + if isinstance(typ, type) and issubclass(typ, Enum): + allowlist.add((typ.__module__, typ.__name__)) + return + + +def _already_seen(typ: Any, seen: set[Any], seen_ids: set[int]) -> bool: + try: + if typ in seen: + return True + seen.add(typ) + return False + except TypeError: + typ_id = id(typ) + if typ_id in seen_ids: + return True + seen_ids.add(typ_id) + return False + + +def _safe_get_type_hints(typ: Any) -> dict[str, Any]: + try: + module = sys.modules.get(getattr(typ, "__module__", "")) + globalns = module.__dict__ if module else None + localns = dict(vars(typ)) if hasattr(typ, "__dict__") else None + return get_type_hints( + typ, globalns=globalns, localns=localns, include_extras=True + ) + except Exception: + return {} + + +def _is_pydantic_model(typ: Any) -> bool: + if not isinstance(typ, type): + return False + if issubclass(typ, BaseModel): + return True + try: + from pydantic.v1 import BaseModel as BaseModelV1 + except Exception: + return False + return issubclass(typ, BaseModelV1) + + +def _pydantic_field_types(typ: type[Any]) -> list[Any]: + if hasattr(typ, "model_fields"): + return [ + field.annotation + for field in typ.model_fields.values() + if getattr(field, "annotation", None) is not None + ] + if hasattr(typ, "__fields__"): + return [ + field.outer_type_ + for field in typ.__fields__.values() + if getattr(field, "outer_type_", None) is not None + ] + return [] diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index d1587077d..c7443e0a2 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -20,6 +20,7 @@ from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.store.base import BaseStore from typing_extensions import Unpack +from langgraph._internal import _serde from langgraph._internal._constants import CACHE_NS_WRITES, PREVIOUS from langgraph._internal._typing import MISSING, DeprecatedKwargs from langgraph.channels.ephemeral_value import EphemeralValue @@ -528,7 +529,7 @@ class entrypoint(Generic[ContextT]): else: output_type = save_type = sig.return_annotation - return Pregel( + graph: Pregel[Any, ContextT, Any, Any] = Pregel( nodes={ func.__name__: PregelNode( bound=bound, @@ -559,5 +560,16 @@ class entrypoint(Generic[ContextT]): cache=self.cache, cache_policy=self.cache_policy, retry_policy=self.retry_policy or (), - context_schema=self.context_schema, # type: ignore[arg-type] + context_schema=self.context_schema, ) + if _serde.STRICT_MSGPACK_ENABLED: + serde_allowlist = _serde.build_serde_allowlist( + schemas=[input_type, output_type, save_type] + + ([self.context_schema] if self.context_schema is not None else []), + channels=graph.channels, + ) + graph._serde_allowlist = serde_allowlist + graph.checkpointer = _serde.apply_checkpointer_allowlist( + graph.checkpointer, serde_allowlist + ) + return graph diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 4d0c90457..a87c17527 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -29,6 +29,7 @@ from langgraph.store.base import BaseStore from pydantic import BaseModel, TypeAdapter from typing_extensions import NotRequired, Required, Self, Unpack, is_typeddict +from langgraph._internal import _serde from langgraph._internal._constants import ( INTERRUPT, NS_END, @@ -1079,6 +1080,28 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): CompiledStateGraph: The compiled `StateGraph`. """ checkpointer = ensure_valid_checkpointer(checkpointer) + serde_allowlist: set[tuple[str, ...]] | None = None + if _serde.STRICT_MSGPACK_ENABLED: + schema_types: list[type[Any]] = [ + self.state_schema, + self.input_schema, + self.output_schema, + ] + if self.context_schema is not None: + schema_types.append(self.context_schema) + for node in self.nodes.values(): + schema_types.append(node.input_schema) + for branches in self.branches.values(): + for branch in branches.values(): + if branch.input_schema is not None: + schema_types.append(branch.input_schema) + serde_allowlist = _serde.build_serde_allowlist( + schemas=schema_types, + channels=self.channels, + ) + checkpointer = _serde.apply_checkpointer_allowlist( + checkpointer, serde_allowlist + ) # assign default values interrupt_before = interrupt_before or [] @@ -1135,6 +1158,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): cache=cache, name=name or "LangGraph", ) + compiled._serde_allowlist = serde_allowlist compiled.attach_node(START, None) for key, node in self.nodes.items(): diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 37e8125f9..ae6a0014d 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -48,6 +48,7 @@ from langgraph.store.base import BaseStore from pydantic import BaseModel, TypeAdapter from typing_extensions import Self, Unpack, deprecated, is_typeddict +from langgraph._internal import _serde from langgraph._internal._config import ( ensure_config, merge_configs, @@ -698,9 +699,17 @@ class Pregel( self.config = config self.trigger_to_nodes = trigger_to_nodes or {} self.name = name + self._serde_allowlist: set[tuple[str, ...]] | None = None if auto_validate: self.validate() + def _apply_checkpointer_allowlist( + self, checkpointer: BaseCheckpointSaver | None + ) -> BaseCheckpointSaver | None: + if not _serde.STRICT_MSGPACK_ENABLED: + return checkpointer + return _serde.apply_checkpointer_allowlist(checkpointer, self._serde_allowlist) + def get_graph( self, config: RunnableConfig | None = None, *, xray: int | bool = False ) -> Graph: @@ -1239,6 +1248,8 @@ class Pregel( checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) + if isinstance(checkpointer, BaseCheckpointSaver): + checkpointer = self._apply_checkpointer_allowlist(checkpointer) if not checkpointer: raise ValueError("No checkpointer set") @@ -1281,6 +1292,8 @@ class Pregel( checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) + if isinstance(checkpointer, BaseCheckpointSaver): + checkpointer = self._apply_checkpointer_allowlist(checkpointer) if not checkpointer: raise ValueError("No checkpointer set") @@ -1329,6 +1342,8 @@ class Pregel( checkpointer: BaseCheckpointSaver | None = config[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) + if isinstance(checkpointer, BaseCheckpointSaver): + checkpointer = self._apply_checkpointer_allowlist(checkpointer) if not checkpointer: raise ValueError("No checkpointer set") @@ -1380,6 +1395,8 @@ class Pregel( checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) + if isinstance(checkpointer, BaseCheckpointSaver): + checkpointer = self._apply_checkpointer_allowlist(checkpointer) if not checkpointer: raise ValueError("No checkpointer set") @@ -1446,6 +1463,8 @@ class Pregel( checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) + if isinstance(checkpointer, BaseCheckpointSaver): + checkpointer = self._apply_checkpointer_allowlist(checkpointer) if not checkpointer: raise ValueError("No checkpointer set") @@ -1890,6 +1909,8 @@ class Pregel( checkpointer: BaseCheckpointSaver | None = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) + if isinstance(checkpointer, BaseCheckpointSaver): + checkpointer = self._apply_checkpointer_allowlist(checkpointer) if not checkpointer: raise ValueError("No checkpointer set") @@ -2378,6 +2399,8 @@ class Pregel( raise RuntimeError("checkpointer=True cannot be used for root graphs.") else: checkpointer = self.checkpointer + if isinstance(checkpointer, BaseCheckpointSaver): + checkpointer = self._apply_checkpointer_allowlist(checkpointer) if checkpointer and not config.get(CONF): raise ValueError( "Checkpointer requires one or more of the following 'configurable' " diff --git a/libs/langgraph/tests/conftest_checkpointer.py b/libs/langgraph/tests/conftest_checkpointer.py index 99c0686fc..b74e55923 100644 --- a/libs/langgraph/tests/conftest_checkpointer.py +++ b/libs/langgraph/tests/conftest_checkpointer.py @@ -1,3 +1,4 @@ +import os from contextlib import asynccontextmanager, contextmanager from uuid import uuid4 @@ -5,6 +6,7 @@ import pytest from langgraph.checkpoint.postgres import PostgresSaver from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver from langgraph.checkpoint.serde.encrypted import EncryptedSerializer +from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer from langgraph.checkpoint.sqlite import SqliteSaver from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver from psycopg import AsyncConnection, Connection @@ -18,30 +20,60 @@ from tests.memory_assert import ( # noqa: E402 ) DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5442/" +STRICT_MSGPACK = os.getenv("LANGGRAPH_STRICT_MSGPACK", "false").lower() in ( + "1", + "true", + "yes", +) + + +def _strict_msgpack_serde() -> JsonPlusSerializer: + return JsonPlusSerializer(allowed_msgpack_modules=None) + + +def _apply_strict_msgpack(checkpointer) -> None: + if not STRICT_MSGPACK: + return + serde = _strict_msgpack_serde() + if hasattr(checkpointer, "serde"): + checkpointer.serde = serde + if hasattr(checkpointer, "saver") and hasattr(checkpointer.saver, "serde"): + checkpointer.saver.serde = serde @contextmanager def _checkpointer_memory(): - yield MemorySaverAssertImmutable() + if STRICT_MSGPACK: + yield MemorySaverAssertImmutable(serde=_strict_msgpack_serde()) + else: + yield MemorySaverAssertImmutable() @contextmanager def _checkpointer_memory_migrate_sends(): - yield MemorySaverNeedsPendingSendsMigration() + checkpointer = MemorySaverNeedsPendingSendsMigration() + _apply_strict_msgpack(checkpointer) + yield checkpointer @contextmanager def _checkpointer_sqlite(): with SqliteSaver.from_conn_string(":memory:") as checkpointer: + _apply_strict_msgpack(checkpointer) yield checkpointer @contextmanager def _checkpointer_sqlite_aes(): with SqliteSaver.from_conn_string(":memory:") as checkpointer: - checkpointer.serde = EncryptedSerializer.from_pycryptodome_aes( - key=b"1234567890123456" - ) + if STRICT_MSGPACK: + checkpointer.serde = EncryptedSerializer.from_pycryptodome_aes( + serde=_strict_msgpack_serde(), key=b"1234567890123456" + ) + else: + checkpointer.serde = EncryptedSerializer.from_pycryptodome_aes( + key=b"1234567890123456" + ) yield checkpointer @@ -57,6 +89,7 @@ def _checkpointer_postgres(): DEFAULT_POSTGRES_URI + database ) as checkpointer: checkpointer.setup() + _apply_strict_msgpack(checkpointer) yield checkpointer finally: # drop unique db @@ -79,6 +112,7 @@ def _checkpointer_postgres_pipe(): # setup can't run inside pipeline because of implicit transaction with checkpointer.conn.pipeline() as pipe: checkpointer.pipe = pipe + _apply_strict_msgpack(checkpointer) yield checkpointer finally: # drop unique db @@ -99,6 +133,7 @@ def _checkpointer_postgres_pool(): ) as pool: checkpointer = PostgresSaver(pool) checkpointer.setup() + _apply_strict_msgpack(checkpointer) yield checkpointer finally: # drop unique db @@ -109,6 +144,7 @@ def _checkpointer_postgres_pool(): @asynccontextmanager async def _checkpointer_sqlite_aio(): async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer: + _apply_strict_msgpack(checkpointer) yield checkpointer @@ -126,6 +162,7 @@ async def _checkpointer_postgres_aio(): DEFAULT_POSTGRES_URI + database ) as checkpointer: await checkpointer.setup() + _apply_strict_msgpack(checkpointer) yield checkpointer finally: # drop unique db @@ -152,6 +189,7 @@ async def _checkpointer_postgres_aio_pipe(): # setup can't run inside pipeline because of implicit transaction async with checkpointer.conn.pipeline() as pipe: checkpointer.pipe = pipe + _apply_strict_msgpack(checkpointer) yield checkpointer finally: # drop unique db @@ -176,6 +214,7 @@ async def _checkpointer_postgres_aio_pool(): ) as pool: checkpointer = AsyncPostgresSaver(pool) await checkpointer.setup() + _apply_strict_msgpack(checkpointer) yield checkpointer finally: # drop unique db diff --git a/libs/langgraph/tests/test_pydantic.py b/libs/langgraph/tests/test_pydantic.py index 87049204f..f49a02871 100644 --- a/libs/langgraph/tests/test_pydantic.py +++ b/libs/langgraph/tests/test_pydantic.py @@ -8,6 +8,7 @@ import uuid from enum import Enum from typing import Annotated, Literal, Optional +from langgraph.checkpoint.base import BaseCheckpointSaver from pydantic import ( BaseModel, ByteSize, @@ -23,7 +24,10 @@ from pydantic import ( from langgraph._internal._pydantic import is_supported_by_pydantic from langgraph.constants import END, START +from langgraph.func import entrypoint, task from langgraph.graph.state import StateGraph +from langgraph.types import Command, Interrupt, interrupt +from tests.any_str import AnyStr def test_is_supported_by_pydantic() -> None: @@ -312,3 +316,47 @@ def test_pydantic_state_field_validator(): g = builder.compile() res = g.invoke(input_state) assert res["text"] == "Hello, Validated John!" + + +class FunctionalState(BaseModel): + a: str + b: str | None = None + + +def test_interrupt_functional_pydantic(sync_checkpointer: BaseCheckpointSaver) -> None: + called_count = 0 + + @task + def foo(state: FunctionalState) -> FunctionalState: + nonlocal called_count + called_count += 1 + return FunctionalState(**{"a": state.a + "foo"}) + + @task + def bar(state: FunctionalState) -> dict: + return {"a": state.a + "bar", "b": state.b} + + @entrypoint(checkpointer=sync_checkpointer) + def graph(inputs: FunctionalState) -> FunctionalState: + fut_foo = foo(inputs) + value = interrupt("Provide value for bar:") + foo_res = fut_foo.result() + assert isinstance(foo_res, FunctionalState) + bar_input = FunctionalState(a=foo_res.a, b=value) + fut_bar = bar(bar_input) + return fut_bar.result() + + config = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + assert graph.invoke(FunctionalState(a=""), config) == { + "__interrupt__": [ + Interrupt( + value="Provide value for bar:", + id=AnyStr(), + ) + ] + } + # Resume with an answer + res = graph.invoke(Command(resume="bar"), config) + assert res == {"a": "foobar", "b": "bar"} + assert called_count == 1 diff --git a/libs/langgraph/tests/test_serde_allowlist.py b/libs/langgraph/tests/test_serde_allowlist.py new file mode 100644 index 000000000..2a90389da --- /dev/null +++ b/libs/langgraph/tests/test_serde_allowlist.py @@ -0,0 +1,159 @@ +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass +from enum import Enum +from typing import TYPE_CHECKING, Annotated, Any, Literal, NewType, Optional, Union + +import pytest +from pydantic import BaseModel +from typing_extensions import NotRequired, Required, TypedDict + +from langgraph._internal._serde import ( + collect_allowlist_from_schemas, + curated_core_allowlist, +) + + +class Color(Enum): + RED = "red" + BLUE = "blue" + + +@dataclass +class InnerDataclass: + value: int + + +class InnerModel(BaseModel): + name: str + + +@dataclass +class Node: + value: int + child: Node | None = None + + +if TYPE_CHECKING: + + class MissingType: + pass + + +@dataclass +class MissingRefDataclass: + payload: MissingType + + +class Payload(TypedDict): + item: InnerDataclass + maybe: NotRequired[InnerModel] + required: Required[str] + + +@dataclass +class NestedDataclass: + inner: InnerDataclass + items: list[InnerModel] + mapping: dict[str, InnerDataclass] + optional: InnerModel | None + union: InnerDataclass | InnerModel + queue: deque[InnerDataclass] + frozen: frozenset[InnerModel] + + +AnnotatedList = Annotated[list[InnerDataclass], "meta"] +UserId = NewType("UserId", int) + + +class DummyChannel: + @property + def ValueType(self) -> type[InnerDataclass]: + return InnerDataclass + + @property + def UpdateType(self) -> type[InnerModel]: + return InnerModel + + +def test_curated_core_allowlist_includes_messages() -> None: + try: + from langchain_core.messages import BaseMessage + except Exception: + pytest.skip("langchain_core not available") + allowlist = curated_core_allowlist() + assert (BaseMessage.__module__, BaseMessage.__name__) in allowlist + + +def test_collect_allowlist_basic_models() -> None: + allowlist = collect_allowlist_from_schemas( + schemas=[InnerDataclass, InnerModel, Color] + ) + assert (InnerDataclass.__module__, InnerDataclass.__name__) in allowlist + assert (InnerModel.__module__, InnerModel.__name__) in allowlist + assert (Color.__module__, Color.__name__) in allowlist + + +def test_collect_allowlist_nested_containers() -> None: + allowlist = collect_allowlist_from_schemas(schemas=[NestedDataclass]) + assert (NestedDataclass.__module__, NestedDataclass.__name__) in allowlist + assert (InnerDataclass.__module__, InnerDataclass.__name__) in allowlist + assert (InnerModel.__module__, InnerModel.__name__) in allowlist + + +def test_collect_allowlist_annotated_and_union() -> None: + allowlist = collect_allowlist_from_schemas( + schemas=[AnnotatedList, InnerModel | None, InnerDataclass | None] + ) + assert (InnerDataclass.__module__, InnerDataclass.__name__) in allowlist + assert (InnerModel.__module__, InnerModel.__name__) in allowlist + + +def test_collect_allowlist_literal_and_any() -> None: + allowlist = collect_allowlist_from_schemas(schemas=[Any, Literal["a"]]) + assert allowlist == set() + + +def test_collect_allowlist_typeddict_fields_only() -> None: + allowlist = collect_allowlist_from_schemas(schemas=[Payload]) + assert (InnerDataclass.__module__, InnerDataclass.__name__) in allowlist + assert (InnerModel.__module__, InnerModel.__name__) in allowlist + assert (Payload.__module__, Payload.__name__) not in allowlist + + +def test_collect_allowlist_forward_refs() -> None: + allowlist = collect_allowlist_from_schemas(schemas=[Node]) + assert (Node.__module__, Node.__name__) in allowlist + + +def test_collect_allowlist_missing_forward_ref() -> None: + allowlist = collect_allowlist_from_schemas(schemas=[MissingRefDataclass]) + assert allowlist == {(MissingRefDataclass.__module__, MissingRefDataclass.__name__)} + + +def test_collect_allowlist_newtype_supertype() -> None: + allowlist = collect_allowlist_from_schemas(schemas=[UserId]) + assert allowlist == set() + + +def test_collect_allowlist_channels() -> None: + channels = {"a": DummyChannel(), "b": DummyChannel()} + allowlist = collect_allowlist_from_schemas(channels=channels) + assert (InnerDataclass.__module__, InnerDataclass.__name__) in allowlist + assert (InnerModel.__module__, InnerModel.__name__) in allowlist + + +def test_collect_allowlist_pep604_union() -> None: + schema = InnerDataclass | InnerModel + allowlist = collect_allowlist_from_schemas(schemas=[schema]) + assert (InnerDataclass.__module__, InnerDataclass.__name__) in allowlist + assert (InnerModel.__module__, InnerModel.__name__) in allowlist + + +def test_collect_allowlist_typing_union_optional() -> None: + typing_optional = Optional[InnerDataclass] # noqa: UP045 + typing_union = Union[InnerDataclass, InnerModel] # noqa: UP007 + allowlist = collect_allowlist_from_schemas(schemas=[typing_optional, typing_union]) + assert (InnerDataclass.__module__, InnerDataclass.__name__) in allowlist + assert (InnerModel.__module__, InnerModel.__name__) in allowlist diff --git a/libs/langgraph/uv.lock b/libs/langgraph/uv.lock index a5186b591..9b75fe18d 100644 --- a/libs/langgraph/uv.lock +++ b/libs/langgraph/uv.lock @@ -1569,6 +1569,7 @@ dev = [ { name = "numpy" }, { name = "pandas" }, { name = "pandas-stubs", specifier = ">=2.2.2.240807" }, + { name = "pycryptodome", specifier = ">=3.23.0" }, { name = "pytest" }, { name = "pytest-asyncio" }, { name = "pytest-mock" }, diff --git a/libs/prebuilt/uv.lock b/libs/prebuilt/uv.lock index 109c99461..351084784 100644 --- a/libs/prebuilt/uv.lock +++ b/libs/prebuilt/uv.lock @@ -373,6 +373,7 @@ dev = [ { name = "numpy" }, { name = "pandas" }, { name = "pandas-stubs", specifier = ">=2.2.2.240807" }, + { name = "pycryptodome", specifier = ">=3.23.0" }, { name = "pytest" }, { name = "pytest-asyncio" }, { name = "pytest-mock" },